Run command in pyflyte-fast-execute in the same process #3029

thomasjpfan · 2025-01-03T20:21:53Z

Why are the changes needed?

With this PR, pyflyte-fast-execute does not run the command in another process, which saves ~7 seconds.

What changes were proposed in this pull request?

With this PR, pyflyte-fast-execute will run pyflyte-map-execute and pyflyte-execute in the same process, which improves startup time. There is a FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS env var that triggers the previous behavior.

How was this patch tested?

I ran:

from flytekit import task


@task(container_image="localhost:30000/flytekit:0.1.3")
def hello(name: str) -> str:
    return f"Hello! {name}"

With pyflyte run --remote wf.py hello --name flyte, it took 8s:

With pyflyte run --remote --env FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS=1 wf.py hello --name flyte, which took 15 s:

Summary by Bito

This PR includes multiple changes: (1) pyflyte-fast-execute optimization with in-process command execution for ~7 second startup time gains, (2) new Environment class for task configurations, (3) VLLM integration and Optuna plugin, (4) improved package management with uv.lock and poetry.lock support, and (5) dependency updates including jinja2 upgrade to 3.1.5 and modified version constraints for grpcio packages.

Unit tests added: True

Estimated effort to review (1-5, lower is better): 5

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot · 2025-01-03T20:22:07Z

Code Review Agent Run #ff52b8

Actionable Suggestions - 3

flytekit/bin/entrypoint.py - 3
- Consider adding error handling for Popen · Line 754-754
- Subprocess call with untrusted input · Line 754-754
- Consider moving static dict to module level · Line 776-776

Review Details

Files reviewed - 2 · Commit Range: c8a5915..c8a5915
- flytekit/bin/entrypoint.py
- flytekit/core/constants.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-01-03T20:25:29Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Feature Improvement - Optimize Command Execution in pyflyte-fast-execute	- `entrypoint.py` - Added in-process command execution capability with fallback to subprocess - `constants.py` - Added new environment variable for controlling process execution mode - `test_remote.py` - Added tests for new in-process execution functionality

flyte-bot · 2025-01-03T20:25:31Z

flytekit/bin/entrypoint.py

+
+    command(args)
+
+


Consider adding error handling for Popen

Consider adding error handling around subprocess.Popen() call to catch potential OSError or ValueError exceptions that could occur during process creation.

Code suggestion

Check the AI-generated fix before applying

Suggested change

try:

p = subprocess.Popen(cmd, env=env)

except (OSError, ValueError) as e:

logger.error(f"Failed to start subprocess: {e}")

exit(1)

Code Review Run #ff52b8

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

flytekit/bin/entrypoint.py

flyte-bot · 2025-01-03T20:25:33Z

flytekit/bin/entrypoint.py

