fix(interactive): pin PORT env var so code-server binds to the correct port

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:

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()

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"