From 15dee9579a000696539b63745d0d036647987c0e Mon Sep 17 00:00:00 2001 From: arbaobao Date: Tue, 24 Sep 2024 06:58:41 +0800 Subject: [PATCH] add dest_dir into pythonpath before loading modules (#2692) Signed-off-by: Nelson Chen --- flytekit/bin/entrypoint.py | 9 ++++++++- tests/flytekit/integration/remote/test_remote.py | 3 +++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index 4b1dec78c6..74069af0a0 100644 --- a/flytekit/bin/entrypoint.py +++ b/flytekit/bin/entrypoint.py @@ -562,7 +562,14 @@ def fast_execute_task_cmd(additional_distribution: str, dest_dir: str, task_exec # 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. - p = subprocess.Popen(cmd) + 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") diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py index 0d9047294c..6607fb65fb 100644 --- a/tests/flytekit/integration/remote/test_remote.py +++ b/tests/flytekit/integration/remote/test_remote.py @@ -32,6 +32,7 @@ PROJECT = "flytesnacks" DOMAIN = "development" VERSION = f"v{os.getpid()}" +DEST_DIR = "/tmp" @pytest.fixture(scope="session") @@ -66,6 +67,8 @@ def run(file_name, wf_name, *args): CONFIG, "run", "--remote", + "--destination-dir", + DEST_DIR, "--image", IMAGE, "--project",