def test_task_template_security_context(sec_ctx): obj = task.TaskTemplate( identifier.Identifier(identifier.ResourceType.TASK, "project", "domain", "name", "version"), "python", parameterizers.LIST_OF_TASK_METADATA[0], parameterizers.LIST_OF_INTERFACES[0], { "a": 1, "b": { "c": 2, "d": 3 } }, container=task.Container( "my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], parameterizers.LIST_OF_RESOURCES[0], {"a": "b"}, {"d": "e"}, ), security_context=sec_ctx, ) assert obj.security_context == sec_ctx expected = obj.security_context if sec_ctx: if sec_ctx.run_as is None and sec_ctx.secrets is None and sec_ctx.tokens is None: expected = None assert task.TaskTemplate.from_flyte_idl( obj.to_flyte_idl()).security_context == expected
def test_task_template(in_tuple): task_metadata, interfaces, resources = in_tuple obj = task.TaskTemplate( identifier.Identifier(identifier.ResourceType.TASK, "project", "domain", "name", "version"), "python", task_metadata, interfaces, { 'a': 1, 'b': { 'c': 2, 'd': 3 } }, container=task.Container("my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {'a': 'b'}, {'d': 'e'})) assert obj.id.resource_type == identifier.ResourceType.TASK assert obj.id.project == "project" assert obj.id.domain == "domain" assert obj.id.name == "name" assert obj.id.version == "version" assert obj.type == "python" assert obj.metadata == task_metadata assert obj.interface == interfaces assert obj.custom == {'a': 1, 'b': {'c': 2, 'd': 3}} assert obj.container.image == "my_image" assert obj.container.resources == resources assert text_format.MessageToString( obj.to_flyte_idl()) == text_format.MessageToString( task.TaskTemplate.from_flyte_idl( obj.to_flyte_idl()).to_flyte_idl())
def test_container(resources): obj = task.Container("my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {'a': 'b'}, {'d': 'e'}) obj.image == "my_image" obj.command == ["this", "is", "a", "cmd"] obj.args == ["this", "is", "an", "arg"] obj.resources == resources obj.env == {'a': 'b'} obj.config == {'d': 'e'} assert obj == task.Container.from_flyte_idl(obj.to_flyte_idl())
def get_sample_container(): """ :rtype: flytekit.models.task.Container """ cpu_resource = _task_model.Resources.ResourceEntry( _task_model.Resources.ResourceName.CPU, "1") resources = _task_model.Resources(requests=[cpu_resource], limits=[cpu_resource]) return _task_model.Container("my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {}, {})
def _get_container_definition( self, environment=None, ): """ :rtype: Container """ args = [] for k, v in _six.iteritems(self.interface.inputs): args.append("--{}".format(k)) args.append("{{{{.Inputs.{}}}}}".format(k)) return _task_models.Container(image=_internal_config.IMAGE.get(), command=[], args=args, resources=_task_models.Resources([], []), env=environment, config={})
def test_task_template_security_context(sec_ctx): obj = task.TaskTemplate( identifier.Identifier(identifier.ResourceType.TASK, "project", "domain", "name", "version"), "python", parameterizers.LIST_OF_TASK_METADATA[0], parameterizers.LIST_OF_INTERFACES[0], {"a": 1, "b": {"c": 2, "d": 3}}, container=task.Container( "my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], parameterizers.LIST_OF_RESOURCES[0], {"a": "b"}, {"d": "e"}, ), security_context=sec_ctx, ) assert obj.security_context == sec_ctx assert text_format.MessageToString(obj.to_flyte_idl()) == text_format.MessageToString( task.TaskTemplate.from_flyte_idl(obj.to_flyte_idl()).to_flyte_idl() )
def test_workflow_closure(): int_type = _types.LiteralType(_types.SimpleType.INTEGER) typed_interface = _interface.TypedInterface( {'a': _interface.Variable(int_type, "description1")}, { 'b': _interface.Variable(int_type, "description2"), 'c': _interface.Variable(int_type, "description3") }) b0 = _literals.Binding( 'a', _literals.BindingData(scalar=_literals.Scalar( primitive=_literals.Primitive(integer=5)))) b1 = _literals.Binding( 'b', _literals.BindingData(promise=_types.OutputReference('my_node', 'b'))) b2 = _literals.Binding( 'b', _literals.BindingData(promise=_types.OutputReference('my_node', 'c'))) node_metadata = _workflow.NodeMetadata(name='node1', timeout=timedelta(seconds=10), retries=_literals.RetryStrategy(0)) task_metadata = _task.TaskMetadata( True, _task.RuntimeMetadata(_task.RuntimeMetadata.RuntimeType.FLYTE_SDK, "1.0.0", "python"), timedelta(days=1), _literals.RetryStrategy(3), "0.1.1b0", "This is deprecated!") cpu_resource = _task.Resources.ResourceEntry( _task.Resources.ResourceName.CPU, "1") resources = _task.Resources(requests=[cpu_resource], limits=[cpu_resource]) task = _task.TaskTemplate( _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "name", "version"), "python", task_metadata, typed_interface, { 'a': 1, 'b': { 'c': 2, 'd': 3 } }, container=_task.Container("my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {}, {})) task_node = _workflow.TaskNode(task.id) node = _workflow.Node(id='my_node', metadata=node_metadata, inputs=[b0], upstream_node_ids=[], output_aliases=[], task_node=task_node) template = _workflow.WorkflowTemplate( id=_identifier.Identifier(_identifier.ResourceType.WORKFLOW, "project", "domain", "name", "version"), metadata=_workflow.WorkflowMetadata(), interface=typed_interface, nodes=[node], outputs=[b1, b2], ) obj = _workflow_closure.WorkflowClosure(workflow=template, tasks=[task]) assert len(obj.tasks) == 1 obj2 = _workflow_closure.WorkflowClosure.from_flyte_idl(obj.to_flyte_idl()) assert obj == obj2
from flytekit.models.core import identifier as _identifier from flytekit.models.core import workflow as _workflow from tests.flytekit.common import parameterizers LIST_OF_DYNAMIC_TASKS = [ _task.TaskTemplate( _identifier.Identifier(_identifier.ResourceType.TASK, "p", "d", "n", "v"), "python", task_metadata, interfaces, _array_job.ArrayJob(2, 2, 2).to_dict(), container=_task.Container( "my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {"a": "b"}, {"d": "e"}, ), ) for task_metadata, interfaces, resources in product( parameterizers.LIST_OF_TASK_METADATA, parameterizers.LIST_OF_INTERFACES, parameterizers.LIST_OF_RESOURCES, ) ] @pytest.mark.parametrize("task", LIST_OF_DYNAMIC_TASKS) def test_future_task_document(task): rs = _literals.RetryStrategy(0) nm = _workflow.NodeMetadata("node-name", _timedelta(minutes=10), rs)
def _get_container_definition( image: str, command: List[str], args: List[str], data_loading_config: Optional[_task_models.DataLoadingConfig] = None, storage_request: Optional[str] = None, ephemeral_storage_request: Optional[str] = None, cpu_request: Optional[str] = None, gpu_request: Optional[str] = None, memory_request: Optional[str] = None, storage_limit: Optional[str] = None, ephemeral_storage_limit: Optional[str] = None, cpu_limit: Optional[str] = None, gpu_limit: Optional[str] = None, memory_limit: Optional[str] = None, environment: Optional[Dict[str, str]] = None, ) -> _task_models.Container: storage_limit = storage_limit storage_request = storage_request ephemeral_storage_limit = ephemeral_storage_limit ephemeral_storage_request = ephemeral_storage_request cpu_limit = cpu_limit cpu_request = cpu_request gpu_limit = gpu_limit gpu_request = gpu_request memory_limit = memory_limit memory_request = memory_request requests = [] if storage_request: requests.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.STORAGE, storage_request)) if ephemeral_storage_request: requests.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.EPHEMERAL_STORAGE, ephemeral_storage_request)) if cpu_request: requests.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.CPU, cpu_request)) if gpu_request: requests.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.GPU, gpu_request)) if memory_request: requests.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.MEMORY, memory_request)) limits = [] if storage_limit: limits.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.STORAGE, storage_limit)) if ephemeral_storage_limit: limits.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.EPHEMERAL_STORAGE, ephemeral_storage_limit)) if cpu_limit: limits.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.CPU, cpu_limit)) if gpu_limit: limits.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.GPU, gpu_limit)) if memory_limit: limits.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.MEMORY, memory_limit)) if environment is None: environment = {} return _task_models.Container( image=image, command=command, args=args, resources=_task_models.Resources(limits=limits, requests=requests), env=environment, config={}, data_loading_config=data_loading_config, )
"version"), "python", task_metadata, interfaces, { 'a': 1, 'b': [1, 2, 3], 'c': 'abc', 'd': { 'x': 1, 'y': 2, 'z': 3 } }, container=task.Container("my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {'a': 'b'}, {'d': 'e'})) for task_metadata, interfaces, resources in product( LIST_OF_TASK_METADATA, LIST_OF_INTERFACES, LIST_OF_RESOURCES) ] LIST_OF_CONTAINERS = [ task.Container("my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {'a': 'b'}, {'d': 'e'}) for resources in LIST_OF_RESOURCES ] LIST_OF_TASK_CLOSURES = [ task.TaskClosure(task.CompiledTask(template)) for template in LIST_OF_TASK_TEMPLATES
task_metadata, interfaces, { "a": 1, "b": [1, 2, 3], "c": "abc", "d": { "x": 1, "y": 2, "z": 3 } }, container=task.Container( "my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {"a": "b"}, {"d": "e"}, ), ) for task_metadata, interfaces, resources in product( LIST_OF_TASK_METADATA, LIST_OF_INTERFACES, LIST_OF_RESOURCES) ] LIST_OF_CONTAINERS = [ task.Container( "my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {"a": "b"}, {"d": "e"},
from flytekit.models import literals as _literals, dynamic_job as _dynamic_job, array_job as _array_job, \ task as _task from flytekit.models.core import workflow as _workflow, identifier as _identifier from tests.flytekit.common import parameterizers LIST_OF_DYNAMIC_TASKS = [ _task.TaskTemplate(_identifier.Identifier(_identifier.ResourceType.TASK, "p", "d", "n", "v"), "python", task_metadata, interfaces, _array_job.ArrayJob(2, 2, 2).to_dict(), container=_task.Container("my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {'a': 'b'}, {'d': 'e'})) for task_metadata, interfaces, resources in product( parameterizers.LIST_OF_TASK_METADATA, parameterizers.LIST_OF_INTERFACES, parameterizers.LIST_OF_RESOURCES) ] @pytest.mark.parametrize("task", LIST_OF_DYNAMIC_TASKS) def test_future_task_document(task): rs = _literals.RetryStrategy(0) nm = _workflow.NodeMetadata('node-name', _timedelta(minutes=10), rs) n = _workflow.Node(id="id", metadata=nm, inputs=[],
def _get_container_definition( image: str, command: List[str], args: List[str], data_loading_config: _task_models.DataLoadingConfig, storage_request: str = None, cpu_request: str = None, gpu_request: str = None, memory_request: str = None, storage_limit: str = None, cpu_limit: str = None, gpu_limit: str = None, memory_limit: str = None, environment: Dict[str, str] = None, ) -> _task_models.Container: storage_limit = storage_limit or _resource_config.DEFAULT_STORAGE_LIMIT.get( ) storage_request = storage_request or _resource_config.DEFAULT_STORAGE_REQUEST.get( ) cpu_limit = cpu_limit or _resource_config.DEFAULT_CPU_LIMIT.get() cpu_request = cpu_request or _resource_config.DEFAULT_CPU_REQUEST.get() gpu_limit = gpu_limit or _resource_config.DEFAULT_GPU_LIMIT.get() gpu_request = gpu_request or _resource_config.DEFAULT_GPU_REQUEST.get() memory_limit = memory_limit or _resource_config.DEFAULT_MEMORY_LIMIT.get() memory_request = memory_request or _resource_config.DEFAULT_MEMORY_REQUEST.get( ) requests = [] if storage_request: requests.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.STORAGE, storage_request)) if cpu_request: requests.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.CPU, cpu_request)) if gpu_request: requests.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.GPU, gpu_request)) if memory_request: requests.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.MEMORY, memory_request)) limits = [] if storage_limit: limits.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.STORAGE, storage_limit)) if cpu_limit: limits.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.CPU, cpu_limit)) if gpu_limit: limits.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.GPU, gpu_limit)) if memory_limit: limits.append( _task_models.Resources.ResourceEntry( _task_models.Resources.ResourceName.MEMORY, memory_limit)) if environment is None: environment = {} return _task_models.Container( image=image, command=command, args=args, resources=_task_models.Resources(limits=limits, requests=requests), env=environment, config={}, data_loading_config=data_loading_config, )
def _op_to_task(self, dag_id, image, op, node_map): """ Generate task given an operator inherited from dsl.ContainerOp. :param airflow.models.BaseOperator op: :param dict(Text, SdkNode) node_map: :rtype: Tuple(base_tasks.SdkTask, SdkNode) """ interface_inputs = {} interface_outputs = {} input_mappings = {} processed_args = None # for key, val in six.iteritems(op.params): # interface_inputs[key] = interface_model.Variable( # _type_helpers.python_std_to_sdk_type(Types.String).to_flyte_literal_type(), # '' # ) # # if param.op_name == '': # binding = promise_common.Input(sdk_type=Types.String, name=param.name) # else: # binding = promise_common.NodeOutput( # sdk_node=node_map[param.op_name], # sdk_type=Types.String, # var=param.name) # input_mappings[param.name] = binding # # for param in op.outputs.values(): # interface_outputs[param.name] = interface_model.Variable( # _type_helpers.python_std_to_sdk_type(Types.String).to_flyte_literal_type(), # '' # ) requests = [] if op.resources: requests.append( task_model.Resources.ResourceEntry( task_model.Resources.ResourceName.Cpu, op.resources.cpus)) requests.append( task_model.Resources.ResourceEntry( task_model.Resources.ResourceName.Memory, op.resources.ram)) requests.append( task_model.Resources.ResourceEntry( task_model.Resources.ResourceName.Gpu, op.resources.gpus)) requests.append( task_model.Resources.ResourceEntry( task_model.Resources.ResourceName.Storage, op.resources.disk)) task_instance = TaskInstance(op, datetime.datetime.now()) command = task_instance.command_as_list(local=True, mark_success=False, ignore_all_deps=True, ignore_depends_on_past=True, ignore_task_deps=True, ignore_ti_state=True, pool=task_instance.pool, pickle_id=dag_id, cfg_path=None) task = base_tasks.SdkTask( op.task_id, SingleStepTask, "airflow_op", task_model.TaskMetadata( False, task_model.RuntimeMetadata( type=task_model.RuntimeMetadata.RuntimeType.Other, version=airflow.version.version, flavor='airflow'), datetime.timedelta(seconds=0), literals_model.RetryStrategy(0), '1', None, ), interface_common.TypedInterface(inputs=interface_inputs, outputs=interface_outputs), custom=None, container=task_model.Container( image=image, command=command, args=[], resources=task_model.Resources(limits=[], requests=requests), env={}, config={}, )) return task, task(**input_mappings).assign_id_and_return(op.task_id)