def compile(self, ctx: FlyteContext, *args, **kwargs):
return create_and_link_node(
ctx,
entity=self,
interface=self.python_interface,
**kwargs,
)
def compile(self, ctx: FlyteContext, *args, **kwargs):
return create_and_link_node(
ctx,
entity=self,
interface=self.python_interface,
timeout=self.metadata.timeout,
retry_strategy=self.metadata.retry_strategy,
**kwargs,
)
def test_create_and_link_node():
@task
def t1(a: typing.Union[int, typing.List[int]]) -> typing.Union[int, typing.List[int]]:
return a
with pytest.raises(FlyteAssertion, match="Cannot create node when not compiling..."):
ctx = context_manager.FlyteContext.current_context()
create_and_link_node(ctx, t1, a=3)
ctx = context_manager.FlyteContext.current_context().with_compilation_state(CompilationState(prefix=""))
p = create_and_link_node(ctx, t1, a=3)
assert p.ref.node_id == "n0"
assert p.ref.var == "o0"
assert len(p.ref.node.bindings) == 1
@task
def t2(a: typing.Optional[int] = None) -> typing.Union[int]:
return a
p = create_and_link_node(ctx, t2)
assert p.ref.var == "o0"
assert len(p.ref.node.bindings) == 0
def __call__(self, *args, **kwargs):
if len(args) > 0:
raise AssertionError("Only Keyword Arguments are supported for launch plan executions")
ctx = FlyteContext.current_context()
if ctx.compilation_state is not None:
inputs = self.saved_inputs
inputs.update(kwargs)
return create_and_link_node(ctx, entity=self, interface=self.workflow._native_interface, **inputs)
else:
# Calling a launch plan should just forward the call to the workflow, nothing more. But let's add in the
# saved inputs.
inputs = self.saved_inputs
inputs.update(kwargs)
return self.workflow(*args, **inputs)
def compile(self, ctx: FlyteContext, *args, **kwargs):
return create_and_link_node(ctx, entity=self, **kwargs)
def __call__(self, *args, **kwargs):
"""
The call pattern for Workflows is close to, but not exactly, the call pattern for Tasks. For local execution,
it goes
__call__ -> _local_execute -> execute
From execute, different things happen for the two Workflow styles. For PythonFunctionWorkflows, the Python
function is run, for the ImperativeWorkflow, each node is run one at a time.
"""
if len(args) > 0:
raise AssertionError("Only Keyword Arguments are supported for Workflow executions")
ctx = FlyteContextManager.current_context()
# Get default agruements and override with kwargs passed in
input_kwargs = self.python_interface.default_inputs_as_kwargs
input_kwargs.update(kwargs)
# The first condition is compilation.
if ctx.compilation_state is not None:
return create_and_link_node(ctx, entity=self, interface=self.python_interface, **input_kwargs)
# This condition is hit when this workflow (self) is being called as part of a parent's workflow local run.
# The context specifying the local workflow execution has already been set.
elif (
ctx.execution_state is not None and ctx.execution_state.mode == ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION
):
if ctx.execution_state.branch_eval_mode == BranchEvalMode.BRANCH_SKIPPED:
if self.python_interface and self.python_interface.output_tuple_name:
variables = [k for k in self.python_interface.outputs.keys()]
output_tuple = collections.namedtuple(self.python_interface.output_tuple_name, variables)
nones = [None for _ in self.python_interface.outputs.keys()]
return output_tuple(*nones)
else:
return None
# We are already in a local execution, just continue the execution context
return self._local_execute(ctx, **input_kwargs)
# Last is starting a local workflow execution
else:
# Run some sanity checks
# Even though the _local_execute call generally expects inputs to be Promises, we don't have to do the
# conversion here in this loop. The reason is because we don't prevent users from specifying inputs
# as direct scalars, which means there's another Promise-generating loop inside _local_execute too
for k, v in input_kwargs.items():
if k not in self.interface.inputs:
raise ValueError(f"Received unexpected keyword argument {k}")
if isinstance(v, Promise):
raise ValueError(f"Received a promise for a workflow call, when expecting a native value for {k}")
result = None
with FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.new_execution_state().with_params(mode=ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION)
)
) as child_ctx:
result = self._local_execute(child_ctx, **input_kwargs)
expected_outputs = len(self.python_interface.outputs)
if expected_outputs == 0:
if result is None or isinstance(result, VoidPromise):
return None
else:
raise Exception(f"Workflow local execution expected 0 outputs but something received {result}")
if (1 < expected_outputs == len(result)) or (result is not None and expected_outputs == 1):
if isinstance(result, Promise):
v = [v for k, v in self.python_interface.outputs.items()][0] # get output native type
return TypeEngine.to_python_value(ctx, result.val, v)
else:
for prom in result:
if not isinstance(prom, Promise):
raise Exception("should be promises")
native_list = [
TypeEngine.to_python_value(ctx, promise.val, self.python_interface.outputs[promise.var])
for promise in result
]
return tuple(native_list)
raise ValueError("expected outputs and actual outputs do not match")
def __call__(self, *args, **kwargs):
if len(args) > 0:
raise AssertionError(
"Only Keyword Arguments are supported for Workflow executions")
ctx = FlyteContext.current_context()
# Handle subworkflows in compilation
if ctx.compilation_state is not None:
input_kwargs = self._native_interface.default_inputs_as_kwargs
input_kwargs.update(kwargs)
return create_and_link_node(ctx,
entity=self,
interface=self._native_interface,
**input_kwargs)
elif (ctx.execution_state is not None and ctx.execution_state.mode
== ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION):
# We are already in a local execution, just continue the execution context
return self._local_execute(ctx, **kwargs)
# When someone wants to run the workflow function locally. Assume that the inputs given are given as Python
# native values. _local_execute will always translate Python native literals to Flyte literals, so no worries
# there, but it'll return Promise objects.
else:
# Run some sanity checks
# Even though the _local_execute call generally expects inputs to be Promises, we don't have to do the
# conversion here in this loop. The reason is because we don't prevent users from specifying inputs
# as direct scalars, which means there's another Promise-generating loop inside _local_execute too
for k, v in kwargs.items():
if k not in self.interface.inputs:
raise ValueError(
f"Received unexpected keyword argument {k}")
if isinstance(v, Promise):
raise ValueError(
f"Received a promise for a workflow call, when expecting a native value for {k}"
)
with ctx.new_execution_context(
mode=ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION) as ctx:
result = self._local_execute(ctx, **kwargs)
expected_outputs = len(self._native_interface.outputs)
if expected_outputs == 0:
if result is None or isinstance(result, VoidPromise):
return None
else:
raise Exception(
f"Workflow local execution expected 0 outputs but something received {result}"
)
if (expected_outputs > 1 and len(result) == expected_outputs) or (
expected_outputs == 1 and result is not None):
if isinstance(result, Promise):
v = [v
for k, v in self._native_interface.outputs.items()][0]
return TypeEngine.to_python_value(ctx, result.val, v)
else:
for prom in result:
if not isinstance(prom, Promise):
raise Exception("should be promises")
native_list = [
TypeEngine.to_python_value(
ctx, promise.val,
self._native_interface.outputs[promise.var])
for promise in result
]
return tuple(native_list)
raise ValueError(
"expected outputs and actual outputs do not match")
def compile(self, ctx: FlyteContext, *args, **kwargs):
"""
Generates a node that encapsulates this task in a workflow definition.
"""
return create_and_link_node(ctx, entity=self, **kwargs)
