diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index 70af22f1e1..0b58a6d14b 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -400,8 +400,12 @@ class FastSerializationSettings(object): """ enabled: bool = False + # This is the location that the code should be copied into. destination_dir: Optional[str] = None + # This is the zip file where the new code was uploaded to. + distribution_location: Optional[str] = None + @dataclass(frozen=True) class SerializationSettings(object): diff --git a/flytekit/core/python_function_task.py b/flytekit/core/python_function_task.py index fa98b7ca89..35016f9571 100644 --- a/flytekit/core/python_function_task.py +++ b/flytekit/core/python_function_task.py @@ -238,18 +238,6 @@ def compile_into_workflow( "Compilation for a dynamic workflow called in fast execution mode but no additional code " "distribution could be retrieved" ) - logger.warn(f"ctx.execution_state.additional_context {ctx.execution_state.additional_context}") - for task_template in tts: - sanitized_args = [] - for arg in task_template.container.args: - if arg == "{{ .remote_package_path }}": - sanitized_args.append(ctx.execution_state.additional_context.get("dynamic_addl_distro")) - elif arg == "{{ .dest_dir }}": - sanitized_args.append(ctx.execution_state.additional_context.get("dynamic_dest_dir", ".")) - else: - sanitized_args.append(arg) - del task_template.container.args[:] - task_template.container.args.extend(sanitized_args) dj_spec = _dynamic_job.DynamicJobSpec( min_successes=len(workflow_spec.template.nodes), @@ -290,7 +278,13 @@ def dynamic_execute(self, task_function: Callable, **kwargs) -> Any: if is_fast_execution: ctx = ctx.with_serialization_settings( ctx.serialization_settings.new_builder() - .with_fast_serialization_settings(FastSerializationSettings(enabled=True)) + .with_fast_serialization_settings( + FastSerializationSettings( + enabled=True, + destination_dir=ctx.execution_state.additional_context.get("dynamic_dest_dir", "."), + distribution_location=ctx.execution_state.additional_context.get("dynamic_addl_distro"), + ) + ) .build() ) diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 83d724e4ce..8c9750e8cd 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -68,19 +68,17 @@ def _fast_serialize_command_fn( ) -> Callable[[SerializationSettings], List[str]]: default_command = task.get_default_command(settings) - dest_dir = ( - settings.fast_serialization_settings.destination_dir if settings.fast_serialization_settings is not None else "" - ) - if dest_dir is None or dest_dir == "": - dest_dir = "{{ .dest_dir }}" - def fn(settings: SerializationSettings) -> List[str]: return [ "pyflyte-fast-execute", "--additional-distribution", - "{{ .remote_package_path }}", + settings.fast_serialization_settings.distribution_location + if settings.fast_serialization_settings and settings.fast_serialization_settings.distribution_location + else "{{ .remote_package_path }}", "--dest-dir", - dest_dir, + settings.fast_serialization_settings.destination_dir + if settings.fast_serialization_settings and settings.fast_serialization_settings.destination_dir + else "{{ .dest_dir }}", "--", *default_command, ] diff --git a/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py b/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py index 6cbf7d57fa..a4872bb321 100644 --- a/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py +++ b/plugins/flytekit-k8s-pod/flytekitplugins/pod/task.py @@ -75,7 +75,7 @@ def _serialize_pod_spec(self, settings: SerializationSettings) -> Dict[str, Any] final_containers = [] for container in containers: # In the case of the primary container, we overwrite specific container attributes with the default values - # used in an SDK runnable task. + # used in the regular Python task. if container.name == self.task_config.primary_container_name: sdk_default_container = super().get_container(settings) diff --git a/plugins/flytekit-k8s-pod/tests/test_pod.py b/plugins/flytekit-k8s-pod/tests/test_pod.py index 9677f51211..85fd2d8041 100644 --- a/plugins/flytekit-k8s-pod/tests/test_pod.py +++ b/plugins/flytekit-k8s-pod/tests/test_pod.py @@ -10,6 +10,7 @@ from flytekit import Resources, TaskMetadata, dynamic, map_task, task from flytekit.core import context_manager from flytekit.core.context_manager import FastSerializationSettings +from flytekit.core.type_engine import TypeEngine from flytekit.extend import ExecutionState, Image, ImageConfig, SerializationSettings from flytekit.tools.translator import get_serializable @@ -391,3 +392,65 @@ def simple_pod_task(i: int): "task-name", "simple_pod_task", ] + + +def test_fast(): + REQUESTS_GPU = Resources(cpu="123m", mem="234Mi", ephemeral_storage="123M", gpu="1") + LIMITS_GPU = Resources(cpu="124M", mem="235Mi", ephemeral_storage="124M", gpu="1") + + def get_minimal_pod_task_config() -> Pod: + primary_container = V1Container(name="flytetask") + pod_spec = V1PodSpec(containers=[primary_container]) + return Pod(pod_spec=pod_spec, primary_container_name="flytetask") + + @task( + task_config=get_minimal_pod_task_config(), + requests=REQUESTS_GPU, + limits=LIMITS_GPU, + ) + def pod_task_with_resources(dummy_input: str) -> str: + return dummy_input + + @dynamic(requests=REQUESTS_GPU, limits=LIMITS_GPU) + def dynamic_task_with_pod_subtask(dummy_input: str) -> str: + pod_task_with_resources(dummy_input=dummy_input) + return dummy_input + + default_img = Image(name="default", fqn="test", tag="tag") + serialization_settings = SerializationSettings( + project="project", + domain="domain", + version="version", + env={"FOO": "baz"}, + image_config=ImageConfig(default_image=default_img, images=[default_img]), + fast_serialization_settings=FastSerializationSettings(enabled=True), + ) + + with context_manager.FlyteContextManager.with_context( + context_manager.FlyteContextManager.current_context().with_serialization_settings(serialization_settings) + ) as ctx: + with context_manager.FlyteContextManager.with_context( + ctx.with_execution_state( + ctx.execution_state.with_params( + mode=ExecutionState.Mode.TASK_EXECUTION, + additional_context={ + "dynamic_addl_distro": "s3://my-s3-bucket/fast/123", + "dynamic_dest_dir": "/User/flyte/workflows", + }, + ) + ) + ) as ctx: + input_literal_map = TypeEngine.dict_to_literal_map(ctx, {"dummy_input": "hi"}) + dynamic_job_spec = dynamic_task_with_pod_subtask.dispatch_execute(ctx, input_literal_map) + # print(dynamic_job_spec) + assert len(dynamic_job_spec._nodes) == 1 + assert len(dynamic_job_spec.tasks) == 1 + args = " ".join(dynamic_job_spec.tasks[0].k8s_pod.pod_spec["containers"][0]["args"]) + assert args.startswith( + "pyflyte-fast-execute --additional-distribution s3://my-s3-bucket/fast/123 " + "--dest-dir /User/flyte/workflows" + ) + assert dynamic_job_spec.tasks[0].k8s_pod.pod_spec["containers"][0]["resources"]["limits"]["cpu"] == "124M" + assert dynamic_job_spec.tasks[0].k8s_pod.pod_spec["containers"][0]["resources"]["requests"]["gpu"] == "1" + + assert context_manager.FlyteContextManager.size() == 1 diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 600e90d74b..a2508ce5de 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -621,7 +621,9 @@ def my_wf(a: int) -> typing.List[str]: ) ) ) as ctx: - dynamic_job_spec = my_subwf.compile_into_workflow(ctx, my_subwf._task_function, a=5) + input_literal_map = TypeEngine.dict_to_literal_map(ctx, {"a": 5}) + + dynamic_job_spec = my_subwf.dispatch_execute(ctx, input_literal_map) assert len(dynamic_job_spec._nodes) == 5 assert len(dynamic_job_spec.tasks) == 1 args = " ".join(dynamic_job_spec.tasks[0].container.args)