Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_dynamic_launch_plan_yielding_of_constant_workflow():
outputs = lp_yield_empty_wf.unit_test()
# TODO: Currently, Flytekit will not return early and not do anything if there are any workflow nodes detected
# in the output of a dynamic task.
dj_spec = outputs[_sdk_constants.FUTURES_FILE_NAME]
assert len(dj_spec.nodes) == 1
assert len(dj_spec.outputs) == 1
assert dj_spec.outputs[0].var == "out"
assert len(outputs.keys()) == 1
def test_dynamic_launch_plan_yielding():
outputs = lp_yield_task.unit_test(num=10)
# TODO: Currently, Flytekit will not return early and not do anything if there are any workflow nodes detected
# in the output of a dynamic task.
dj_spec = outputs[_sdk_constants.FUTURES_FILE_NAME]
assert dj_spec.min_successes == 1
launch_plan_node = dj_spec.nodes[0]
node_id = launch_plan_node.id
assert "models-test-dynamic-wfs-id-lp" in node_id
assert node_id.endswith("-0")
# Assert that the output of the dynamic job spec is bound to the single node in the spec, the workflow node
# containing the launch plan
assert dj_spec.outputs[0].var == "out"
assert dj_spec.outputs[0].binding.promise.node_id == node_id
assert dj_spec.outputs[0].binding.promise.var == "task_output"
{'a': 9}, _type_map_from_variable_map(_task_defs.add_one.interface.inputs))
input_file = os.path.join(input_dir.name, "inputs.pb")
_utils.write_proto_to_file(literal_map.to_flyte_idl(), input_file)
with _utils.AutoDeletingTempDir("out") as output_dir:
_execute_task(
_task_defs.add_one.task_module,
_task_defs.add_one.task_function_name,
input_file,
output_dir.name,
False
)
p = _utils.load_proto_from_file(
_literals_pb2.LiteralMap,
os.path.join(output_dir.name, _constants.OUTPUT_FILE_NAME)
)
raw_map = _type_helpers.unpack_literal_map_to_sdk_python_std(
_literal_models.LiteralMap.from_flyte_idl(p),
_type_map_from_variable_map(_task_defs.add_one.interface.outputs)
)
assert raw_map['b'] == 10
assert len(raw_map) == 1
def test_default_python_task():
assert isinstance(default_task, _sdk_runnable.SdkRunnableTask)
assert default_task.interface.inputs['in1'].description == ''
assert default_task.interface.inputs['in1'].type == \
_type_models.LiteralType(simple=_type_models.SimpleType.INTEGER)
assert default_task.interface.outputs['out1'].description == ''
assert default_task.interface.outputs['out1'].type == \
_type_models.LiteralType(simple=_type_models.SimpleType.STRING)
assert default_task.type == _common_constants.SdkTaskType.PYTHON_TASK
assert default_task.task_function_name == 'default_task'
assert default_task.task_module == __name__
assert default_task.metadata.timeout == _datetime.timedelta(seconds=0)
assert default_task.metadata.deprecated_error_message == ''
assert default_task.metadata.discoverable is False
assert default_task.metadata.discovery_version == ''
assert default_task.metadata.retries.retries == 0
assert len(default_task.container.resources.limits) == 0
assert len(default_task.container.resources.requests) == 0
),
"__implicit_schema": _interface_model.Variable(
type=_types.LiteralType(simple=_types.SimpleType.STRING),
description="The schema set as an implicit input"
)
},
{
# Set the schema for the Presto query as an output
"results": _interface_model.Variable(
type=_types.LiteralType(schema=output_schema.schema_type),
description="The schema for the Presto query"
)
})
super(SdkPrestoTask, self).__init__(
_constants.SdkTaskType.PRESTO_TASK,
metadata,
i,
_MessageToDict(presto_query.to_flyte_idl()),
)
# Set user provided inputs
task_inputs(self)
outputs = ' '.join(p['outputs'][0]['data']['text/plain'])
if outputs is not None:
dict = _literal_models._literals_pb2.LiteralMap()
_text_format.Parse(outputs, dict)
# Add output_notebook as an output to the task.
output_notebook = _task_output.OutputReference(
_type_helpers.get_sdk_type_from_literal_type(_Types.Blob.to_flyte_literal_type()))
output_notebook.set(output_notebook_path)
output_literal_map = _literal_models.LiteralMap.from_flyte_idl(dict)
output_literal_map.literals[OUTPUT_NOTEBOOK] = output_notebook.sdk_value
return {
_constants.OUTPUT_FILE_NAME: output_literal_map
}
_task_models.RuntimeMetadata(
_task_models.RuntimeMetadata.RuntimeType.FLYTE_SDK, __version__,
"python"),
timeout or _datetime.timedelta(seconds=0),
_literals.RetryStrategy(retries),
interruptible,
discovery_version,
None
)
# The interface is defined using the inputs and outputs
i = _interface.TypedInterface(inputs=types_to_variable(inputs), outputs=types_to_variable(outputs))
# This sets the base SDKTask with container etc
super(SdkRawContainerTask, self).__init__(
_constants.SdkTaskType.RAW_CONTAINER_TASK,
metadata,
i,
None,
container=_get_container_definition(
image=image,
args=args,
command=command,
data_loading_config=self._data_loading_config,
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,
def wrapper(fn):
return (SensorTask or cls)(
task_function=fn,
task_type=_common_constants.SdkTaskType.SENSOR_TASK,
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,
timeout=timeout,
environment=environment,
custom={},
discovery_version='',
discoverable=False,
tags={
'exec_project': _internal_config.EXECUTION_PROJECT.get(),
'exec_domain': _internal_config.EXECUTION_DOMAIN.get(),
'exec_workflow': _internal_config.EXECUTION_WORKFLOW.get(),
'exec_launchplan': _internal_config.EXECUTION_LAUNCHPLAN.get(),
'api_version': _api_version
}
),
logging=_logging,
tmp_dir=task_dir
),
inputs
)
except _exception_scopes.FlyteScopedException as e:
_logging.error("!!! Begin Error Captured by Flyte !!!")
output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument(
_error_models.ContainerError(
e.error_code,
e.verbose_message,
e.kind
)
)
_logging.error(e.verbose_message)
_logging.error("!!! End Error Captured by Flyte !!!")
except Exception:
_logging.error("!!! Begin Unknown System Error Captured by Flyte !!!")
exc_str = _traceback.format_exc()
output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument(
_error_models.ContainerError(
"SYSTEM:Unknown",
exc_str,
_error_models.ContainerError.Kind.RECOVERABLE
def wrapper(fn):
# Just saving everything as a hash for now, will figure out what to do with this in the future.
task_obj = {}
task_obj['task_type'] = _common_constants.SdkTaskType.PYTHON_TASK,
task_obj['retries'] = retries,
task_obj['storage_request'] = storage_request,
task_obj['cpu_request'] = cpu_request,
task_obj['gpu_request'] = gpu_request,
task_obj['memory_request'] = memory_request,
task_obj['storage_limit'] = storage_limit,
task_obj['cpu_limit'] = cpu_limit,
task_obj['gpu_limit'] = gpu_limit,
task_obj['memory_limit'] = memory_limit,
task_obj['environment'] = environment,
task_obj['custom'] = {}
metadata = _task_model.TaskMetadata(
cache,
_task_model.RuntimeMetadata(
_task_model.RuntimeMetadata.RuntimeType.FLYTE_SDK,