Support worker env vars before importing the server

miles/utils/workers/serving/serve.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+from miles.utils.function_registry import load_function
+from miles.utils.workers.serving.utils import split_worker_argv
+
+
+def main() -> None:
+    own_argv, worker_argv = split_worker_argv(sys.argv[1:])
+
+    parser = argparse.ArgumentParser(description="Compute worker env vars, then exec into the rpc server")
+    parser.add_argument("--env-var-fn", default=None, help="Env var computation function as 'package.module.callable'")
+    args, inner_own_argv = parser.parse_known_args(own_argv)
+    _log(f"start own_argv={own_argv} worker_argv={worker_argv}")
+
+    env = dict(os.environ)
+    if args.env_var_fn is not None:
+        computed_env_vars: dict[str, str] = load_function(args.env_var_fn)(worker_argv)
+        _log(f"env_var_fn={args.env_var_fn} computed={computed_env_vars}")
+        env.update(computed_env_vars)
+
+    inner_argv = [
+        sys.executable,
+        "-m",
+        "miles.utils.workers.serving.serve_inner",
+        *inner_own_argv,
+        "--",
+        *worker_argv,
+    ]
+    _log(f"exec {inner_argv}")
+    os.execve(sys.executable, inner_argv, env)
+
+
+def _log(message: str) -> None:
+    print(f"[serve] {message}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

tests/fast/utils/workers/e2e/env_var_hooks.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from tests.fast.utils.workers.import_probe import report_imported_top_level_modules
+
+IMPORTED_MODULES_ENV_VAR = "MILES_E2E_IMPORTED_MODULES"
+ENV_VAR_FN_FAILURE_MESSAGE = "env var hook refuses to run"
+
+
+def compute_env_vars(argv: list[str]) -> dict[str, str]:
+    return {IMPORTED_MODULES_ENV_VAR: report_imported_top_level_modules()}
+
+
+def raise_env_var_error(argv: list[str]) -> dict[str, str]:
+    raise RuntimeError(ENV_VAR_FN_FAILURE_MESSAGE)

tests/fast/utils/workers/e2e/harness.py

@@ -15,6 +15,7 @@
 
 REPO_ROOT = Path(__file__).resolve().parents[5]
 WORKER_PATH = "tests.fast.utils.workers.e2e.e2e_worker.make_worker"
+ENV_FN_PATH = "tests.fast.utils.workers.e2e.e2e_worker.compute_env_vars"
 
 READY_TIMEOUT_SECONDS = 60.0
 STOP_TIMEOUT_SECONDS = 15.0
@@ -75,6 +76,7 @@ def spawn_server(
     log_path: Path,
     port: int | None = None,
     worker_argv: list[str] | None = None,
+    env_var_fn: bool = True,
     extra_env: dict[str, str] | None = None,
     worker_path: str = WORKER_PATH,
 ) -> ServerProcess:
@@ -85,7 +87,9 @@ def spawn_server(
     env["PYTHONUNBUFFERED"] = "1"
     env.update(extra_env or {})
 
-    argv = [sys.executable, "-m", "miles.utils.workers.serving.serve_inner", "--worker", worker_path]
+    argv = [sys.executable, "-m", "miles.utils.workers.serving.serve", "--worker", worker_path]
+    if env_var_fn:
+        argv += ["--env-var-fn", ENV_FN_PATH]
     argv += ["--host", "127.0.0.1", "--port", str(port)]
     argv += ["--", "--state-dir", str(state_dir)]
     argv += worker_argv or []

tests/fast/utils/workers/e2e/test_serve_entrypoint.py

@@ -1,11 +1,188 @@
+import contextlib
+import os
 import socket
+import subprocess
+import sys
+from collections.abc import Callable, Iterator
+from pathlib import Path
 
+import pytest
 from tests.fast.utils.workers.e2e.e2e_worker import WORKER_FACTORY_ERROR
+from tests.fast.utils.workers.e2e.env_var_hooks import ENV_VAR_FN_FAILURE_MESSAGE, IMPORTED_MODULES_ENV_VAR
+from tests.fast.utils.workers.e2e.harness import (
+    READY_TIMEOUT_SECONDS,
+    REPO_ROOT,
+    WORKER_PATH,
+    ServerProcess,
+    port_is_refused,
+    reserve_port,
+    wait_until_serving,
+)
+from tests.fast.utils.workers.import_probe import unexpected_light_entrypoint_imports
 
+SMOKE_MODULE = "tests.fast.utils.workers.e2e.env_var_hooks"
+SMOKE_ENV_FN_PATH = f"{SMOKE_MODULE}.compute_env_vars"
+SMOKE_RAISING_ENV_FN_PATH = f"{SMOKE_MODULE}.raise_env_var_error"
 RAISING_WORKER_PATH = "tests.fast.utils.workers.e2e.e2e_worker.make_raising_worker"
 EXIT_TIMEOUT_SECONDS = 60.0
 
 
+@pytest.fixture
+def spawn_with_env_var_fn(state_dir: Path, tmp_path: Path) -> Iterator[Callable[..., ServerProcess]]:
+    started: list[ServerProcess] = []
+
+    def start(env_var_fn: str, *, worker_path: str = WORKER_PATH) -> ServerProcess:
+        port = reserve_port()
+        log_path = tmp_path / f"env-var-fn-server-{len(started)}.log"
+
+        env = dict(os.environ)
+        env["PYTHONPATH"] = f"{REPO_ROOT}{os.pathsep}{env.get('PYTHONPATH', '')}"
+        env["PYTHONUNBUFFERED"] = "1"
+
+        argv = [sys.executable, "-m", "miles.utils.workers.serving.serve", "--worker", worker_path]
+        argv += ["--env-var-fn", env_var_fn, "--host", "127.0.0.1", "--port", str(port)]
+        argv += ["--", "--state-dir", str(state_dir)]
+
+        with log_path.open("w") as log_file:
+            process = subprocess.Popen(
+                argv, cwd=REPO_ROOT, env=env, stdout=log_file, stderr=subprocess.STDOUT, start_new_session=True
+            )
+
+        server = ServerProcess(port=port, process=process, log_path=log_path)
+        started.append(server)
+        return server
+
+    yield start
+
+    for server in started:
+        server.stop()
+        server.kill()
+
+
+class TestExecChain:
+    async def test_the_served_process_is_the_spawned_one(self, handle, server):
+        """execve keeps the pid, so terminating the spawned process really stops the server."""
+        assert await handle.report_pid() == server.process.pid
+
+    async def test_worker_argv_reaches_the_factory(self, handle):
+        """Everything after -- is handed to the worker factory."""
+        argv = await handle.report_argv()
+        assert "--state-dir" in argv
+
+    async def test_worker_argv_keeps_its_own_separator(self, spawn, make_handle):
+        """Only the first -- splits, so worker argv may contain further separators."""
+        server = spawn(worker_argv=["--flag", "--", "--inner"])
+        handle = make_handle(server)
+        await handle.wait_ready(timeout=READY_TIMEOUT_SECONDS)
+
+        argv = await handle.report_argv()
+        assert argv[-3:] == ["--flag", "--", "--inner"]
+
+    async def test_env_var_hook_receives_worker_argv(self, handle):
+        """The env-var hook is called with the worker argv, not the entrypoint argv."""
+        recorded = await handle.report_env(name="MILES_E2E_ARGV")
+        assert "--state-dir" in recorded
+
+    async def test_only_allowlisted_modules_are_imported_before_the_hook(self, spawn_with_env_var_fn, make_handle):
+        """When the hook runs, the light entrypoint has imported no top-level module outside the allowlist."""
+        server = spawn_with_env_var_fn(SMOKE_ENV_FN_PATH)
+        wait_until_serving(server)
+        handle = make_handle(server)
+        await handle.wait_ready(timeout=READY_TIMEOUT_SECONDS)
+
+        reported = await handle.report_env(name=IMPORTED_MODULES_ENV_VAR)
+        assert unexpected_light_entrypoint_imports(reported) == []
+
+    async def test_parent_environment_is_inherited(self, spawn, make_handle):
+        """Environment from the launcher reaches the worker."""
+        server = spawn(extra_env={"MILES_E2E_MARKER": "inherited"})
+        handle = make_handle(server)
+        await handle.wait_ready(timeout=READY_TIMEOUT_SECONDS)
+
+        assert await handle.report_env(name="MILES_E2E_MARKER") == "inherited"
+
+    async def test_env_var_hook_overrides_an_inherited_value(self, spawn, make_handle):
+        """A computed variable replaces the same-named value the launcher exported."""
+        server = spawn(extra_env={"MILES_E2E_ARGV": "from-parent"})
+        handle = make_handle(server)
+        await handle.wait_ready(timeout=READY_TIMEOUT_SECONDS)
+
+        recorded = await handle.report_env(name="MILES_E2E_ARGV")
+        assert recorded != "from-parent"
+        assert "--state-dir" in recorded
+
+    async def test_env_var_hook_is_optional(self, spawn, make_handle):
+        """Serving without the hook still works."""
+        server = spawn(env_var_fn=False)
+        handle = make_handle(server)
+        await handle.wait_ready(timeout=READY_TIMEOUT_SECONDS)
+
+        assert await handle.demo_sync(a=1, b=1) == 2
+        assert await handle.report_env(name="MILES_E2E_ARGV") is None
+
+
+class TestStartupFailures:
+    async def test_unknown_worker_path_fails_fast(self, spawn):
+        """A worker path that cannot be imported exits instead of serving."""
+        server = spawn(worker_path="no.such.module.make_worker", wait=False)
+        assert server.wait(timeout=30.0) not in (None, 0)
+        assert port_is_refused(server.port)
+
+    async def test_missing_worker_argument_is_a_usage_error(self, spawn):
+        """argparse rejects a missing --worker with its usage exit code."""
+        env = dict(os.environ)
+        env["PYTHONPATH"] = f"{REPO_ROOT}{os.pathsep}{env.get('PYTHONPATH', '')}"
+        result = subprocess.run(
+            [sys.executable, "-m", "miles.utils.workers.serving.serve", "--host", "127.0.0.1"],
+            cwd=REPO_ROOT,
+            env=env,
+            capture_output=True,
+            timeout=60,
+        )
+
+        assert result.returncode == 2
+        assert b"usage" in result.stderr.lower()
+
+    async def test_port_conflict_fails_fast(self, spawn, server):
+        """A second server on a taken port exits without disturbing the first."""
+        conflicting = spawn(port=server.port, wait=False)
+        assert conflicting.wait(timeout=30.0) not in (None, 0)
+        assert server.is_running()
+
+    @pytest.mark.parametrize("bad_path", ["no_colon_module", "miles.utils.workers.serving.serve.no_such_attr"])
+    async def test_bad_factory_paths_fail_fast(self, spawn, bad_path):
+        """Malformed or missing factory paths exit rather than serving a broken worker."""
+        server = spawn(worker_path=bad_path, wait=False)
+        assert server.wait(timeout=30.0) not in (None, 0)
+
+    async def test_unknown_env_var_fn_module_fails_fast(self, spawn_with_env_var_fn):
+        """An env-var hook whose module cannot be imported exits instead of serving."""
+        server = spawn_with_env_var_fn("no.such.module.compute_env_vars")
+        assert server.wait(timeout=30.0) not in (None, 0)
+        assert port_is_refused(server.port)
+        assert "ModuleNotFoundError" in server.logs()
+
+    async def test_missing_env_var_fn_attribute_fails_fast(self, spawn_with_env_var_fn):
+        """An env-var hook naming an attribute the module lacks exits instead of serving."""
+        server = spawn_with_env_var_fn(f"{SMOKE_MODULE}.no_such_attr")
+        assert server.wait(timeout=30.0) not in (None, 0)
+        assert port_is_refused(server.port)
+
+        logs = server.logs()
+        assert "AttributeError" in logs
+        assert "no_such_attr" in logs
+
+    async def test_raising_env_var_fn_fails_fast(self, spawn_with_env_var_fn):
+        """An env-var hook that raises when called exits instead of serving, and reports its own error."""
+        server = spawn_with_env_var_fn(SMOKE_RAISING_ENV_FN_PATH)
+        assert server.wait(timeout=30.0) not in (None, 0)
+        assert port_is_refused(server.port)
+
+        logs = server.logs()
+        assert "RuntimeError" in logs
+        assert ENV_VAR_FN_FAILURE_MESSAGE in logs
+
+
 class TestWorkerFactoryFailure:
     def test_raising_worker_factory_fails_before_binding_the_port(self, spawn) -> None:
         """A worker factory that raises fails startup before the port is bound and reports its own error."""
@@ -18,3 +195,50 @@ def test_raising_worker_factory_fails_before_binding_the_port(self, spawn) -> No
         ), f"startup failed before reaching the worker factory:\n{server.logs()}"
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
             sock.bind(("127.0.0.1", server.port))
+
+
+class TestPortBinding:
+    async def test_server_is_unreachable_on_non_loopback_addresses(self, server):
+        """Binding 127.0.0.1 keeps the port reachable on loopback and refused on the machine's other address."""
+        address = _non_loopback_ipv4_address()
+        if address is None:
+            pytest.skip("no non-loopback ipv4 address on this machine")
+
+        assert not port_is_refused(server.port)
+        assert _connection_is_refused(address, server.port)
+
+
+def _non_loopback_ipv4_address() -> str | None:
+    candidates: list[str] = []
+
+    with contextlib.suppress(OSError):
+        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe:
+            probe.connect(("8.8.8.8", 53))
+            candidates.append(probe.getsockname()[0])
+
+    with contextlib.suppress(OSError):
+        for info in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET, socket.SOCK_STREAM):
+            candidates.append(info[4][0])
+
+    for address in candidates:
+        if not address.startswith("127.") and address != "0.0.0.0" and _is_local_address(address):
+            return address
+    return None
+
+
+def _is_local_address(address: str) -> bool:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        try:
+            sock.bind((address, 0))
+        except OSError:
+            return False
+    return True
+
+
+def _connection_is_refused(address: str, port: int) -> bool:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.settimeout(2.0)
+        try:
+            return sock.connect_ex((address, port)) != 0
+        except OSError:
+            return True

