diff --git a/src/flyte/_debug/utils.py b/src/flyte/_debug/utils.py index 8eb2ec480..24ef6a16d 100644 --- a/src/flyte/_debug/utils.py +++ b/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: diff --git a/src/flyte/_debug/vscode.py b/src/flyte/_debug/vscode.py index 068b59f9c..7469227b2 100644 --- a/src/flyte/_debug/vscode.py +++ b/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() diff --git a/tests/flyte/test_debug.py b/tests/flyte/test_debug.py index 07c4dc527..3a584a0e2 100644 --- a/tests/flyte/test_debug.py +++ b/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"])