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)
