def test_basic_unit_test():
def add_one(wf_params, value_in, value_out):
value_out.set(value_in + 1)
t = sdk_runnable.SdkRunnableTask(
add_one,
_common_constants.SdkTaskType.PYTHON_TASK,
"1",
1,
None,
None,
None,
None,
None,
None,
None,
None,
None,
False,
None,
{},
None,
)
t.add_inputs({'value_in': interface.Variable(primitives.Integer.to_flyte_literal_type(), "")})
t.add_outputs({'value_out': interface.Variable(primitives.Integer.to_flyte_literal_type(), "")})
out = t.unit_test(value_in=1)
assert out['value_out'] == 2
def test_basic_unit_test():
def add_one(wf_params, value_in, value_out):
value_out.set(value_in + 1)
t = sdk_runnable.SdkRunnableTask(
add_one,
_common_constants.SdkTaskType.PYTHON_TASK,
"1",
1,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
False,
None,
{},
None,
)
t.add_inputs({'value_in': interface.Variable(primitives.Integer.to_flyte_literal_type(), "")})
t.add_outputs({'value_out': interface.Variable(primitives.Integer.to_flyte_literal_type(), "")})
out = t.unit_test(value_in=1)
assert out['value_out'] == 2
with _pytest.raises(_user_exceptions.FlyteAssertion) as e:
t()
assert "value_in" in str(e.value)
assert "INTEGER" in str(e.value)
def get_sdk_node(
self,
pipeline_context,
instance,
pipeline_run,
step_key,
task_type=constants.SdkTaskType.PYTHON_TASK,
cache_version="",
retries=0,
interruptible=False,
deprecated="",
storage_request=None,
cpu_request=None,
gpu_request=None,
memory_request=None,
storage_limit=None,
cpu_limit=None,
gpu_limit=None,
memory_limit=None,
cache=False,
timeout=datetime.timedelta(seconds=0),
environment=None,
):
execution_step = self.execution_plan.get_step_by_key(step_key)
flyte_inputs = self.flyte_inputs(execution_step.step_input_dict,
execution_step.solid_name)
flyte_outputs = self.flyte_outputs(execution_step.step_output_dict,
execution_step.solid_name)
def wrapper(wf_params, *args, **kwargs): # pylint: disable=unused-argument
# TODO: We can't update config values via inputs from Flyte, because they are immutable
plan = self.execution_plan.build_subset_plan([step_key])
for param, arg in kwargs.items():
self.inject_intermediates(pipeline_context, execution_step,
param, arg)
results = list(
execute_plan(
plan,
instance,
run_config=self.run_config,
pipeline_run=pipeline_run,
))
for result in results:
step_context = pipeline_context.for_step(execution_step)
self.output_value(step_context, step_key, result,
execution_step, kwargs)
# This will take the wrapper definition and re-create it with explicit parameters as keyword argumentss
wrapper = forge.sign(forge.arg("wf_params"),
*map(forge.arg, flyte_inputs.keys()),
*map(forge.arg, flyte_outputs.keys()))(wrapper)
# flytekit uses this name for an internal representation, make it unique to the step key
wrapper.__name__ = execution_step.solid_name
task = sdk_runnable.SdkRunnableTask(
task_function=wrapper,
task_type=task_type,
discovery_version=cache_version,
retries=retries,
interruptible=interruptible,
deprecated=deprecated,
storage_request=storage_request,
cpu_request=cpu_request,
gpu_request=gpu_request,
memory_request=memory_request,
storage_limit=storage_limit,
cpu_limit=cpu_limit,
gpu_limit=gpu_limit,
memory_limit=memory_limit,
discoverable=cache,
timeout=timeout,
environment=environment,
custom={},
)
if flyte_inputs:
task = inputs(task, **flyte_inputs)
if flyte_outputs:
task = outputs(task, **flyte_outputs)
return task