From 1b562210eeec1397ee5bc0f901ae51b4e4a5edd2 Mon Sep 17 00:00:00 2001 From: Ketan Umare Date: Thu, 6 Feb 2025 16:05:22 -0800 Subject: [PATCH 1/3] Uses fixed and default inputs for launchplans on the cli Signed-off-by: Ketan Umare --- .gitignore | 1 + flytekit/clis/sdk_in_container/run.py | 18 +++++++++++++++++- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index ac4cf37b06..0db8768ef2 100644 --- a/.gitignore +++ b/.gitignore @@ -39,3 +39,4 @@ coverage.xml # Version file is auto-generated by setuptools_scm flytekit/_version.py +testing diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index 44141a5cc1..7a08ef31af 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -1052,9 +1052,25 @@ def _create_command( r = run_level_params.remote_instance() flyte_ctx = r.context + final_inputs_with_defaults = loaded_entity.python_interface.inputs_with_defaults + if isinstance(loaded_entity, LaunchPlan): + # For LaunchPlans it is essential to handle fixed inputs and default inputs in a special way + # Fixed inputs are inputs that are always passed to the launch plan and cannot be overridden + # Default inputs are inputs that are optional and have a default value + # The final inputs to the launch plan are a combination of the fixed inputs and the default inputs + all_inputs = loaded_entity.python_interface.inputs_with_defaults + default_inputs = loaded_entity.saved_inputs + pmap = loaded_entity.parameters + final_inputs_with_defaults = {} + for name, _ in pmap.parameters.items(): + _type, v = all_inputs[name] + if name in default_inputs: + v = default_inputs[name] + final_inputs_with_defaults[name] = _type, v + # Add options for each of the workflow inputs params = [] - for input_name, input_type_val in loaded_entity.python_interface.inputs_with_defaults.items(): + for input_name, input_type_val in final_inputs_with_defaults.items(): literal_var = loaded_entity.interface.inputs.get(input_name) python_type, default_val = input_type_val required = type(None) not in get_args(python_type) and default_val is None From c0bcb2730b5bfe85a4de5cfabaa46761d7ebbfa2 Mon Sep 17 00:00:00 2001 From: Ketan Umare Date: Thu, 6 Feb 2025 17:07:11 -0800 Subject: [PATCH 2/3] Unit tests Signed-off-by: Ketan Umare --- tests/flytekit/unit/cli/pyflyte/test_run.py | 88 +++++++++---------- .../flytekit/unit/cli/pyflyte/test_run_lps.py | 72 +++++++++++++++ tests/flytekit/unit/cli/pyflyte/workflow.py | 2 +- 3 files changed, 117 insertions(+), 45 deletions(-) create mode 100644 tests/flytekit/unit/cli/pyflyte/test_run_lps.py diff --git a/tests/flytekit/unit/cli/pyflyte/test_run.py b/tests/flytekit/unit/cli/pyflyte/test_run.py index 848dbbf6e1..e4ab4145dc 100644 --- a/tests/flytekit/unit/cli/pyflyte/test_run.py +++ b/tests/flytekit/unit/cli/pyflyte/test_run.py @@ -16,7 +16,7 @@ from flytekit.clis.sdk_in_container.run import ( RunLevelParams, get_entities_in_file, - run_command, + run_command, WorkflowCommand, ) from flytekit.configuration import Config, Image, ImageConfig from flytekit.core.task import task @@ -28,8 +28,7 @@ from flytekit.remote import FlyteRemote from typing import Iterator, List from flytekit.types.iterator import JSON -from flytekit import workflow - +from flytekit import workflow, LaunchPlan pytest.importorskip("pandas") @@ -205,7 +204,7 @@ def test_pyflyte_run_cli(workflow_file): "--s", json.dumps({"x": {"i": 1, "a": ["h", "e"]}}), "--t", - json.dumps({"i": [{"i":1,"a":["h","e"]}]}), + json.dumps({"i": [{"i": 1, "a": ["h", "e"]}]}), ], catch_exceptions=False, ) @@ -293,7 +292,8 @@ def test_all_types_with_yaml_input(): result = runner.invoke( pyflyte.main, - ["run", os.path.join(DIR_NAME, "workflow.py"), "my_wf", "--inputs-file", os.path.join(os.path.dirname(os.path.realpath(__file__)), "my_wf_input.yaml")], + ["run", os.path.join(DIR_NAME, "workflow.py"), "my_wf", "--inputs-file", + os.path.join(os.path.dirname(os.path.realpath(__file__)), "my_wf_input.yaml")], catch_exceptions=False, ) assert result.exit_code == 0, result.stdout @@ -301,7 +301,7 @@ def test_all_types_with_yaml_input(): def test_all_types_with_pipe_input(monkeypatch): runner = CliRunner() - input= str(json.load(open(os.path.join(os.path.dirname(os.path.realpath(__file__)), "my_wf_input.json"),"r"))) + input = str(json.load(open(os.path.join(os.path.dirname(os.path.realpath(__file__)), "my_wf_input.json"), "r"))) monkeypatch.setattr("sys.stdin", io.StringIO(input)) result = runner.invoke( pyflyte.main, @@ -321,18 +321,18 @@ def test_all_types_with_pipe_input(monkeypatch): "pipe_input, option_input", [ ( - str( - json.load( - open( - os.path.join( - os.path.dirname(os.path.realpath(__file__)), - "my_wf_input.json", - ), - "r", + str( + json.load( + open( + os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "my_wf_input.json", + ), + "r", + ) ) - ) - ), - "GREEN", + ), + "GREEN", ) ], ) @@ -579,11 +579,11 @@ def test_list_default_arguments(wf_path): reason="Github macos-latest image does not have docker installed as per https://github.com/orgs/community/discussions/25777", ) def test_pyflyte_run_run( - mock_image, - image_string, - leaf_configuration_file_name, - final_image_config, - mock_image_spec_builder, + mock_image, + image_string, + leaf_configuration_file_name, + final_image_config, + mock_image_spec_builder, ): mock_image.return_value = "cr.flyte.org/flyteorg/flytekit:py3.9-latest" ImageBuildEngine.register("test", mock_image_spec_builder) @@ -597,10 +597,10 @@ def tk(): ... image_config = ImageConfig.validate_image(None, "", image_tuple) pp = ( - pathlib.Path(__file__).parent.parent.parent - / "configuration" - / "configs" - / leaf_configuration_file_name + pathlib.Path(__file__).parent.parent.parent + / "configuration" + / "configs" + / leaf_configuration_file_name ) obj = RunLevelParams( @@ -641,7 +641,7 @@ def jsons(): @mock.patch("flytekit.configuration.default_images.DefaultImages.default_image") def test_pyflyte_run_with_iterator_json_type( - mock_image, mock_image_spec_builder, caplog + mock_image, mock_image_spec_builder, caplog ): mock_image.return_value = "cr.flyte.org/flyteorg/flytekit:py3.9-latest" ImageBuildEngine.register( @@ -679,10 +679,10 @@ def tk_simple_iterator(x: Iterator[int] = iter([1, 2, 3])) -> Iterator[int]: image_config = ImageConfig.validate_image(None, "", image_tuple) pp = ( - pathlib.Path(__file__).parent.parent.parent - / "configuration" - / "configs" - / "no_images.yaml" + pathlib.Path(__file__).parent.parent.parent + / "configuration" + / "configs" + / "no_images.yaml" ) obj = RunLevelParams( @@ -796,9 +796,9 @@ def test_pyflyte_run_with_none(a_val, workflow_file): [ (["--env", "MY_ENV_VAR=hello"], '["MY_ENV_VAR"]', "hello"), ( - ["--env", "MY_ENV_VAR=hello", "--env", "ABC=42"], - '["MY_ENV_VAR","ABC"]', - "hello,42", + ["--env", "MY_ENV_VAR=hello", "--env", "ABC=42"], + '["MY_ENV_VAR","ABC"]', + "hello,42", ), ], ) @@ -813,16 +813,16 @@ def test_pyflyte_run_with_none(a_val, workflow_file): def test_envvar_local_execution(envs, envs_argument, expected_output, workflow_file): runner = CliRunner() args = ( - [ - "run", - ] - + envs - + [ - workflow_file, - "wf_with_env_vars", - "--env_vars", - ] - + [envs_argument] + [ + "run", + ] + + envs + + [ + workflow_file, + "wf_with_env_vars", + "--env_vars", + ] + + [envs_argument] ) result = runner.invoke( pyflyte.main, diff --git a/tests/flytekit/unit/cli/pyflyte/test_run_lps.py b/tests/flytekit/unit/cli/pyflyte/test_run_lps.py new file mode 100644 index 0000000000..30211a5534 --- /dev/null +++ b/tests/flytekit/unit/cli/pyflyte/test_run_lps.py @@ -0,0 +1,72 @@ +import click + +from flytekit import task, workflow, LaunchPlan +from flytekit.clis.sdk_in_container.run import WorkflowCommand, RunLevelParams + +import mock +import pytest + + +@task +def two_inputs(x: int, y: str) -> str: + return f"{x},{y}" + +@workflow +def two_inputs_wf(x: int, y: str) -> str: + return two_inputs(x, y) + +lp_fixed_y_default_x = LaunchPlan.get_or_create( + workflow=two_inputs_wf, + name="fixed-default-inputs", + fixed_inputs={"y": "hello"}, + default_inputs={"x": 1} +) + +lp_fixed_y = LaunchPlan.get_or_create( + workflow=two_inputs_wf, + name="fixed-y", + fixed_inputs={"y": "hello"}, +) + +lp_fixed_x = LaunchPlan.get_or_create( + workflow=two_inputs_wf, + name="fixed-x", + fixed_inputs={"x": 1}, +) + +lp_fixed_all = LaunchPlan.get_or_create( + workflow=two_inputs_wf, + name="fixed-all", + fixed_inputs={"x": 1, "y": "test"}, +) + +lp_default_x = LaunchPlan.get_or_create( + name="default-inputs", + workflow=two_inputs_wf, + default_inputs={"x": 1} +) + +lp_simple = LaunchPlan.get_or_create( + workflow=two_inputs_wf, + name="no-fixed-default", +) + +@pytest.mark.parametrize("lp_execs", [ + (lp_fixed_y_default_x, {"x": 1}), + (lp_fixed_y, {"x": None}), + (lp_fixed_x, {"y": None}), + (lp_fixed_all, {}), + (lp_default_x, {"y": None, "x": 1}), + (lp_simple, {"x": None, "y": None}), +]) +def test_workflowcommand_create_command(lp_execs): + cmd = WorkflowCommand("testfile.py") + rp = RunLevelParams() + ctx = click.Context(cmd, obj=rp) + lp, exp_opts = lp_execs + opts = cmd._create_command(ctx, "test_entity", rp, lp, "launch plan").params + for o in opts: + if "input" in o.name: + continue + assert o.name in exp_opts + assert o.default == exp_opts[o.name] diff --git a/tests/flytekit/unit/cli/pyflyte/workflow.py b/tests/flytekit/unit/cli/pyflyte/workflow.py index 96351485af..7502b9b993 100644 --- a/tests/flytekit/unit/cli/pyflyte/workflow.py +++ b/tests/flytekit/unit/cli/pyflyte/workflow.py @@ -8,7 +8,7 @@ from dataclasses_json import DataClassJsonMixin from typing_extensions import Annotated -from flytekit import kwtypes, task, workflow +from flytekit import kwtypes, task, workflow, LaunchPlan from flytekit.types.directory import FlyteDirectory from flytekit.types.file import FlyteFile from flytekit.types.structured import StructuredDataset From 16463d29ec3043062c541096c8d376f5c853a681 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Wed, 12 Feb 2025 12:18:37 -0800 Subject: [PATCH 3/3] Apply suggestions from code review --- tests/flytekit/unit/cli/pyflyte/workflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit/unit/cli/pyflyte/workflow.py b/tests/flytekit/unit/cli/pyflyte/workflow.py index 7502b9b993..96351485af 100644 --- a/tests/flytekit/unit/cli/pyflyte/workflow.py +++ b/tests/flytekit/unit/cli/pyflyte/workflow.py @@ -8,7 +8,7 @@ from dataclasses_json import DataClassJsonMixin from typing_extensions import Annotated -from flytekit import kwtypes, task, workflow, LaunchPlan +from flytekit import kwtypes, task, workflow from flytekit.types.directory import FlyteDirectory from flytekit.types.file import FlyteFile from flytekit.types.structured import StructuredDataset