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

flytekit/bin/entrypoint.py

@@ -30,7 +30,7 @@
 from flytekit.core import utils
 from flytekit.core.base_task import IgnoreOutputs, PythonTask
 from flytekit.core.checkpointer import SyncCheckpoint
-from flytekit.core.constants import FLYTE_FAIL_ON_ERROR
+from flytekit.core.constants import FLYTE_FAIL_ON_ERROR, FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS
 from flytekit.core.context_manager import (
     ExecutionParameters,
     ExecutionState,
@@ -721,6 +721,37 @@ def execute_task_cmd(
     )
 
 
+def _run_cmd_in_new_process(cmd, dest_dir):
+    """Run cmd in a new process."""
+    env = os.environ.copy()
+    if dest_dir is not None:
+        dest_dir_resolved = os.path.realpath(os.path.expanduser(dest_dir))
+        if "PYTHONPATH" in env:
+            env["PYTHONPATH"] += os.pathsep + dest_dir_resolved
+        else:
+            env["PYTHONPATH"] = dest_dir_resolved
+    p = subprocess.Popen(cmd, env=env)
+
+    def handle_sigterm(signum, frame):
+        logger.info(f"passing signum {signum} [frame={frame}] to subprocess")
+        p.send_signal(signum)
+
+    signal.signal(signal.SIGTERM, handle_sigterm)
+    returncode = p.wait()
+    exit(returncode)
+
+
+def _run_cmd_in_current_process(command: click.Command, args: List[str], dest_dir: Optional[str]):
+    """Run command with args in the same process."""
+
+    if dest_dir is not None:
+        dest_dir_resolved = os.path.realpath(os.path.expanduser(dest_dir))
+        if all(os.path.realpath(path) != dest_dir_resolved for path in sys.path):
+            sys.path.append(dest_dir_resolved)
+
+    command(args)
+
+
 @_pass_through.command("pyflyte-fast-execute")
 @click.option("--additional-distribution", required=False)
 @click.option("--dest-dir", required=False)
@@ -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]}
+
     # Use the commandline to run the task execute command rather than calling it directly in python code
     # since the current runtime bytecode references the older user code, rather than the downloaded distribution.
-    env = os.environ.copy()
-    if dest_dir is not None:
-        dest_dir_resolved = os.path.realpath(os.path.expanduser(dest_dir))
-        if "PYTHONPATH" in env:
-            env["PYTHONPATH"] += os.pathsep + dest_dir_resolved
-        else:
-            env["PYTHONPATH"] = dest_dir_resolved
-    p = subprocess.Popen(cmd, env=env)
-
-    def handle_sigterm(signum, frame):
-        logger.info(f"passing signum {signum} [frame={frame}] to subprocess")
-        p.send_signal(signum)
-
-    signal.signal(signal.SIGTERM, handle_sigterm)
-    returncode = p.wait()
-    exit(returncode)
+    if str2bool(os.getenv(FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS)) or cmd[0] not in commands_to_run_in_process:
+        logger.debug(f"Running {cmd[0]} in a new process")
+        _run_cmd_in_new_process(cmd, dest_dir)
+    else:
+        logger.debug(f"Running {cmd[0]} in the same process")
+        _run_cmd_in_current_process(commands_to_run_in_process[cmd[0]], cmd[1:], dest_dir)
 
 
 @_pass_through.command("pyflyte-map-execute")

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"
+
 # Executions launched by the current eager task will be tagged with this key:current_eager_exec_name
 EAGER_TAG_KEY = "eager-exec"
 

tests/flytekit/integration/remote/test_remote.py

@@ -65,15 +65,18 @@ def register():
     assert out.returncode == 0
 
 
-def run(file_name, wf_name, *args) -> str:
+def run(file_name, wf_name, *wf_args, run_args=None) -> str:
     # Copy the environment and set the environment variable
+    if run_args is None:
+        run_args = []
     out = subprocess.run(
         [
             "pyflyte",
             "--verbose",
             "-c",
             CONFIG,
             "run",
+            *run_args,
             "--remote",
             "--destination-dir",
             DEST_DIR,
@@ -85,7 +88,7 @@ def run(file_name, wf_name, *args) -> str:
             DOMAIN,
             MODULE_PATH / file_name,
             wf_name,
-            *args,
+            *wf_args,
         ],
         capture_output=True,  # Capture the output streams
         text=True,  # Return outputs as strings (not bytes)
@@ -110,6 +113,14 @@ def test_remote_run():
     run("default_lp.py", "my_wf")
 
 
+def test_remote_run_in_new_process():
+    # child_workflow.parent_wf asynchronously register a parent wf1 with child lp from another wf2.
+    run("child_workflow.py", "parent_wf", "--a", "3", run_args=["--env", "FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS=1"])
+
+    # run twice to make sure it will register a new version of the workflow.
+    run("default_lp.py", "my_wf", run_args=["--env", "FLYTE_FAST_EXECUTE_CMD_IN_NEW_PROCESS=1"])
+
+
 def test_remote_eager_run():
     # child_workflow.parent_wf asynchronously register a parent wf1 with child lp from another wf2.
     run("eager_example.py", "simple_eager_workflow", "--x", "3")