diff --git a/.claude/skills/benchmarking/SKILL.md b/.claude/skills/benchmarking/SKILL.md index 99365390c9..a07f05e1f3 100644 --- a/.claude/skills/benchmarking/SKILL.md +++ b/.claude/skills/benchmarking/SKILL.md @@ -69,7 +69,7 @@ pgrep -a -u "$USER" -f 'multiprocessing' pgrep -a -u "$USER" -f 'torch/_inductor/compile_worker' ``` -Then confirm with `nvidia-smi` that the memory actually came back. Tens of GB in use with no server of yours running means orphans. Leftover IPC socket files under `/tmp/mstar_$USER/` (or `--socket-path-prefix`) are harmless and persist between runs — what breaks a new server is a *live* process still holding one. +Then confirm with `nvidia-smi` that the memory actually came back. Tens of GB in use with no server of yours running means orphans. Leftover IPC socket files under `/tmp/mstar_${USER}_/` (or `--socket-path-prefix`) are harmless and persist between runs — what breaks a new server is a *live* process still holding one. **`benchmark.worker_phases.server` starts the server in its own process group**, separate from the wrapper's. Signalling the wrapper's group leaves the server running; the next arm's wrapper then refuses to start, and a client that doesn't check keeps benchmarking the previous arm's server. In a scripted A/B, stop both groups, check the port is free before each arm, and log which checkout each arm imported (`python -c 'import mstar; print(mstar.__file__)'` with the arm's environment) so you can verify every number came from the build you think it did. @@ -203,5 +203,5 @@ How many repeats, and how large a delta counts as real, depend on the model and - **Don't pipe a multi-cell sweep through `tail`.** A twelve-cell run piped through `tail -100` lost five of them; they were recoverable only because each cell also tee'd to its own `--output-dir`. Analyze the per-cell files, not the console scrollback. - **Size a readiness wait for the slowest model's cold compile, not the one you developed against.** A 15-minute `/health` poll was ample for orpheus and expired midway through qwen3-omni's first Inductor compile, which then left the server running and holding its GPUs after the harness had given up on it. - **`--ignore-eos` isn't safe for every model.** It's the right default for plain text output, but Qwen3-Omni's talker has undefined behavior past EOS and crashes on stale SHM reads, which looks like a regression. -- **Two servers on one machine need distinct `--socket-path-prefix`.** `--port` alone doesn't isolate the ZMQ IPC sockets, and one server's handshake gets taken by the other. +- **Two servers on one machine need distinct socket prefixes.** The default prefix is keyed on the port (`/tmp/mstar__/`), so distinct ports are enough when `--socket-path-prefix` is left unset. Two servers given the same explicit prefix share the ZMQ IPC sockets, and one server's handshake gets taken by the other. - **A stale Rust extension** starts up healthy and fails every request with unexpected-keyword `TypeError`s. Rebuild after merging or rebasing onto a `main` that touched `rust/`. diff --git a/docs/serving.rst b/docs/serving.rst index 52b93a6097..622fad8df2 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -48,8 +48,9 @@ mstar serve - ``SHM`` - Tensor transport: ``SHM`` (safe single-node default), ``TCP``, or ``RDMA``. * - ``--socket-path-prefix`` - - ``/tmp/mstar_/`` - - ZMQ IPC socket prefix. + - ``/tmp/mstar__/`` + - ZMQ IPC socket prefix; one per server, so servers on other ports do + not share sockets. * - ``--upload-dir`` - ``/tmp/mstar_uploads_/`` - Temp directory for uploaded media. @@ -123,8 +124,10 @@ mstar-serve - HF default - HuggingFace weight cache directory. * - ``--socket-path-prefix`` - - ``/tmp/mstar`` - - ZMQ IPC socket prefix (shared with conductor/workers). + - ``/tmp/mstar__/`` + - ZMQ IPC socket prefix (shared with conductor/workers). The Dynamo + worker (``python -m mstar.integrations.dynamo``) has no HTTP port, so + it defaults to a new directory per launch and removes it at exit. * - ``--upload-dir`` - ``/tmp/mstar_uploads`` - Temp directory for uploaded media. diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 8ab6f3c866..0bb4caf684 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -37,6 +37,7 @@ from mstar.utils.gc_freeze import freeze_after_setup from mstar.utils.logging_config import quiet_noisy_loggers from mstar.utils.orphan import watch_parent +from mstar.utils.socket_prefix import per_port_socket_prefix logger = logging.getLogger(__name__) @@ -1591,8 +1592,9 @@ def main(argv: list[str] | None = None): parser.add_argument("--port", type=int, default=8000) parser.add_argument("--mooncake-port", type=int, default=8080) parser.add_argument( - "--socket-path-prefix", type=str, default="/tmp/mstar", - help="ZMQ IPC socket path prefix (shared with conductor/workers)", + "--socket-path-prefix", type=str, default=None, + help="ZMQ IPC socket path prefix (shared with conductor/workers); " + "default /tmp/mstar__/", ) parser.add_argument( "--upload-dir", type=str, default="/tmp/mstar_uploads", @@ -1645,6 +1647,8 @@ def main(argv: list[str] | None = None): "$PATH, then rust/server/target/release)", ) args = parser.parse_args(argv) + if args.socket_path_prefix is None: + args.socket_path_prefix = per_port_socket_prefix(args.port) logging.basicConfig( level=getattr(logging, args.log_level), diff --git a/mstar/cli/main.py b/mstar/cli/main.py index 4236d23856..d0d8a0d6a0 100644 --- a/mstar/cli/main.py +++ b/mstar/cli/main.py @@ -20,6 +20,7 @@ from pathlib import Path import mstar +from mstar.utils.socket_prefix import per_port_socket_prefix # Model name -> default config (relative to the repo's configs/). DEFAULT_CONFIGS: dict[str, str] = { @@ -207,7 +208,7 @@ def _serve(args: argparse.Namespace) -> None: os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus who = getpass.getuser() if hasattr(getpass, "getuser") else "mstar" - socket_prefix = args.socket_path_prefix or f"/tmp/mstar_{who}/" + socket_prefix = args.socket_path_prefix or per_port_socket_prefix(args.port) upload_dir = args.upload_dir or f"/tmp/mstar_uploads_{who}/" argv = [ diff --git a/mstar/integrations/dynamo/worker.py b/mstar/integrations/dynamo/worker.py index 85176329aa..856a66c926 100644 --- a/mstar/integrations/dynamo/worker.py +++ b/mstar/integrations/dynamo/worker.py @@ -20,6 +20,7 @@ import contextlib import logging import multiprocessing as mp +import shutil import signal from typing import TYPE_CHECKING, Any, Callable, NamedTuple @@ -32,6 +33,7 @@ TensorBridge, get_tensor_spec, ) +from mstar.utils.socket_prefix import per_launch_socket_prefix if TYPE_CHECKING: from mstar.api_server.entrypoint import APIServer @@ -206,7 +208,9 @@ def main(argv: list[str] | None = None) -> None: help="HF snapshot dir backing the frontend-side model card") parser.add_argument("--served-model-name", default=None) parser.add_argument("--cache-dir", default=None) - parser.add_argument("--socket-path-prefix", default="/tmp/mstar") + parser.add_argument("--socket-path-prefix", default=None, + help="ZMQ IPC socket prefix; default a new /tmp/mstar__* " + "directory per launch, removed at exit") parser.add_argument("--upload-dir", default="/tmp/mstar_uploads") parser.add_argument("--tensor-comm-protocol", default="RDMA", help="Tensor transfer protocol: RDMA, TCP, or SHM (shared memory)") @@ -220,10 +224,6 @@ def main(argv: list[str] | None = None) -> None: ) args = parser.parse_args(argv) - from mstar.api_server.entrypoint import ( - DeadConductorError, - _shutdown_conductor_process, - ) from mstar.utils.logging_config import quiet_noisy_loggers logging.basicConfig( @@ -232,6 +232,25 @@ def main(argv: list[str] | None = None) -> None: ) quiet_noisy_loggers() + # No HTTP port to key the prefix on, and replicas of one endpoint + # (same namespace, component and endpoint) commonly share a host. + owned_prefix = None + if args.socket_path_prefix is None: + owned_prefix = args.socket_path_prefix = per_launch_socket_prefix() + logger.info("IPC socket prefix: %s", owned_prefix) + try: + _serve_until_shutdown(args) + finally: + if owned_prefix is not None: + shutil.rmtree(owned_prefix, ignore_errors=True) + + +def _serve_until_shutdown(args) -> None: + from mstar.api_server.entrypoint import ( + DeadConductorError, + _shutdown_conductor_process, + ) + server, conductor_proc, model_name = build_server(args) exit_code = 0 try: diff --git a/mstar/utils/socket_prefix.py b/mstar/utils/socket_prefix.py new file mode 100644 index 0000000000..84df11baf6 --- /dev/null +++ b/mstar/utils/socket_prefix.py @@ -0,0 +1,37 @@ +"""Default IPC socket prefixes for the server entry points. + +Every process of a deployment binds fixed names under its prefix +(``conductor.ipc``, ``api_server.ipc``, ``worker_0.ipc``, ...), and libzmq +replaces an existing socket file on ``bind``. Two servers on one prefix +therefore take over each other's inboxes, so each default here is unique to +one server on the host. +""" + +import getpass +import os +import re +import tempfile + + +def _user() -> str: + """The user name, or the uid when there is none or it is not a short + ASCII name: a socket path holds at most 107 bytes, and the longest + entity name under the prefix needs about 50 of them.""" + try: + name = getpass.getuser() + except (KeyError, OSError): + # No passwd entry for this uid, as in containers run with an arbitrary uid. + name = "" + return name if re.fullmatch(r"[A-Za-z0-9._-]{1,32}", name) else str(os.getuid()) + + +def per_port_socket_prefix(port: int) -> str: + """``/tmp/mstar__/`` for an HTTP server: the port is already + unique to one server on the host.""" + return f"/tmp/mstar_{_user()}_{port}/" + + +def per_launch_socket_prefix() -> str: + """A new ``/tmp/mstar__`` directory (mode 0700) for a server + with no port of its own; the caller removes it at exit.""" + return tempfile.mkdtemp(prefix=f"mstar_{_user()}_", dir="/tmp") diff --git a/test/cpu-core.txt b/test/cpu-core.txt index 127bfb8ec8..15fe022c09 100644 --- a/test/cpu-core.txt +++ b/test/cpu-core.txt @@ -13,6 +13,7 @@ test/modular/test_rid_handles.py test/modular/test_tensor_uuid.py test/modular/test_shm_tensor_comm.py test/modular/test_deployment_lock.py +test/modular/test_default_socket_prefix.py test/modular/test_waypoint_components.py test/modular/test_waypoint_dit.py test/modular/test_waypoint_weight_loader.py diff --git a/test/modular/test_default_socket_prefix.py b/test/modular/test_default_socket_prefix.py new file mode 100644 index 0000000000..21f6e4db0a --- /dev/null +++ b/test/modular/test_default_socket_prefix.py @@ -0,0 +1,148 @@ +"""Each server entry point defaults to an IPC socket prefix of its own (#490). + +Every process of a deployment binds fixed names under its prefix, and libzmq +replaces an existing socket file on bind. With one shared default, a second +server's API server binds over the first one's ``api_server.ipc`` before the +deployment lock refuses its conductor, and the first server's replies land in +the second one's inbox. +""" + +import os +import shutil +import stat +import sys +import time +import types + +import pytest + +from mstar.communication.communicator import ZMQCommunicator +from mstar.utils import socket_prefix + + +class _StopError(Exception): + pass + + +def _entrypoint_prefix(monkeypatch, tmp_path, *argv) -> str: + """The prefix ``mstar-serve`` hands the API server, without starting one.""" + from mstar.api_server import entrypoint + from mstar.model import registry + + seen = {} + + def fake_api_server(**kwargs): + seen.update(kwargs) + raise _StopError + + class _Model: + PREPROCESS_TORCH_THREADS = None + + def __init__(self, **kwargs): + pass + + cfg = tmp_path / "cfg.yaml" + cfg.write_text("model: fake\n") + monkeypatch.setattr(registry, "get_model_class", lambda name: _Model) + monkeypatch.setattr(entrypoint, "model_init_kwargs", lambda name: {}) + monkeypatch.setattr(entrypoint, "APIServer", fake_api_server) + with pytest.raises(_StopError): + entrypoint.main(["--config", str(cfg), *argv]) + return seen["socket_path_prefix"] + + +def test_mstar_serve_entrypoint_keys_the_default_prefix_on_the_port(monkeypatch, tmp_path): + a = _entrypoint_prefix(monkeypatch, tmp_path, "--port", "8000") + b = _entrypoint_prefix(monkeypatch, tmp_path, "--port", "8001") + assert a != b + assert a == socket_prefix.per_port_socket_prefix(8000) + + +def test_mstar_serve_entrypoint_keeps_an_explicit_prefix(monkeypatch, tmp_path): + got = _entrypoint_prefix(monkeypatch, tmp_path, "--socket-path-prefix", "/x/y/") + assert got == "/x/y/" + + +def test_mstar_serve_cli_uses_the_same_rule(monkeypatch): + from mstar.cli import main as cli + + seen = [] + monkeypatch.setitem( + sys.modules, "mstar.api_server.entrypoint", + types.SimpleNamespace(main=seen.append), + ) + cli.main(["serve", "kokoro", "--port", "9123"]) + argv = seen[0] + assert argv[argv.index("--socket-path-prefix") + 1] == socket_prefix.per_port_socket_prefix(9123) + + +def _dynamo_prefix(monkeypatch, *argv) -> dict: + """The prefix the Dynamo worker builds its server on, and whether that + directory still exists after ``main`` returns.""" + from mstar.integrations.dynamo import worker + + seen = {} + + def fake_build_server(args): + seen["prefix"] = args.socket_path_prefix + seen["mode"] = stat.S_IMODE(os.stat(args.socket_path_prefix).st_mode) + raise _StopError + + monkeypatch.setattr(worker, "build_server", fake_build_server) + with pytest.raises(_StopError): + worker.main(["--config", "c.yaml", "--model-path", "/m", *argv]) + return seen + + +def test_dynamo_workers_get_a_private_prefix_per_launch_removed_at_exit(monkeypatch): + a = _dynamo_prefix(monkeypatch) + b = _dynamo_prefix(monkeypatch) + assert a["prefix"] != b["prefix"] + assert a["mode"] == 0o700 + assert not os.path.exists(a["prefix"]) and not os.path.exists(b["prefix"]) + + +def test_dynamo_worker_keeps_and_does_not_remove_an_explicit_prefix(monkeypatch, tmp_path): + got = _dynamo_prefix(monkeypatch, "--socket-path-prefix", str(tmp_path)) + assert got["prefix"] == str(tmp_path) + assert tmp_path.is_dir() + + +def test_two_servers_on_their_default_prefixes_keep_their_own_inboxes(): + """The #490 repro: on a shared prefix, the second server's api_server binds + over the first one's, so the first conductor's reply reaches the second.""" + a = socket_prefix.per_launch_socket_prefix() + b = socket_prefix.per_launch_socket_prefix() + comms = [] + try: + for prefix in (a, b): + comms.append(ZMQCommunicator( + my_id="api_server", push_ids=["conductor"], ipc_socket_path_prefix=prefix)) + comms.append(ZMQCommunicator( + my_id="conductor", push_ids=["api_server"], ipc_socket_path_prefix=prefix)) + a_api, a_cond, b_api, _ = comms + a_cond.send("api_server", {"reply": "a"}) + got = [] + deadline = time.monotonic() + 5 + while not got and time.monotonic() < deadline: + got = a_api.get_all_new_messages() + time.sleep(0.01) + assert got == [{"reply": "a"}] + assert b_api.get_all_new_messages() == [] + finally: + for c in comms: + c.pull_socket.close() + for s in c.push_sockets.values(): + s.close() + for prefix in (a, b): + shutil.rmtree(prefix, ignore_errors=True) + + +@pytest.mark.parametrize("name, kept", [("alice", True), ("a" * 32, True), ("a" * 33, False), + ("\u00e9" * 30, False), ("a b", False), ("", False)]) +def test_the_prefix_holds_the_user_name_only_when_it_is_short_and_ascii(monkeypatch, name, kept): + monkeypatch.setattr(socket_prefix.getpass, "getuser", lambda: name) + prefix = socket_prefix.per_port_socket_prefix(65535) + assert prefix == f"/tmp/mstar_{name if kept else os.getuid()}_65535/" + # a 45-byte socket name, longer than any a deployment binds, still fits a socket path + assert len((prefix + "x" * 41 + ".ipc").encode()) <= 107 diff --git a/test/modular/test_launch_script_socket_prefix.py b/test/modular/test_launch_script_socket_prefix.py index f5e4504d66..b9b6c0b9ea 100644 --- a/test/modular/test_launch_script_socket_prefix.py +++ b/test/modular/test_launch_script_socket_prefix.py @@ -1,8 +1,9 @@ """Every server binds the same IPC names under its ``--socket-path-prefix``, so two servers on one host need two prefixes. A launch script whose prefix is -keyed only on the user, or that leaves the entrypoint's ``/tmp/mstar`` default, -gives every server it starts the same sockets: the second one takes over the -first one's API inboxes before the deployment lock refuses it.""" +keyed only on the user, or that relied on the entrypoint's old shared +``/tmp/mstar`` default, gives every server it starts the same sockets: the +second one takes over the first one's API inboxes before the deployment lock +refuses it.""" import re from pathlib import Path @@ -54,7 +55,7 @@ def test_launch_script_passes_a_socket_prefix(script): for command in _code(script).splitlines(): if ENTRYPOINT.search(command): assert "--socket-path-prefix" in command, ( - f"{script.relative_to(REPO)} starts the entrypoint on its shared /tmp/mstar default: {command.strip()}" + f"{script.relative_to(REPO)} starts the entrypoint without --socket-path-prefix: {command.strip()}" )