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)
