def _local_execute(self, ctx: FlyteContext, **kwargs) -> Union[Tuple[Promise], Promise, VoidPromise]:
"""
This code is used only in the case when we want to dispatch_execute with outputs from a previous node
For regular execution, dispatch_execute is invoked directly.
"""
# Unwrap the kwargs values. After this, we essentially have a LiteralMap
# The reason why we need to do this is because the inputs during local execute can be of 2 types
# - Promises or native constants
# Promises as essentially inputs from previous task executions
# native constants are just bound to this specific task (default values for a task input)
# Also alongwith promises and constants, there could be dictionary or list of promises or constants
kwargs = translate_inputs_to_literals(
ctx, input_kwargs=kwargs, interface=self.interface, native_input_types=self.get_input_types()
)
input_literal_map = _literal_models.LiteralMap(literals=kwargs)
outputs_literal_map = self.dispatch_execute(ctx, input_literal_map)
outputs_literals = outputs_literal_map.literals
# TODO maybe this is the part that should be done for local execution, we pass the outputs to some special
# location, otherwise we dont really need to right? The higher level execute could just handle literalMap
# After running, we again have to wrap the outputs, if any, back into Promise objects
output_names = list(self.interface.outputs.keys())
if len(output_names) != len(outputs_literals):
# Length check, clean up exception
raise AssertionError(f"Length difference {len(output_names)} {len(outputs_literals)}")
# Tasks that don't return anything still return a VoidPromise
if len(output_names) == 0:
return VoidPromise(self.name)
vals = [Promise(var, outputs_literals[var]) for var in output_names]
return create_task_output(vals, self.python_interface)
def _local_execute(self, ctx: FlyteContext, **kwargs) -> Union[Tuple[Promise], Promise, VoidPromise]:
"""
Please see the _local_execute comments in the main task.
"""
# Unwrap the kwargs values. After this, we essentially have a LiteralMap
# The reason why we need to do this is because the inputs during local execute can be of 2 types
# - Promises or native constants
# Promises as essentially inputs from previous task executions
# native constants are just bound to this specific task (default values for a task input)
# Also alongwith promises and constants, there could be dictionary or list of promises or constants
kwargs = translate_inputs_to_literals(
ctx,
incoming_values=kwargs,
flyte_interface_types=self.interface.inputs,
native_types=self.python_interface.inputs,
)
input_literal_map = _literal_models.LiteralMap(literals=kwargs)
outputs_literal_map = self.unwrap_literal_map_and_execute(ctx, input_literal_map)
# After running, we again have to wrap the outputs, if any, back into Promise objects
outputs_literals = outputs_literal_map.literals
output_names = list(self.python_interface.outputs.keys())
if len(output_names) != len(outputs_literals):
# Length check, clean up exception
raise AssertionError(f"Length difference {len(output_names)} {len(outputs_literals)}")
# Tasks that don't return anything still return a VoidPromise
if len(output_names) == 0:
return VoidPromise(self.name)
vals = [Promise(var, outputs_literals[var]) for var in output_names]
return create_task_output(vals, self.python_interface)
def _local_execute(self, ctx: FlyteContext, **kwargs) -> Union[Tuple[Promise], Promise, VoidPromise]:
# This is done to support the invariant that Workflow local executions always work with Promise objects
# holding Flyte literal values. Even in a wf, a user can call a sub-workflow with a Python native value.
for k, v in kwargs.items():
if not isinstance(v, Promise):
t = self.python_interface.inputs[k]
kwargs[k] = Promise(var=k, val=TypeEngine.to_literal(ctx, v, t, self.interface.inputs[k].type))
# The output of this will always be a combination of Python native values and Promises containing Flyte
# Literals.
function_outputs = self.execute(**kwargs)
# First handle the empty return case.
# A workflow function may return a task that doesn't return anything
# def wf():
# return t1()
# or it may not return at all
# def wf():
# t1()
# In the former case we get the task's VoidPromise, in the latter we get None
if isinstance(function_outputs, VoidPromise) or function_outputs is None:
if len(self.python_interface.outputs) != 0:
raise FlyteValueException(
function_outputs,
f"{function_outputs} received but interface has {len(self.python_interface.outputs)} outputs.",
)
return VoidPromise(self.name)
# Because we should've already returned in the above check, we just raise an error here.
if len(self.python_interface.outputs) == 0:
raise FlyteValueException(
function_outputs, f"{function_outputs} received but should've been VoidPromise or None."
)
expected_output_names = list(self.python_interface.outputs.keys())
if len(expected_output_names) == 1:
# Here we have to handle the fact that the wf could've been declared with a typing.NamedTuple of
# length one. That convention is used for naming outputs - and single-length-NamedTuples are
# particularly troublesome but elegant handling of them is not a high priority
# Again, we're using the output_tuple_name as a proxy.
if self.python_interface.output_tuple_name and isinstance(function_outputs, tuple):
wf_outputs_as_map = {expected_output_names[0]: function_outputs[0]}
else:
wf_outputs_as_map = {expected_output_names[0]: function_outputs}
else:
wf_outputs_as_map = {expected_output_names[i]: function_outputs[i] for i, _ in enumerate(function_outputs)}
# Basically we need to repackage the promises coming from the tasks into Promises that match the workflow's
# interface. We do that by extracting out the literals, and creating new Promises
wf_outputs_as_literal_dict = translate_inputs_to_literals(
ctx,
wf_outputs_as_map,
flyte_interface_types=self.interface.outputs,
native_types=self.python_interface.outputs,
)
# Recreate new promises that use the workflow's output names.
new_promises = [Promise(var, wf_outputs_as_literal_dict[var]) for var in expected_output_names]
return create_task_output(new_promises, self.python_interface)
def _compute_outputs(
self,
n: Node) -> Optional[Union[Promise, Tuple[Promise], VoidPromise]]:
curr = self.compute_output_vars()
if curr is None or len(curr) == 0:
return VoidPromise(n.id, NodeOutput(node=n, var="placeholder"))
promises = [
Promise(var=x, val=NodeOutput(node=n, var=x)) for x in curr
]
# TODO: Is there a way to add the Python interface here? Currently, it's an optional arg.
return create_task_output(promises)
def end_branch(
self
) -> Optional[Union[Condition, Tuple[Promise], Promise, VoidPromise]]:
"""
This should be invoked after every branch has been visited
"""
if self._last_case:
FlyteContextManager.pop_context()
curr = self.compute_output_vars()
if curr is None:
return VoidPromise(self.name)
promises = [Promise(var=x, val=None) for x in curr]
return create_task_output(promises)
return self._condition
def _compute_outputs(
self, selected_output_promise
) -> Optional[Union[Tuple[Promise], Promise, VoidPromise]]:
"""
For the local execution case only returns the least common set of outputs
"""
curr = self.compute_output_vars()
if curr is None:
return VoidPromise(self.name)
if not isinstance(selected_output_promise, tuple):
selected_output_promise = (selected_output_promise, )
promises = [
Promise(var=x, val=v.val)
for x, v in zip(curr, selected_output_promise)
]
return create_task_output(promises)
def _compute_outputs(self, n: Node) -> Union[Promise, Tuple[Promise], VoidPromise]:
output_var_sets: typing.List[typing.Set[str]] = []
for c in self._cases:
if c.output_promise is None and c.err is None:
# One node returns a void output and no error, we will default to a
# Void output
return VoidPromise(n.id)
if c.output_promise is not None:
if isinstance(c.output_promise, tuple):
output_var_sets.append(set([i.var for i in c.output_promise]))
else:
output_var_sets.append(set([c.output_promise.var]))
curr = output_var_sets[0]
if len(output_var_sets) > 1:
for x in output_var_sets[1:]:
curr = curr.intersection(x)
promises = [Promise(var=x, val=NodeOutput(node=n, var=x)) for x in curr]
# TODO: Is there a way to add the Python interface here? Currently, it's an optional arg.
return create_task_output(promises)
def _local_execute(
self, ctx: FlyteContext,
**kwargs) -> Union[Tuple[Promise], Promise, VoidPromise]:
"""
Performs local execution of a workflow. kwargs are expected to be Promises for the most part (unless,
someone has hardcoded in my_wf(input_1=5) or something).
:param ctx: The FlyteContext
:param kwargs: parameters for the workflow itself
"""
logger.info(
f"Executing Workflow {self._name}, ctx{ctx.execution_state.Mode}")
# This is done to support the invariant that Workflow local executions always work with Promise objects
# holding Flyte literal values. Even in a wf, a user can call a sub-workflow with a Python native value.
for k, v in kwargs.items():
if not isinstance(v, Promise):
t = self._native_interface.inputs[k]
kwargs[k] = Promise(var=k,
val=TypeEngine.to_literal(
ctx, v, t,
self.interface.inputs[k].type))
function_outputs = self.execute(**kwargs)
if (isinstance(function_outputs, VoidPromise)
or function_outputs is None
or len(self.python_interface.outputs) == 0):
# The reason this is here is because a workflow function may return a task that doesn't return anything
# def wf():
# return t1()
# or it may not return at all
# def wf():
# t1()
# In the former case we get the task's VoidPromise, in the latter we get None
return VoidPromise(self.name)
# TODO: Can we refactor the task code to be similar to what's in this function?
promises = _workflow_fn_outputs_to_promise(
ctx, self._native_interface.outputs, self.interface.outputs,
function_outputs)
# TODO: With the native interface, create_task_output should be able to derive the typed interface, and it
# should be able to do the conversion of the output of the execute() call directly.
return create_task_output(promises, self._native_interface)
def local_execute(self, ctx: FlyteContext,
**kwargs) -> Union[Tuple[Promise], Promise, VoidPromise]:
"""
This function is used only in the local execution path and is responsible for calling dispatch execute.
Use this function when calling a task with native values (or Promises containing Flyte literals derived from
Python native values).
"""
# Unwrap the kwargs values. After this, we essentially have a LiteralMap
# The reason why we need to do this is because the inputs during local execute can be of 2 types
# - Promises or native constants
# Promises as essentially inputs from previous task executions
# native constants are just bound to this specific task (default values for a task input)
# Also along with promises and constants, there could be dictionary or list of promises or constants
kwargs = translate_inputs_to_literals(
ctx,
incoming_values=kwargs,
flyte_interface_types=self.interface.inputs, # type: ignore
native_types=self.get_input_types(),
)
input_literal_map = _literal_models.LiteralMap(literals=kwargs)
# if metadata.cache is set, check memoized version
if self.metadata.cache:
# TODO: how to get a nice `native_inputs` here?
logger.info(
f"Checking cache for task named {self.name}, cache version {self.metadata.cache_version} "
f"and inputs: {input_literal_map}")
outputs_literal_map = LocalTaskCache.get(
self.name, self.metadata.cache_version, input_literal_map)
# The cache returns None iff the key does not exist in the cache
if outputs_literal_map is None:
logger.info("Cache miss, task will be executed now")
outputs_literal_map = self.dispatch_execute(
ctx, input_literal_map)
# TODO: need `native_inputs`
LocalTaskCache.set(self.name, self.metadata.cache_version,
input_literal_map, outputs_literal_map)
logger.info(
f"Cache set for task named {self.name}, cache version {self.metadata.cache_version} "
f"and inputs: {input_literal_map}")
else:
logger.info("Cache hit")
else:
es = ctx.execution_state
b = es.user_space_params.with_task_sandbox()
ctx = ctx.current_context().with_execution_state(
es.with_params(user_space_params=b.build())).build()
outputs_literal_map = self.dispatch_execute(ctx, input_literal_map)
outputs_literals = outputs_literal_map.literals
# TODO maybe this is the part that should be done for local execution, we pass the outputs to some special
# location, otherwise we dont really need to right? The higher level execute could just handle literalMap
# After running, we again have to wrap the outputs, if any, back into Promise objects
output_names = list(self.interface.outputs.keys()) # type: ignore
if len(output_names) != len(outputs_literals):
# Length check, clean up exception
raise AssertionError(
f"Length difference {len(output_names)} {len(outputs_literals)}"
)
# Tasks that don't return anything still return a VoidPromise
if len(output_names) == 0:
return VoidPromise(self.name)
vals = [Promise(var, outputs_literals[var]) for var in output_names]
return create_task_output(vals, self.python_interface)