tests/fast/utils/workers/import_probe.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import sys
+
+IMPORTED_MODULES_SEPARATOR = ","
+
+# Everything here lands in sys.modules through site initialization, before the
+# entrypoint runs a line of its own, so none of it says anything about what the
+# entrypoint imports. nvidia_cutlass_dsl arrives via a .pth in dist-packages.
+ALLOWED_LIGHT_ENTRYPOINT_IMPORTS = frozenset(
+    {
+        "__main__",
+        "miles",
+        "tests",
+        "sitecustomize",
+        "usercustomize",
+        "_distutils_hack",
+        "_virtualenv",
+        "nvidia_cutlass_dsl",
+    }
+)
+
+
+_INSTALLER_SHIM_PREFIX = "__editable__"
+
+
+def imported_top_level_modules() -> set[str]:
+    top_level_names = {name.partition(".")[0] for name in sys.modules}
+    return {
+        name
+        for name in top_level_names
+        if name not in sys.stdlib_module_names and not name.startswith(_INSTALLER_SHIM_PREFIX)
+    }
+
+
+def report_imported_top_level_modules() -> str:
+    return IMPORTED_MODULES_SEPARATOR.join(sorted(imported_top_level_modules()))
+
+
+def unexpected_light_entrypoint_imports(reported: str) -> list[str]:
+    imported = {name for name in reported.split(IMPORTED_MODULES_SEPARATOR) if name}
+    return sorted(imported - ALLOWED_LIGHT_ENTRYPOINT_IMPORTS)

tests/fast/utils/workers/serving/serve_smoke_worker.py

@@ -0,0 +1,30 @@
+import os
+
+from tests.fast.utils.workers.import_probe import report_imported_top_level_modules
+
+IMPORTED_MODULES_ENV_VAR = "MILES_SERVE_SMOKE_IMPORTED_MODULES"
+
+
+class SmokeWorker:
+    def __init__(self, argv: list[str]):
+        self._argv = argv
+
+    def demo_sync(self, a: int, b: int) -> int:
+        return a + b
+
+    def report_argv(self) -> list[str]:
+        return self._argv
+
+    def report_env(self, name: str) -> str | None:
+        return os.environ.get(name)
+
+
+def make_worker(argv: list[str]) -> SmokeWorker:
+    return SmokeWorker(argv)
+
+
+def compute_env_vars(argv: list[str]) -> dict[str, str]:
+    return {
+        "MILES_SERVE_SMOKE_ENV": ",".join(argv),
+        IMPORTED_MODULES_ENV_VAR: report_imported_top_level_modules(),
+    }