Skip to content

Commit fd83ca0

Browse files
committed
fix(interactive): pin PORT env var so code-server binds to the correct port
code-server reads the PORT environment variable as one of its bind-address sources. When a pod/image sets PORT (common in Kubernetes), it overrode the explicit --bind-addr flag and code-server bound to the wrong port. Pass env={"PORT": str(self.port)} to the code-server subprocess so the inherited PORT can't override --bind-addr. execute_command now accepts an optional env dict that is merged on top of os.environ. Signed-off-by: Jeev B <jeevb@users.noreply.github.com>
1 parent 9527091 commit fd83ca0

3 files changed

Lines changed: 66 additions & 3 deletions

File tree

‎flytekit/interactive/utils.py‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,14 +63,26 @@ def get_task_inputs(task_module_name, task_name, context_working_dir):
6363
return native_inputs
6464

6565

66-
def execute_command(cmd):
66+
def execute_command(cmd, env=None):
6767
"""
6868
Execute a command in the shell.
69+
70+
Args:
71+
cmd (str): The command to execute.
72+
env (dict, optional): Environment variables to set for the subprocess. These are merged on top of the
73+
current process environment, so callers can override specific variables (e.g. ``PORT``) without
74+
dropping the rest of the inherited environment.
6975
"""
7076

7177
logger = flytekit.current_context().logging
7278

73-
process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
79+
subprocess_env = None
80+
if env is not None:
81+
subprocess_env = {**os.environ, **env}
82+
83+
process = subprocess.Popen(
84+
cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=subprocess_env
85+
)
7486
logger.info(f"cmd: {cmd}")
7587
stdout, stderr = process.communicate()
7688
if process.returncode != EXIT_CODE_SUCCESS:

‎flytekit/interactive/vscode_lib/decorator.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -443,7 +443,10 @@ def execute(self, *args, **kwargs):
443443
child_process = multiprocessing.Process(
444444
target=execute_command,
445445
kwargs={
446-
"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}"
446+
"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}",
447+
# code-server also reads the PORT env var when resolving its bind address. Explicitly pin it to
448+
# self.port so an inherited PORT (e.g. injected by Kubernetes) can't override --bind-addr.
449+
"env": {"PORT": str(self.port)},
447450
},
448451
)
449452
child_process.start()

‎tests/flytekit/unit/interactive/test_flyteinteractive_vscode.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import os
12
from collections import OrderedDict
23

34
import mock
@@ -13,6 +14,7 @@
1314
from flytekit.interactive.constants import (
1415
EXIT_CODE_SUCCESS,
1516
)
17+
from flytekit.interactive.utils import execute_command
1618
from flytekit.interactive.vscode_lib.decorator import (
1719
get_code_server_info,
1820
get_installed_extensions,
@@ -363,3 +365,49 @@ def test_get_installed_extensions_failed(mock_run):
363365

364366
expected_extensions = []
365367
assert installed_extensions == expected_extensions
368+
369+
370+
def test_vscode_passes_port_env_to_child_process(vscode_patches, mock_remote_execution):
371+
"""code-server must bind to self.port regardless of a PORT env var inherited from the pod/image."""
372+
(
373+
mock_process,
374+
_,
375+
_,
376+
_,
377+
_,
378+
_,
379+
_,
380+
) = vscode_patches
381+
382+
@task
383+
@vscode(port=1234)
384+
def t():
385+
return
386+
387+
@workflow
388+
def wf():
389+
t()
390+
391+
with mock.patch.dict(os.environ, {"PORT": "9999"}):
392+
wf()
393+
394+
mock_process.assert_called_once()
395+
_, kwargs = mock_process.call_args
396+
assert kwargs["kwargs"]["env"] == {"PORT": "1234"}
397+
398+
399+
@mock.patch("subprocess.Popen")
400+
def test_execute_command_passes_env_to_subprocess(mock_popen):
401+
"""execute_command should merge the provided env over os.environ when launching the subprocess."""
402+
mock_proc = mock.Mock()
403+
mock_proc.returncode = EXIT_CODE_SUCCESS
404+
mock_proc.communicate.return_value = (b"", b"")
405+
mock_popen.return_value = mock_proc
406+
407+
with mock.patch.dict(os.environ, {"PORT": "9999", "EXISTING": "keep"}):
408+
execute_command("code-server --bind-addr 0.0.0.0:1234", env={"PORT": "1234"})
409+
410+
_, kwargs = mock_popen.call_args
411+
assert kwargs["env"]["PORT"] == "1234"
412+
# The inherited environment is preserved for other variables.
413+
assert kwargs["env"]["EXISTING"] == "keep"

0 commit comments

Comments
 (0)