From 38f140a9fe9b7edd940798c1ca45d9a09221473e Mon Sep 17 00:00:00 2001 From: Ketan Umare <16888709+kumare3@users.noreply.github.com> Date: Mon, 24 Jun 2024 18:02:24 -0700 Subject: [PATCH 001/192] Run flytekit tasks on remote with local default values passed correctly (#2525) * Run flytekit tasks on remote with local default values passed correctly ```python @task def default_inputs(i: int = 0, f: float = 10.0, e: Color = Color.RED, b: bool = True, j: typing.Optional[int] = None): print(i, f, e, b, j) ``` Running this on remote will now work correctly ```bash pyflyte run --remote exhaustive.py default_inputs ``` Signed-off-by: Ketan Umare * fixing lint Signed-off-by: Ketan Umare --------- Signed-off-by: Ketan Umare Co-authored-by: Ketan Umare --- flytekit/clis/sdk_in_container/run.py | 18 ++++++++++++++--- flytekit/core/interface.py | 2 +- .../default_arguments/task_defaults.py | 14 +++++++++++++ tests/flytekit/unit/cli/pyflyte/test_run.py | 20 +++++++++++++++++++ .../unit/core/test_python_function_task.py | 20 +++++++++++++++++++ 5 files changed, 70 insertions(+), 4 deletions(-) create mode 100644 tests/flytekit/unit/cli/pyflyte/default_arguments/task_defaults.py diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index 323988b340..e99f114818 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -509,6 +509,13 @@ def _update_flyte_context(params: RunLevelParams) -> FlyteContext.Builder: return ctx_builder.with_file_access(file_access) +def is_optional(_type): + """ + Checks if the given type is Optional Type + """ + return typing.get_origin(_type) is typing.Union and type(None) in typing.get_args(_type) + + def run_command(ctx: click.Context, entity: typing.Union[PythonFunctionWorkflow, PythonTask]): """ Returns a function that is used to implement WorkflowCommand and execute a flyte workflow. @@ -526,8 +533,13 @@ def _run(*args, **kwargs): click.secho(f"Running Execution on {'Remote' if run_level_params.is_remote else 'local'}.", fg="cyan") try: inputs = {} - for input_name, _ in entity.python_interface.inputs.items(): + for input_name, v in entity.python_interface.inputs_with_defaults.items(): processed_click_value = kwargs.get(input_name) + optional_v = False + if processed_click_value is None and isinstance(v, typing.Tuple): + optional_v = is_optional(v[0]) + if len(v) == 2: + processed_click_value = v[1] if isinstance(processed_click_value, ArtifactQuery): if run_level_params.is_remote: click.secho( @@ -542,7 +554,7 @@ def _run(*args, **kwargs): raise click.UsageError( f"Default for '{input_name}' is a query, which must be specified when running locally." ) - if processed_click_value is not None: + if processed_click_value is not None or optional_v: inputs[input_name] = processed_click_value if not run_level_params.is_remote: @@ -780,7 +792,7 @@ def _create_command( ctx: click.Context, entity_name: str, run_level_params: RunLevelParams, - loaded_entity: typing.Any, + loaded_entity: [PythonTask, WorkflowBase], is_workflow: bool, ): """ diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 13b6af2d4b..7f0af5f8e6 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -25,7 +25,7 @@ def repr_kv(k: str, v: Union[Type, Tuple[Type, Any]]) -> str: if isinstance(v, tuple): - if v[1]: + if v[1] is not None: return f"{k}: {v[0]}={v[1]}" return f"{k}: {v[0]}" return f"{k}: {v}" diff --git a/tests/flytekit/unit/cli/pyflyte/default_arguments/task_defaults.py b/tests/flytekit/unit/cli/pyflyte/default_arguments/task_defaults.py new file mode 100644 index 0000000000..45c6c7d14c --- /dev/null +++ b/tests/flytekit/unit/cli/pyflyte/default_arguments/task_defaults.py @@ -0,0 +1,14 @@ +import enum + +from flytekit import task + + +class Color(enum.Enum): + RED = 'red' + GREEN = 'green' + BLUE = 'blue' + + +@task +def foo(i: int = 0, j: str = "Hello", c: Color = Color.RED): + print(i, j, c) diff --git a/tests/flytekit/unit/cli/pyflyte/test_run.py b/tests/flytekit/unit/cli/pyflyte/test_run.py index 6957828743..3bb7697d47 100644 --- a/tests/flytekit/unit/cli/pyflyte/test_run.py +++ b/tests/flytekit/unit/cli/pyflyte/test_run.py @@ -516,3 +516,23 @@ def test_envvar_local_execution(envs, envs_argument, expected_output, workflow_f ) output = result.stdout.strip().split("\n")[-1].strip() assert output == expected_output + + +@pytest.mark.parametrize( + "task_path", + [("task_defaults.py")], +) +def test_list_default_arguments(task_path): + runner = CliRunner() + dir_name = os.path.dirname(os.path.realpath(__file__)) + result = runner.invoke( + pyflyte.main, + [ + "run", + os.path.join(dir_name, "default_arguments", task_path), + "foo", + ], + catch_exceptions=False, + ) + assert result.exit_code == 0 + assert result.stdout == "Running Execution on local.\n0 Hello Color.RED\n\n" diff --git a/tests/flytekit/unit/core/test_python_function_task.py b/tests/flytekit/unit/core/test_python_function_task.py index 8ba875b664..e775cb922a 100644 --- a/tests/flytekit/unit/core/test_python_function_task.py +++ b/tests/flytekit/unit/core/test_python_function_task.py @@ -264,3 +264,23 @@ def t1(i: str): @task(node_dependency_hints=[t1]) def t2(i: str): pass + + +def test_default_inputs(): + @task + def foo(x: int = 0, y: str = "Hello") -> int: + return x + + assert foo.python_interface.default_inputs_as_kwargs == {"x": 0, "y": "Hello"} + + @task + def foo2(x: int, y: str = "Hello") -> int: + return x + + assert foo2.python_interface.inputs_with_defaults == {"x": (int, None), "y": (str, "Hello")} + + @task + def foo3(x: int, y: str) -> int: + return x + + assert foo3.python_interface.inputs_with_defaults == {"x": (int, None), "y": (str, None)} From 487ac7b19dc85e56ce3fad1d9babad9baaeb59c6 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Tue, 25 Jun 2024 03:30:26 -0400 Subject: [PATCH 002/192] Adds support for wandb backlinks through FLYTE_EXECUTION_URL (#2497) Signed-off-by: Thomas J. Fan --- .../flytekit-wandb/flytekitplugins/wandb/tracking.py | 11 ++++++++++- plugins/flytekit-wandb/tests/test_wandb_init.py | 12 +++++++++--- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/plugins/flytekit-wandb/flytekitplugins/wandb/tracking.py b/plugins/flytekit-wandb/flytekitplugins/wandb/tracking.py index 3d0a4ac894..8bf033f2ab 100644 --- a/plugins/flytekit-wandb/flytekitplugins/wandb/tracking.py +++ b/plugins/flytekit-wandb/flytekitplugins/wandb/tracking.py @@ -94,7 +94,16 @@ def execute(self, *args, **kwargs): else: wand_id = self.id - wandb.init(project=self.project, entity=self.entity, id=wand_id, **self.init_kwargs) + run = wandb.init(project=self.project, entity=self.entity, id=wand_id, **self.init_kwargs) + + # If FLYTE_EXECUTION_URL is defined, inject it into wandb to link back to the execution. + execution_url = os.getenv("FLYTE_EXECUTION_URL") + if execution_url is not None: + notes_list = [f"[Execution URL]({execution_url})"] + if run.notes: + notes_list.append(run.notes) + run.notes = os.linesep.join(notes_list) + output = self.task_function(*args, **kwargs) wandb.finish() return output diff --git a/plugins/flytekit-wandb/tests/test_wandb_init.py b/plugins/flytekit-wandb/tests/test_wandb_init.py index 664e4a77ac..1db9f5c0fd 100644 --- a/plugins/flytekit-wandb/tests/test_wandb_init.py +++ b/plugins/flytekit-wandb/tests/test_wandb_init.py @@ -1,3 +1,4 @@ +import os from unittest.mock import Mock, patch import pytest @@ -57,10 +58,9 @@ def test_local_execution_with_id(wandb_mock): wandb_mock.init.assert_called_with(project="abc", entity="xyz", id="1234", tags=["my_tag"]) -@patch("flytekitplugins.wandb.tracking.os") @patch("flytekitplugins.wandb.tracking.FlyteContextManager") @patch("flytekitplugins.wandb.tracking.wandb") -def test_non_local_execution(wandb_mock, manager_mock, os_mock): +def test_non_local_execution(wandb_mock, manager_mock, monkeypatch): # Pretend that the execution is remote ctx_mock = Mock() ctx_mock.execution_state.is_local_execution.return_value = False @@ -69,13 +69,19 @@ def test_non_local_execution(wandb_mock, manager_mock, os_mock): ctx_mock.user_space_params.execution_id.name = "my_execution_id" manager_mock.current_context.return_value = ctx_mock - os_mock.environ = {} + execution_url = "http://execution_url.com/afsdfsafafasdfs" + monkeypatch.setattr("flytekitplugins.wandb.tracking.os.environ", {"FLYTE_EXECUTION_URL": execution_url}) + + run_mock = Mock() + run_mock.notes = "" + wandb_mock.init.return_value = run_mock train_model() wandb_mock.init.assert_called_with(project="abc", entity="xyz", id="my_execution_id", tags=["my_tag"]) ctx_mock.user_space_params.secrets.get.assert_called_with(key="abc", group="xyz") wandb_mock.login.assert_called_with(key="this_is_the_secret", host="https://api.wandb.ai") + assert run_mock.notes == f"[Execution URL]({execution_url})" def test_errors(): From 907a4cb3228bf1a25c430a6a1cbe3f797bea24b3 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 26 Jun 2024 01:03:31 +0800 Subject: [PATCH 003/192] Fix flaky test_initialise_azure_file_provider_with_default_credential (#2530) Signed-off-by: Kevin Su --- tests/flytekit/unit/core/test_data_persistence.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/flytekit/unit/core/test_data_persistence.py b/tests/flytekit/unit/core/test_data_persistence.py index b2852e58f3..159214fe43 100644 --- a/tests/flytekit/unit/core/test_data_persistence.py +++ b/tests/flytekit/unit/core/test_data_persistence.py @@ -165,8 +165,14 @@ def test_initialise_azure_file_provider_with_service_principal(): assert fp.get_filesystem().tenant_id == "tenantid" -@mock.patch.dict(os.environ, {"FLYTE_AZURE_STORAGE_ACCOUNT_NAME": "accountname", "AZURE_STORAGE_ANON": "false"}) def test_initialise_azure_file_provider_with_default_credential(): - fp = FileAccessProvider("/tmp", "abfs://container/path/within/container") - assert fp.get_filesystem().account_name == "accountname" - assert isinstance(fp.get_filesystem().sync_credential, DefaultAzureCredential) + with mock.patch.dict( + os.environ, + { + "FLYTE_AZURE_STORAGE_ACCOUNT_NAME": "accountname", + "AZURE_STORAGE_ANON": "false", + }, + ): + fp = FileAccessProvider("/tmp", "abfs://container/path/within/container") + assert fp.get_filesystem().account_name == "accountname" + assert isinstance(fp.get_filesystem().sync_credential, DefaultAzureCredential) From 1654e86a16a89e090c8322b868294f2a70ed3c44 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 26 Jun 2024 01:45:30 +0800 Subject: [PATCH 004/192] Install numpy<2.0.0 for pandera plugin (#2526) Signed-off-by: Kevin Su --- .github/workflows/pythonbuild.yml | 2 +- plugins/flytekit-envd/tests/test_image_spec.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index dfa83fca20..53ab938369 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -427,7 +427,7 @@ jobs: # TODO: move to protobuf>=5. Github issue: https://github.com/flyteorg/flyte/issues/5448 uv pip install --system -U $GITHUB_WORKSPACE "protobuf<5" # TODO: remove this when numpy v2 in onnx has been resolved - if [[ ${{ matrix.plugin-names }} == *"onnx"* || ${{ matrix.plugin-names }} == "flytekit-sqlalchemy" ]]; then + if [[ ${{ matrix.plugin-names }} == *"onnx"* || ${{ matrix.plugin-names }} == "flytekit-sqlalchemy" || ${{ matrix.plugin-names }} == "flytekit-pandera" ]]; then uv pip install --system "numpy<2.0.0" fi uv pip freeze diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index 5b7b73f755..5f252f6357 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/plugins/flytekit-envd/tests/test_image_spec.py @@ -101,7 +101,7 @@ def build(): def test_image_spec_extra_index_url(): image_spec = ImageSpec( - packages=["-U --pre pandas", "torch", "torchvision"], + packages=["-U pandas", "torch", "torchvision"], base_image="cr.flyte.org/flyteorg/flytekit:py3.9-latest", pip_extra_index_url=[ "https://download.pytorch.org/whl/cpu", @@ -120,7 +120,7 @@ def test_image_spec_extra_index_url(): def build(): base(image="cr.flyte.org/flyteorg/flytekit:py3.9-latest", dev=False) run(commands=[]) - install.python_packages(name=["-U --pre pandas", "torch", "torchvision"]) + install.python_packages(name=["-U pandas", "torch", "torchvision"]) install.apt_packages(name=[]) runtime.environ(env={{'PYTHONPATH': '/root:', '_F_IMG_ID': '{image_name}'}}, extra_path=['/root']) config.pip_index(url="https://pypi.org/simple", extra_url="https://download.pytorch.org/whl/cpu https://pypi.anaconda.org/scientific-python-nightly-wheels/simple") From 117560c9c766ad260a3760bf2a765259f70b139e Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 26 Jun 2024 06:56:45 +0800 Subject: [PATCH 005/192] refactor(core): Improve dataclass attribute serialization (#2531) Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 8 +++---- tests/flytekit/unit/core/test_type_engine.py | 24 ++++++++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 4937703ef4..b938902e4f 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -591,10 +591,10 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]) -> typing.A else: return python_val else: - for v in dataclasses.fields(python_type): - val = python_val.__getattribute__(v.name) - field_type = v.type - python_val.__setattr__(v.name, self._serialize_flyte_type(val, field_type)) + dataclass_attributes = typing.get_type_hints(python_type) + for n, t in dataclass_attributes.items(): + val = python_val.__getattribute__(n) + python_val.__setattr__(n, self._serialize_flyte_type(val, t)) return python_val def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) -> Optional[T]: diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index abdc314f29..d077c75396 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -920,6 +920,30 @@ class TestStructD(DataClassJsonMixin): assert ot == o +@mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data") +def test_dataclass_with_postponed_annotation(mock_put_data): + remote_path = "s3://tmp/file" + mock_put_data.return_value = remote_path + + @dataclass + class Data: + a: int + f: "FlyteFile" + + ctx = FlyteContext.current_context() + tf = DataclassTransformer() + t = tf.get_literal_type(Data) + assert t.simple == SimpleType.STRUCT + with tempfile.TemporaryDirectory() as tmp: + test_file = os.path.join(tmp, "abc.txt") + with open(test_file, "w") as f: + f.write("123") + + pv = Data(a=1, f=FlyteFile(test_file, remote_path=remote_path)) + lt = tf.to_literal(ctx, pv, Data, t) + assert lt.scalar.generic.fields["f"].struct_value.fields["path"].string_value == remote_path + + @mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data") def test_optional_flytefile_in_dataclass(mock_upload_dir): mock_upload_dir.return_value = True From a3bb3ed75601810809fab84ce8c497ac10421d75 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 26 Jun 2024 10:29:43 +0800 Subject: [PATCH 006/192] Add filter support to recent_executions (#2524) Signed-off-by: Kevin Su --- flytekit/remote/remote.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 70549c5b30..301ce4b5fb 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -576,12 +576,14 @@ def recent_executions( project: typing.Optional[str] = None, domain: typing.Optional[str] = None, limit: typing.Optional[int] = 100, + filters: typing.Optional[typing.List[filter_models.Filter]] = None, ) -> typing.List[FlyteWorkflowExecution]: # Ignore token for now exec_models, _ = self.client.list_executions_paginated( project or self.default_project, domain or self.default_domain, limit, + filters=filters, sort_by=MOST_RECENT_FIRST, ) return [FlyteWorkflowExecution.promote_from_model(e) for e in exec_models] From 647b0711492a07c001e367cde3f377ecedbc9892 Mon Sep 17 00:00:00 2001 From: Samhita Alla Date: Wed, 26 Jun 2024 15:16:12 +0530 Subject: [PATCH 007/192] modify the name of the dict transformer to typed dict (#2534) Signed-off-by: Samhita Alla --- flytekit/core/type_engine.py | 2 +- tests/flytekit/unit/core/test_type_engine.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index b938902e4f..2a087cb4ca 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -1677,7 +1677,7 @@ class DictTransformer(TypeTransformer[dict]): """ def __init__(self): - super().__init__("Python Dictionary", dict) + super().__init__("Typed Dict", dict) @staticmethod def extract_types_or_metadata(t: Optional[Type[dict]]) -> typing.Tuple: diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index d077c75396..e20b870c93 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -1810,7 +1810,7 @@ def test_union_containers(): lv = TypeEngine.to_literal(ctx, list_of_maps_of_list_ints, pt, lt) assert lv.scalar.union.stored_type.structure.tag == "Typed List" lv = TypeEngine.to_literal(ctx, map_of_list_ints, pt, lt) - assert lv.scalar.union.stored_type.structure.tag == "Python Dictionary" + assert lv.scalar.union.stored_type.structure.tag == "Typed Dict" @pytest.mark.skipif(sys.version_info < (3, 10), reason="PEP604 requires >=3.10.") From 553ec2066385f81cec37595a05966f2a0e748031 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 27 Jun 2024 02:04:52 +0800 Subject: [PATCH 008/192] Ensure that PythonAutoContainerTask has a container_image attribute (#2527) Signed-off-by: Kevin Su --- .github/workflows/monodocs_build.yml | 2 +- .github/workflows/pythonbuild.yml | 2 +- flytekit/core/python_auto_container.py | 21 +++++++++++++++------ 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/.github/workflows/monodocs_build.yml b/.github/workflows/monodocs_build.yml index 7f11de452c..7585c464fe 100644 --- a/.github/workflows/monodocs_build.yml +++ b/.github/workflows/monodocs_build.yml @@ -1,7 +1,7 @@ name: Monodocs Build concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number }} cancel-in-progress: true on: diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index 53ab938369..8d450ffc87 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -14,7 +14,7 @@ env: FLYTE_SDK_LOGGING_LEVEL: 10 # debug concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number }} cancel-in-progress: true jobs: diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 2c4703cdd3..93d832d4e2 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -87,19 +87,28 @@ def __init__( kwargs["metadata"] = kwargs["metadata"] if "metadata" in kwargs else TaskMetadata() kwargs["metadata"].pod_template_name = pod_template_name + self._container_image = container_image + # TODO(katrogan): Implement resource overrides + self._resources = ResourceSpec( + requests=requests if requests else Resources(), limits=limits if limits else Resources() + ) + + # The serialization of the other tasks (Task -> protobuf), as well as the initialization of the current task, may occur simultaneously. + # We should make sure super().__init__ is being called after setting _container_image because PythonAutoContainerTask + # is added to the FlyteEntities in super().__init__, and the translator will iterate over + # FlyteEntities and call entity.container_image(). + # Therefore, we need to ensure the _container_image attribute is set + # before appending the task to FlyteEntities. + # https://github.com/flyteorg/flytekit/blob/876877abd8d6f4f54dec2738a0ca07a12e9115b1/flytekit/tools/translator.py#L181 + super().__init__( task_type=task_type, name=name, task_config=task_config, security_ctx=sec_ctx, + environment=environment, **kwargs, ) - self._container_image = container_image - # TODO(katrogan): Implement resource overrides - self._resources = ResourceSpec( - requests=requests if requests else Resources(), limits=limits if limits else Resources() - ) - self._environment = environment or {} compilation_state = FlyteContextManager.current_context().compilation_state if compilation_state and compilation_state.task_resolver: From 6d35b75bf4e9f6d2f0e390910384fae6951ca72c Mon Sep 17 00:00:00 2001 From: pryce-turner <31577879+pryce-turner@users.noreply.github.com> Date: Thu, 27 Jun 2024 10:19:43 -0700 Subject: [PATCH 009/192] Catch silent subproc_execute failures (#2404) * Made outfile ephemeral Signed-off-by: pryce-turner * Changed error handling to warn log for pipe with shell commands Signed-off-by: pryce-turner --------- Signed-off-by: pryce-turner --- flytekit/extras/tasks/shell.py | 8 ++++++++ tests/flytekit/unit/extras/tasks/test_shell.py | 14 ++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/flytekit/extras/tasks/shell.py b/flytekit/extras/tasks/shell.py index 03f6b61ebc..ef9cd0c0e1 100644 --- a/flytekit/extras/tasks/shell.py +++ b/flytekit/extras/tasks/shell.py @@ -79,6 +79,14 @@ def subproc_execute(command: typing.Union[List[str], str], **kwargs) -> ProcessR try: # Execute the command and capture stdout and stderr result = subprocess.run(command, **kwargs) + print(result.check_returncode()) + + if "|" in command and kwargs.get("shell"): + logger.warning( + """Found a pipe in the command and shell=True. + This can lead to silent failures if subsequent commands + succeed despite previous failures.""" + ) # Access the stdout and stderr output return ProcessResult(result.returncode, result.stdout, result.stderr) diff --git a/tests/flytekit/unit/extras/tasks/test_shell.py b/tests/flytekit/unit/extras/tasks/test_shell.py index bb06f1a5dc..65a7a50e39 100644 --- a/tests/flytekit/unit/extras/tasks/test_shell.py +++ b/tests/flytekit/unit/extras/tasks/test_shell.py @@ -370,3 +370,17 @@ def test_subproc_execute_with_shell(): subproc_execute(cmd, shell=True) cont = open(opth).read() assert "hello" in cont + + +def test_subproc_execute_missing_dep(): + cmd = ["non-existent", "blah"] + with pytest.raises(Exception) as e: + subproc_execute(cmd) + assert "executable could not be found" in str(e.value) + + +def test_subproc_execute_error(): + cmd = ["ls", "--banana"] + with pytest.raises(Exception) as e: + subproc_execute(cmd) + assert "Failed with return code" in str(e.value) From bd4c7085e8cc6068900144bb55548e7c6c56f7c3 Mon Sep 17 00:00:00 2001 From: Greg Gydush <35151789+ggydush@users.noreply.github.com> Date: Thu, 27 Jun 2024 12:08:43 -0700 Subject: [PATCH 010/192] fix: Prevent local files from being moved when using FlyteFile on local executions (#2476) * fix: Do not copy local files when using FlyteFile Signed-off-by: ggydush * fix: Prevent copying of local files when running local execution Signed-off-by: ggydush * fix: Revert Signed-off-by: ggydush * fix: Fix another location of should upload Signed-off-by: ggydush * test: Fix failing test cases Signed-off-by: ggydush * fix: Fix to still handle uploads Signed-off-by: ggydush --------- Signed-off-by: ggydush Signed-off-by: Eduardo Apolinario Co-authored-by: Eduardo Apolinario --- flytekit/types/file/file.py | 5 +++++ tests/flytekit/unit/core/test_flyte_file.py | 9 ++------- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 15113f8f01..567b29de57 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -440,6 +440,9 @@ def to_literal( # Set the remote destination if one was given instead of triggering a random one below remote_path = python_val.remote_path or None + if ctx.execution_state.is_local_execution and python_val.remote_path is None: + should_upload = False + elif isinstance(python_val, pathlib.Path) or isinstance(python_val, str): source_path = str(python_val) if issubclass(python_type, FlyteFile): @@ -455,6 +458,8 @@ def to_literal( p = pathlib.Path(python_val) if not p.is_file(): raise TypeTransformerFailedError(f"Error converting {python_val} because it's not a file.") + if ctx.execution_state.is_local_execution: + should_upload = False # python_type must be os.PathLike - see check at beginning of function else: should_upload = False diff --git a/tests/flytekit/unit/core/test_flyte_file.py b/tests/flytekit/unit/core/test_flyte_file.py index ff038c164b..b1331adea7 100644 --- a/tests/flytekit/unit/core/test_flyte_file.py +++ b/tests/flytekit/unit/core/test_flyte_file.py @@ -205,7 +205,7 @@ def my_wf(fname: os.PathLike = SAMPLE_DATA) -> int: assert sample_lp.parameters.parameters["fname"].default.scalar.blob.uri == SAMPLE_DATA -def test_file_handling_local_file_gets_copied(): +def test_file_handling_local_file_does_not_get_copied(): @task def t1() -> FlyteFile: # Use this test file itself, since we know it exists. @@ -223,12 +223,7 @@ def my_wf() -> FlyteFile: assert len(top_level_files) == 1 # the flytekit_local folder x = my_wf() - - # After running, this test file should've been copied to the mock remote location. - mock_remote_files = os.listdir(os.path.join(random_dir, "mock_remote")) - assert len(mock_remote_files) == 1 # the file - # File should've been copied to the mock remote folder - assert x.path.startswith(random_dir) + assert x.path == __file__ def test_file_handling_local_file_gets_force_no_copy(): From 2abf095274a3e2ce470410cde2690ec52dc698e8 Mon Sep 17 00:00:00 2001 From: ddl-rliu <140021987+ddl-rliu@users.noreply.github.com> Date: Thu, 27 Jun 2024 18:08:14 -0700 Subject: [PATCH 011/192] [fix] Validate interface variable names are python-valid (#2538) * Validate interface variable names Signed-off-by: ddl-rliu * Remove unused import Signed-off-by: ddl-rliu * Fix lint error Signed-off-by: ddl-rliu --------- Signed-off-by: ddl-rliu --- flytekit/core/interface.py | 2 ++ tests/flytekit/unit/core/test_interface.py | 10 ++++++++++ 2 files changed, 12 insertions(+) diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 7f0af5f8e6..7faf541661 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -70,6 +70,8 @@ def __init__( self._inputs: Union[Dict[str, Tuple[Type, Any]], Dict[str, Type]] = {} # type: ignore if inputs: for k, v in inputs.items(): + if not k.isidentifier(): + raise ValueError(f"Input name must be valid Python identifier: {k!r}") if type(v) is tuple and len(cast(Tuple, v)) > 1: self._inputs[k] = v # type: ignore else: diff --git a/tests/flytekit/unit/core/test_interface.py b/tests/flytekit/unit/core/test_interface.py index ecf25fb6e0..e860729a83 100644 --- a/tests/flytekit/unit/core/test_interface.py +++ b/tests/flytekit/unit/core/test_interface.py @@ -320,6 +320,16 @@ def z(a: int, b: str) -> typing.NamedTuple("NT", x_str=str, y_int=int): assert typed_interface.outputs.get("y_int").description == "description for y_int" +def test_init_interface_with_invalid_parameters(): + from flytekit.core.interface import Interface + + with pytest.raises(ValueError, match=r"Input name must be valid Python identifier:"): + _ = Interface({"my.input": int}, {}) + + with pytest.raises(ValueError, match=r"Type names and field names must be valid identifiers:"): + _ = Interface({}, {"my.output": int}) + + def test_parameter_change_to_pickle_type(): ctx = context_manager.FlyteContext.current_context() From 2719b349cd6ac9b18cad858ff2fc6d6f79986b1e Mon Sep 17 00:00:00 2001 From: Chi-Sheng Liu Date: Sat, 29 Jun 2024 01:53:51 +0800 Subject: [PATCH 012/192] [Feature] Support positional arguments (#2522) - Change the `inputs` and `outputs` attributes in the `Interface` class to `OrderedDict` to preserve the order. - Write values in positional arguments to `kwargs`. Resolves: flyteorg/flyte#5320 Signed-off-by: Chi-Sheng Liu --- flytekit/core/promise.py | 40 +++--- .../flytekit/unit/core/test_serialization.py | 123 ++++++++++++++++++ 2 files changed, 143 insertions(+), 20 deletions(-) diff --git a/flytekit/core/promise.py b/flytekit/core/promise.py index 557d621dd4..c4f71eb2d6 100644 --- a/flytekit/core/promise.py +++ b/flytekit/core/promise.py @@ -1202,19 +1202,22 @@ def flyte_entity_call_handler( #. Start a local execution - This means that we're not already in a local workflow execution, which means that we should expect inputs to be native Python values and that we should return Python native values. """ - # Sanity checks - # Only keyword args allowed - if len(args) > 0: - raise _user_exceptions.FlyteAssertion( - f"When calling tasks, only keyword args are supported. " - f"Aborting execution as detected {len(args)} positional args {args}" - ) # Make sure arguments are part of interface for k, v in kwargs.items(): - if k not in cast(SupportsNodeCreation, entity).python_interface.inputs: - raise AssertionError( - f"Received unexpected keyword argument '{k}' in function '{cast(SupportsNodeCreation, entity).name}'" - ) + if k not in entity.python_interface.inputs: + raise AssertionError(f"Received unexpected keyword argument '{k}' in function '{entity.name}'") + + # Check if we have more arguments than expected + if len(args) > len(entity.python_interface.inputs): + raise AssertionError( + f"Received more arguments than expected in function '{entity.name}'. Expected {len(entity.python_interface.inputs)} but got {len(args)}" + ) + + # Convert args to kwargs + for arg, input_name in zip(args, entity.python_interface.inputs.keys()): + if input_name in kwargs: + raise AssertionError(f"Got multiple values for argument '{input_name}' in function '{entity.name}'") + kwargs[input_name] = arg ctx = FlyteContextManager.current_context() if ctx.execution_state and ( @@ -1234,15 +1237,12 @@ def flyte_entity_call_handler( child_ctx.execution_state and child_ctx.execution_state.branch_eval_mode == BranchEvalMode.BRANCH_SKIPPED ): - if ( - len(cast(SupportsNodeCreation, entity).python_interface.inputs) > 0 - or len(cast(SupportsNodeCreation, entity).python_interface.outputs) > 0 - ): - output_names = list(cast(SupportsNodeCreation, entity).python_interface.outputs.keys()) + if len(entity.python_interface.inputs) > 0 or len(entity.python_interface.outputs) > 0: + output_names = list(entity.python_interface.outputs.keys()) if len(output_names) == 0: return VoidPromise(entity.name) vals = [Promise(var, None) for var in output_names] - return create_task_output(vals, cast(SupportsNodeCreation, entity).python_interface) + return create_task_output(vals, entity.python_interface) else: return None return cast(LocallyExecutable, entity).local_execute(ctx, **kwargs) @@ -1255,7 +1255,7 @@ def flyte_entity_call_handler( cast(ExecutionParameters, child_ctx.user_space_params)._decks = [] result = cast(LocallyExecutable, entity).local_execute(child_ctx, **kwargs) - expected_outputs = len(cast(SupportsNodeCreation, entity).python_interface.outputs) + expected_outputs = len(entity.python_interface.outputs) if expected_outputs == 0: if result is None or isinstance(result, VoidPromise): return None @@ -1268,10 +1268,10 @@ def flyte_entity_call_handler( if (1 < expected_outputs == len(cast(Tuple[Promise], result))) or ( result is not None and expected_outputs == 1 ): - return create_native_named_tuple(ctx, result, cast(SupportsNodeCreation, entity).python_interface) + return create_native_named_tuple(ctx, result, entity.python_interface) raise AssertionError( f"Expected outputs and actual outputs do not match." f"Result {result}. " - f"Python interface: {cast(SupportsNodeCreation, entity).python_interface}" + f"Python interface: {entity.python_interface}" ) diff --git a/tests/flytekit/unit/core/test_serialization.py b/tests/flytekit/unit/core/test_serialization.py index 2fcf8bbd94..725e3e14fc 100644 --- a/tests/flytekit/unit/core/test_serialization.py +++ b/tests/flytekit/unit/core/test_serialization.py @@ -943,3 +943,126 @@ def wf_with_input() -> typing.Optional[typing.List[int]]: ) assert wf_with_input() == input_val + +def test_positional_args_task(): + arg1 = 5 + arg2 = 6 + ret = 17 + + @task + def t1(x: int, y: int) -> int: + return x + y * 2 + + @workflow + def wf_pure_positional_args() -> int: + return t1(arg1, arg2) + + @workflow + def wf_mixed_positional_and_keyword_args() -> int: + return t1(arg1, y=arg2) + + wf_pure_positional_args_spec = get_serializable(OrderedDict(), serialization_settings, wf_pure_positional_args) + wf_mixed_positional_and_keyword_args_spec = get_serializable(OrderedDict(), serialization_settings, wf_mixed_positional_and_keyword_args) + + arg1_binding = Scalar(primitive=Primitive(integer=arg1)) + arg2_binding = Scalar(primitive=Primitive(integer=arg2)) + output_type = LiteralType(simple=SimpleType.INTEGER) + + assert wf_pure_positional_args_spec.template.nodes[0].inputs[0].binding.value == arg1_binding + assert wf_pure_positional_args_spec.template.nodes[0].inputs[1].binding.value == arg2_binding + assert wf_pure_positional_args_spec.template.interface.outputs["o0"].type == output_type + + + assert wf_mixed_positional_and_keyword_args_spec.template.nodes[0].inputs[0].binding.value == arg1_binding + assert wf_mixed_positional_and_keyword_args_spec.template.nodes[0].inputs[1].binding.value == arg2_binding + assert wf_mixed_positional_and_keyword_args_spec.template.interface.outputs["o0"].type == output_type + + assert wf_pure_positional_args() == ret + assert wf_mixed_positional_and_keyword_args() == ret + +def test_positional_args_workflow(): + arg1 = 5 + arg2 = 6 + ret = 17 + + @task + def t1(x: int, y: int) -> int: + return x + y * 2 + + @workflow + def sub_wf(x: int, y: int) -> int: + return t1(x=x, y=y) + + @workflow + def wf_pure_positional_args() -> int: + return sub_wf(arg1, arg2) + + @workflow + def wf_mixed_positional_and_keyword_args() -> int: + return sub_wf(arg1, y=arg2) + + wf_pure_positional_args_spec = get_serializable(OrderedDict(), serialization_settings, wf_pure_positional_args) + wf_mixed_positional_and_keyword_args_spec = get_serializable(OrderedDict(), serialization_settings, wf_mixed_positional_and_keyword_args) + + arg1_binding = Scalar(primitive=Primitive(integer=arg1)) + arg2_binding = Scalar(primitive=Primitive(integer=arg2)) + output_type = LiteralType(simple=SimpleType.INTEGER) + + assert wf_pure_positional_args_spec.template.nodes[0].inputs[0].binding.value == arg1_binding + assert wf_pure_positional_args_spec.template.nodes[0].inputs[1].binding.value == arg2_binding + assert wf_pure_positional_args_spec.template.interface.outputs["o0"].type == output_type + + assert wf_mixed_positional_and_keyword_args_spec.template.nodes[0].inputs[0].binding.value == arg1_binding + assert wf_mixed_positional_and_keyword_args_spec.template.nodes[0].inputs[1].binding.value == arg2_binding + assert wf_mixed_positional_and_keyword_args_spec.template.interface.outputs["o0"].type == output_type + + assert wf_pure_positional_args() == ret + assert wf_mixed_positional_and_keyword_args() == ret + +def test_positional_args_chained_tasks(): + @task + def t1(x: int, y: int) -> int: + return x + y * 2 + + @workflow + def wf() -> int: + x = t1(2, y = 3) + y = t1(3, 4) + return t1(x, y = y) + + assert wf() == 30 + +def test_positional_args_task_inputs_from_workflow_args(): + @task + def t1(x: int, y: int, z: int) -> int: + return x + y * 2 + z * 3 + + @workflow + def wf(x: int, y: int) -> int: + return t1(x, y=y, z=3) + + assert wf(1, 2) == 14 + +def test_unexpected_kwargs_task_raises_error(): + @task + def t1(a: int) -> int: + return a + + with pytest.raises(AssertionError, match="Received unexpected keyword argument"): + t1(b=6) + +def test_too_many_positional_args_task_raises_error(): + @task + def t1(a: int) -> int: + return a + + with pytest.raises(AssertionError, match="Received more arguments than expected"): + t1(1, 2) + +def test_both_positional_and_keyword_args_task_raises_error(): + @task + def t1(a: int) -> int: + return a + + with pytest.raises(AssertionError, match="Got multiple values for argument"): + t1(1, a=2) From 1085725bb1f2fe6160737fcf6b21de7ebb040d9a Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Mon, 1 Jul 2024 13:25:33 -0400 Subject: [PATCH 013/192] Adds a Default ImageSpec image builder (#2346) Signed-off-by: Thomas J. Fan Co-authored-by: Kevin Su --- flytekit/image_spec/__init__.py | 6 +- flytekit/image_spec/default_builder.py | 275 ++++++++++++++++++ .../flytekit-envd/tests/test_image_spec.py | 2 +- .../core/image_spec/test_default_builder.py | 124 ++++++++ 4 files changed, 405 insertions(+), 2 deletions(-) create mode 100644 flytekit/image_spec/default_builder.py create mode 100644 tests/flytekit/unit/core/image_spec/test_default_builder.py diff --git a/flytekit/image_spec/__init__.py b/flytekit/image_spec/__init__.py index e47f7f159c..f5a8992bac 100644 --- a/flytekit/image_spec/__init__.py +++ b/flytekit/image_spec/__init__.py @@ -15,4 +15,8 @@ ImageSpec """ -from .image_spec import ImageSpec +from .default_builder import DefaultImageBuilder +from .image_spec import ImageBuildEngine, ImageSpec + +# Set this to a lower priority compared to `envd` to maintain backward compatibility +ImageBuildEngine.register("default", DefaultImageBuilder(), priority=1) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py new file mode 100644 index 0000000000..7e91eebf10 --- /dev/null +++ b/flytekit/image_spec/default_builder.py @@ -0,0 +1,275 @@ +import os +import shutil +import subprocess +import sys +import tempfile +import warnings +from pathlib import Path +from string import Template +from typing import ClassVar + +import click + +from flytekit.image_spec.image_spec import ( + _F_IMG_ID, + ImageSpec, + ImageSpecBuilder, +) +from flytekit.tools.ignore import DockerIgnore, GitIgnore, IgnoreGroup, StandardIgnore + +UV_PYTHON_INSTALL_COMMAND_TEMPLATE = Template("""\ +RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \ + --mount=from=uv,source=/uv,target=/usr/bin/uv \ + --mount=type=bind,target=requirements_uv.txt,src=requirements_uv.txt \ + /usr/bin/uv \ + pip install --python /root/micromamba/envs/dev/bin/python $PIP_EXTRA \ + --requirement requirements_uv.txt +""") + +PIP_PYTHON_INSTALL_COMMAND_TEMPLATE = Template("""\ +RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/pip,id=pip \ + --mount=type=bind,target=requirements_pip.txt,src=requirements_pip.txt \ + /root/micromamba/envs/dev/bin/python -m pip install $PIP_EXTRA \ + --requirement requirements_pip.txt +""") + +APT_INSTALL_COMMAND_TEMPLATE = Template( + """\ +RUN --mount=type=cache,sharing=locked,mode=0777,target=/var/cache/apt,id=apt \ + apt-get update && apt-get install -y --no-install-recommends \ + $APT_PACKAGES +""" +) + +DOCKER_FILE_TEMPLATE = Template( + """\ +#syntax=docker/dockerfile:1.5 +FROM ghcr.io/astral-sh/uv:0.2.13 as uv +FROM mambaorg/micromamba:1.5.8-bookworm-slim as micromamba + +FROM $BASE_IMAGE + +USER root +$APT_INSTALL_COMMAND +RUN update-ca-certificates + +RUN id -u flytekit || useradd --create-home --shell /bin/bash flytekit +RUN chown -R flytekit /root && chown -R flytekit /home + +RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/micromamba/pkgs,\ +id=micromamba \ + --mount=from=micromamba,source=/usr/bin/micromamba,target=/usr/bin/micromamba \ + /usr/bin/micromamba create -n dev -c conda-forge $CONDA_CHANNELS \ + python=$PYTHON_VERSION $CONDA_PACKAGES + +$UV_PYTHON_INSTALL_COMMAND +$PIP_PYTHON_INSTALL_COMMAND + +# Configure user space +ENV PATH="/root/micromamba/envs/dev/bin:$$PATH" +ENV FLYTE_SDK_RICH_TRACEBACKS=0 SSL_CERT_DIR=/etc/ssl/certs $ENV + +# Adds nvidia just in case it exists +ENV PATH="$$PATH:/usr/local/nvidia/bin:/usr/local/cuda/bin" \ + LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$$LD_LIBRARY_PATH" + +$COPY_COMMAND_RUNTIME +RUN $RUN_COMMANDS + +WORKDIR /root +SHELL ["/bin/bash", "-c"] + +USER flytekit +RUN echo "export PATH=$$PATH" >> $$HOME/.profile +""" +) + + +def get_flytekit_for_pypi(): + """Get flytekit version on PyPI.""" + from flytekit import __version__ + + if not __version__ or "dev" in __version__: + return "flytekit" + else: + return f"flytekit=={__version__}" + + +def create_docker_context(image_spec: ImageSpec, tmp_dir: Path): + """Populate tmp_dir with Dockerfile as specified by the `image_spec`.""" + base_image = image_spec.base_image or "debian:bookworm-slim" + + requirements = [get_flytekit_for_pypi()] + + if image_spec.cuda is not None or image_spec.cudnn is not None: + msg = ( + "cuda and cudnn do not need to be specified. If you are installed " + "a GPU accelerated library on PyPI, then it likely will install cuda " + "from PyPI." + "With conda you can installed cuda from the `nvidia` channel by adding `nvidia` to " + "ImageSpec.conda_channels and adding packages from " + "https://anaconda.org/nvidia into ImageSpec.conda_packages. If you require " + "cuda for non-python dependencies, you can set a `base_image` with cuda " + "preinstalled." + ) + raise ValueError(msg) + + if image_spec.requirements: + with open(image_spec.requirements) as f: + requirements.extend([line.strip() for line in f.readlines()]) + + if image_spec.packages: + requirements.extend(image_spec.packages) + + uv_requirements = [] + + # uv does not support git + subdirectory, so we use pip to install them instead + pip_requirements = [] + + for requirement in requirements: + if "git" in requirement and "subdirectory" in requirement: + pip_requirements.append(requirement) + else: + uv_requirements.append(requirement) + + requirements_uv_path = tmp_dir / "requirements_uv.txt" + requirements_uv_path.write_text("\n".join(uv_requirements)) + + pip_extra = f"--index-url {image_spec.pip_index}" if image_spec.pip_index else "" + uv_python_install_command = UV_PYTHON_INSTALL_COMMAND_TEMPLATE.substitute(PIP_EXTRA=pip_extra) + + if pip_requirements: + requirements_uv_path = tmp_dir / "requirements_pip.txt" + requirements_uv_path.write_text(os.linesep.join(pip_requirements)) + + pip_python_install_command = PIP_PYTHON_INSTALL_COMMAND_TEMPLATE.substitute(PIP_EXTRA=pip_extra) + else: + pip_python_install_command = "" + + env_dict = {"PYTHONPATH": "/root", _F_IMG_ID: image_spec.image_name()} + + if image_spec.env: + env_dict.update(image_spec.env) + + env = " ".join(f"{k}={v}" for k, v in env_dict.items()) + + apt_packages = ["ca-certificates"] + if image_spec.apt_packages: + apt_packages.extend(image_spec.apt_packages) + + apt_install_command = APT_INSTALL_COMMAND_TEMPLATE.substitute(APT_PACKAGES=" ".join(apt_packages)) + + if image_spec.source_root: + source_path = tmp_dir / "src" + + ignore = IgnoreGroup(image_spec.source_root, [GitIgnore, DockerIgnore, StandardIgnore]) + shutil.copytree( + image_spec.source_root, + source_path, + ignore=shutil.ignore_patterns(*ignore.list_ignored()), + dirs_exist_ok=True, + ) + copy_command_runtime = "COPY --chown=flytekit ./src /root" + else: + copy_command_runtime = "" + + conda_packages = image_spec.conda_packages or [] + conda_channels = image_spec.conda_channels or [] + + if conda_packages: + conda_packages_concat = " ".join(conda_packages) + else: + conda_packages_concat = "" + + if conda_channels: + conda_channels_concat = " ".join(f"-c {channel}" for channel in conda_channels) + else: + conda_channels_concat = "" + + if image_spec.python_version: + python_version = image_spec.python_version + else: + python_version = f"{sys.version_info.major}.{sys.version_info.minor}" + + if image_spec.commands: + run_commands = " && ".join(image_spec.commands) + else: + run_commands = "" + + docker_content = DOCKER_FILE_TEMPLATE.substitute( + PYTHON_VERSION=python_version, + UV_PYTHON_INSTALL_COMMAND=uv_python_install_command, + PIP_PYTHON_INSTALL_COMMAND=pip_python_install_command, + CONDA_PACKAGES=conda_packages_concat, + CONDA_CHANNELS=conda_channels_concat, + APT_INSTALL_COMMAND=apt_install_command, + BASE_IMAGE=base_image, + ENV=env, + COPY_COMMAND_RUNTIME=copy_command_runtime, + RUN_COMMANDS=run_commands, + ) + + dockerfile_path = tmp_dir / "Dockerfile" + dockerfile_path.write_text(docker_content) + + +class DefaultImageBuilder(ImageSpecBuilder): + """Image builder using Docker and buildkit.""" + + _SUPPORTED_IMAGE_SPEC_PARAMETERS: ClassVar[set] = { + "name", + "python_version", + "builder", + "source_root", + "env", + "registry", + "packages", + "conda_packages", + "conda_channels", + "requirements", + "apt_packages", + "platform", + "cuda", + "cudnn", + "base_image", + "pip_index", + # "registry_config", + "commands", + } + + def build_image(self, image_spec: ImageSpec) -> str: + return self._build_image(image_spec) + + def _build_image(self, image_spec: ImageSpec, *, push: bool = True) -> str: + # For testing, set `push=False`` to just build the image locally and not push to + # registry + unsupported_parameters = [ + name + for name, value in vars(image_spec).items() + if value is not None and name not in self._SUPPORTED_IMAGE_SPEC_PARAMETERS and not name.startswith("_") + ] + if unsupported_parameters: + msg = f"The following parameters are unsupported and ignored: " f"{unsupported_parameters}" + warnings.warn(msg, UserWarning, stacklevel=2) + + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_path = Path(tmp_dir) + create_docker_context(image_spec, tmp_path) + + command = [ + "docker", + "image", + "build", + "--tag", + f"{image_spec.image_name()}", + "--platform", + image_spec.platform, + ] + + if image_spec.registry and push: + command.append("--push") + command.append(tmp_dir) + + concat_command = " ".join(command) + click.secho(f"Run command: {concat_command} ", fg="blue") + subprocess.run(command, check=True) diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index 5f252f6357..31cd92effe 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/plugins/flytekit-envd/tests/test_image_spec.py @@ -11,7 +11,7 @@ @pytest.fixture(scope="module", autouse=True) def register_envd_higher_priority(): # Register a new envd platform with the highest priority so the test in this file uses envd - highest_priority_builder = max(ImageBuildEngine._REGISTRY, key=ImageBuildEngine._REGISTRY.get) + highest_priority_builder = max(ImageBuildEngine._REGISTRY, key=lambda name: ImageBuildEngine._REGISTRY[name][1]) highest_priority = ImageBuildEngine._REGISTRY[highest_priority_builder][1] yield ImageBuildEngine.register( "envd_high_priority", diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py new file mode 100644 index 0000000000..bd47b4558e --- /dev/null +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -0,0 +1,124 @@ +import os + +import pytest + +from flytekit.image_spec import ImageSpec +from flytekit.image_spec.default_builder import DefaultImageBuilder, create_docker_context + + +def test_create_docker_context(tmp_path): + docker_context_path = tmp_path / "builder_root" + docker_context_path.mkdir() + + source_root = tmp_path / "other_files" + source_root.mkdir() + hello_world_path = source_root / "hello_world.txt" + hello_world_path.write_text("hello") + + other_requirements_path = tmp_path / "requirements.txt" + other_requirements_path.write_text("threadpoolctl\n") + + image_spec = ImageSpec( + name="FLYTEKIT", + python_version="3.12", + env={"MY_ENV": "MY_VALUE"}, + apt_packages=["curl"], + conda_packages=["scipy==1.13.0", "numpy"], + packages=["pandas==2.2.1"], + requirements=os.fspath(other_requirements_path), + source_root=os.fspath(source_root), + commands=["mkdir my_dir"], + ) + + create_docker_context(image_spec, docker_context_path) + + dockerfile_path = docker_context_path / "Dockerfile" + assert dockerfile_path.exists() + dockerfile_content = dockerfile_path.read_text() + + assert "curl" in dockerfile_content + assert "scipy==1.13.0 numpy" in dockerfile_content + assert "python=3.12" in dockerfile_content + assert "--requirement requirements_uv.txt" in dockerfile_content + assert "COPY --chown=flytekit ./src /root" in dockerfile_content + assert "RUN mkdir my_dir" in dockerfile_content + + requirements_path = docker_context_path / "requirements_uv.txt" + assert requirements_path.exists() + + requirements_content = requirements_path.read_text() + assert "pandas==2.2.1" in requirements_content + assert "threadpoolctl" in requirements_content + + tmp_hello_world = docker_context_path / "src" / "hello_world.txt" + assert tmp_hello_world.exists() + assert tmp_hello_world.read_text() == "hello" + + +def test_create_docker_context_with_git_subfolder(tmp_path): + # uv's pip install errors with git and subdirectory + # In this case, we go back to pip instead + docker_context_path = tmp_path / "builder_root" + docker_context_path.mkdir() + + image_spec = ImageSpec( + name="FLYTEKIT", + python_version="3.12", + apt_packages=["git"], + packages=["git+https://github.com/flyteorg/flytekit.git@master#subdirectory=plugins/flytekit-wandb"], + ) + + create_docker_context(image_spec, docker_context_path) + + dockerfile_path = docker_context_path / "Dockerfile" + assert dockerfile_path.exists() + dockerfile_content = dockerfile_path.read_text() + + assert "--requirement requirements_pip.txt" in dockerfile_content + requirements_path = docker_context_path / "requirements_pip.txt" + assert requirements_path.exists() + + +def test_create_docker_context_cuda(tmp_path): + docker_context_path = tmp_path / "builder_root" + docker_context_path.mkdir() + + image_spec = ImageSpec(cuda="12.4.1", cudnn="8") + + msg = "cuda and cudnn do not need to be specified. If you are installed" + + with pytest.raises(ValueError, match=msg): + create_docker_context(image_spec, docker_context_path) + + +@pytest.mark.skipif( + os.environ.get("_FLYTEKIT_TEST_DEFAULT_BUILDER", "0") == "0", + reason="Set _FLYTEKIT_TEST_DEFAULT_BUILDER=1 to run this test", +) +def test_build(tmp_path): + docker_context_path = tmp_path / "builder_root" + docker_context_path.mkdir() + + source_root = tmp_path / "other_files" + source_root.mkdir() + hello_world_path = source_root / "hello_world.txt" + hello_world_path.write_text("hello") + + other_requirements_path = tmp_path / "requirements.txt" + other_requirements_path.write_text("threadpoolctl\n") + + image_spec = ImageSpec( + name="FLYTEKIT", + python_version="3.12", + env={"MY_ENV": "MY_VALUE"}, + apt_packages=["curl"], + conda_packages=["scipy==1.13.0", "numpy"], + packages=["pandas==2.2.1"], + requirements=os.fspath(other_requirements_path), + source_root=os.fspath(source_root), + commands=["mkdir my_dir"], + ) + + builder = DefaultImageBuilder() + + builder.build_image(image_spec) From 308982abbb993568cc16c19071763a04d036d3af Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Mon, 1 Jul 2024 18:31:24 -0400 Subject: [PATCH 014/192] Fix pandera plugin for 0.20 (#2545) Signed-off-by: Thomas J. Fan --- plugins/flytekit-pandera/tests/test_plugin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/flytekit-pandera/tests/test_plugin.py b/plugins/flytekit-pandera/tests/test_plugin.py index a3e7c82565..1357fdf135 100644 --- a/plugins/flytekit-pandera/tests/test_plugin.py +++ b/plugins/flytekit-pandera/tests/test_plugin.py @@ -7,7 +7,7 @@ def test_pandera_dataframe_type_hints(): - class InSchema(pandera.SchemaModel): + class InSchema(pandera.DataFrameModel): col1: pandera.typing.Series[int] col2: pandera.typing.Series[float] From 47248eb3184c04e15a7cc7eca6fb988765284a86 Mon Sep 17 00:00:00 2001 From: Future-Outlier Date: Mon, 1 Jul 2024 15:55:31 -0700 Subject: [PATCH 015/192] Fix `Flytefile` Remote Path Error (#2544) * print messages for sandbox execute and local execute Signed-off-by: Future-Outlier * remove PR 2476 change Signed-off-by: Future-Outlier * print execution mode in flytefile Signed-off-by: Future-Outlier * add should upload Signed-off-by: Future-Outlier * add condition ctx.execution_state.mode != ExecutionState.Mode.TASK_EXECUTION Signed-off-by: Future-Outlier * bug found, need to call the function Signed-off-by: Future-Outlier * remove comparison Signed-off-by: Future-Outlier * use ctx.execution_state.is_local_execution() Signed-off-by: Future-Outlier * lint Signed-off-by: Future-Outlier * fix bug Signed-off-by: Future-Outlier * remove comments Signed-off-by: Future-Outlier * revert is_local_execution Signed-off-by: Future-Outlier --------- Signed-off-by: Future-Outlier --- flytekit/core/context_manager.py | 2 +- flytekit/types/file/file.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index c51b60c1c9..506e55c829 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -547,7 +547,7 @@ def with_params( user_space_params=user_space_params if user_space_params else self.user_space_params, ) - def is_local_execution(self): + def is_local_execution(self) -> bool: return ( self.mode == ExecutionState.Mode.LOCAL_TASK_EXECUTION or self.mode == ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 567b29de57..5304fd21ed 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -440,7 +440,7 @@ def to_literal( # Set the remote destination if one was given instead of triggering a random one below remote_path = python_val.remote_path or None - if ctx.execution_state.is_local_execution and python_val.remote_path is None: + if ctx.execution_state.is_local_execution() and python_val.remote_path is None: should_upload = False elif isinstance(python_val, pathlib.Path) or isinstance(python_val, str): @@ -458,7 +458,7 @@ def to_literal( p = pathlib.Path(python_val) if not p.is_file(): raise TypeTransformerFailedError(f"Error converting {python_val} because it's not a file.") - if ctx.execution_state.is_local_execution: + if ctx.execution_state.is_local_execution(): should_upload = False # python_type must be os.PathLike - see check at beginning of function else: From 214be0cebbef3aac78b6f84390c8cb180531d8d6 Mon Sep 17 00:00:00 2001 From: novahow <58504997+novahow@users.noreply.github.com> Date: Tue, 2 Jul 2024 10:00:10 +0800 Subject: [PATCH 016/192] Core/enable deck (#2314) * add additional_decks support Signed-off-by: novahow modified: flytekit/core/base_task.py modified: flytekit/core/python_function_task.py modified: flytekit/core/task.py modified: flytekit/deck/deck.py modified: tests/flytekit/unit/core/test_flyte_file.py * add tests and remove confusing fields Signed-off-by: novahow modified: flytekit/core/base_task.py modified: flytekit/deck/deck.py modified: tests/flytekit/unit/deck/test_deck.py * add deckselector Signed-off-by: novahow modified: flytekit/core/base_task.py modified: flytekit/core/context_manager.py modified: flytekit/core/task.py modified: flytekit/deck/deck.py modified: tests/flytekit/unit/deck/test_deck.py * make deck_selector to tuple Signed-off-by: novahow * fix remote deck bug Signed-off-by: novahow * fix timelinedeck and remove rendered_deck param Signed-off-by: novahow * fix UI Signed-off-by: novahow * fix timelinedeck test multiple time_info Signed-off-by: novahow * nit Signed-off-by: novahow * nit with enum Signed-off-by: novahow * nit deck_fields Signed-off-by: novahow * enable all decks, remove plotly dep Signed-off-by: novahow * kevin's update Signed-off-by: Kevin Su * nit Signed-off-by: Kevin Su * remove chart Signed-off-by: Kevin Su --------- Signed-off-by: novahow Signed-off-by: Kevin Su Co-authored-by: Kevin Su --- flytekit/core/base_task.py | 60 ++++++++++++++++----- flytekit/core/context_manager.py | 12 ++++- flytekit/core/python_function_task.py | 16 +++--- flytekit/core/task.py | 12 +++++ flytekit/deck/__init__.py | 2 +- flytekit/deck/deck.py | 71 ++++++++++++++++-------- tests/flytekit/unit/core/test_utils.py | 3 ++ tests/flytekit/unit/deck/test_deck.py | 75 ++++++++++++++++++++++---- 8 files changed, 195 insertions(+), 56 deletions(-) diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index e7c4c44296..8b7f0190ef 100644 --- a/flytekit/core/base_task.py +++ b/flytekit/core/base_task.py @@ -70,6 +70,7 @@ from flytekit.core.tracker import TrackedInstance from flytekit.core.type_engine import TypeEngine, TypeTransformerFailedError from flytekit.core.utils import timeit +from flytekit.deck import DeckField from flytekit.loggers import logger from flytekit.models import dynamic_job as _dynamic_job from flytekit.models import interface as _interface_models @@ -464,6 +465,13 @@ def __init__( environment: Optional[Dict[str, str]] = None, disable_deck: Optional[bool] = None, enable_deck: Optional[bool] = None, + deck_fields: Optional[Tuple[DeckField, ...]] = ( + DeckField.SOURCE_CODE, + DeckField.DEPENDENCIES, + DeckField.TIMELINE, + DeckField.INPUT, + DeckField.OUTPUT, + ), **kwargs, ): """ @@ -479,6 +487,8 @@ def __init__( execution of the task. Supplied as a dictionary of key/value pairs disable_deck (bool): (deprecated) If true, this task will not output deck html file enable_deck (bool): If true, this task will output deck html file + deck_fields (Tuple[DeckField]): Tuple of decks to be + generated for this task. Valid values can be selected from fields of ``flytekit.deck.DeckField`` enum """ super().__init__( task_type=task_type, @@ -490,22 +500,35 @@ def __init__( self._environment = environment if environment else {} self._task_config = task_config + # first we resolve the conflict between params regarding decks, if any two of [disable_deck, enable_deck] + # are set, we raise an error + configured_deck_params = [disable_deck is not None, enable_deck is not None] + if sum(configured_deck_params) > 1: + raise ValueError("only one of [disable_deck, enable_deck] can be set") + if disable_deck is not None: warnings.warn( "disable_deck was deprecated in 1.10.0, please use enable_deck instead", FutureWarning, ) - # Confirm that disable_deck and enable_deck do not contradict each other - if disable_deck is not None and enable_deck is not None: - raise ValueError("disable_deck and enable_deck cannot both be set at the same time") - if enable_deck is not None: self._disable_deck = not enable_deck elif disable_deck is not None: self._disable_deck = disable_deck else: self._disable_deck = True + + self._deck_fields = list(deck_fields) if (deck_fields is not None and self.disable_deck is False) else [] + + deck_members = set([_field for _field in DeckField]) + # enumerate additional decks, check if any of them are invalid + for deck in self._deck_fields: + if deck not in deck_members: + raise ValueError( + f"Element {deck} from deck_fields param is not a valid deck field. Please use one of {deck_members}" + ) + if self._python_interface.docstring: if self.docs is None: self._docs = Documentation( @@ -645,18 +668,20 @@ def _output_to_literal_map(self, native_outputs: Dict[int, Any], ctx: FlyteConte def _write_decks(self, native_inputs, native_outputs_as_map, ctx, new_user_params): if self._disable_deck is False: - from flytekit.deck.deck import Deck, _output_deck + from flytekit.deck.deck import Deck, DeckField, _output_deck - INPUT = "Inputs" - OUTPUT = "Outputs" + INPUT = DeckField.INPUT + OUTPUT = DeckField.OUTPUT - input_deck = Deck(INPUT) - for k, v in native_inputs.items(): - input_deck.append(TypeEngine.to_html(ctx, v, self.get_type_for_input_var(k, v))) + if DeckField.INPUT in self.deck_fields: + input_deck = Deck(INPUT.value) + for k, v in native_inputs.items(): + input_deck.append(TypeEngine.to_html(ctx, v, self.get_type_for_input_var(k, v))) - output_deck = Deck(OUTPUT) - for k, v in native_outputs_as_map.items(): - output_deck.append(TypeEngine.to_html(ctx, v, self.get_type_for_output_var(k, v))) + if DeckField.OUTPUT in self.deck_fields: + output_deck = Deck(OUTPUT.value) + for k, v in native_outputs_as_map.items(): + output_deck.append(TypeEngine.to_html(ctx, v, self.get_type_for_output_var(k, v))) if ctx.execution_state and ctx.execution_state.is_local_execution(): # When we run the workflow remotely, flytekit outputs decks at the end of _dispatch_execute @@ -681,6 +706,8 @@ def dispatch_execute( may be none * ``DynamicJobSpec`` is returned when a dynamic workflow is executed """ + if DeckField.TIMELINE.value in self.deck_fields and ctx.user_space_params is not None: + ctx.user_space_params.decks.append(ctx.user_space_params.timeline_deck) # Invoked before the task is executed new_user_params = self.pre_execute(ctx.user_space_params) @@ -791,6 +818,13 @@ def disable_deck(self) -> bool: """ return self._disable_deck + @property + def deck_fields(self) -> List[DeckField]: + """ + If not empty, this task will output deck html file for the specified decks + """ + return self._deck_fields + class TaskResolverMixin(object): """ diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index 506e55c829..b30d454a50 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -88,6 +88,7 @@ class Builder(object): execution_date: typing.Optional[datetime] = None logging: Optional[_logging.Logger] = None task_id: typing.Optional[_identifier.Identifier] = None + output_metadata_prefix: Optional[str] = None def __init__(self, current: typing.Optional[ExecutionParameters] = None): self.stats = current.stats if current else None @@ -100,6 +101,7 @@ def __init__(self, current: typing.Optional[ExecutionParameters] = None): self.attrs = current._attrs if current else {} self.raw_output_prefix = current.raw_output_prefix if current else None self.task_id = current.task_id if current else None + self.output_metadata_prefix = current.output_metadata_prefix if current else None def add_attr(self, key: str, v: typing.Any) -> ExecutionParameters.Builder: self.attrs[key] = v @@ -118,6 +120,7 @@ def build(self) -> ExecutionParameters: decks=self.decks, raw_output_prefix=self.raw_output_prefix, task_id=self.task_id, + output_metadata_prefix=self.output_metadata_prefix, **self.attrs, ) @@ -181,6 +184,7 @@ def __init__( self._checkpoint = checkpoint self._decks = decks self._task_id = task_id + self._timeline_deck = None @property def stats(self) -> taggable.TaggableStats: @@ -273,7 +277,7 @@ def default_deck(self) -> Deck: @property def timeline_deck(self) -> "TimeLineDeck": # type: ignore - from flytekit.deck.deck import TimeLineDeck + from flytekit.deck.deck import DeckField, TimeLineDeck time_line_deck = None for deck in self.decks: @@ -281,8 +285,12 @@ def timeline_deck(self) -> "TimeLineDeck": # type: ignore time_line_deck = deck break if time_line_deck is None: - time_line_deck = TimeLineDeck("Timeline") + if self._timeline_deck is not None: + time_line_deck = self._timeline_deck + else: + time_line_deck = TimeLineDeck(DeckField.TIMELINE.value, auto_add_to_deck=False) + self._timeline_deck = time_line_deck return time_line_deck def __getattr__(self, attr_name: str) -> typing.Any: diff --git a/flytekit/core/python_function_task.py b/flytekit/core/python_function_task.py index f2f354fa27..c3464e053d 100644 --- a/flytekit/core/python_function_task.py +++ b/flytekit/core/python_function_task.py @@ -352,7 +352,7 @@ def dynamic_execute(self, task_function: Callable, **kwargs) -> Any: def _write_decks(self, native_inputs, native_outputs_as_map, ctx, new_user_params): if self._disable_deck is False: - from flytekit.deck import Deck + from flytekit.deck import Deck, DeckField from flytekit.deck.renderer import PythonDependencyRenderer # These errors are raised if the source code can not be retrieved @@ -360,12 +360,14 @@ def _write_decks(self, native_inputs, native_outputs_as_map, ctx, new_user_param source_code = inspect.getsource(self._task_function) from flytekit.deck.renderer import SourceCodeRenderer - source_code_deck = Deck("Source Code") - renderer = SourceCodeRenderer() - source_code_deck.append(renderer.to_html(source_code)) + if DeckField.SOURCE_CODE in self.deck_fields: + source_code_deck = Deck(DeckField.SOURCE_CODE.value) + renderer = SourceCodeRenderer() + source_code_deck.append(renderer.to_html(source_code)) - python_dependencies_deck = Deck("Dependencies") - renderer = PythonDependencyRenderer() - python_dependencies_deck.append(renderer.to_html()) + if DeckField.DEPENDENCIES in self.deck_fields: + python_dependencies_deck = Deck(DeckField.DEPENDENCIES.value) + renderer = PythonDependencyRenderer() + python_dependencies_deck.append(renderer.to_html()) return super()._write_decks(native_inputs, native_outputs_as_map, ctx, new_user_params) diff --git a/flytekit/core/task.py b/flytekit/core/task.py index d30947509d..7e420269d3 100644 --- a/flytekit/core/task.py +++ b/flytekit/core/task.py @@ -12,6 +12,7 @@ from flytekit.core.python_function_task import PythonFunctionTask from flytekit.core.reference_entity import ReferenceEntity, TaskReference from flytekit.core.resources import Resources +from flytekit.deck import DeckField from flytekit.extras.accelerators import BaseAccelerator from flytekit.image_spec.image_spec import ImageSpec from flytekit.models.documentation import Documentation @@ -114,6 +115,7 @@ def task( docs: Optional[Documentation] = ..., disable_deck: Optional[bool] = ..., enable_deck: Optional[bool] = ..., + deck_fields: Optional[Tuple[DeckField, ...]] = ..., pod_template: Optional["PodTemplate"] = ..., pod_template_name: Optional[str] = ..., accelerator: Optional[BaseAccelerator] = ..., @@ -151,6 +153,7 @@ def task( docs: Optional[Documentation] = ..., disable_deck: Optional[bool] = ..., enable_deck: Optional[bool] = ..., + deck_fields: Optional[Tuple[DeckField, ...]] = ..., pod_template: Optional["PodTemplate"] = ..., pod_template_name: Optional[str] = ..., accelerator: Optional[BaseAccelerator] = ..., @@ -187,6 +190,13 @@ def task( docs: Optional[Documentation] = None, disable_deck: Optional[bool] = None, enable_deck: Optional[bool] = None, + deck_fields: Optional[Tuple[DeckField, ...]] = ( + DeckField.SOURCE_CODE, + DeckField.DEPENDENCIES, + DeckField.TIMELINE, + DeckField.INPUT, + DeckField.OUTPUT, + ), pod_template: Optional["PodTemplate"] = None, pod_template_name: Optional[str] = None, accelerator: Optional[BaseAccelerator] = None, @@ -307,6 +317,7 @@ def launch_dynamically(): :param task_resolver: Provide a custom task resolver. :param disable_deck: (deprecated) If true, this task will not output deck html file :param enable_deck: If true, this task will output deck html file + :param deck_fields: If specified and enble_deck is True, this task will output deck html file with the fields specified in the tuple :param docs: Documentation about this task :param pod_template: Custom PodTemplate for this task. :param pod_template_name: The name of the existing PodTemplate resource which will be used in this task. @@ -339,6 +350,7 @@ def wrapper(fn: Callable[..., Any]) -> PythonFunctionTask[T]: task_resolver=task_resolver, disable_deck=disable_deck, enable_deck=enable_deck, + deck_fields=deck_fields, docs=docs, pod_template=pod_template, pod_template_name=pod_template_name, diff --git a/flytekit/deck/__init__.py b/flytekit/deck/__init__.py index 5250ad0adc..58da56cf64 100644 --- a/flytekit/deck/__init__.py +++ b/flytekit/deck/__init__.py @@ -18,5 +18,5 @@ SourceCodeRenderer """ -from .deck import Deck +from .deck import Deck, DeckField from .renderer import MarkdownRenderer, SourceCodeRenderer, TopFrameRenderer diff --git a/flytekit/deck/deck.py b/flytekit/deck/deck.py index 3ce9d058a4..fbb398ef49 100644 --- a/flytekit/deck/deck.py +++ b/flytekit/deck/deck.py @@ -1,3 +1,4 @@ +import enum import os import typing from typing import Optional @@ -10,6 +11,18 @@ DECK_FILE_NAME = "deck.html" +class DeckField(str, enum.Enum): + """ + DeckField is used to specify the fields that will be rendered in the deck. + """ + + INPUT = "Input" + OUTPUT = "Output" + SOURCE_CODE = "Source Code" + TIMELINE = "Timeline" + DEPENDENCIES = "Dependencies" + + class Deck: """ Deck enable users to get customizable and default visibility into their tasks. @@ -52,10 +65,11 @@ def t2() -> Annotated[pd.DataFrame, TopFrameRenderer(10)]: return iris_df """ - def __init__(self, name: str, html: Optional[str] = ""): + def __init__(self, name: str, html: Optional[str] = "", auto_add_to_deck: bool = True): self._name = name self._html = html - FlyteContextManager.current_context().user_space_params.decks.append(self) + if auto_add_to_deck: + FlyteContextManager.current_context().user_space_params.decks.append(self) def append(self, html: str) -> "Deck": assert isinstance(html, str) @@ -79,8 +93,8 @@ class TimeLineDeck(Deck): Instead, the complete data set is used to create a comprehensive visualization of the execution time of each part of the task. """ - def __init__(self, name: str, html: Optional[str] = ""): - super().__init__(name, html) + def __init__(self, name: str, html: Optional[str] = "", auto_add_to_deck: bool = True): + super().__init__(name, html, auto_add_to_deck) self.time_info = [] def append_time_info(self, info: dict): @@ -89,19 +103,9 @@ def append_time_info(self, info: dict): @property def html(self) -> str: - try: - from flytekitplugins.deck.renderer import GanttChartRenderer, TableRenderer - except ImportError: - warning_info = "Plugin 'flytekit-deck-standard' is not installed. To display time line, install the plugin in the image." - logger.warning(warning_info) - return warning_info - if len(self.time_info) == 0: return "" - import pandas - - df = pandas.DataFrame(self.time_info) note = """

Note:

    @@ -109,16 +113,36 @@ def html(self) -> str:
  1. For accurate execution time measurements, users should refer to wall time and process time.
""" - # set the accuracy to microsecond - df["ProcessTime"] = df["ProcessTime"].apply(lambda time: "{:.6f}".format(time)) - df["WallTime"] = df["WallTime"].apply(lambda time: "{:.6f}".format(time)) - gantt_chart_html = GanttChartRenderer().to_html(df) - time_table_html = TableRenderer().to_html( - df[["Name", "WallTime", "ProcessTime"]], - header_labels=["Name", "Wall Time(s)", "Process Time(s)"], - ) - return gantt_chart_html + time_table_html + note + return generate_time_table(self.time_info) + note + + +def generate_time_table(data: dict) -> str: + html = [ + '', + """ + + + + + + + + """, + "", + ] + + # Add table rows + for row in data: + html.append("") + html.append(f"") + html.append(f"") + html.append(f"") + html.append("") + html.append("") + + html.append("
NameWall Time(s)Process Time(s)
{row['Name']}{row['WallTime']:.6f}{row['ProcessTime']:.6f}
") + return "".join(html) def _get_deck( @@ -129,6 +153,7 @@ def _get_deck( If ignore_jupyter is set to True, then it will return a str even in a jupyter environment. """ deck_map = {deck.name: deck.html for deck in new_user_params.decks} + raw_html = get_deck_template().render(metadata=deck_map) if not ignore_jupyter and ipython_check(): try: diff --git a/tests/flytekit/unit/core/test_utils.py b/tests/flytekit/unit/core/test_utils.py index ca0d07565d..3e9c42dba0 100644 --- a/tests/flytekit/unit/core/test_utils.py +++ b/tests/flytekit/unit/core/test_utils.py @@ -33,9 +33,12 @@ def test_timeit(): ctx = FlyteContextManager.current_context() ctx.user_space_params._decks = [] + from flytekit.deck import DeckField + with timeit("Set disable_deck to False"): kwargs = {} kwargs["disable_deck"] = False + kwargs["deck_fields"] = (DeckField.TIMELINE.value,) ctx = FlyteContextManager.current_context() time_info_list = ctx.user_space_params.timeline_deck.time_info diff --git a/tests/flytekit/unit/deck/test_deck.py b/tests/flytekit/unit/deck/test_deck.py index 9cf497bccd..ce07317a94 100644 --- a/tests/flytekit/unit/deck/test_deck.py +++ b/tests/flytekit/unit/deck/test_deck.py @@ -7,7 +7,7 @@ import flytekit from flytekit import Deck, FlyteContextManager, task -from flytekit.deck import MarkdownRenderer, SourceCodeRenderer, TopFrameRenderer +from flytekit.deck import DeckField, MarkdownRenderer, SourceCodeRenderer, TopFrameRenderer from flytekit.deck.deck import _output_deck from flytekit.deck.renderer import PythonDependencyRenderer @@ -40,20 +40,21 @@ def test_timeline_deck(): ) ctx = FlyteContextManager.current_context() ctx.user_space_params._decks = [] + ctx.user_space_params._timeline_deck = None timeline_deck = ctx.user_space_params.timeline_deck timeline_deck.append_time_info(time_info) assert timeline_deck.name == "Timeline" assert len(timeline_deck.time_info) == 1 assert timeline_deck.time_info[0] == time_info - assert len(ctx.user_space_params.decks) == 1 + assert len(ctx.user_space_params.decks) == 0 @pytest.mark.parametrize( "disable_deck,expected_decks", [ - (None, 1), # time line deck - (False, 5), # time line deck + source code deck + python dependency deck + input and output decks - (True, 1), # time line deck + (None, 0), + (False, 5), # source code + dependency + input + output + timeline decks + (True, 0), ], ) def test_deck_for_task(disable_deck, expected_decks): @@ -71,26 +72,80 @@ def t1(a: int) -> str: assert len(ctx.user_space_params.decks) == expected_decks +@pytest.mark.parametrize( + "deck_fields,enable_deck,expected_decks", + [ + ((), True, 0), + ((DeckField.INPUT.value), False, 0), + ( + (DeckField.OUTPUT.value, DeckField.INPUT.value, DeckField.TIMELINE.value, DeckField.DEPENDENCIES.value), + True, + 4, # time line deck + dependency + input and output decks + ), + (None, True, 5), # source code + dependency + input + output + timeline decks + ], +) +@mock.patch("flytekit.deck.deck._output_deck") +def test_additional_deck_for_task(_output_deck, deck_fields, enable_deck, expected_decks): + ctx = FlyteContextManager.current_context() + + kwargs = {} + if deck_fields is not None: + kwargs["deck_fields"] = deck_fields + if enable_deck is not None: + kwargs["enable_deck"] = enable_deck + + @task(**kwargs) + def t1(a: int) -> str: + return str(a) + + t1(a=3) + assert len(ctx.user_space_params.decks) == expected_decks + + +@pytest.mark.parametrize( + "deck_fields,enable_deck,disable_deck", + [ + (None, True, False), + (("WrongDeck", DeckField.INPUT.value, DeckField.OUTPUT.value), True, None), # WrongDeck is not a valid field + ], +) +def test_invalid_deck_params(deck_fields, enable_deck, disable_deck): + kwargs = {} + if deck_fields is not None: + kwargs["deck_fields"] = deck_fields + if enable_deck is not None: + kwargs["enable_deck"] = enable_deck + if disable_deck is not None: + kwargs["disable_deck"] = disable_deck + + with pytest.raises(ValueError): + + @task(**kwargs) + def t1(a: int) -> str: + return str(a) + + @pytest.mark.skipif("pandas" not in sys.modules, reason="Pandas is not installed.") @pytest.mark.filterwarnings("ignore:disable_deck was deprecated") @pytest.mark.parametrize( "enable_deck,disable_deck, expected_decks, expect_error", [ - (None, None, 2, False), # default deck and time line deck + (None, None, 1, False), # default deck ( None, False, 6, False, - ), # default deck and time line deck + source code deck + python dependency deck + input and output decks - (None, True, 2, False), # default deck and time line deck + ), # default deck + source code + dependency + input + output + timeline decks + (None, True, 1, False), # default deck ( True, None, 6, False, - ), # default deck and time line deck + source code deck + python dependency deck + input and output decks - (False, None, 2, False), # default deck and time line deck + ), # default deck + source code + dependency + input + output + timeline decks + (False, None, 1, False), # default deck (True, True, -1, True), # Set both disable_deck and enable_deck to True and confirm that it fails (False, False, -1, True), # Set both disable_deck and enable_deck to False and confirm that it fails ], From d419b660b3c34f30fca7fa737fb7dce269ea7a77 Mon Sep 17 00:00:00 2001 From: Future-Outlier Date: Mon, 1 Jul 2024 20:49:28 -0700 Subject: [PATCH 017/192] Add thonmas and future outlier to CODEOWNERS (#2548) --- CODEOWNERS | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CODEOWNERS b/CODEOWNERS index dafddb874b..917fe185c0 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -1,3 +1,3 @@ # These owners will be the default owners for everything in # the repo. Unless a later match takes precedence. -* @wild-endeavor @kumare3 @eapolinario @pingsutw @cosmicBboy @samhita-alla +* @wild-endeavor @kumare3 @eapolinario @pingsutw @cosmicBboy @samhita-alla @thomasjpfan @future-outlier From bd01c3d0db180af02959e81365aaf4201617c52e Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 2 Jul 2024 14:04:19 +0800 Subject: [PATCH 018/192] Improve error message for ImageSpec (#2498) Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 28 +++++++++++-------- .../flytekitplugins/envd/image_builder.py | 16 +++++++++-- .../unit/core/image_spec/test_image_spec.py | 2 +- 3 files changed, 32 insertions(+), 14 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 05abacdbf9..a17a2e272e 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -104,9 +104,11 @@ def is_container(self) -> bool: return os.environ.get(_F_IMG_ID) == self.image_name() return True - def exist(self) -> bool: + def exist(self) -> Optional[bool]: """ Check if the image exists in the registry. + Return True if the image exists in the registry, False otherwise. + Return None if failed to check if the image exists due to the permission issue or other reasons. """ import docker from docker.errors import APIError, ImageNotFound @@ -121,7 +123,9 @@ def exist(self) -> bool: except APIError as e: if e.response.status_code == 404: return False - return True + + click.secho(f"Failed to check if the image exists with error:\n {e}", fg="red") + return None except ImageNotFound: return False except Exception as e: @@ -153,9 +157,8 @@ def exist(self) -> bool: f" pip install --upgrade docker" ) - click.secho(f"Failed to check if the image exists with error : {e}", fg="red") - click.secho("Flytekit assumes that the image already exists.", fg="blue") - return True + click.secho(f"Failed to check if the image exists with error:\n {e}", fg="red") + return None def __hash__(self): return hash(asdict(self).__str__()) @@ -240,14 +243,17 @@ def should_build(self, image_spec: ImageSpec) -> bool: True if the image should be built, otherwise it returns False. """ img_name = image_spec.image_name() - if not image_spec.exist(): + exist = image_spec.exist() + if exist is False: click.secho(f"Image {img_name} not found. building...", fg="blue") return True - if image_spec._is_force_push: - click.secho(f"Image {img_name} found but overwriting existing image.", fg="blue") - return True - - click.secho(f"Image {img_name} found. Skip building.", fg="blue") + elif exist is True: + if image_spec._is_force_push: + click.secho(f"Overwriting existing image {img_name}.", fg="blue") + return True + click.secho(f"Image {img_name} found. Skip building.", fg="blue") + else: + click.secho(f"Flytekit assumes the image {img_name} already exists.", fg="blue") return False diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 8a8d893cfd..af5c32ec6d 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -2,10 +2,13 @@ import pathlib import shutil import subprocess +from dataclasses import asdict from importlib import metadata import click from packaging.version import Version +from rich import print +from rich.pretty import Pretty from flytekit.configuration import DefaultImages from flytekit.core import context_manager @@ -35,7 +38,16 @@ def build_image(self, image_spec: ImageSpec): else: build_command += f" --tag {image_spec.image_name()}" envd_context_switch(image_spec.registry) - execute_command(build_command) + try: + execute_command(build_command) + except Exception as e: + click.secho("❌ Failed to build image spec:", fg="red") + print( + Pretty( + asdict(image_spec, dict_factory=lambda x: {k: v for (k, v) in x if v is not None}), indent_size=2 + ) + ) + raise e from None def envd_context_switch(registry: str): @@ -72,7 +84,7 @@ def execute_command(command: str): if p.returncode != 0: _, stderr = p.communicate() - raise Exception(f"failed to run command {command} with error {stderr}") + raise RuntimeError(f"failed to run command {command} with error:\n {stderr.decode()}") return result diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py index 011828d4ce..02c734b119 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -66,7 +66,7 @@ def test_image_spec(mock_image_spec_builder): assert "dummy" in ImageBuildEngine._REGISTRY assert calculate_hash_from_image_spec(image_spec) == tag - assert image_spec.exist() is True + assert image_spec.exist() is None # Remove the dummy builder, and build the image again # The image has already been built, so it shouldn't fail. From a9c130800168df54149e3f33518d462e5a13a3ec Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 2 Jul 2024 14:12:47 +0800 Subject: [PATCH 019/192] Improve error message for pyflyte run (#2472) Signed-off-by: Kevin Su --- flytekit/clis/sdk_in_container/run.py | 3 +- flytekit/clis/sdk_in_container/serve.py | 24 ++++-- flytekit/clis/sdk_in_container/utils.py | 4 +- flytekit/configuration/file.py | 1 - flytekit/core/base_task.py | 1 - flytekit/core/context_manager.py | 6 +- flytekit/core/interface.py | 4 +- flytekit/core/reference_entity.py | 1 - flytekit/core/shim_task.py | 1 - flytekit/core/tracker.py | 6 +- flytekit/core/utils.py | 8 +- flytekit/extend/backend/base_agent.py | 8 +- flytekit/loggers.py | 82 +++++++++++++------ flytekit/remote/remote.py | 24 +++--- .../types/structured/structured_dataset.py | 10 +-- 15 files changed, 106 insertions(+), 77 deletions(-) diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index e99f114818..1518d592f7 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -528,7 +528,8 @@ def _run(*args, **kwargs): # By the time we get to this function, all the loading has already happened run_level_params: RunLevelParams = ctx.obj - logger.debug(f"Running {entity.name} with {kwargs} and run_level_params {run_level_params}") + entity_type = "workflow" if isinstance(entity, PythonFunctionWorkflow) else "task" + logger.debug(f"Running {entity_type} {entity.name} with input {kwargs}") click.secho(f"Running Execution on {'Remote' if run_level_params.is_remote else 'local'}.", fg="cyan") try: diff --git a/flytekit/clis/sdk_in_container/serve.py b/flytekit/clis/sdk_in_container/serve.py index efe7086126..d6ceab54fc 100644 --- a/flytekit/clis/sdk_in_container/serve.py +++ b/flytekit/clis/sdk_in_container/serve.py @@ -8,6 +8,8 @@ add_AsyncAgentServiceServicer_to_server, add_SyncAgentServiceServicer_to_server, ) +from rich.console import Console +from rich.table import Table @click.group("serve") @@ -55,8 +57,8 @@ def agent(_: click.Context, port, worker, timeout): async def _start_grpc_server(port: int, worker: int, timeout: int): from flytekit.extend.backend.agent_service import AgentMetadataService, AsyncAgentService, SyncAgentService + click.secho("🚀 Starting the agent service...") _start_http_server() - click.secho("Starting the agent service...", fg="blue") print_agents_metadata() server = grpc.aio.server(futures.ThreadPoolExecutor(max_workers=worker)) @@ -75,7 +77,7 @@ def _start_http_server(): try: from prometheus_client import start_http_server - click.secho("Starting up the server to expose the prometheus metrics...", fg="blue") + click.secho("Starting up the server to expose the prometheus metrics...") start_http_server(9090) except ImportError as e: click.secho(f"Failed to start the prometheus server with error {e}", fg="red") @@ -104,7 +106,17 @@ def print_agents_metadata(): from flytekit.extend.backend.base_agent import AgentRegistry agents = AgentRegistry.list_agents() - for agent in agents: - name = agent.name - metadata = [category.name for category in agent.supported_task_categories] - click.secho(f"Starting {name} that supports task categories {metadata}", fg="blue") + + table = Table(title="Agent Metadata") + table.add_column("Agent Name", style="cyan", no_wrap=True) + table.add_column("Support Task Types", style="cyan") + table.add_column("Is Sync", style="green") + + for a in agents: + categories = "" + for c in a.supported_task_categories: + categories += f"{c.name} (v{c.version}) " + table.add_row(a.name, categories, str(a.is_sync)) + + console = Console() + console.print(table) diff --git a/flytekit/clis/sdk_in_container/utils.py b/flytekit/clis/sdk_in_container/utils.py index 942fb1f53f..056ac2638b 100644 --- a/flytekit/clis/sdk_in_container/utils.py +++ b/flytekit/clis/sdk_in_container/utils.py @@ -9,7 +9,7 @@ from flytekit.exceptions.base import FlyteException from flytekit.exceptions.user import FlyteInvalidInputException -from flytekit.loggers import get_level_from_cli_verbosity, logger, upgrade_to_rich_logging +from flytekit.loggers import get_level_from_cli_verbosity, logger project_option = click.Option( param_decls=["-p", "--project"], @@ -137,7 +137,7 @@ class ErrorHandlingCommand(click.RichGroup): def invoke(self, ctx: click.Context) -> typing.Any: verbose = ctx.params["verbose"] log_level = get_level_from_cli_verbosity(verbose) - upgrade_to_rich_logging(log_level=log_level) + logger.setLevel(log_level) try: return super().invoke(ctx) except Exception as e: diff --git a/flytekit/configuration/file.py b/flytekit/configuration/file.py index 80bb8fd62b..521bc72f61 100644 --- a/flytekit/configuration/file.py +++ b/flytekit/configuration/file.py @@ -219,7 +219,6 @@ def _get_from_yaml(self, c: YamlConfigEntry) -> typing.Any: d = d[k] return d except KeyError: - logger.debug(f"Switch {c.switch} could not be found in yaml config") return None def get(self, c: typing.Union[LegacyConfigEntry, YamlConfigEntry]) -> typing.Any: diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index 8b7f0190ef..c161df296e 100644 --- a/flytekit/core/base_task.py +++ b/flytekit/core/base_task.py @@ -756,7 +756,6 @@ def dispatch_execute( return self._async_execute(native_inputs, native_outputs, ctx, exec_ctx, new_user_params) - logger.debug("Task executed successfully in user level") # Lets run the post_execute method. This may result in a IgnoreOutputs Exception, which is # bubbled up to be handled at the callee layer. native_outputs = self.post_execute(new_user_params, native_outputs) diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index b30d454a50..340046e941 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -33,7 +33,7 @@ from flytekit.core.node import Node from flytekit.interfaces.cli_identifiers import WorkflowExecutionIdentifier from flytekit.interfaces.stats import taggable -from flytekit.loggers import logger, user_space_logger +from flytekit.loggers import developer_logger, user_space_logger from flytekit.models.core import identifier as _identifier if typing.TYPE_CHECKING: @@ -867,7 +867,7 @@ def push_context(ctx: FlyteContext, f: Optional[traceback.FrameSummary] = None) context_list.append(ctx) flyte_context_Var.set(context_list) t = "\t" - logger.debug( + developer_logger.debug( f"{t * ctx.level}[{len(flyte_context_Var.get())}] Pushing context - {'compile' if ctx.compilation_state else 'execute'}, branch[{ctx.in_a_condition}], {ctx.get_origin_stackframe_repr()}" ) return ctx @@ -878,7 +878,7 @@ def pop_context() -> FlyteContext: ctx = context_list.pop() flyte_context_Var.set(context_list) t = "\t" - logger.debug( + developer_logger.debug( f"{t * ctx.level}[{len(flyte_context_Var.get()) + 1}] Popping context - {'compile' if ctx.compilation_state else 'execute'}, branch[{ctx.in_a_condition}], {ctx.get_origin_stackframe_repr()}" ) if len(flyte_context_Var.get()) == 0: diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 7faf541661..42bc87242e 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -16,7 +16,7 @@ from flytekit.core.sentinel import DYNAMIC_INPUT_BINDING from flytekit.core.type_engine import TypeEngine, UnionTransformer from flytekit.exceptions.user import FlyteValidationException -from flytekit.loggers import logger +from flytekit.loggers import developer_logger, logger from flytekit.models import interface as _interface_models from flytekit.models.literals import Literal, Scalar, Void @@ -512,7 +512,7 @@ def t(a: int, b: str) -> Dict[str, int]: ... else: # Handle all other single return types - logger.debug(f"Task returns unnamed native tuple {return_annotation}") + developer_logger.debug(f"Task returns unnamed native tuple {return_annotation}") return {default_output_name(): cast(Type, return_annotation)} diff --git a/flytekit/core/reference_entity.py b/flytekit/core/reference_entity.py index b54c4d67f6..1c33bbedaa 100644 --- a/flytekit/core/reference_entity.py +++ b/flytekit/core/reference_entity.py @@ -124,7 +124,6 @@ def unwrap_literal_map_and_execute( except Exception as e: logger.exception(f"Exception when executing {e}") raise e - logger.debug("Task executed successfully in user level") expected_output_names = list(self.python_interface.outputs.keys()) if len(expected_output_names) == 1: diff --git a/flytekit/core/shim_task.py b/flytekit/core/shim_task.py index f96db3e49c..b205bbab08 100644 --- a/flytekit/core/shim_task.py +++ b/flytekit/core/shim_task.py @@ -108,7 +108,6 @@ def dispatch_execute( logger.exception(f"Exception when executing {e}") raise e - logger.debug("Task executed successfully in user level") # Lets run the post_execute method. This may result in a IgnoreOutputs Exception, which is # bubbled up to be handled at the callee layer. native_outputs = self.post_execute(new_user_params, native_outputs) diff --git a/flytekit/core/tracker.py b/flytekit/core/tracker.py index 2c45d926e8..2d7c0360ed 100644 --- a/flytekit/core/tracker.py +++ b/flytekit/core/tracker.py @@ -9,7 +9,7 @@ from flytekit.configuration.feature_flags import FeatureFlags from flytekit.exceptions import system as _system_exceptions -from flytekit.loggers import logger +from flytekit.loggers import developer_logger, logger def import_module_from_file(module_name, file): @@ -125,12 +125,12 @@ def find_lhs(self) -> str: if self._instantiated_in is None or self._instantiated_in == "": raise _system_exceptions.FlyteSystemException(f"Object {self} does not have an _instantiated in") - logger.debug(f"Looking for LHS for {self} from {self._instantiated_in}") + developer_logger.debug(f"Looking for LHS for {self} from {self._instantiated_in}") m = importlib.import_module(self._instantiated_in) for k in dir(m): try: if getattr(m, k) is self: - logger.debug(f"Found LHS for {self}, {k}") + developer_logger.debug(f"Found LHS for {self}, {k}") self._lhs = k return k except ValueError as err: diff --git a/flytekit/core/utils.py b/flytekit/core/utils.py index 3106b3294e..4c064e8f34 100644 --- a/flytekit/core/utils.py +++ b/flytekit/core/utils.py @@ -335,13 +335,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): ) ) - logger.info( - "{}. [Wall Time: {}s, Process Time: {}s]".format( - self._name, - end_wall_time - self._start_wall_time, - end_process_time - self._start_process_time, - ) - ) + logger.info(f"{self._name}. [Time: {end_wall_time - self._start_wall_time:.6f}s]") class ClassDecorator(ABC): diff --git a/flytekit/extend/backend/base_agent.py b/flytekit/extend/backend/base_agent.py index 1630bf71aa..9d42910070 100644 --- a/flytekit/extend/backend/base_agent.py +++ b/flytekit/extend/backend/base_agent.py @@ -18,7 +18,7 @@ from rich.logging import RichHandler from rich.progress import Progress -from flytekit import FlyteContext, PythonFunctionTask, logger +from flytekit import FlyteContext, PythonFunctionTask from flytekit.configuration import ImageConfig, SerializationSettings from flytekit.core import utils from flytekit.core.base_task import PythonTask @@ -210,8 +210,6 @@ def register(agent: Union[AsyncAgentBase, SyncAgentBase], override: bool = False ) AgentRegistry._METADATA[agent.name] = agent_metadata - logger.info(f"Registering {agent.name} for task type: {agent.task_category}") - @staticmethod def get_agent(task_type_name: str, task_type_version: int = 0) -> Union[SyncAgentBase, AsyncAgentBase]: task_category = TaskCategory(name=task_type_name, version=task_type_version) @@ -272,8 +270,8 @@ async def _do( return await mirror_async_methods( agent.do, task_template=template, inputs=literal_map, output_prefix=output_prefix ) - except Exception as error_message: - raise FlyteUserException(f"Failed to run the task {self.name} with error: {error_message}") + except Exception as e: + raise FlyteUserException(f"Failed to run the task {self.name} with error: {e}") from None class AsyncAgentExecutorMixin: diff --git a/flytekit/loggers.py b/flytekit/loggers.py index 1a0165f007..12b3e90423 100644 --- a/flytekit/loggers.py +++ b/flytekit/loggers.py @@ -13,12 +13,14 @@ # For now, assume this is the environment variable whose usage will remain unchanged and controls output for all # loggers defined in this file. LOGGING_ENV_VAR = "FLYTE_SDK_LOGGING_LEVEL" +LOGGING_DEV_ENV_VAR = "FLYTE_SDK_DEV_LOGGING_LEVEL" LOGGING_FMT_ENV_VAR = "FLYTE_SDK_LOGGING_FORMAT" LOGGING_RICH_FMT_ENV_VAR = "FLYTE_SDK_RICH_TRACEBACKS" # By default, the root flytekit logger to debug so everything is logged, but enable fine-tuning logger = logging.getLogger("flytekit") user_space_logger = logging.getLogger("user_space") +developer_logger = logging.getLogger("developer") # Stop propagation so that configuration is isolated to this file (so that it doesn't matter what the # global Python root logger is set to). @@ -71,6 +73,27 @@ def set_user_logger_properties( user_space_logger.setLevel(level) +def set_developer_properties( + handler: typing.Optional[logging.Handler] = None, + filter: typing.Optional[logging.Filter] = None, + level: typing.Optional[int] = None, +): + """ + developer logger is only used for debugging. It is possible to selectively tune the logging for the developer. + + :param handler: logging.Handler to add to the user_space_logger + :param filter: logging.Filter to add to the user_space_logger + :param level: logging level to set the user_space_logger to + """ + global developer_logger + if handler is not None: + developer_logger.addHandler(handler) + if filter is not None: + developer_logger.addFilter(filter) + if level is not None: + developer_logger.setLevel(level) + + def _get_env_logging_level(default_level: int = logging.WARNING) -> int: """ Returns the logging level set in the environment variable, or logging.WARNING if the environment variable is not @@ -83,6 +106,17 @@ def initialize_global_loggers(): """ Initializes the global loggers to the default configuration. """ + # Use Rich logging while running in the local execution or jupyter notebook. + if ( + os.environ.get("FLYTE_INTERNAL_EXECUTION_ID", None) is None or interactive.ipython_check() + ) and is_rich_logging_enabled(): + try: + upgrade_to_rich_logging() + return + except OSError as e: + logger.warning(f"Failed to initialize rich logging: {e}") + pass + handler = logging.StreamHandler() handler.setLevel(logging.DEBUG) formatter = logging.Formatter(fmt="[%(name)s] %(message)s") @@ -92,10 +126,7 @@ def initialize_global_loggers(): set_flytekit_log_properties(handler, None, _get_env_logging_level()) set_user_logger_properties(handler, None, logging.INFO) - - # Use Rich logging while running in the local execution - if os.environ.get("FLYTE_INTERNAL_EXECUTION_ID", None) is None or interactive.ipython_check(): - upgrade_to_rich_logging() + set_developer_properties(handler, None, logging.INFO) def is_rich_logging_enabled() -> bool: @@ -103,29 +134,26 @@ def is_rich_logging_enabled() -> bool: def upgrade_to_rich_logging(log_level: typing.Optional[int] = logging.WARNING): - if not is_rich_logging_enabled(): - return - try: - import click - from rich.console import Console - from rich.logging import RichHandler - - import flytekit - - handler = RichHandler( - tracebacks_suppress=[click, flytekit], - rich_tracebacks=True, - omit_repeated_times=False, - log_time_format="%H:%M:%S.%f", - console=Console(width=os.get_terminal_size().columns), - ) - formatter = logging.Formatter(fmt="%(message)s") - handler.setFormatter(formatter) - set_flytekit_log_properties(handler, None, _get_env_logging_level(default_level=log_level)) - set_user_logger_properties(handler, None, logging.INFO) - except OSError as e: - logger.debug(f"Failed to initialize rich logging: {e}") - pass + import click + from rich.console import Console + from rich.logging import RichHandler + + import flytekit + + handler = RichHandler( + tracebacks_suppress=[click, flytekit], + rich_tracebacks=True, + omit_repeated_times=False, + show_path=False, + log_time_format="%H:%M:%S.%f", + console=Console(width=os.get_terminal_size().columns), + ) + + formatter = logging.Formatter(fmt="%(filename)s:%(lineno)d - %(message)s") + handler.setFormatter(formatter) + set_flytekit_log_properties(handler, None, _get_env_logging_level(default_level=log_level)) + set_user_logger_properties(handler, None, logging.INFO) + set_developer_properties(handler, None, logging.INFO) def get_level_from_cli_verbosity(verbosity: int) -> int: diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 301ce4b5fb..66b1ae54b6 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -52,7 +52,7 @@ FlyteEntityNotExistException, FlyteValueException, ) -from flytekit.loggers import logger +from flytekit.loggers import developer_logger, logger from flytekit.models import common as common_models from flytekit.models import filters as filter_models from flytekit.models import launch_plan as launch_plan_models @@ -153,12 +153,12 @@ def _get_entity_identifier( ) -def _get_git_repo_url(source_path): +def _get_git_repo_url(source_path: str): """ Get git repo URL from remote.origin.url """ try: - git_config = source_path / ".git" / "config" + git_config = pathlib.Path(source_path) / ".git" / "config" if not git_config.exists(): raise ValueError(f"{source_path} is not a git repo") @@ -182,7 +182,7 @@ def _get_git_repo_url(source_path): raise ValueError("Unable to parse url") except Exception as e: - logger.debug(str(e)) + logger.debug(f"unable to find the git config in {source_path} with error: {str(e)}") return "" @@ -669,7 +669,6 @@ def raw_register( workflow_model.TaskNode, ), ): - logger.debug("Ignoring nodes for registration.") return None elif isinstance(cp_entity, ReferenceSpec): @@ -683,7 +682,7 @@ def raw_register( try: self.client.create_task(task_identifer=ident, task_spec=cp_entity) except FlyteEntityAlreadyExistsException: - logger.info(f" {ident} Already Exists!") + logger.debug(f" {ident} Already Exists!") return ident if isinstance(cp_entity, admin_workflow_models.WorkflowSpec): @@ -693,7 +692,7 @@ def raw_register( try: self.client.create_workflow(workflow_identifier=ident, workflow_spec=cp_entity) except FlyteEntityAlreadyExistsException: - logger.info(f" {ident} Already Exists!") + logger.debug(f" {ident} Already Exists!") if create_default_launchplan: if not og_entity: @@ -717,7 +716,7 @@ def raw_register( try: self.client.create_launch_plan(lp_entity.id, lp_entity.spec) except FlyteEntityAlreadyExistsException: - logger.info(f" {lp_entity.id} Already Exists!") + logger.debug(f" {lp_entity.id} Already Exists!") return ident if isinstance(cp_entity, launch_plan_models.LaunchPlan): @@ -725,7 +724,7 @@ def raw_register( try: self.client.create_launch_plan(launch_plan_identifer=ident, launch_plan_spec=cp_entity.spec) except FlyteEntityAlreadyExistsException: - logger.info(f" {ident} Already Exists!") + logger.debug(f" {ident} Already Exists!") return ident raise AssertionError(f"Unknown entity of type {type(cp_entity)}") @@ -889,7 +888,6 @@ def upload_file( if not to_upload.is_file(): raise ValueError(f"{to_upload} is not a single file, upload arg must be a single file.") md5_bytes, str_digest, _ = hash_file(to_upload) - logger.debug(f"Text hash of file to upload is {str_digest}") upload_location = self.client.get_upload_signed_url( project=project or self.default_project, @@ -923,7 +921,9 @@ def upload_file( f"Request to send data {upload_location.signed_url} failed.\nResponse: {rsp.text}", ) - logger.debug(f"Uploading {to_upload} to {upload_location.signed_url} native url {upload_location.native_url}") + developer_logger.debug( + f"Uploading {to_upload} to {upload_location.signed_url} native url {upload_location.native_url}" + ) return md5_bytes, upload_location.native_url @@ -1616,7 +1616,7 @@ def execute_reference_workflow( try: flyte_lp = self.fetch_launch_plan(**resolved_identifiers_dict) except FlyteEntityNotExistException: - remote_logger.info("Try to register default launch plan because it wasn't found in Flyte Admin!") + logger.info("Try to register default launch plan because it wasn't found in Flyte Admin!") default_lp = LaunchPlan.get_default_launch_plan(self.context, entity) self.register_launch_plan( default_lp, diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index 18e9eeb09a..f8c62febf1 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -18,7 +18,7 @@ from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import TypeEngine, TypeTransformer from flytekit.deck.renderer import Renderable -from flytekit.loggers import logger +from flytekit.loggers import developer_logger, logger from flytekit.models import literals from flytekit.models import types as type_models from flytekit.models.literals import Literal, Scalar, StructuredDatasetMetadata @@ -515,7 +515,9 @@ def register_for_protocol( f"Already registered a handler for {(h.python_type, protocol, h.supported_format)}" ) lowest_level[h.supported_format] = h - logger.debug(f"Registered {h} as handler for {h.python_type}, protocol {protocol}, fmt {h.supported_format}") + developer_logger.debug( + f"Registered {h} as handler for {h.python_type}, protocol {protocol}, fmt {h.supported_format}" + ) if (default_format_for_type or default_for_type) and h.supported_format != GENERIC_FORMAT: if h.python_type in cls.DEFAULT_FORMATS and not override: @@ -524,9 +526,7 @@ def register_for_protocol( f"Not using handler {h} with format {h.supported_format} as default for {h.python_type}, {cls.DEFAULT_FORMATS[h.python_type]} already specified." ) else: - logger.debug( - f"Setting format {h.supported_format} for dataframes of type {h.python_type} from handler {h}" - ) + logger.debug(f"Use {type(h).__name__} as default handler for {h.python_type}.") cls.DEFAULT_FORMATS[h.python_type] = h.supported_format if default_storage_for_type or default_for_type: if h.protocol in cls.DEFAULT_PROTOCOLS and not override: From 825a50e9e92c359264fb5cb610b1e2a4c5678009 Mon Sep 17 00:00:00 2001 From: Future-Outlier Date: Tue, 2 Jul 2024 09:54:58 -0700 Subject: [PATCH 020/192] Support s3 path for click Dir param type (#2547) Signed-off-by: Future-Outlier --- flytekit/interaction/click_types.py | 10 ++++++---- tests/flytekit/unit/interaction/test_click_types.py | 8 ++++++++ 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/flytekit/interaction/click_types.py b/flytekit/interaction/click_types.py index 5728dce3d0..f0426ac7f2 100644 --- a/flytekit/interaction/click_types.py +++ b/flytekit/interaction/click_types.py @@ -80,13 +80,15 @@ def convert( ) -> typing.Any: if isinstance(value, ArtifactQuery): return value - p = pathlib.Path(value) + # set remote_directory to false if running pyflyte run locally. This makes sure that the original # directory is used and not a random one. remote_directory = None if getattr(ctx.obj, "is_remote", False) else False - if p.exists() and p.is_dir(): - return FlyteDirectory(path=value, remote_directory=remote_directory) - raise click.BadParameter(f"parameter should be a valid directory path, {value}") + if not FileAccessProvider.is_remote(value): + p = pathlib.Path(value) + if not p.exists() or not p.is_dir(): + raise click.BadParameter(f"parameter should be a valid flytedirectory path, {value}") + return FlyteDirectory(path=value, remote_directory=remote_directory) class StructuredDatasetParamType(click.ParamType): diff --git a/tests/flytekit/unit/interaction/test_click_types.py b/tests/flytekit/unit/interaction/test_click_types.py index cb32982916..e21a283271 100644 --- a/tests/flytekit/unit/interaction/test_click_types.py +++ b/tests/flytekit/unit/interaction/test_click_types.py @@ -28,6 +28,14 @@ dummy_param = click.Option(["--dummy"], type=click.STRING, default="dummy") +def test_dir_param(): + import os + m = mock.MagicMock() + current_file_directory = os.path.dirname(os.path.abspath(__file__)) + l = DirParamType().convert(current_file_directory, m, m) + assert l.path == current_file_directory + r = DirParamType().convert("https://tmp/dir", m, m) + assert r.path == "https://tmp/dir" def test_file_param(): m = mock.MagicMock() From a6a86518abbd316693708abf2a12712e480901c4 Mon Sep 17 00:00:00 2001 From: Chun-Mao Lai <72752478+Mecoli1219@users.noreply.github.com> Date: Thu, 4 Jul 2024 03:07:32 +0800 Subject: [PATCH 021/192] remove upper bound of plugin dependencies for flytekit-greatexpectations (#2506) Signed-off-by: Mecoli1219 --- plugins/flytekit-greatexpectations/setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/plugins/flytekit-greatexpectations/setup.py b/plugins/flytekit-greatexpectations/setup.py index 506dd4853b..4ade555296 100644 --- a/plugins/flytekit-greatexpectations/setup.py +++ b/plugins/flytekit-greatexpectations/setup.py @@ -6,8 +6,8 @@ plugin_requires = [ "flytekit>=1.5.0", - "great-expectations>=0.13.30,<=0.18.8", - "sqlalchemy>=1.4.23,<2.0.0", + "great-expectations>=0.13.30", + "sqlalchemy>=1.4.23", "pyspark==3.3.1", "s3fs<2023.6.0", ] From 153b0df2f63d72e7b7c2d82a66e52c521ffc4033 Mon Sep 17 00:00:00 2001 From: "Fabio M. Graetz, Ph.D" Date: Wed, 3 Jul 2024 23:33:16 +0200 Subject: [PATCH 022/192] Add myself to code owners of flytekit-kf-pytorch (#2556) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Fabio Grätz Co-authored-by: Fabio Grätz --- CODEOWNERS | 1 + 1 file changed, 1 insertion(+) diff --git a/CODEOWNERS b/CODEOWNERS index 917fe185c0..a1c6d04c0f 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -1,3 +1,4 @@ # These owners will be the default owners for everything in # the repo. Unless a later match takes precedence. * @wild-endeavor @kumare3 @eapolinario @pingsutw @cosmicBboy @samhita-alla @thomasjpfan @future-outlier +plugins/flytekit-kf-pytorch @fg91 @wild-endeavor @kumare3 @eapolinario @pingsutw @cosmicBboy @samhita-alla @thomasjpfan @future-outlier From b82a25f755c99fc009331827e06eb81f4d9ecb9b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bu=C4=9Fra=20Gedik?= Date: Wed, 3 Jul 2024 14:34:05 -0700 Subject: [PATCH 023/192] Add run policy (#2555) Signed-off-by: bugra.gedik Co-authored-by: bugra.gedik --- .../flytekitplugins/kfpytorch/task.py | 29 ++++++++++++------- .../tests/test_elastic_task.py | 27 ++++++++++++++++- 2 files changed, 45 insertions(+), 11 deletions(-) diff --git a/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py b/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py index 94b575e2a9..95477726e1 100644 --- a/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py +++ b/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py @@ -107,7 +107,7 @@ class PyTorch(object): master: Master = field(default_factory=lambda: Master()) worker: Worker = field(default_factory=lambda: Worker()) - run_policy: Optional[RunPolicy] = field(default_factory=lambda: None) + run_policy: Optional[RunPolicy] = None # Support v0 config for backwards compatibility num_workers: Optional[int] = None @@ -130,6 +130,7 @@ class Elastic(object): max_restarts (int): Maximum number of worker group restarts before failing. rdzv_configs (Dict[str, Any]): Additional rendezvous configs to pass to torch elastic, e.g. `{"timeout": 1200, "join_timeout": 900}`. See `torch.distributed.launcher.api.LaunchConfig` and `torch.distributed.elastic.rendezvous.dynamic_rendezvous.create_handler`. + run_policy: Configuration for the run policy. """ nnodes: Union[int, str] = 1 @@ -138,6 +139,7 @@ class Elastic(object): monitor_interval: int = 5 max_restarts: int = 0 rdzv_configs: Dict[str, Any] = field(default_factory=dict) + run_policy: Optional[RunPolicy] = None class PyTorchFunctionTask(PythonFunctionTask[PyTorch]): @@ -181,21 +183,15 @@ def _convert_replica_spec( restart_policy=replica_config.restart_policy.value if replica_config.restart_policy else None, ) - def _convert_run_policy(self, run_policy: RunPolicy) -> kubeflow_common.RunPolicy: - return kubeflow_common.RunPolicy( - clean_pod_policy=run_policy.clean_pod_policy.value if run_policy.clean_pod_policy else None, - ttl_seconds_after_finished=run_policy.ttl_seconds_after_finished, - active_deadline_seconds=run_policy.active_deadline_seconds, - backoff_limit=run_policy.backoff_limit, - ) - def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: worker = self._convert_replica_spec(self.task_config.worker) # support v0 config for backwards compatibility if self.task_config.num_workers: worker.replicas = self.task_config.num_workers - run_policy = self._convert_run_policy(self.task_config.run_policy) if self.task_config.run_policy else None + run_policy = ( + _convert_run_policy_to_flyte_idl(self.task_config.run_policy) if self.task_config.run_policy else None + ) pytorch_job = pytorch_task.DistributedPyTorchTrainingTask( worker_replicas=worker, master_replicas=self._convert_replica_spec(self.task_config.master), @@ -263,6 +259,15 @@ def spawn_helper( return ElasticWorkerResult(return_value=return_val, decks=flytekit.current_context().decks) +def _convert_run_policy_to_flyte_idl(run_policy: RunPolicy) -> kubeflow_common.RunPolicy: + return kubeflow_common.RunPolicy( + clean_pod_policy=run_policy.clean_pod_policy.value if run_policy.clean_pod_policy else None, + ttl_seconds_after_finished=run_policy.ttl_seconds_after_finished, + active_deadline_seconds=run_policy.active_deadline_seconds, + backoff_limit=run_policy.backoff_limit, + ) + + class PytorchElasticFunctionTask(PythonFunctionTask[Elastic]): """ Plugin for distributed training with torch elastic/torchrun (see @@ -445,11 +450,15 @@ def get_custom(self, settings: SerializationSettings) -> Optional[Dict[str, Any] nproc_per_node=self.task_config.nproc_per_node, max_restarts=self.task_config.max_restarts, ) + run_policy = ( + _convert_run_policy_to_flyte_idl(self.task_config.run_policy) if self.task_config.run_policy else None + ) job = pytorch_task.DistributedPyTorchTrainingTask( worker_replicas=pytorch_task.DistributedPyTorchTrainingReplicaSpec( replicas=self.max_nodes, ), elastic_config=elastic_config, + run_policy=run_policy, ) return MessageToDict(job) diff --git a/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py b/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py index fd13a39659..b79192f3f1 100644 --- a/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py +++ b/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py @@ -7,10 +7,11 @@ import torch import torch.distributed as dist from dataclasses_json import DataClassJsonMixin -from flytekitplugins.kfpytorch.task import Elastic +from flytekitplugins.kfpytorch.task import CleanPodPolicy, Elastic, RunPolicy import flytekit from flytekit import task, workflow +from flytekit.configuration import SerializationSettings from flytekit.exceptions.user import FlyteRecoverableException @@ -187,3 +188,27 @@ def wf(recoverable: bool): else: with pytest.raises(RuntimeError): wf(recoverable=recoverable) + + +def test_run_policy() -> None: + """Test that run policy is propagated to custom spec.""" + + run_policy = RunPolicy( + clean_pod_policy=CleanPodPolicy.ALL, + ttl_seconds_after_finished=10 * 60, + active_deadline_seconds=36000, + backoff_limit=None, + ) + + # nnodes must be > 1 to get pytorchjob spec + @task(task_config=Elastic(nnodes=2, nproc_per_node=2, run_policy=run_policy)) + def test_task(): + pass + + spec = test_task.get_custom(SerializationSettings(image_config=None)) + + assert spec["runPolicy"] == { + "cleanPodPolicy": "CLEANPOD_POLICY_ALL", + "ttlSecondsAfterFinished": 600, + "activeDeadlineSeconds": 36000, + } From 3cd587683fbfc12592622d596391ee6e07548acb Mon Sep 17 00:00:00 2001 From: pryce-turner <31577879+pryce-turner@users.noreply.github.com> Date: Wed, 3 Jul 2024 16:43:42 -0700 Subject: [PATCH 024/192] Added entrypoint to imagespec and default builder (#2553) * Added entrypoint to imagespec and default builder Signed-off-by: pryce-turner * Added docstring Signed-off-by: pryce-turner * Revert reqs.txt Signed-off-by: pryce-turner * Fixed falsey behavior of empty list, moved ENTRYPOINT directive into optional str to avoid entrypoint reset behavior Signed-off-by: pryce-turner --------- Signed-off-by: pryce-turner --- flytekit/image_spec/default_builder.py | 9 +++++++++ flytekit/image_spec/image_spec.py | 2 ++ .../core/image_spec/test_default_builder.py | 20 +++++++++++++++++++ .../unit/core/image_spec/test_image_spec.py | 2 ++ 4 files changed, 33 insertions(+) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index 7e91eebf10..a09f445ba9 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -1,3 +1,4 @@ +import json import os import shutil import subprocess @@ -73,6 +74,8 @@ ENV PATH="$$PATH:/usr/local/nvidia/bin:/usr/local/cuda/bin" \ LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$$LD_LIBRARY_PATH" +$ENTRYPOINT + $COPY_COMMAND_RUNTIME RUN $RUN_COMMANDS @@ -191,6 +194,11 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path): else: python_version = f"{sys.version_info.major}.{sys.version_info.minor}" + if image_spec.entrypoint is None: + entrypoint = "" + else: + entrypoint = f"ENTRYPOINT {json.dumps(image_spec.entrypoint)}" + if image_spec.commands: run_commands = " && ".join(image_spec.commands) else: @@ -206,6 +214,7 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path): BASE_IMAGE=base_image, ENV=env, COPY_COMMAND_RUNTIME=copy_command_runtime, + ENTRYPOINT=entrypoint, RUN_COMMANDS=run_commands, ) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index a17a2e272e..7cde9fa70a 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -45,6 +45,7 @@ class ImageSpec: pip_index: Specify the custom pip index url pip_extra_index_url: Specify one or more pip index urls as a list registry_config: Specify the path to a JSON registry config file + entrypoint: List of strings to overwrite the entrypoint of the base image with, set to [] to remove the entrypoint. commands: Command to run during the building process tag_format: Custom string format for image tag. The ImageSpec hash passed in as `spec_hash`. For example, to add a "dev" suffix to the image tag, set `tag_format="{spec_hash}-dev"` @@ -68,6 +69,7 @@ class ImageSpec: pip_index: Optional[str] = None pip_extra_index_url: Optional[List[str]] = None registry_config: Optional[str] = None + entrypoint: Optional[List[str]] = None commands: Optional[List[str]] = None tag_format: Optional[str] = None diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index bd47b4558e..8941da81b9 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -28,6 +28,7 @@ def test_create_docker_context(tmp_path): requirements=os.fspath(other_requirements_path), source_root=os.fspath(source_root), commands=["mkdir my_dir"], + entrypoint=["/bin/bash"], ) create_docker_context(image_spec, docker_context_path) @@ -42,6 +43,7 @@ def test_create_docker_context(tmp_path): assert "--requirement requirements_uv.txt" in dockerfile_content assert "COPY --chown=flytekit ./src /root" in dockerfile_content assert "RUN mkdir my_dir" in dockerfile_content + assert "ENTRYPOINT [\"/bin/bash\"]" in dockerfile_content requirements_path = docker_context_path / "requirements_uv.txt" assert requirements_path.exists() @@ -79,6 +81,24 @@ def test_create_docker_context_with_git_subfolder(tmp_path): assert requirements_path.exists() +def test_create_docker_context_with_null_entrypoint(tmp_path): + docker_context_path = tmp_path / "builder_root" + docker_context_path.mkdir() + + image_spec = ImageSpec( + name="FLYTEKIT", + python_version="3.12", + entrypoint=[], + ) + + create_docker_context(image_spec, docker_context_path) + + dockerfile_path = docker_context_path / "Dockerfile" + assert dockerfile_path.exists() + dockerfile_content = dockerfile_path.read_text() + assert "ENTRYPOINT []" in dockerfile_content + + def test_create_docker_context_cuda(tmp_path): docker_context_path = tmp_path / "builder_root" docker_context_path.mkdir() diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py index 02c734b119..c1e52953bb 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -25,6 +25,7 @@ def test_image_spec(mock_image_spec_builder): cudnn="8", requirements=REQUIREMENT_FILE, registry_config=REGISTRY_CONFIG_FILE, + entrypoint=["/bin/bash"], ) assert image_spec._is_force_push is False @@ -50,6 +51,7 @@ def test_image_spec(mock_image_spec_builder): assert image_spec.is_container() is True assert image_spec.commands == ["echo hello"] assert image_spec._is_force_push is True + assert image_spec.entrypoint == ["/bin/bash"] tag = calculate_hash_from_image_spec(image_spec) assert "=" != tag[-1] From 55bb44522ea805576daf93ed0bc45aa5d03a5788 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 4 Jul 2024 08:45:41 +0800 Subject: [PATCH 025/192] Set default width if fail to get terminal size (#2558) * Set default width if fail to get terminal size Signed-off-by: Kevin Su * lint Signed-off-by: Kevin Su * test Signed-off-by: Kevin Su --------- Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 2 +- flytekit/loggers.py | 8 +++++++- tests/flytekit/loggers.py | 11 +++++++++++ 3 files changed, 19 insertions(+), 2 deletions(-) create mode 100644 tests/flytekit/loggers.py diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 2a087cb4ca..92a4373839 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -1650,7 +1650,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: res_tag = trans.name found_res = True except Exception as e: - logger.debug(f"Failed to convert from {lv} to {v}", e) + logger.debug(f"Failed to convert from {lv} to {v} with error: {e}") if is_ambiguous: raise TypeError( diff --git a/flytekit/loggers.py b/flytekit/loggers.py index 12b3e90423..8e38c2b4ee 100644 --- a/flytekit/loggers.py +++ b/flytekit/loggers.py @@ -140,13 +140,19 @@ def upgrade_to_rich_logging(log_level: typing.Optional[int] = logging.WARNING): import flytekit + try: + width = os.get_terminal_size().columns + except Exception as e: + logger.debug(f"Failed to get terminal size: {e}") + width = 80 + handler = RichHandler( tracebacks_suppress=[click, flytekit], rich_tracebacks=True, omit_repeated_times=False, show_path=False, log_time_format="%H:%M:%S.%f", - console=Console(width=os.get_terminal_size().columns), + console=Console(width=width), ) formatter = logging.Formatter(fmt="%(filename)s:%(lineno)d - %(message)s") diff --git a/tests/flytekit/loggers.py b/tests/flytekit/loggers.py new file mode 100644 index 0000000000..3b94259049 --- /dev/null +++ b/tests/flytekit/loggers.py @@ -0,0 +1,11 @@ +import sys +from subprocess import run + + +def test_error(tmp_path): + EXAMPLE = "import flytekit" + + script_path = tmp_path / "script.py" + script_path.write_text(EXAMPLE) + out = run([sys.executable, script_path], text=True, capture_output=True) + assert out.stderr == "" From d4bebfa9136ecafa4b36f854e65d58d31700a460 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 4 Jul 2024 11:22:33 +0800 Subject: [PATCH 026/192] Disable rich traceback for papermill plugin (#2559) Signed-off-by: Kevin Su --- plugins/flytekit-papermill/tests/conftest.py | 8 ++++++++ 1 file changed, 8 insertions(+) create mode 100644 plugins/flytekit-papermill/tests/conftest.py diff --git a/plugins/flytekit-papermill/tests/conftest.py b/plugins/flytekit-papermill/tests/conftest.py new file mode 100644 index 0000000000..04564a0752 --- /dev/null +++ b/plugins/flytekit-papermill/tests/conftest.py @@ -0,0 +1,8 @@ +import os + +import pytest + + +@pytest.fixture(scope="module", autouse=True) +def set_default_envs(): + os.environ["FLYTE_SDK_RICH_TRACEBACKS"] = "0" From a7285067bfd90d53c373163c7db961fe09e3231d Mon Sep 17 00:00:00 2001 From: novahow <58504997+novahow@users.noreply.github.com> Date: Fri, 5 Jul 2024 06:02:43 +0800 Subject: [PATCH 027/192] fix datetime in eager workflow (#2541) Signed-off-by: novahow --- flytekit/experimental/eager_function.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/flytekit/experimental/eager_function.py b/flytekit/experimental/eager_function.py index d47a2baef2..7eec791726 100644 --- a/flytekit/experimental/eager_function.py +++ b/flytekit/experimental/eager_function.py @@ -179,7 +179,11 @@ async def __call__(self, **kwargs): self.async_stack.set_node(node) poll_interval = self._poll_interval or timedelta(seconds=30) - time_to_give_up = datetime.max if self._timeout is None else datetime.now(timezone.utc) + self._timeout + time_to_give_up = ( + (datetime.max.replace(tzinfo=timezone.utc)) + if self._timeout is None + else datetime.now(timezone.utc) + self._timeout + ) while datetime.now(timezone.utc) < time_to_give_up: execution = self.remote.sync(execution) @@ -208,7 +212,11 @@ async def terminate(self): ) poll_interval = self._poll_interval or timedelta(seconds=6) - time_to_give_up = datetime.max if self._timeout is None else datetime.now(timezone.utc) + self._timeout + time_to_give_up = ( + (datetime.max.replace(tzinfo=timezone.utc)) + if self._timeout is None + else datetime.now(timezone.utc) + self._timeout + ) while datetime.now(timezone.utc) < time_to_give_up: execution = self.remote.sync(execution) From ecbfe51d415edf5a57c9d89d3d3ddc8713c688cc Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Fri, 5 Jul 2024 18:07:18 -0400 Subject: [PATCH 028/192] Fix CSS with Flyte decks (#2565) Signed-off-by: Thomas J. Fan --- flytekit/deck/html/template.html | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/flytekit/deck/html/template.html b/flytekit/deck/html/template.html index 19e0256880..a58e8a3252 100644 --- a/flytekit/deck/html/template.html +++ b/flytekit/deck/html/template.html @@ -61,9 +61,8 @@ } #flyte-frame-container > div.active { - display: Block; - padding: 2rem 4rem; - width: 100%; + display: block; + padding: 2rem 2rem; } From 05f4ae671c64f1ab1e948aed00ccc2524d7d4e51 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Fri, 5 Jul 2024 23:10:03 +0100 Subject: [PATCH 029/192] Fix `FlyteDirectory` on Azure (#2564) Signed-off-by: Thomas Newton --- flytekit/remote/remote_fs.py | 4 ++-- tests/flytekit/unit/remote/test_fs_remote.py | 14 +++++++++----- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/flytekit/remote/remote_fs.py b/flytekit/remote/remote_fs.py index 10131f63fa..c85d93959f 100644 --- a/flytekit/remote/remote_fs.py +++ b/flytekit/remote/remote_fs.py @@ -167,8 +167,8 @@ def extract_common(native_urls: typing.List[str]) -> str: else: break - fs = fsspec.filesystem(get_protocol(native_urls[0])) - sep = fs.sep + fs_class = fsspec.get_filesystem_class(get_protocol(native_urls[0])) + sep = fs_class.sep # split the common prefix on the last separator so we don't get any trailing characters. common_prefix = common_prefix.rsplit(sep, 1)[0] logger.debug(f"Returning {common_prefix} from {native_urls}") diff --git a/tests/flytekit/unit/remote/test_fs_remote.py b/tests/flytekit/unit/remote/test_fs_remote.py index dd3473e44f..efc6e94e8b 100644 --- a/tests/flytekit/unit/remote/test_fs_remote.py +++ b/tests/flytekit/unit/remote/test_fs_remote.py @@ -116,14 +116,18 @@ def test_remote_upload_with_data_persistence(sandbox_remote): assert res == f"s3://my-s3-bucket/flytesnacks/development/O55F24U7RMLDOUI3LZ6SL4ZBEI======/{only_file}" -def test_common_matching(): +@pytest.mark.parametrize("url_prefix", ["s3://my-s3-bucket", "abfs://my-azure-container", "abfss://my-azure-container", "gcs://my-gcs-bucket"]) +def test_common_matching(url_prefix): urls = [ - "s3://my-s3-bucket/flytesnacks/development/ABCYZWMPACZAJ2MABGMOZ6CCPY======/source/empty.md", - "s3://my-s3-bucket/flytesnacks/development/ABCXKL5ZZWXY3PDLM3OONUHHME======/source/nested/more.txt", - "s3://my-s3-bucket/flytesnacks/development/ABCXBAPBKONMADXVW5Q3J6YBWM======/source/original.txt", + url_prefix + url_suffix + for url_suffix in [ + "/flytesnacks/development/ABCYZWMPACZAJ2MABGMOZ6CCPY======/source/empty.md", + "/flytesnacks/development/ABCXKL5ZZWXY3PDLM3OONUHHME======/source/nested/more.txt", + "/flytesnacks/development/ABCXBAPBKONMADXVW5Q3J6YBWM======/source/original.txt", + ] ] - assert FlyteFS.extract_common(urls) == "s3://my-s3-bucket/flytesnacks/development" + assert FlyteFS.extract_common(urls) == url_prefix + "/flytesnacks/development" def test_hashing(sandbox_remote, source_folder): From 4be4e33c3299961f9d085063a9740b8eaffc4588 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Fri, 5 Jul 2024 18:10:44 -0400 Subject: [PATCH 030/192] Makes home dir if it does not exists (#2562) Signed-off-by: Thomas J. Fan --- flytekit/image_spec/default_builder.py | 3 ++- tests/flytekit/unit/core/image_spec/test_default_builder.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index a09f445ba9..23b356f77b 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -83,7 +83,8 @@ SHELL ["/bin/bash", "-c"] USER flytekit -RUN echo "export PATH=$$PATH" >> $$HOME/.profile +RUN mkdir -p $$HOME && \ + echo "export PATH=$$PATH" >> $$HOME/.profile """ ) diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index 8941da81b9..c15f6b343a 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -44,6 +44,7 @@ def test_create_docker_context(tmp_path): assert "COPY --chown=flytekit ./src /root" in dockerfile_content assert "RUN mkdir my_dir" in dockerfile_content assert "ENTRYPOINT [\"/bin/bash\"]" in dockerfile_content + assert "mkdir -p $HOME" in dockerfile_content requirements_path = docker_context_path / "requirements_uv.txt" assert requirements_path.exists() From e742be36a0beba095fb1ee20ce1db867b7f0b90c Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Mon, 8 Jul 2024 14:17:53 -0400 Subject: [PATCH 031/192] Removes isodate requirement (#2568) Signed-off-by: Thomas J. Fan --- dev-requirements.txt | 1 - pyproject.toml | 1 - 2 files changed, 2 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index d3279adcac..98c56efe69 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -183,7 +183,6 @@ ipython==8.25.0 isodate==0.6.1 # via # azure-storage-blob - # flytekit jaraco-classes==3.4.0 # via # keyring diff --git a/pyproject.toml b/pyproject.toml index 22a8fa1a40..a3f9d8710b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,6 @@ dependencies = [ "grpcio", "grpcio-status", "importlib-metadata", - "isodate", "jinja2", "joblib", "jsonlines", From 5ff9eb3e7d258f7bd73dfeb79018169072a31e30 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 9 Jul 2024 02:43:35 +0800 Subject: [PATCH 032/192] Read FLYTE_SDK_DEV_LOGGING_LEVEL from env (#2571) Signed-off-by: Kevin Su --- flytekit/loggers.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/flytekit/loggers.py b/flytekit/loggers.py index 8e38c2b4ee..07bd6f388b 100644 --- a/flytekit/loggers.py +++ b/flytekit/loggers.py @@ -13,6 +13,7 @@ # For now, assume this is the environment variable whose usage will remain unchanged and controls output for all # loggers defined in this file. LOGGING_ENV_VAR = "FLYTE_SDK_LOGGING_LEVEL" +# The environment variable controls the logging level for the developer logger. LOGGING_DEV_ENV_VAR = "FLYTE_SDK_DEV_LOGGING_LEVEL" LOGGING_FMT_ENV_VAR = "FLYTE_SDK_LOGGING_FORMAT" LOGGING_RICH_FMT_ENV_VAR = "FLYTE_SDK_RICH_TRACEBACKS" @@ -102,6 +103,14 @@ def _get_env_logging_level(default_level: int = logging.WARNING) -> int: return int(os.getenv(LOGGING_ENV_VAR, default_level)) +def _get_dev_env_logging_level(default_level: int = logging.INFO) -> int: + """ + Returns the logging level set in the environment variable, or logging.INFO if the environment variable is not + set. + """ + return int(os.getenv(LOGGING_DEV_ENV_VAR, default_level)) + + def initialize_global_loggers(): """ Initializes the global loggers to the default configuration. @@ -159,7 +168,7 @@ def upgrade_to_rich_logging(log_level: typing.Optional[int] = logging.WARNING): handler.setFormatter(formatter) set_flytekit_log_properties(handler, None, _get_env_logging_level(default_level=log_level)) set_user_logger_properties(handler, None, logging.INFO) - set_developer_properties(handler, None, logging.INFO) + set_developer_properties(handler, None, _get_dev_env_logging_level()) def get_level_from_cli_verbosity(verbosity: int) -> int: From caf3139584914dc532585cf272313f685c093fe2 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 9 Jul 2024 02:46:34 +0800 Subject: [PATCH 033/192] print native input/output when cache hit (#2567) Signed-off-by: Kevin Su --- flytekit/core/base_task.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index c161df296e..88c2b39c02 100644 --- a/flytekit/core/base_task.py +++ b/flytekit/core/base_task.py @@ -283,7 +283,7 @@ def local_execute( # native constants are just bound to this specific task (default values for a task input) # Also along with promises and constants, there could be dictionary or list of promises or constants try: - kwargs = translate_inputs_to_literals( + literals = translate_inputs_to_literals( ctx, incoming_values=kwargs, flyte_interface_types=self.interface.inputs, @@ -293,20 +293,19 @@ def local_execute( msg = f"Failed to convert inputs of task '{self.name}':\n {exc}" logger.error(msg) raise TypeError(msg) from exc - input_literal_map = _literal_models.LiteralMap(literals=kwargs) + input_literal_map = _literal_models.LiteralMap(literals=literals) # if metadata.cache is set, check memoized version local_config = LocalConfig.auto() if self.metadata.cache and local_config.cache_enabled: - # TODO: how to get a nice `native_inputs` here? - logger.info( - f"Checking cache for task named {self.name}, cache version {self.metadata.cache_version} " - f", inputs: {input_literal_map}, and ignore input vars: {self.metadata.cache_ignore_input_vars}" - ) if local_config.cache_overwrite: outputs_literal_map = None logger.info("Cache overwrite, task will be executed now") else: + logger.info( + f"Checking cache for task named {self.name}, cache version {self.metadata.cache_version} " + f", inputs: {kwargs}, and ignore input vars: {self.metadata.cache_ignore_input_vars}" + ) outputs_literal_map = LocalTaskCache.get( self.name, self.metadata.cache_version, input_literal_map, self.metadata.cache_ignore_input_vars ) @@ -327,7 +326,7 @@ def local_execute( ) logger.info( f"Cache set for task named {self.name}, cache version {self.metadata.cache_version} " - f", inputs: {input_literal_map}, and ignore input vars: {self.metadata.cache_ignore_input_vars}" + f", inputs: {kwargs}, and ignore input vars: {self.metadata.cache_ignore_input_vars}" ) else: # This code should mirror the call to `sandbox_execute` in the above cache case. From 89e5461f281c0f5e5c38500ef7da1a38ffcb09c8 Mon Sep 17 00:00:00 2001 From: Dennis Keck <26092524+fellhorn@users.noreply.github.com> Date: Mon, 8 Jul 2024 22:51:26 +0200 Subject: [PATCH 034/192] Fix: Set OMP_NUM_THREADS by default in Elastic (#2569) Signed-off-by: Dennis Keck <26092524+fellhorn@users.noreply.github.com> --- .../flytekitplugins/kfpytorch/task.py | 20 ++++++++++++++++ .../tests/test_elastic_task.py | 24 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py b/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py index 95477726e1..73d5e07c1b 100644 --- a/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py +++ b/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py @@ -122,6 +122,10 @@ class Elastic(object): Single-node elastic training is executed in a k8s pod when `nnodes` is set to 1. Multi-node training is executed otherwise using a `Pytorch Job `_. + Like `torchrun`, this plugin sets the environment variable `OMP_NUM_THREADS` to 1 if it is not set. + Please see https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html for potential performance improvements. + To change `OMP_NUM_THREADS`, specify it in the environment dict of the flytekit task decorator or via `pyflyte run --env`. + Args: nnodes (Union[int, str]): Number of nodes, or the range of nodes in form :. nproc_per_node (str): Number of workers per node. @@ -332,6 +336,22 @@ def _execute(self, **kwargs) -> Any: ) ) + # If OMP_NUM_THREADS is not set, set it to 1 to avoid overloading the system. + # Doing so to copy the default behavior of torchrun. + # See https://github.com/pytorch/pytorch/blob/eea4ece256d74c6f25c1f4eab37b3f2f4aeefd4d/torch/distributed/run.py#L791 + if "OMP_NUM_THREADS" not in os.environ and self.task_config.nproc_per_node > 1: + omp_num_threads = 1 + logger.warning( + "\n*****************************************\n" + "Setting OMP_NUM_THREADS environment variable for each process to be " + "%s in default, to avoid your system being overloaded, " + "please further tune the variable for optimal performance in " + "your application as needed. \n" + "*****************************************", + omp_num_threads, + ) + os.environ["OMP_NUM_THREADS"] = str(omp_num_threads) + config = LaunchConfig( run_id=flytekit.current_context().execution_id.name, min_nodes=self.min_nodes, diff --git a/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py b/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py index b79192f3f1..9cb62b993c 100644 --- a/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py +++ b/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py @@ -14,6 +14,12 @@ from flytekit.configuration import SerializationSettings from flytekit.exceptions.user import FlyteRecoverableException +@pytest.fixture(autouse=True, scope="function") +def restore_env(): + original_env = os.environ.copy() + yield + os.environ.clear() + os.environ.update(original_env) @dataclass class Config(DataClassJsonMixin): @@ -212,3 +218,21 @@ def test_task(): "ttlSecondsAfterFinished": 600, "activeDeadlineSeconds": 36000, } + +@pytest.mark.parametrize("start_method", ["spawn", "fork"]) +def test_omp_num_threads(start_method: str) -> None: + """Test that the env var OMP_NUM_THREADS is set by default and not overwritten if set.""" + + @task(task_config=Elastic(nnodes=1, nproc_per_node=2, start_method=start_method)) + def test_task_omp_default(): + assert os.environ["OMP_NUM_THREADS"] == "1" + + test_task_omp_default() + + os.environ["OMP_NUM_THREADS"] = "42" + + @task(task_config=Elastic(nnodes=1, nproc_per_node=2, start_method=start_method)) + def test_task_omp_set(): + assert os.environ["OMP_NUM_THREADS"] == "42" + + test_task_omp_set() From b0c581ed5ed72c789de8f8e9e10ce6a610d993e5 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 9 Jul 2024 07:29:49 +0800 Subject: [PATCH 035/192] Use logging level in env for default dev logger (#2572) Signed-off-by: Kevin Su --- flytekit/loggers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/loggers.py b/flytekit/loggers.py index 07bd6f388b..41d1bd3ad8 100644 --- a/flytekit/loggers.py +++ b/flytekit/loggers.py @@ -135,7 +135,7 @@ def initialize_global_loggers(): set_flytekit_log_properties(handler, None, _get_env_logging_level()) set_user_logger_properties(handler, None, logging.INFO) - set_developer_properties(handler, None, logging.INFO) + set_developer_properties(handler, None, _get_dev_env_logging_level()) def is_rich_logging_enabled() -> bool: From 5b3e7256c8cb14f19dd727e964b26be346d491b5 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 8 Jul 2024 16:37:27 -0700 Subject: [PATCH 036/192] Bump certifi (#2566) Bumps [certifi](https://github.com/certifi/python-certifi) from 2023.7.22 to 2024.7.4. - [Commits](https://github.com/certifi/python-certifi/compare/2023.07.22...2024.07.04) --- updated-dependencies: - dependency-name: certifi dependency-type: direct:production ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .../remote/mock_flyte_repo/workflows/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt index 164b9db345..cffbcd071c 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/requirements.txt @@ -41,7 +41,7 @@ botocore==1.29.161 # via aiobotocore cachetools==5.3.1 # via google-auth -certifi==2023.7.22 +certifi==2024.7.4 # via # kubernetes # requests From 1f37fa46317e6f43c88280afe4fae0692835cb9c Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Tue, 9 Jul 2024 12:04:41 -0400 Subject: [PATCH 037/192] Allow for flytekit version to be specified in default image builder (#2563) * Allow for flytekit version to be specified Signed-off-by: Thomas J. Fan * REV Revert Signed-off-by: Thomas J. Fan * REV Revert Signed-off-by: Thomas J. Fan * Adds more comments Signed-off-by: Thomas J. Fan --------- Signed-off-by: Thomas J. Fan --- flytekit/image_spec/default_builder.py | 21 ++++++++++- .../core/image_spec/test_default_builder.py | 36 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index 23b356f77b..57399cf07f 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -1,5 +1,6 @@ import json import os +import re import shutil import subprocess import sys @@ -99,11 +100,25 @@ def get_flytekit_for_pypi(): return f"flytekit=={__version__}" +_PACKAGE_NAME_RE = re.compile(r"^[\w-]+") + + +def _is_flytekit(package: str) -> bool: + """Return True if `package` is flytekit. `package` is expected to be a valid version + spec. i.e. `flytekit==1.12.3`, `flytekit`, `flytekit~=1.12.3`. + """ + m = _PACKAGE_NAME_RE.match(package) + if not m: + return False + name = m.group() + return name == "flytekit" + + def create_docker_context(image_spec: ImageSpec, tmp_dir: Path): """Populate tmp_dir with Dockerfile as specified by the `image_spec`.""" base_image = image_spec.base_image or "debian:bookworm-slim" - requirements = [get_flytekit_for_pypi()] + requirements = [] if image_spec.cuda is not None or image_spec.cudnn is not None: msg = ( @@ -125,6 +140,10 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path): if image_spec.packages: requirements.extend(image_spec.packages) + # Adds flytekit if it is not specified + if not any(_is_flytekit(package) for package in requirements): + requirements.append(get_flytekit_for_pypi()) + uv_requirements = [] # uv does not support git + subdirectory, so we use pip to install them instead diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index c15f6b343a..e42337b567 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -2,6 +2,7 @@ import pytest +import flytekit from flytekit.image_spec import ImageSpec from flytekit.image_spec.default_builder import DefaultImageBuilder, create_docker_context @@ -100,6 +101,41 @@ def test_create_docker_context_with_null_entrypoint(tmp_path): assert "ENTRYPOINT []" in dockerfile_content +@pytest.mark.parametrize("flytekit_spec", [None, "flytekit>=1.12.3", "flytekit==1.12.3"]) +def test_create_docker_context_with_flytekit(tmp_path, flytekit_spec, monkeypatch): + + # pretend version is 1.13.0 + mock_version = "1.13.0" + monkeypatch.setattr(flytekit, "__version__", mock_version) + + docker_context_path = tmp_path / "builder_root" + docker_context_path.mkdir() + + if flytekit_spec: + packages = [flytekit_spec] + else: + packages = [] + + image_spec = ImageSpec( + name="FLYTEKIT", packages=packages + ) + + create_docker_context(image_spec, docker_context_path) + + dockerfile_path = docker_context_path / "Dockerfile" + assert dockerfile_path.exists() + + requirements_path = docker_context_path / "requirements_uv.txt" + assert requirements_path.exists() + + requirements_content = requirements_path.read_text() + if flytekit_spec: + flytekit_spec in requirements_content + assert f"flytekit=={mock_version}" not in requirements_content + else: + assert f"flytekit=={mock_version}" in requirements_content + + def test_create_docker_context_cuda(tmp_path): docker_context_path = tmp_path / "builder_root" docker_context_path.mkdir() From ccade5a277793204057325071c64d6e451e29e53 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Tue, 9 Jul 2024 14:21:46 -0400 Subject: [PATCH 038/192] Removes jinja2 dependency (#2570) Signed-off-by: Thomas J. Fan --- dev-requirements.in | 1 - dev-requirements.txt | 6 ------ flytekit/deck/deck.py | 31 ++++++++++++++++--------------- flytekit/deck/html/template.html | 10 ++-------- pyproject.toml | 14 +++++--------- 5 files changed, 23 insertions(+), 39 deletions(-) diff --git a/dev-requirements.in b/dev-requirements.in index ca37177df7..2c91767a01 100644 --- a/dev-requirements.in +++ b/dev-requirements.in @@ -20,7 +20,6 @@ IPython keyrings.alt setuptools_scm pytest-icdiff -jinja2 # Tensorflow is not available for python 3.12 yet: https://github.com/tensorflow/tensorflow/issues/62003 tensorflow; python_version<'3.12' diff --git a/dev-requirements.txt b/dev-requirements.txt index 98c56efe69..41525cf1ad 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -195,10 +195,6 @@ jaraco-functools==4.0.1 # via keyring jedi==0.19.1 # via ipython -jinja2==3.1.4 - # via - # -r dev-requirements.in - # flytekit jmespath==1.0.1 # via botocore joblib==1.4.2 @@ -220,8 +216,6 @@ markdown-it-py==3.0.0 # via # flytekit # rich -markupsafe==2.1.5 - # via jinja2 marshmallow==3.21.2 # via # dataclasses-json diff --git a/flytekit/deck/deck.py b/flytekit/deck/deck.py index fbb398ef49..025306d47b 100644 --- a/flytekit/deck/deck.py +++ b/flytekit/deck/deck.py @@ -1,6 +1,8 @@ import enum import os import typing +from html import escape +from string import Template from typing import Optional from flytekit.core.context_manager import ExecutionParameters, ExecutionState, FlyteContext, FlyteContextManager @@ -153,8 +155,16 @@ def _get_deck( If ignore_jupyter is set to True, then it will return a str even in a jupyter environment. """ deck_map = {deck.name: deck.html for deck in new_user_params.decks} + nav_htmls = [] + body_htmls = [] - raw_html = get_deck_template().render(metadata=deck_map) + for key, value in deck_map.items(): + nav_htmls.append(f'
  • {escape(key)}
    • ') + # Can not escape here because this is HTML. Escaping it will present the HTML as text. + # The renderer must ensure that the HTML is safe. + body_htmls.append(f"
    {value}
    ") + + raw_html = get_deck_template().substitute(NAV_HTML="".join(nav_htmls), BODY_HTML="".join(body_htmls)) if not ignore_jupyter and ipython_check(): try: from IPython.core.display import HTML @@ -184,18 +194,9 @@ def _output_deck(task_name: str, new_user_params: ExecutionParameters): logger.error(f"Failed to write flyte deck html with error {e}.") -def get_deck_template() -> "Template": - from jinja2 import Environment, FileSystemLoader, select_autoescape - +def get_deck_template() -> Template: root = os.path.dirname(os.path.abspath(__file__)) - templates_dir = os.path.join(root, "html") - env = Environment( - loader=FileSystemLoader(templates_dir), - # 🔥 include autoescaping for security purposes - # sources: - # - https://jinja.palletsprojects.com/en/3.0.x/api/#autoescaping - # - https://stackoverflow.com/a/38642558/8474894 (see in comments) - # - https://stackoverflow.com/a/68826578/8474894 - autoescape=select_autoescape(enabled_extensions=("html",)), - ) - return env.get_template("template.html") + templates_dir = os.path.join(root, "html", "template.html") + with open(templates_dir, "r") as f: + template_content = f.read() + return Template(template_content) diff --git a/flytekit/deck/html/template.html b/flytekit/deck/html/template.html index a58e8a3252..4a560b7930 100644 --- a/flytekit/deck/html/template.html +++ b/flytekit/deck/html/template.html @@ -69,20 +69,14 @@ -
    - {% for key, value in metadata.items() %} -
    {{ value | safe }}
    - {% endfor %} + $BODY_HTML
    -