def test_single_step_entrypoint_in_proc():
with _TemporaryConfiguration(os.path.join(os.path.dirname(__file__), 'fake.config'),
internal_overrides={
'project': 'test',
'domain': 'development'
}):
with _utils.AutoDeletingTempDir("in") as input_dir:
literal_map = _type_helpers.pack_python_std_map_to_literal_map(
{'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_single_step_entrypoint_out_of_proc():
with _TemporaryConfiguration(os.path.join(os.path.dirname(__file__), 'fake.config'),
internal_overrides={
'project': 'test',
'domain': 'development'
}):
with _utils.AutoDeletingTempDir("in") as input_dir:
literal_map = _type_helpers.pack_python_std_map_to_literal_map({'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:
cmd = []
cmd.extend(["--task-module", _task_defs.add_one.task_module])
cmd.extend(["--task-name", _task_defs.add_one.task_function_name])
cmd.extend(["--inputs", input_file])
cmd.extend(["--output-prefix", output_dir.name])
result = CliRunner().invoke(execute_task_cmd, cmd)
assert result.exit_code == 0
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_arrayjob_entrypoint_in_proc():
with _TemporaryConfiguration(os.path.join(os.path.dirname(__file__),
'fake.config'),
internal_overrides={
'project': 'test',
'domain': 'development'
}):
with _utils.AutoDeletingTempDir("dir") as dir:
literal_map = _type_helpers.pack_python_std_map_to_literal_map(
{'a': 9},
_type_map_from_variable_map(
_task_defs.add_one.interface.inputs))
input_dir = os.path.join(dir.name, "1")
os.mkdir(
input_dir) # auto cleanup will take this subdir into account
input_file = os.path.join(input_dir, "inputs.pb")
_utils.write_proto_to_file(literal_map.to_flyte_idl(), input_file)
# construct indexlookup.pb which has array: [1]
mapped_index = _literals.Literal(
_literals.Scalar(primitive=_literals.Primitive(integer=1)))
index_lookup_collection = _literals.LiteralCollection(
[mapped_index])
index_lookup_file = os.path.join(dir.name, "indexlookup.pb")
_utils.write_proto_to_file(index_lookup_collection.to_flyte_idl(),
index_lookup_file)
# fake arrayjob task by setting environment variables
orig_env_index_var_name = os.environ.get(
'BATCH_JOB_ARRAY_INDEX_VAR_NAME')
orig_env_array_index = os.environ.get('AWS_BATCH_JOB_ARRAY_INDEX')
os.environ[
'BATCH_JOB_ARRAY_INDEX_VAR_NAME'] = 'AWS_BATCH_JOB_ARRAY_INDEX'
os.environ['AWS_BATCH_JOB_ARRAY_INDEX'] = '0'
execute_task(_task_defs.add_one.task_module,
_task_defs.add_one.task_function_name, dir.name,
dir.name, False)
raw_map = _type_helpers.unpack_literal_map_to_sdk_python_std(
_literal_models.LiteralMap.from_flyte_idl(
_utils.load_proto_from_file(
_literals_pb2.LiteralMap,
os.path.join(input_dir, _constants.OUTPUT_FILE_NAME))),
_type_map_from_variable_map(
_task_defs.add_one.interface.outputs))
assert raw_map['b'] == 10
assert len(raw_map) == 1
# reset the env vars
if orig_env_index_var_name:
os.environ[
'BATCH_JOB_ARRAY_INDEX_VAR_NAME'] = orig_env_index_var_name
if orig_env_array_index:
os.environ['AWS_BATCH_JOB_ARRAY_INDEX'] = orig_env_array_index
def _python_std_input_map_to_literal_map(self, inputs):
"""
:param dict[Text,Any] inputs: A dictionary of Python standard inputs that will be type-checked and compiled
to a LiteralMap
:rtype: flytekit.models.literals.LiteralMap
"""
return _type_helpers.pack_python_std_map_to_literal_map(inputs, {
k: _type_helpers.get_sdk_type_from_literal_type(v.type)
for k, v in _six.iteritems(self.interface.inputs)
})
def _python_std_input_map_to_literal_map(self, inputs):
"""
:param dict[Text,Any] inputs: A dictionary of Python standard inputs that will be type-checked and compiled
to a LiteralMap
:rtype: flytekit.models.literals.LiteralMap
"""
return _type_helpers.pack_python_std_map_to_literal_map(
inputs,
{k: user_input.sdk_type for k, user_input in _six.iteritems(self.default_inputs.parameters) if k in inputs},
)
def unit_test(self, **input_map):
"""
:param dict[Text, T] input_map: Python Std input from users. We will cast these to the appropriate Flyte
literals.
:returns: Depends on the behavior of the specific task in the unit engine.
"""
return _engine_loader.get_engine('unit').get_task(self).execute(
_type_helpers.pack_python_std_map_to_literal_map(
input_map, {
k: _type_helpers.get_sdk_type_from_literal_type(v.type)
for k, v in _six.iteritems(self.interface.inputs)
}))
def local_execute(self, **input_map):
"""
:param dict[Text, T] input_map: Python Std input from users. We will cast these to the appropriate Flyte
literals.
:rtype: dict[Text, T]
:returns: The output produced by this task in Python standard format.
"""
return _engine_loader.get_engine('local').get_task(self).execute(
_type_helpers.pack_python_std_map_to_literal_map(
input_map, {
k: _type_helpers.get_sdk_type_from_literal_type(v.type)
for k, v in _six.iteritems(self.interface.inputs)
}))
def __init__(
self,
sdk_workflow,
default_inputs=None,
fixed_inputs=None,
role=None,
schedule=None,
notifications=None,
labels=None,
annotations=None,
auth_role=None,
raw_output_data_config=None,
):
"""
:param flytekit.common.local_workflow.SdkRunnableWorkflow sdk_workflow:
:param dict[Text,flytekit.common.promise.Input] default_inputs:
:param dict[Text,Any] fixed_inputs: These inputs will be fixed and not need to be set when executing this
launch plan.
:param Text role: Deprecated. IAM role to execute this launch plan with.
:param flytekit.models.schedule.Schedule: Schedule to apply to this workflow.
:param list[flytekit.models.common.Notification]: List of notifications to apply to this launch plan.
:param flytekit.models.common.Labels labels: Any custom kubernetes labels to apply to workflows executed by this
launch plan.
:param flytekit.models.common.Annotations annotations: Any custom kubernetes annotations to apply to workflows
executed by this launch plan.
Any custom kubernetes annotations to apply to workflows executed by this launch plan.
:param flytekit.models.common.Authrole auth_role: The auth method with which to execute the workflow.
:param flytekit.models.common.RawOutputDataConfig raw_output_data_config: Config for offloading data
"""
if role and auth_role:
raise ValueError(
"Cannot set both role and auth. Role is deprecated, use auth instead."
)
fixed_inputs = fixed_inputs or {}
default_inputs = default_inputs or {}
if role:
auth_role = _common_models.AuthRole(assumable_iam_role=role)
# The constructor for SdkLaunchPlan sets the id to None anyways so we don't bother passing in an ID. The ID
# should be set in one of three places,
# 1) When the object is registered (in the code above)
# 2) By the dynamic task code after this runnable object has already been __call__'ed. The SdkNode produced
# maintains a link to this object and will set the ID according to the configuration variables present.
# 3) When SdkLaunchPlan.fetch() is run
super(SdkRunnableLaunchPlan, self).__init__(
None,
_launch_plan_models.LaunchPlanMetadata(
schedule=schedule or _schedule_model.Schedule(""),
notifications=notifications or [],
),
_interface_models.ParameterMap(default_inputs),
_type_helpers.pack_python_std_map_to_literal_map(
fixed_inputs,
{
k: _type_helpers.get_sdk_type_from_literal_type(var.type)
for k, var in _six.iteritems(sdk_workflow.interface.inputs)
if k in fixed_inputs
},
),
labels or _common_models.Labels({}),
annotations or _common_models.Annotations({}),
auth_role,
raw_output_data_config or _common_models.RawOutputDataConfig(""),
)
self._interface = _interface.TypedInterface(
{k: v.var
for k, v in _six.iteritems(default_inputs)},
sdk_workflow.interface.outputs,
)
self._upstream_entities = {sdk_workflow}
self._sdk_workflow = sdk_workflow
def __init__(
self,
sdk_workflow,
default_inputs=None,
fixed_inputs=None,
role=None,
schedule=None,
notifications=None,
labels=None,
annotations=None,
auth=None,
):
"""
:param flytekit.common.workflow.SdkWorkflow sdk_workflow:
:param dict[Text,flytekit.common.promise.Input] default_inputs:
:param dict[Text,Any] fixed_inputs: These inputs will be fixed and not need to be set when executing this
launch plan.
:param Text role: Deprecated. IAM role to execute this launch plan with.
:param flytekit.models.schedule.Schedule: Schedule to apply to this workflow.
:param list[flytekit.models.common.Notification]: List of notifications to apply to this launch plan.
:param flytekit.models.common.Labels labels: Any custom kubernetes labels to apply to workflows executed by this
launch plan.
:param flytekit.models.common.Annotations annotations: Any custom kubernetes annotations to apply to workflows
executed by this launch plan.
Any custom kubernetes annotations to apply to workflows executed by this launch plan.
:param flytekit.models.launch_plan.Auth auth: The auth method with which to execute the workflow.
"""
if role and auth:
raise ValueError(
"Cannot set both role and auth. Role is deprecated, use auth instead."
)
fixed_inputs = fixed_inputs or {}
default_inputs = default_inputs or {}
if role:
auth = _launch_plan_models.Auth(assumable_iam_role=role)
super(SdkRunnableLaunchPlan, self).__init__(
_identifier.Identifier(_identifier_model.ResourceType.WORKFLOW,
_internal_config.PROJECT.get(),
_internal_config.DOMAIN.get(),
_uuid.uuid4().hex,
_internal_config.VERSION.get()),
_launch_plan_models.LaunchPlanMetadata(
schedule=schedule or _schedule_model.Schedule(''),
notifications=notifications or []),
_interface_models.ParameterMap(default_inputs),
_type_helpers.pack_python_std_map_to_literal_map(
fixed_inputs, {
k: _type_helpers.get_sdk_type_from_literal_type(var.type)
for k, var in _six.iteritems(sdk_workflow.interface.inputs)
if k in fixed_inputs
}),
labels or _common_models.Labels({}),
annotations or _common_models.Annotations({}),
auth,
)
self._interface = _interface.TypedInterface(
{k: v.var
for k, v in _six.iteritems(default_inputs)},
sdk_workflow.interface.outputs)
self._upstream_entities = {sdk_workflow}
self._sdk_workflow = sdk_workflow
