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 combines multiple enhancements: (1) pyflyte-fast-execute optimization with in-process command execution for ~7 second startup time gains (with FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS flag for legacy behavior), (2) new Environment class for task configurations, (3) VLLM integration and Optuna plugin for model serving and hyperparameter optimization, and (4) improved package management with uv.lock and poetry.lock support

Unit tests added: True

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

Signed-off-by: Thomas J. Fan <[email protected]>

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 support for running commands in same process with fallback to new process - `constants.py` - Added new environment variable FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS - `test_remote.py` - Added tests for running commands in new process mode

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.01%. Comparing base (f3996f6) to head (c8a5915).

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #3029       +/-   ##
===========================================
+ Coverage   51.28%   79.01%   +27.72%     
===========================================
  Files         201      220       +19     
  Lines       21281    22215      +934     
  Branches     2734     2734               
===========================================
+ Hits        10915    17554     +6639     
+ Misses       9767     3923     -5844     
- Partials      599      738      +139

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

Signed-off-by: Thomas J. Fan <[email protected]>

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 <[email protected]>

Signed-off-by: Thomas J. Fan <[email protected]>

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 <[email protected]>

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

Run pyflyte-fast-execute in the same process

c8a5915

Signed-off-by: Thomas J. Fan <[email protected]>

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 <[email protected]>

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 <[email protected]>

Add FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS to integration test

9e86728

Signed-off-by: Thomas J. Fan <[email protected]>

Add run_args

eda9c8c

Signed-off-by: Thomas J. Fan <[email protected]>

Fix tests

b3b54a9

Signed-off-by: Thomas J. Fan <[email protected]>

flyte-bot reviewed Jan 9, 2025

View reviewed changes

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

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

+    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

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