@@ -742,24 +773,16 @@ def fast_execute_task_cmd(additional_distribution: str, dest_dir: str, task_exec
            cmd.extend(["--dynamic-addl-distro", additional_distribution, "--dynamic-dest-dir", dest_dir])
        cmd.append(arg)

+    commands_to_run_in_process = {cmd.name: cmd for cmd in [map_execute_task_cmd, execute_task_cmd]}


Consider moving static dict to module level

Consider moving the commands_to_run_in_process dictionary definition outside the function to avoid recreating it on every call. This dictionary is static and could be defined at module level.

Code suggestion

Check the AI-generated fix before applying

@@ -1,1 +1,3 @@ +_commands_to_run_in_process = {cmd.name: cmd for cmd in [map_execute_task_cmd, execute_task_cmd]} + def fast_execute_task_cmd(additional_distribution: str, dest_dir: str, task_execute_cmd: List[str]): @@ -776,1 +778,0 @@ - commands_to_run_in_process = {cmd.name: cmd for cmd in [map_execute_task_cmd, execute_task_cmd]}

Code Review Run #ff52b8

Is this a valid issue, or was it incorrectly flagged by the Agent?

it was incorrectly flagged

codecov · 2025-01-03T20:28:46Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 79.58%. Comparing base (dfa8f04) to head (1fe153b).

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #3029       +/-   ##
===========================================
+ Coverage   47.23%   79.58%   +32.34%     
===========================================
  Files         202      202               
  Lines       21355    21356        +1     
  Branches     2744     2744               
===========================================
+ Hits        10088    16996     +6908     
+ Misses      10776     3593     -7183     
- Partials      491      767      +276

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot · 2025-01-03T21:48:58Z

Code Review Agent Run #33c19c

Actionable Suggestions - 1

flytekit/bin/entrypoint.py - 1
- Consider handling None from os.getenv · Line 780-780

Review Details

Files reviewed - 1 · Commit Range: c8a5915..00abba6
- flytekit/bin/entrypoint.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flytekit/bin/entrypoint.py

pingsutw · 2025-01-06T22:49:34Z

flytekit/core/constants.py

@@ -23,6 +23,9 @@
 # Set this environment variable to true to force the task to return non-zero exit code on failure.
 FLYTE_FAIL_ON_ERROR = "FLYTE_FAIL_ON_ERROR"

+# Set this environment variable to true to force pyflyte-fast-execute to run task-execute-cmd in a separate process
+FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS = "FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS"
+


Could we run a workflow in the integration tests with FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS enabled? Just want to make sure the old behavior is still working.

flytekit/tests/flytekit/integration/remote/test_remote.py

Line 105 in 99f6d9f

def test_remote_run():

…r_flyte_fast_execute Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot · 2025-01-09T16:17:43Z

Code Review Agent Run #cfcbee

Actionable Suggestions - 0

Additional Suggestions - 10

plugins/flytekit-spark/tests/test_environment.py - 1
- Consider parameterizing Spark app name assertion · Line 17-19
tests/flytekit/integration/remote/utils.py - 1
- Consider making parquet file path configurable · Line 89-89
plugins/flytekit-optuna/setup.py - 1
- Consider adding Python 3.11 support · Line 21-21
tests/flytekit/unit/core/test_workflows.py - 1
- Consider expanding workflow lhs test coverage · Line 360-363
tests/flytekit/integration/remote/test_remote.py - 2
- Consider adding type hints for parameters · Line 68-68
- Consider configurable timeout for execution wait · Line 873-873
plugins/flytekit-inference/tests/test_vllm.py - 1
- Consider breaking down long assertion statement · Line 41-41
plugins/flytekit-inference/flytekitplugins/inference/vllm/serve.py - 1
- Consider consolidating hf_secret validation checks · Line 44-47
flytekit/image_spec/default_builder.py - 1
- Consider moving poetry arg to template · Line 196-197
tests/flytekit/unit/core/image_spec/test_default_builder.py - 1
- Consider using complete error message pattern · Line 271-271

Review Details

Files reviewed - 26 · Commit Range: 00abba6..eda9c8c
- flytekit/__init__.py
- flytekit/core/array_node_map_task.py
- flytekit/core/environment.py
- flytekit/core/workflow.py
- flytekit/image_spec/default_builder.py
- flytekit/remote/remote.py
- flytekit/types/directory/types.py
- flytekit/types/file/file.py
- plugins/flytekit-inference/flytekitplugins/inference/__init__.py
- plugins/flytekit-inference/flytekitplugins/inference/vllm/serve.py
- plugins/flytekit-inference/setup.py
- plugins/flytekit-inference/tests/test_vllm.py
- plugins/flytekit-optuna/flytekitplugins/optuna/__init__.py
- plugins/flytekit-optuna/flytekitplugins/optuna/optimizer.py
- plugins/flytekit-optuna/setup.py
- plugins/flytekit-optuna/tests/test_optimizer.py
- plugins/flytekit-spark/tests/test_environment.py
- pyproject.toml
- tests/flytekit/integration/remote/test_remote.py
- tests/flytekit/integration/remote/utils.py
- tests/flytekit/integration/remote/workflows/basic/attr_access_sd.py
- tests/flytekit/unit/core/image_spec/test_default_builder.py
- tests/flytekit/unit/core/test_environment.py
- tests/flytekit/unit/core/test_flyte_directory.py
- tests/flytekit/unit/core/test_flyte_file.py
- tests/flytekit/unit/core/test_workflows.py
Files skipped - 3
- .github/workflows/pythonbuild.yml - Reason: Filter setting
- plugins/flytekit-inference/README.md - Reason: Filter setting
- plugins/flytekit-optuna/README.md - Reason: Filter setting
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot · 2025-01-09T18:47:55Z

Code Review Agent Run #53b9a2

Actionable Suggestions - 1

tests/flytekit/integration/remote/test_remote.py - 1
- Parameter reordering may break API compatibility · Line 68-68

Review Details

Files reviewed - 1 · Commit Range: eda9c8c..b3b54a9
- tests/flytekit/integration/remote/test_remote.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

tests/flytekit/integration/remote/test_remote.py

…r_flyte_fast_execute Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot · 2025-01-10T18:52:17Z

Code Review Agent Run #52ec9b

Actionable Suggestions - 0

Review Details

Files reviewed - 2 · Commit Range: b3b54a9..1fe153b
- plugins/flytekit-onnx-pytorch/dev-requirements.txt
- pyproject.toml
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

Run pyflyte-fast-execute in the same process

c8a5915

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

thomasjpfan requested review from wild-endeavor, kumare3, eapolinario, pingsutw, cosmicBboy, samhita-alla and Future-Outlier as code owners January 3, 2025 20:21

flyte-bot reviewed Jan 3, 2025

View reviewed changes

flytekit/bin/entrypoint.py Show resolved Hide resolved

flyte-bot reviewed Jan 3, 2025

View reviewed changes

Use str2bool

00abba6

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot reviewed Jan 3, 2025

View reviewed changes

flytekit/bin/entrypoint.py Show resolved Hide resolved

pingsutw reviewed Jan 6, 2025

View reviewed changes

thomasjpfan added 3 commits January 9, 2025 10:07

Merge remote-tracking branch 'upstream/master' into thomasjpfan/faste…

e9b0f6c

…r_flyte_fast_execute Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

Add FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS to integration test

9e86728

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

Add run_args

eda9c8c

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

Fix tests

b3b54a9

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flyte-bot reviewed Jan 9, 2025

View reviewed changes

tests/flytekit/integration/remote/test_remote.py Show resolved Hide resolved

Merge remote-tracking branch 'upstream/master' into thomasjpfan/faste…

1fe153b

…r_flyte_fast_execute Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Run command in pyflyte-fast-execute in the same process #3029

Run command in pyflyte-fast-execute in the same process #3029

thomasjpfan commented Jan 3, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Jan 3, 2025 •

edited

Loading

Code Review Agent Run #ff52b8

flyte-bot commented Jan 3, 2025 •

edited

Loading

Changelist by Bito

flyte-bot Jan 3, 2025

flyte-bot Jan 3, 2025 •

edited by thomasjpfan

Loading

codecov bot commented Jan 3, 2025 •

edited

Loading

flyte-bot commented Jan 3, 2025 •

edited

Loading

Code Review Agent Run #33c19c

pingsutw Jan 6, 2025

flyte-bot commented Jan 9, 2025 •

edited

Loading

Code Review Agent Run #cfcbee

flyte-bot commented Jan 9, 2025 •

edited

Loading

Code Review Agent Run #53b9a2

flyte-bot commented Jan 10, 2025 •

edited

Loading

Code Review Agent Run #52ec9b

+    try:
+        p = subprocess.Popen(cmd, env=env)
+    except (OSError, ValueError) as e:
+        logger.error(f"Failed to start subprocess: {e}")
+        exit(1)

Run command in pyflyte-fast-execute in the same process #3029

Are you sure you want to change the base?

Run command in pyflyte-fast-execute in the same process #3029

Conversation

thomasjpfan commented Jan 3, 2025 • edited by flyte-bot Loading

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Summary by Bito

flyte-bot commented Jan 3, 2025 • edited Loading

Code Review Agent Run #ff52b8

flyte-bot commented Jan 3, 2025 • edited Loading

Changelist by Bito

flyte-bot Jan 3, 2025

Choose a reason for hiding this comment

flyte-bot Jan 3, 2025 • edited by thomasjpfan Loading

Choose a reason for hiding this comment

codecov bot commented Jan 3, 2025 • edited Loading

Codecov Report

flyte-bot commented Jan 3, 2025 • edited Loading

Code Review Agent Run #33c19c

pingsutw Jan 6, 2025

Choose a reason for hiding this comment

flyte-bot commented Jan 9, 2025 • edited Loading

Code Review Agent Run #cfcbee

flyte-bot commented Jan 9, 2025 • edited Loading

Code Review Agent Run #53b9a2

flyte-bot commented Jan 10, 2025 • edited Loading

Code Review Agent Run #52ec9b

thomasjpfan commented Jan 3, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Jan 3, 2025 •

edited

Loading

flyte-bot commented Jan 3, 2025 •

edited

Loading

flyte-bot Jan 3, 2025 •

edited by thomasjpfan

Loading

codecov bot commented Jan 3, 2025 •

edited

Loading

flyte-bot commented Jan 3, 2025 •

edited

Loading

flyte-bot commented Jan 9, 2025 •

edited

Loading

flyte-bot commented Jan 9, 2025 •

edited

Loading

flyte-bot commented Jan 10, 2025 •

edited

Loading