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

src/flyte/_debug/utils.py

@@ -1,14 +1,31 @@
 import asyncio
+import os
+from typing import Dict, Optional
 
 from flyte._debug.constants import EXIT_CODE_SUCCESS
 from flyte._logging import logger
 
 
-async def execute_command(cmd: str):
+async 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.
     """
-    process = await asyncio.create_subprocess_shell(cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
+    subprocess_env = None
+    if env is not None:
+        subprocess_env = {**os.environ, **env}
+
+    process = await asyncio.create_subprocess_shell(
+        cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env=subprocess_env
+    )
     logger.info(f"cmd: {cmd}")
     stdout, stderr = await process.communicate()
     if process.returncode != EXIT_CODE_SUCCESS:

src/flyte/_debug/vscode.py

@@ -9,7 +9,7 @@
 import tarfile
 import time
 from pathlib import Path
-from typing import List
+from typing import Dict, List, Optional
 
 import aiofiles
 import click
@@ -26,6 +26,7 @@
     FLYTE_ENABLE_VSCODE_KEY,
     HEARTBEAT_PATH,
     MAX_IDLE_SECONDS,
+    VSCODE_PORT,
 )
 from flyte._debug.utils import (
     execute_command,
@@ -277,6 +278,16 @@ def prepare_launch_json(ctx: click.Context, pid: int):
         json.dump(settings_json, file, indent=4)
 
 
+def _run_code_server(cmd: str, env: Optional[Dict[str, str]] = None) -> None:
+    """
+    Entrypoint for the code-server child process.
+
+    Defined at module level (rather than as a lambda) so it is picklable: the ``spawn`` and ``forkserver``
+    multiprocessing start methods pickle the target, and ``forkserver`` is the default on Linux for Python 3.14+.
+    """
+    asyncio.run(execute_command(cmd, env=env))
+
+
 async def _start_vscode_server(ctx: click.Context):
     if ctx.params["tgz"] is None:
         await download_vscode()
@@ -285,11 +296,18 @@ async def _start_vscode_server(ctx: click.Context):
             download_tgz(ctx.params["dest"], ctx.params["version"], ctx.params["tgz"]), download_vscode()
         )
     code_server_idle_timeout_seconds = os.getenv("CODE_SERVER_IDLE_TIMEOUT_SECONDS", str(MAX_IDLE_SECONDS))
+    cmd = (
+        f"code-server --bind-addr 0.0.0.0:{VSCODE_PORT}"
+        f" --idle-timeout-seconds {code_server_idle_timeout_seconds}"
+        f" --disable-workspace-trust --auth none {os.getcwd()}"
+    )
     child_process = multiprocessing.Process(
-        target=lambda cmd: asyncio.run(asyncio.run(execute_command(cmd))),
+        target=_run_code_server,
         kwargs={
-            "cmd": f"code-server --bind-addr 0.0.0.0:6060 --idle-timeout-seconds {code_server_idle_timeout_seconds}"
-            f" --disable-workspace-trust --auth none {os.getcwd()}"
+            "cmd": cmd,
+            # code-server also reads the PORT env var when resolving its bind address. Explicitly pin it to
+            # VSCODE_PORT so an inherited PORT (e.g. injected by Kubernetes) can't override --bind-addr.
+            "env": {"PORT": str(VSCODE_PORT)},
         },
     )
     child_process.start()

tests/flyte/test_debug.py

@@ -2,13 +2,17 @@
 
 from __future__ import annotations
 
+import os
 from types import SimpleNamespace
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
+import pytest
 from click.testing import CliRunner
 from flyteidl2.core.execution_pb2 import TaskLog
 
 from flyte._debug.client import _build_full_url, _extract_vscode_uri
+from flyte._debug.constants import EXIT_CODE_SUCCESS
+from flyte._debug.utils import execute_command
 from flyte._run import _Runner
 
 # ---------------------------------------------------------------------------
@@ -198,3 +202,79 @@ def test_caches_result(self):
             assert run.get_debug_url() == expected
             assert run.get_debug_url() == expected
             mock_watch.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# execute_command: env override pins PORT without dropping inherited vars
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_execute_command_passes_env_to_subprocess():
+    """execute_command should merge the provided env over os.environ when launching the subprocess."""
+    mock_proc = MagicMock()
+    mock_proc.returncode = EXIT_CODE_SUCCESS
+    mock_proc.communicate = AsyncMock(return_value=(b"", b""))
+
+    with (
+        patch("asyncio.create_subprocess_shell", new=AsyncMock(return_value=mock_proc)) as mock_create,
+        patch.dict(os.environ, {"PORT": "9999", "EXISTING": "keep"}),
+    ):
+        await execute_command("code-server --bind-addr 0.0.0.0:6060", env={"PORT": "6060"})
+
+    _, kwargs = mock_create.call_args
+    # The provided env overrides the inherited PORT so code-server binds to the right port.
+    assert kwargs["env"]["PORT"] == "6060"
+    # The inherited environment is preserved for other variables.
+    assert kwargs["env"]["EXISTING"] == "keep"
+
+
+@pytest.mark.asyncio
+async def test_execute_command_without_env_inherits_process_env():
+    """When no env override is given, the subprocess inherits the parent environment (env=None)."""
+    mock_proc = MagicMock()
+    mock_proc.returncode = EXIT_CODE_SUCCESS
+    mock_proc.communicate = AsyncMock(return_value=(b"", b""))
+
+    with patch("asyncio.create_subprocess_shell", new=AsyncMock(return_value=mock_proc)) as mock_create:
+        await execute_command("echo hello")
+
+    _, kwargs = mock_create.call_args
+    assert kwargs["env"] is None
+
+
+# ---------------------------------------------------------------------------
+# _start_vscode_server: pins PORT for code-server with a picklable target
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_start_vscode_server_pins_port_env_with_picklable_target():
+    """_start_vscode_server must launch code-server with PORT pinned to VSCODE_PORT via a picklable target."""
+    import pickle
+
+    from flyte._debug import vscode
+    from flyte._debug.constants import VSCODE_PORT
+
+    mock_proc = MagicMock()
+    mock_proc.pid = 4321
+    # is_alive() False so the heartbeat loop is skipped and the coroutine returns immediately.
+    mock_proc.is_alive.return_value = False
+
+    ctx = MagicMock()
+    ctx.params = {"tgz": None}
+
+    with (
+        patch.object(vscode, "download_vscode", new=AsyncMock()),
+        patch.object(vscode, "prepare_launch_json"),
+        patch.object(vscode.multiprocessing, "Process", return_value=mock_proc) as mock_process_cls,
+    ):
+        await vscode._start_vscode_server(ctx)
+
+    mock_process_cls.assert_called_once()
+    _, kwargs = mock_process_cls.call_args
+    # PORT is pinned so an inherited value can't override --bind-addr.
+    assert kwargs["kwargs"]["env"] == {"PORT": str(VSCODE_PORT)}
+    assert f"--bind-addr 0.0.0.0:{VSCODE_PORT}" in kwargs["kwargs"]["cmd"]
+    # The target must be picklable (lambdas are not) so the spawn/forkserver start methods work.
+    pickle.dumps(kwargs["target"])