def test_resources(resource_list): obj = task.Resources(resource_list, resource_list) obj1 = task.Resources([], resource_list) obj2 = task.Resources(resource_list, []) assert obj.requests == obj2.requests assert obj.limits == obj1.limits assert obj == task.Resources.from_flyte_idl(obj.to_flyte_idl())
def _get_container_definition(self, environment=None, **kwargs): """ :rtype: flytekit.models.task.Container """ return _spark_task.SdkRunnableSparkContainer( command=[], args=[], resources=_task_models.Resources(limits=[], requests=[]), env=environment or {}, config={}, )
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_resources( cls, storage_request=None, cpu_request=None, gpu_request=None, memory_request=None, storage_limit=None, cpu_limit=None, gpu_limit=None, memory_limit=None, ): """ :param Text storage_request: :param Text cpu_request: :param Text gpu_request: :param Text memory_request: :param Text storage_limit: :param Text cpu_limit: :param Text gpu_limit: :param Text memory_limit: """ 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) ) return _task_models.Resources(limits=limits, requests=requests)
def _get_container_definition(self, environment=None, **kwargs): """ :rtype: flytekit.models.task.Container """ return SdkRunnableSparkContainer( command=[], args=[ "execute_spark_task", "--task-module", self.task_module, "--task-name", self.task_function_name, "--inputs", "{{.input}}", "--output-prefix", "{{.outputPrefix}}" ], resources=_task_models.Resources(limits=[], requests=[]), env=environment or {}, config={})
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_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
def _get_container_definition( self, storage_request=None, cpu_request=None, gpu_request=None, memory_request=None, storage_limit=None, cpu_limit=None, gpu_limit=None, memory_limit=None, environment=None, cls=None, ): """ :param Text storage_request: :param Text cpu_request: :param Text gpu_request: :param Text memory_request: :param Text storage_limit: :param Text cpu_limit: :param Text gpu_limit: :param Text memory_limit: :param dict[Text,Text] environment: :param cls Optional[type]: Type of container to instantiate. Generally should subclass SdkRunnableContainer. :rtype: flytekit.models.task.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)) return (cls or SdkRunnableContainer)( command=[], args=[ "pyflyte-execute", "--task-module", self.task_module, "--task-name", self.task_function_name, "--inputs", "{{.input}}", "--output-prefix", "{{.outputPrefix}}" ], resources=_task_models.Resources(limits=limits, requests=requests), env=environment, config={})
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, )
interface.TypedInterface({'a': interface.Variable(t, "description 1")}, {'b': interface.Variable(t, "description 2")}) for t in LIST_OF_ALL_LITERAL_TYPES ] LIST_OF_RESOURCE_ENTRIES = [ task.Resources.ResourceEntry(task.Resources.ResourceName.CPU, "1"), task.Resources.ResourceEntry(task.Resources.ResourceName.GPU, "1"), task.Resources.ResourceEntry(task.Resources.ResourceName.MEMORY, "1G"), task.Resources.ResourceEntry(task.Resources.ResourceName.STORAGE, "1G") ] LIST_OF_RESOURCE_ENTRY_LISTS = [LIST_OF_RESOURCE_ENTRIES] LIST_OF_RESOURCES = [ task.Resources(request, limit) for request, limit in product( LIST_OF_RESOURCE_ENTRY_LISTS, LIST_OF_RESOURCE_ENTRY_LISTS) ] LIST_OF_RUNTIME_METADATA = [ task.RuntimeMetadata(task.RuntimeMetadata.RuntimeType.OTHER, "1.0.0", "python"), task.RuntimeMetadata(task.RuntimeMetadata.RuntimeType.FLYTE_SDK, "1.0.0b0", "golang") ] LIST_OF_RETRY_POLICIES = [ literals.RetryStrategy(retries=i) for i in [0, 1, 3, 100] ] LIST_OF_INTERRUPTIBLE = [None, True, False]
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 _get_container_definition(self, storage_request=None, cpu_request=None, gpu_request=None, memory_request=None, storage_limit=None, cpu_limit=None, gpu_limit=None, memory_limit=None, environment=None, **kwargs): """ :param Text storage_request: :param Text cpu_request: :param Text gpu_request: :param Text memory_request: :param Text storage_limit: :param Text cpu_limit: :param Text gpu_limit: :param Text memory_limit: :param dict[Text, Text] environment: :rtype: flytekit.models.task.Container """ storage_limit = storage_limit or storage_request cpu_limit = cpu_limit or cpu_request gpu_limit = gpu_limit or gpu_request memory_limit = memory_limit or memory_request 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)) return SdkRunnableContainer( command=[], args=[ "pyflyte-execute", "--task-module", self.task_module, "--task-name", self.task_function_name, "--inputs", "{{.input}}", "--output-prefix", "{{.outputPrefix}}" ], resources=_task_models.Resources(limits=limits, requests=requests), env=environment, 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)