diff --git a/flytekit/interactive/utils.py b/flytekit/interactive/utils.py index 6d39032756..ee0c1a082e 100644 --- a/flytekit/interactive/utils.py +++ b/flytekit/interactive/utils.py @@ -2,6 +2,7 @@ import os import subprocess import sys +from typing import Dict, Optional from flyteidl.core import literals_pb2 as _literals_pb2 @@ -63,14 +64,27 @@ def get_task_inputs(task_module_name, task_name, context_working_dir): return native_inputs -def execute_command(cmd): +def execute_command(cmd: str, env: Optional[Dict[str, str]] = None): """ Execute a command in the shell. + + Args: + cmd (str): The command to execute. + env (Optional[Dict[str, str]]): Environment variables to set for the subprocess. These are merged on top of + the current process environment, so callers can override specific variables (e.g. ``PORT``) without + dropping the rest of the inherited environment. + + Raises: + RuntimeError: If the command exits with a non-zero return code. """ logger = flytekit.current_context().logging - process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + subprocess_env = None + if env is not None: + subprocess_env = {**os.environ, **env} + + process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=subprocess_env) logger.info(f"cmd: {cmd}") stdout, stderr = process.communicate() if process.returncode != EXIT_CODE_SUCCESS: diff --git a/flytekit/interactive/vscode_lib/decorator.py b/flytekit/interactive/vscode_lib/decorator.py index f709e26bb2..d1f24b8081 100644 --- a/flytekit/interactive/vscode_lib/decorator.py +++ b/flytekit/interactive/vscode_lib/decorator.py @@ -443,7 +443,10 @@ def execute(self, *args, **kwargs): child_process = multiprocessing.Process( target=execute_command, kwargs={ - "cmd": f"code-server --bind-addr 0.0.0.0:{self.port} --idle-timeout-seconds {code_server_idle_timeout_seconds} --disable-workspace-trust --auth none {task_function_source_dir}" + "cmd": f"code-server --bind-addr 0.0.0.0:{self.port} --idle-timeout-seconds {code_server_idle_timeout_seconds} --disable-workspace-trust --auth none {task_function_source_dir}", + # code-server also reads the PORT env var when resolving its bind address. Explicitly pin it to + # self.port so an inherited PORT (e.g. injected by Kubernetes) can't override --bind-addr. + "env": {"PORT": str(self.port)}, }, ) child_process.start() diff --git a/tests/flytekit/unit/interactive/test_flyteinteractive_vscode.py b/tests/flytekit/unit/interactive/test_flyteinteractive_vscode.py index 792e030018..2d93036696 100644 --- a/tests/flytekit/unit/interactive/test_flyteinteractive_vscode.py +++ b/tests/flytekit/unit/interactive/test_flyteinteractive_vscode.py @@ -1,3 +1,4 @@ +import os from collections import OrderedDict import mock @@ -13,6 +14,7 @@ from flytekit.interactive.constants import ( EXIT_CODE_SUCCESS, ) +from flytekit.interactive.utils import execute_command from flytekit.interactive.vscode_lib.decorator import ( get_code_server_info, get_installed_extensions, @@ -363,3 +365,49 @@ def test_get_installed_extensions_failed(mock_run): expected_extensions = [] assert installed_extensions == expected_extensions + + +def test_vscode_passes_port_env_to_child_process(vscode_patches, mock_remote_execution): + """code-server must bind to self.port regardless of a PORT env var inherited from the pod/image.""" + ( + mock_process, + _, + _, + _, + _, + _, + _, + ) = vscode_patches + + @task + @vscode(port=1234) + def t(): + return + + @workflow + def wf(): + t() + + with mock.patch.dict(os.environ, {"PORT": "9999"}): + wf() + + mock_process.assert_called_once() + _, kwargs = mock_process.call_args + assert kwargs["kwargs"]["env"] == {"PORT": "1234"} + + +@mock.patch("subprocess.Popen") +def test_execute_command_passes_env_to_subprocess(mock_popen): + """execute_command should merge the provided env over os.environ when launching the subprocess.""" + mock_proc = mock.Mock() + mock_proc.returncode = EXIT_CODE_SUCCESS + mock_proc.communicate.return_value = (b"", b"") + mock_popen.return_value = mock_proc + + with mock.patch.dict(os.environ, {"PORT": "9999", "EXISTING": "keep"}): + execute_command("code-server --bind-addr 0.0.0.0:1234", env={"PORT": "1234"}) + + _, kwargs = mock_popen.call_args + assert kwargs["env"]["PORT"] == "1234" + # The inherited environment is preserved for other variables. + assert kwargs["env"]["EXISTING"] == "keep"