diff --git a/miles/utils/workers/serving/serve.py b/miles/utils/workers/serving/serve.py new file mode 100644 index 00000000000..21441e2ecdb --- /dev/null +++ b/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() diff --git a/tests/fast/utils/workers/e2e/env_var_hooks.py b/tests/fast/utils/workers/e2e/env_var_hooks.py new file mode 100644 index 00000000000..cdd4b78020c --- /dev/null +++ b/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) diff --git a/tests/fast/utils/workers/e2e/harness.py b/tests/fast/utils/workers/e2e/harness.py index d796f4cc08a..e154c44bbf4 100644 --- a/tests/fast/utils/workers/e2e/harness.py +++ b/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 [] diff --git a/tests/fast/utils/workers/e2e/test_serve_entrypoint.py b/tests/fast/utils/workers/e2e/test_serve_entrypoint.py index 4f7a7522479..d43803d64e7 100644 --- a/tests/fast/utils/workers/e2e/test_serve_entrypoint.py +++ b/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 diff --git a/tests/fast/utils/workers/import_probe.py b/tests/fast/utils/workers/import_probe.py new file mode 100644 index 00000000000..52385cbd443 --- /dev/null +++ b/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) diff --git a/tests/fast/utils/workers/serving/serve_smoke_worker.py b/tests/fast/utils/workers/serving/serve_smoke_worker.py new file mode 100644 index 00000000000..d47b44332b2 --- /dev/null +++ b/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(), + } diff --git a/tests/fast/utils/workers/serving/test_serve.py b/tests/fast/utils/workers/serving/test_serve.py new file mode 100644 index 00000000000..190958321f5 --- /dev/null +++ b/tests/fast/utils/workers/serving/test_serve.py @@ -0,0 +1,235 @@ +import contextlib +import os +import subprocess +import sys +import time +from pathlib import Path + +import httpx +import pytest +from tests.fast.utils.workers.import_probe import unexpected_light_entrypoint_imports +from tests.fast.utils.workers.serving.serve_smoke_worker import IMPORTED_MODULES_ENV_VAR, SmokeWorker + +from miles.utils.http_utils import find_available_port +from miles.utils.workers.rpc.client.handle import RpcWorkerHandle +from miles.utils.workers.rpc.client.misc import ServerRestartedError +from miles.utils.workers.serving import serve as serve_module +from miles.utils.workers.serving.utils import split_worker_argv +from miles.utils.workers.worker_handle import WorkerUnreachableError + +_REPO_ROOT = Path(__file__).resolve().parents[5] +_WORKER_PATH = "tests.fast.utils.workers.serving.serve_smoke_worker.make_worker" +_ENV_FN_PATH = "tests.fast.utils.workers.serving.serve_smoke_worker.compute_env_vars" + + +class TestSplitWorkerArgv: + def test_splits_on_the_separator(self): + """Everything after -- is worker argv, everything before belongs to the entrypoint.""" + assert split_worker_argv(["--worker", "m:f", "--", "--greeting", "hi"]) == ( + ["--worker", "m:f"], + ["--greeting", "hi"], + ) + + def test_no_separator_means_no_worker_argv(self): + """Without a separator the whole argv belongs to the entrypoint.""" + assert split_worker_argv(["--worker", "m:f"]) == (["--worker", "m:f"], []) + + def test_later_separators_stay_with_the_worker(self): + """Only the first separator splits, so worker argv may contain its own --.""" + assert split_worker_argv(["--", "a", "--", "b"]) == ([], ["a", "--", "b"]) + + +class TestOuterServeForwarding: + def test_only_env_hook_is_consumed_before_exec( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """The outer serve layer forwards all non-hook own arguments verbatim.""" + captured: dict[str, object] = {} + + def fake_execve(path: str, argv: list[str], env: dict[str, str]) -> None: + captured.update(path=path, argv=argv, env=env) + + monkeypatch.setattr(serve_module.os, "execve", fake_execve) + monkeypatch.setattr( + serve_module, + "load_function", + lambda path: lambda worker_argv: {"MILES_TEST_ENV": ",".join(worker_argv)}, + ) + + monkeypatch.setattr( + sys, + "argv", + [ + "serve.py", + "--worker", + "package.worker", + "--host", + "127.0.0.1", + "--env-var-fn", + "package.env", + "--port", + "9000", + "--", + "--flag", + "value", + ], + ) + + serve_module.main() + + assert captured["argv"] == [ + sys.executable, + "-m", + "miles.utils.workers.serving.serve_inner", + "--worker", + "package.worker", + "--host", + "127.0.0.1", + "--port", + "9000", + "--", + "--flag", + "value", + ] + assert captured["env"]["MILES_TEST_ENV"] == "--flag,value" + + def test_env_var_hook_overrides_same_named_parent_variable( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A computed variable wins over a same-named variable inherited from the parent environment.""" + captured: dict[str, object] = {} + + def fake_execve(path: str, argv: list[str], env: dict[str, str]) -> None: + captured.update(path=path, argv=argv, env=env) + + monkeypatch.setenv("MILES_TEST_ENV", "from-parent") + monkeypatch.setenv("MILES_TEST_UNTOUCHED", "kept") + monkeypatch.setattr(serve_module.os, "execve", fake_execve) + monkeypatch.setattr( + serve_module, + "load_function", + lambda path: lambda worker_argv: {"MILES_TEST_ENV": ",".join(worker_argv)}, + ) + + monkeypatch.setattr( + sys, + "argv", + [ + "serve.py", + "--worker", + "package.worker", + "--env-var-fn", + "package.env", + "--", + "--flag", + "value", + ], + ) + + serve_module.main() + + assert captured["env"]["MILES_TEST_ENV"] == "--flag,value" + assert captured["env"]["MILES_TEST_UNTOUCHED"] == "kept" + + +def _spawn_serve(port: int) -> subprocess.Popen: + env = dict(os.environ) + env["PYTHONPATH"] = f"{_REPO_ROOT}{os.pathsep}{env.get('PYTHONPATH', '')}" + + return subprocess.Popen( + [ + sys.executable, + "-m", + "miles.utils.workers.serving.serve", + "--worker", + _WORKER_PATH, + "--env-var-fn", + _ENV_FN_PATH, + "--host", + "127.0.0.1", + "--port", + str(port), + "--", + "--greeting", + "hello", + ], + cwd=_REPO_ROOT, + env=env, + ) + + +async def _wait_ready_or_die(handle, process: subprocess.Popen, timeout: float = 60.0) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + assert process.poll() is None, f"serve subprocess exited early with code {process.returncode}" + with contextlib.suppress(WorkerUnreachableError): + await handle.wait_ready(timeout=2.0) + return + raise AssertionError("serve subprocess never became ready") + + +def _stop(process: subprocess.Popen) -> None: + process.terminate() + try: + process.wait(timeout=10.0) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10.0) + + +class TestServeEndToEnd: + async def test_serve_call_argv_and_env(self): + """serve.py boots a worker subprocess that answers typed calls with argv and env applied.""" + port = find_available_port(20000 + os.getpid() % 10000) + process = _spawn_serve(port) + + async with httpx.AsyncClient(trust_env=False) as client: + handle = RpcWorkerHandle( + SmokeWorker, + server_url=f"http://127.0.0.1:{port}", + http_client=client, + ) + try: + await _wait_ready_or_die(handle, process) + assert await handle.demo_sync(a=3, b=4) == 7 + assert await handle.report_argv() == ["--greeting", "hello"] + assert await handle.report_env(name="MILES_SERVE_SMOKE_ENV") == "--greeting,hello" + reported = await handle.report_env(name=IMPORTED_MODULES_ENV_VAR) + assert unexpected_light_entrypoint_imports(reported) == [] + finally: + _stop(process) + + async def test_restart_detected_by_stable_boot_uuid_client(self): + """A stable-boot-uuid client notices when the serve subprocess is restarted.""" + port = find_available_port(21000 + os.getpid() % 10000) + first_process = _spawn_serve(port) + second_process: subprocess.Popen | None = None + + async with httpx.AsyncClient(trust_env=False) as client: + handle = RpcWorkerHandle( + SmokeWorker, + server_url=f"http://127.0.0.1:{port}", + require_stable_boot_uuid=True, + http_client=client, + ) + try: + await _wait_ready_or_die(handle, first_process) + assert await handle.demo_sync(a=1, b=1) == 2 + + _stop(first_process) + second_process = _spawn_serve(port) + fresh_handle = RpcWorkerHandle( + SmokeWorker, + server_url=f"http://127.0.0.1:{port}", + http_client=client, + ) + await _wait_ready_or_die(fresh_handle, second_process) + + with pytest.raises(ServerRestartedError): + await handle.demo_sync(a=1, b=1) + finally: + if second_process is not None: + _stop(second_process) + _stop(first_process)