def unwrap_literal_map_and_execute(
self, ctx: FlyteContext, input_literal_map: _literal_models.LiteralMap
) -> Union[VoidPromise, _literal_models.LiteralMap,
_dynamic_job.DynamicJobSpec]:
"""
Please see the implementation of the dispatch_execute function in the real task.
"""
# Invoked before the task is executed
# Translate the input literals to Python native
native_inputs = TypeEngine.literal_map_to_kwargs(
ctx, input_literal_map, self.python_interface.inputs)
logger.info(f"Invoking {self.name} with inputs: {native_inputs}")
try:
native_outputs = self.execute(**native_inputs)
except Exception as e:
logger.exception(f"Exception when executing {e}")
raise e
logger.info(
f"Task executed successfully in user level, outputs: {native_outputs}"
)
expected_output_names = list(self.python_interface.outputs.keys())
if len(expected_output_names) == 1:
native_outputs_as_map = {expected_output_names[0]: native_outputs}
elif len(expected_output_names) == 0:
native_outputs_as_map = {}
else:
native_outputs_as_map = {
expected_output_names[i]: native_outputs[i]
for i, _ in enumerate(native_outputs)
}
# We manually construct a LiteralMap here because task inputs and outputs actually violate the assumption
# built into the IDL that all the values of a literal map are of the same type.
literals = {}
for k, v in native_outputs_as_map.items():
literal_type = self.interface.outputs[k].type
py_type = self.python_interface.outputs[k]
if isinstance(v, tuple):
raise AssertionError(
f"Output({k}) in task{self.name} received a tuple {v}, instead of {py_type}"
)
literals[k] = TypeEngine.to_literal(ctx, v, py_type, literal_type)
outputs_literal_map = _literal_models.LiteralMap(literals=literals)
# After the execute has been successfully completed
return outputs_literal_map
def dispatch_execute(
self, ctx: FlyteContext, input_literal_map: _literal_models.LiteralMap
) -> Union[_literal_models.LiteralMap, _dynamic_job.DynamicJobSpec]:
"""
This method translates Flyte's Type system based input values and invokes the actual call to the executor
This method is also invoked during runtime.
* ``VoidPromise`` is returned in the case when the task itself declares no outputs.
* ``Literal Map`` is returned when the task returns either one more outputs in the declaration. Individual outputs
may be none
* ``DynamicJobSpec`` is returned when a dynamic workflow is executed
"""
# Invoked before the task is executed
new_user_params = self.pre_execute(ctx.user_space_params)
# Create another execution context with the new user params, but let's keep the same working dir
with ctx.new_execution_context(
mode=ctx.execution_state.mode,
execution_params=new_user_params,
working_dir=ctx.execution_state.working_dir,
) as exec_ctx:
# TODO We could support default values here too - but not part of the plan right now
# Translate the input literals to Python native
native_inputs = TypeEngine.literal_map_to_kwargs(
exec_ctx, input_literal_map, self.python_interface.inputs)
# TODO: Logger should auto inject the current context information to indicate if the task is running within
# a workflow or a subworkflow etc
logger.info(f"Invoking {self.name} with inputs: {native_inputs}")
try:
native_outputs = self.execute(**native_inputs)
except Exception as e:
logger.exception(f"Exception when executing {e}")
raise e
logger.info(
f"Task executed successfully in user level, outputs: {native_outputs}"
)
# Lets run the post_execute method. This may result in a IgnoreOutputs Exception, which is
# bubbled up to be handled at the callee layer.
native_outputs = self.post_execute(new_user_params, native_outputs)
# Short circuit the translation to literal map because what's returned may be a dj spec (or an
# already-constructed LiteralMap if the dynamic task was a no-op), not python native values
if isinstance(native_outputs,
_literal_models.LiteralMap) or isinstance(
native_outputs, _dynamic_job.DynamicJobSpec):
return native_outputs
expected_output_names = list(self._outputs_interface.keys())
if len(expected_output_names) == 1:
# Here we have to handle the fact that the task 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(
native_outputs, tuple):
native_outputs_as_map = {
expected_output_names[0]: native_outputs[0]
}
else:
native_outputs_as_map = {
expected_output_names[0]: native_outputs
}
elif len(expected_output_names) == 0:
native_outputs_as_map = {}
else:
native_outputs_as_map = {
expected_output_names[i]: native_outputs[i]
for i, _ in enumerate(native_outputs)
}
# We manually construct a LiteralMap here because task inputs and outputs actually violate the assumption
# built into the IDL that all the values of a literal map are of the same type.
literals = {}
for k, v in native_outputs_as_map.items():
literal_type = self._outputs_interface[k].type
py_type = self.get_type_for_output_var(k, v)
if isinstance(v, tuple):
raise AssertionError(
f"Output({k}) in task{self.name} received a tuple {v}, instead of {py_type}"
)
try:
literals[k] = TypeEngine.to_literal(
exec_ctx, v, py_type, literal_type)
except Exception as e:
raise AssertionError(
f"failed to convert return value for var {k}") from e
outputs_literal_map = _literal_models.LiteralMap(literals=literals)
# After the execute has been successfully completed
return outputs_literal_map
def dispatch_execute(
self, ctx: FlyteContext, input_literal_map: _literal_models.LiteralMap
) -> Union[_literal_models.LiteralMap, _dynamic_job.DynamicJobSpec]:
"""
This function is largely similar to the base PythonTask, with the exception that we have to infer the Python
interface before executing. Also, we refer to ``self.task_template`` rather than just ``self`` similar to task
classes that derive from the base ``PythonTask``.
"""
# Invoked before the task is executed
new_user_params = self.pre_execute(ctx.user_space_params)
# Create another execution context with the new user params, but let's keep the same working dir
with FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.execution_state.with_params(
user_space_params=new_user_params))) as exec_ctx:
# Added: Have to reverse the Python interface from the task template Flyte interface
# See docstring for more details.
guessed_python_input_types = TypeEngine.guess_python_types(
self.task_template.interface.inputs)
native_inputs = TypeEngine.literal_map_to_kwargs(
exec_ctx, input_literal_map, guessed_python_input_types)
logger.info(
f"Invoking FlyteTask executor {self.task_template.id.name} with inputs: {native_inputs}"
)
try:
native_outputs = self.execute(**native_inputs)
except Exception as e:
logger.exception(f"Exception when executing {e}")
raise e
logger.debug("Task executed successfully in user level")
# Lets run the post_execute method. This may result in a IgnoreOutputs Exception, which is
# bubbled up to be handled at the callee layer.
native_outputs = self.post_execute(new_user_params, native_outputs)
# Short circuit the translation to literal map because what's returned may be a dj spec (or an
# already-constructed LiteralMap if the dynamic task was a no-op), not python native values
if isinstance(native_outputs,
_literal_models.LiteralMap) or isinstance(
native_outputs, _dynamic_job.DynamicJobSpec):
return native_outputs
expected_output_names = list(
self.task_template.interface.outputs.keys())
if len(expected_output_names) == 1:
# Here we have to handle the fact that the task 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.
# Deleted some stuff
native_outputs_as_map = {
expected_output_names[0]: native_outputs
}
elif len(expected_output_names) == 0:
native_outputs_as_map = {}
else:
native_outputs_as_map = {
expected_output_names[i]: native_outputs[i]
for i, _ in enumerate(native_outputs)
}
# We manually construct a LiteralMap here because task inputs and outputs actually violate the assumption
# built into the IDL that all the values of a literal map are of the same type.
literals = {}
for k, v in native_outputs_as_map.items():
literal_type = self.task_template.interface.outputs[k].type
py_type = type(v)
if isinstance(v, tuple):
raise AssertionError(
f"Output({k}) in task{self.task_template.id.name} received a tuple {v}, instead of {py_type}"
)
try:
literals[k] = TypeEngine.to_literal(
exec_ctx, v, py_type, literal_type)
except Exception as e:
raise AssertionError(
f"failed to convert return value for var {k}") from e
outputs_literal_map = _literal_models.LiteralMap(literals=literals)
# After the execute has been successfully completed
return outputs_literal_map