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