diff --git a/DEPENDENCIES.md b/DEPENDENCIES.md index d7adc9203..8f5410362 100644 --- a/DEPENDENCIES.md +++ b/DEPENDENCIES.md @@ -169,7 +169,9 @@ xr-ai-vllm (utils/xr-ai-vllm/) `--stop` flow. Besides `serve` / `stop_persistent_servers`, exposes the shared wrapper helpers `resolve_model_cache`, `load_config`, `setup_hf_env`, and `gpu_compute_major` (all stdlib-only; pyyaml is imported function-locally - inside `load_config` so the `--stop` path stays dependency-free). + inside `load_config` so the `--stop` path stays dependency-free). Docker + containers carry a deterministic launch fingerprint; containers created + with stale model, image, GPU, or vLLM arguments are replaced before reuse. xr-ai-vad (utils/xr-ai-vad/) └── numpy >=1.24 diff --git a/docs/source/components/ai-services.md b/docs/source/components/ai-services.md index fcaac340c..87c88d466 100644 --- a/docs/source/components/ai-services.md +++ b/docs/source/components/ai-services.md @@ -258,15 +258,31 @@ The persistent vLLM-backed servers (`vlm_server`, `llama_nemotron_llm_server`, **survive stack restarts by design**. Each persistent wrapper script checks its health endpoint before spawning vLLM: -- **Already running** → touch the ready file immediately, then idle. Stack is - ready in seconds; no model reload. -- **Not running** → spawn vLLM normally, wait for `/health`, touch ready file. +- **Already running with a matching launch fingerprint** → touch the ready + file immediately, then idle. Stack is ready in seconds; no model reload. +- **Matching container still starting** → attach to its lifecycle and keep + waiting for `/health` instead of issuing a conflicting second `docker run`. +- **Already running with changed or legacy configuration** → stop, remove, and + recreate the repository-owned container from the current YAML. +- **Healthy endpoint without the expected running container** → fail without + stopping the unowned listener. +- **Stopped Docker container with matching launch fingerprint** → restart it, + wait for `/health`, then touch the ready file. +- **Stopped Docker container with changed or legacy configuration** → remove + and recreate it from the current YAML before waiting for `/health`. +- **Not running** → spawn vLLM normally, wait for `/health`, then touch the + ready file. In pip mode, vLLM is spawned with `start_new_session=True` so the launcher's -`killpg()` does not reach it on shutdown. In docker mode, the container is -launched detached (`docker run -d --name xr-ai-vllm-`) so it -similarly outlives the wrapper. Either way the wrapper exits cleanly and -vLLM keeps running. +`killpg()` does not reach it on shutdown. In docker mode, Docker owns the +container while the foreground `docker run` client uses its own session. +Either way vLLM keeps running after the orchestrator exits. + +Docker containers carry a fingerprint of their image, GPU assignment, model +cache, environment, bootstrap packages, complete vLLM command, and a versioned +launcher-controlled Docker contract. This prevents a failed container created +by one sample profile—or by older launcher behavior—from being restarted later +with stale memory limits, entrypoint, setup commands, or model arguments. **Stopping the persisted servers** — run from the sample directory: diff --git a/tests/test_vllm_docker.py b/tests/test_vllm_docker.py index 9d85d2b0e..ee6695cfd 100644 --- a/tests/test_vllm_docker.py +++ b/tests/test_vllm_docker.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Unit tests for xr_ai_vllm._docker pure helpers.""" + from __future__ import annotations import json @@ -10,12 +11,16 @@ from pathlib import Path from unittest.mock import MagicMock, patch +import pytest from xr_ai_vllm._docker import ( + _CONFIG_LABEL, _already_logged_in, + _launch_fingerprint, _LogStreamer, _registry_for, build_run_argv, container_exists, + container_label, container_running, pid_on_port, run, @@ -99,9 +104,37 @@ def test_container_name_present(self, tmp_path): def test_port_label_set(self, tmp_path): argv = build_run_argv(**self._base_kwargs(tmp_path)) - assert "--label" in argv - idx = argv.index("--label") - assert argv[idx + 1] == "xr-ai-vllm.port=8100" + labels = [argv[index + 1] for index, value in enumerate(argv) if value == "--label"] + assert "xr-ai-vllm.port=8100" in labels + assert any(label.startswith(f"{_CONFIG_LABEL}=") for label in labels) + + def test_configuration_fingerprint_changes_with_vllm_arguments(self, tmp_path): + kwargs = self._base_kwargs(tmp_path) + first = build_run_argv(**kwargs) + kwargs["vllm_argv"] = [*kwargs["vllm_argv"], "--gpu-memory-utilization", "0.78"] + second = build_run_argv(**kwargs) + + def fingerprint(argv): + labels = [argv[index + 1] for index, value in enumerate(argv) if value == "--label"] + return next(label for label in labels if label.startswith(f"{_CONFIG_LABEL}=")) + + assert fingerprint(first) != fingerprint(second) + + def test_configuration_fingerprint_changes_with_contract_version( + self, + tmp_path, + monkeypatch, + ): + kwargs = self._base_kwargs(tmp_path) + first = build_run_argv(**kwargs) + monkeypatch.setattr("xr_ai_vllm._docker._LAUNCH_CONTRACT_VERSION", 2) + second = build_run_argv(**kwargs) + + def fingerprint(argv): + labels = [argv[index + 1] for index, value in enumerate(argv) if value == "--label"] + return next(label for label in labels if label.startswith(f"{_CONFIG_LABEL}=")) + + assert fingerprint(first) != fingerprint(second) def test_network_host(self, tmp_path): argv = build_run_argv(**self._base_kwargs(tmp_path)) @@ -326,6 +359,13 @@ def test_container_running_false_when_docker_missing(self): ): assert not container_running("some-name") + def test_container_label_returns_inspected_value(self): + with patch( + "xr_ai_vllm._docker.subprocess.check_output", + return_value="abc123\n", + ): + assert container_label("some-name", _CONFIG_LABEL) == "abc123" + def test_pid_on_port_returns_none_when_tools_missing(self): with patch( "xr_ai_vllm._docker.subprocess.check_output", @@ -334,18 +374,107 @@ def test_pid_on_port_returns_none_when_tools_missing(self): assert pid_on_port(8100) is None +def _run_kwargs(tmp_path): + return dict( + image="vllm/vllm-openai:v0.20.0", + container_name="xr-ai-vllm-test", + log_prefix="test", + vllm_argv=["vllm", "serve", "model", "--gpu-memory-utilization", "0.78"], + host="0.0.0.0", + port=8107, + model_cache=tmp_path / "models", + hf_token=None, + cuda_visible_devices="1", + extra_env=None, + extra_pip=None, + ready_file=None, + ) + + +def _expected_fingerprint(kwargs): + return _launch_fingerprint( + image=kwargs["image"], + port=kwargs["port"], + model_cache=kwargs["model_cache"], + cuda_visible_devices=kwargs["cuda_visible_devices"], + extra_env=kwargs["extra_env"], + extra_pip=kwargs["extra_pip"], + vllm_argv=kwargs["vllm_argv"], + ) + + class TestRun: - def test_stopped_container_is_removed_and_relaunched(self, tmp_path): + def test_healthy_unowned_listener_is_rejected(self, tmp_path): + kwargs = _run_kwargs(tmp_path) + with ( + patch("xr_ai_vllm._docker._docker_available", return_value=True), + patch("xr_ai_vllm._docker._lifecycle.health_ok", return_value=True), + patch("xr_ai_vllm._docker.container_exists", return_value=False), + patch("xr_ai_vllm._docker.stop_container") as stop, + patch("xr_ai_vllm._docker.remove_container") as remove, + patch("xr_ai_vllm._docker.subprocess.Popen") as popen, + patch("xr_ai_vllm._docker.signal.getsignal", return_value=None), + patch("xr_ai_vllm._docker.signal.signal"), + pytest.raises(SystemExit, match="1"), + ): + run(**kwargs) + + stop.assert_not_called() + remove.assert_not_called() + popen.assert_not_called() + + def test_healthy_stale_container_is_recreated(self, tmp_path): + kwargs = _run_kwargs(tmp_path) + state = {"exists": True} + process = MagicMock() + process.poll.return_value = None argv = ["docker", "run", "fresh-container"] + + def remove(_name): + state["exists"] = False + return True + + with ( + patch("xr_ai_vllm._docker._docker_available", return_value=True), + patch("xr_ai_vllm._docker._lifecycle.health_ok", return_value=True), + patch("xr_ai_vllm._docker.container_exists", side_effect=lambda _name: state["exists"]), + patch("xr_ai_vllm._docker.container_running", return_value=True), + patch("xr_ai_vllm._docker.container_label", return_value="stale"), + patch("xr_ai_vllm._docker.stop_container", return_value=True) as stop, + patch("xr_ai_vllm._docker.remove_container", side_effect=remove) as remove_mock, + patch("xr_ai_vllm._docker._maybe_ngc_login"), + patch("xr_ai_vllm._docker.build_run_argv", return_value=argv), + patch("xr_ai_vllm._docker.subprocess.Popen", return_value=process) as popen, + patch("xr_ai_vllm._docker._LogStreamer", return_value=MagicMock()), + patch("xr_ai_vllm._docker._lifecycle.wait_until_healthy"), + patch("xr_ai_vllm._docker._lifecycle.idle_until_stopped"), + patch("xr_ai_vllm._docker.signal.getsignal", return_value=None), + patch("xr_ai_vllm._docker.signal.signal"), + ): + run(**kwargs) + + stop.assert_called_once_with("xr-ai-vllm-test") + remove_mock.assert_called_once_with("xr-ai-vllm-test") + popen.assert_called_once_with(argv, start_new_session=True) + + def test_stale_stopped_container_is_recreated(self, tmp_path): + kwargs = _run_kwargs(tmp_path) + state = {"exists": True} process = MagicMock() process.poll.return_value = None + argv = ["docker", "run", "fresh-container"] + + def remove(_name): + state["exists"] = False + return True with ( patch("xr_ai_vllm._docker._docker_available", return_value=True), patch("xr_ai_vllm._docker._lifecycle.health_ok", return_value=False), - patch("xr_ai_vllm._docker.container_exists", return_value=True), + patch("xr_ai_vllm._docker.container_exists", side_effect=lambda _name: state["exists"]), patch("xr_ai_vllm._docker.container_running", return_value=False), - patch("xr_ai_vllm._docker.remove_container", return_value=True) as remove, + patch("xr_ai_vllm._docker.container_label", return_value="stale"), + patch("xr_ai_vllm._docker.remove_container", side_effect=remove) as remove_mock, patch("xr_ai_vllm._docker._maybe_ngc_login"), patch("xr_ai_vllm._docker.build_run_argv", return_value=argv), patch("xr_ai_vllm._docker.subprocess.Popen", return_value=process) as popen, @@ -355,20 +484,59 @@ def test_stopped_container_is_removed_and_relaunched(self, tmp_path): patch("xr_ai_vllm._docker.signal.getsignal", return_value=None), patch("xr_ai_vllm._docker.signal.signal"), ): - run( - image="vllm/vllm-openai:v0.20.0", - container_name="xr-ai-vllm-omni", - log_prefix="omni", - vllm_argv=["vllm", "serve", "model"], - host="0.0.0.0", - port=8108, - model_cache=tmp_path, - hf_token=None, - cuda_visible_devices="0", - extra_env=None, - extra_pip=["mamba-ssm"], - ready_file=None, - ) - - remove.assert_called_once_with("xr-ai-vllm-omni") + run(**kwargs) + + remove_mock.assert_called_once_with("xr-ai-vllm-test") popen.assert_called_once_with(argv, start_new_session=True) + + def test_matching_stopped_container_is_restarted(self, tmp_path): + kwargs = _run_kwargs(tmp_path) + wait_handle = MagicMock() + wait_handle.poll.return_value = None + + with ( + patch("xr_ai_vllm._docker._docker_available", return_value=True), + patch("xr_ai_vllm._docker._lifecycle.health_ok", return_value=False), + patch("xr_ai_vllm._docker.container_exists", return_value=True), + patch("xr_ai_vllm._docker.container_running", return_value=False), + patch("xr_ai_vllm._docker.container_label", return_value=_expected_fingerprint(kwargs)), + patch("xr_ai_vllm._docker.start_container", return_value=True) as start, + patch("xr_ai_vllm._docker._wait_for_container", return_value=wait_handle) as wait, + patch("xr_ai_vllm._docker.build_run_argv") as build, + patch("xr_ai_vllm._docker._LogStreamer", return_value=MagicMock()), + patch("xr_ai_vllm._docker._lifecycle.wait_until_healthy"), + patch("xr_ai_vllm._docker._lifecycle.idle_until_stopped"), + patch("xr_ai_vllm._docker.signal.getsignal", return_value=None), + patch("xr_ai_vllm._docker.signal.signal"), + ): + run(**kwargs) + + start.assert_called_once_with("xr-ai-vllm-test") + wait.assert_called_once_with("xr-ai-vllm-test") + build.assert_not_called() + + def test_matching_running_container_continues_startup(self, tmp_path): + kwargs = _run_kwargs(tmp_path) + wait_handle = MagicMock() + wait_handle.poll.return_value = None + + with ( + patch("xr_ai_vllm._docker._docker_available", return_value=True), + patch("xr_ai_vllm._docker._lifecycle.health_ok", return_value=False), + patch("xr_ai_vllm._docker.container_exists", return_value=True), + patch("xr_ai_vllm._docker.container_running", return_value=True), + patch("xr_ai_vllm._docker.container_label", return_value=_expected_fingerprint(kwargs)), + patch("xr_ai_vllm._docker._wait_for_container", return_value=wait_handle) as wait, + patch("xr_ai_vllm._docker.start_container") as start, + patch("xr_ai_vllm._docker.build_run_argv") as build, + patch("xr_ai_vllm._docker._LogStreamer", return_value=MagicMock()), + patch("xr_ai_vllm._docker._lifecycle.wait_until_healthy"), + patch("xr_ai_vllm._docker._lifecycle.idle_until_stopped"), + patch("xr_ai_vllm._docker.signal.getsignal", return_value=None), + patch("xr_ai_vllm._docker.signal.signal"), + ): + run(**kwargs) + + wait.assert_called_once_with("xr-ai-vllm-test") + start.assert_not_called() + build.assert_not_called() diff --git a/utils/xr-ai-vllm/xr_ai_vllm/_docker.py b/utils/xr-ai-vllm/xr_ai_vllm/_docker.py index 00893dd77..52d52e5ef 100644 --- a/utils/xr-ai-vllm/xr_ai_vllm/_docker.py +++ b/utils/xr-ai-vllm/xr_ai_vllm/_docker.py @@ -15,8 +15,10 @@ pull can proceed. Existing `~/.docker/config.json` entries take priority and are not overwritten. """ + from __future__ import annotations +import hashlib import json import logging import os @@ -36,11 +38,41 @@ _DOCKER_CONFIG = Path.home() / ".docker" / "config.json" _LOGIN_DONE: set[str] = set() +_CONFIG_LABEL = "xr-ai-vllm.config" +_LAUNCH_CONTRACT_VERSION = 1 # ── docker run argv builder ────────────────────────────────────────────────── +def _launch_fingerprint( + *, + image: str, + port: int, + model_cache: Path, + cuda_visible_devices: str | None, + extra_env: dict[str, str] | None, + extra_pip: list[str] | None, + vllm_argv: list[str], +) -> str: + payload = { + "launch_contract_version": _LAUNCH_CONTRACT_VERSION, + "image": image, + "port": port, + "model_cache": str(model_cache), + "cuda_visible_devices": cuda_visible_devices, + "extra_env": extra_env or {}, + "extra_pip": extra_pip or [], + "vllm_argv": vllm_argv, + } + encoded = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest()[:20] + + def build_run_argv( *, image: str, @@ -67,6 +99,16 @@ def build_run_argv( # caller needing to know the container name — implementation detail stays # inside this module. argv += ["--label", f"xr-ai-vllm.port={port}"] + fingerprint = _launch_fingerprint( + image=image, + port=port, + model_cache=model_cache, + cuda_visible_devices=cuda_visible_devices, + extra_env=extra_env, + extra_pip=extra_pip, + vllm_argv=vllm_argv, + ) + argv += ["--label", f"{_CONFIG_LABEL}={fingerprint}"] argv += ["--network", "host"] # vLLM workers communicate via /dev/shm; the default 64 MiB tmpfs is too # small for the KV cache shards. --ipc host gives them the host's larger @@ -163,6 +205,26 @@ def container_running(name: str) -> bool: return False +def container_label(name: str, label: str) -> str | None: + """Return one Docker container label, or ``None`` when unavailable.""" + + try: + raw = subprocess.check_output( + [ + "docker", + "inspect", + "--format", + f'{{{{ index .Config.Labels "{label}" }}}}', + name, + ], + text=True, + stderr=subprocess.DEVNULL, + ).strip() + return raw or None + except (FileNotFoundError, subprocess.CalledProcessError): + return None + + def remove_container(name: str) -> bool: """``docker rm`` *name* if it exists; return True if the container was removed. @@ -218,6 +280,30 @@ def stop_container(name: str, timeout_s: int = 20) -> bool: return False +def start_container(name: str) -> bool: + """Start one stopped container without attaching its output.""" + try: + subprocess.run( + ["docker", "start", name], + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + ) + return True + except (FileNotFoundError, subprocess.CalledProcessError): + return False + + +def _wait_for_container(name: str) -> subprocess.Popen: + """Return a liveness handle that exits when the container stops.""" + return subprocess.Popen( + ["docker", "wait", name], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + + # ── NGC auth ──────────────────────────────────────────────────────────────── @@ -466,6 +552,15 @@ def run( sys.exit(2) health_url = _lifecycle.health_url(host, port) + fingerprint = _launch_fingerprint( + image=image, + port=port, + model_cache=model_cache, + cuda_visible_devices=cuda_visible_devices, + extra_env=extra_env, + extra_pip=extra_pip, + vllm_argv=vllm_argv, + ) # On abort (Ctrl-C during model-servers startup) the launcher passes # no_kill=set() and SIGTERMs every wrapper's process group. Without a @@ -505,48 +600,102 @@ def _on_signal(_sig, _frame): signal.signal(signal.SIGINT, _on_signal) signal.signal(signal.SIGTERM, _on_signal) - # Reuse a container that survived a wrapper restart (weight persistence). + proc: subprocess.Popen | None = None + existing = container_exists(container_name) + running = existing and container_running(container_name) + + # A healthy endpoint is reusable only when the expected Docker container + # owns it and its complete launch contract matches. if _lifecycle.health_ok(health_url): + if not running: + log.error( + "Port %d is healthy, but expected container %s is not running; " + "refusing to reuse an unowned listener", + port, + container_name, + ) + sys.exit(1) + if container_label(container_name, _CONFIG_LABEL) == fingerprint: + print( + f"[{log_prefix}] vLLM already running on port {port} — reusing", + flush=True, + ) + if ready_file: + ready_file.touch() + signal.signal(signal.SIGINT, orig_int) + signal.signal(signal.SIGTERM, orig_term) + _lifecycle.idle_until_stopped(health_url, log_prefix) + return + print( - f"[{log_prefix}] vLLM already running on port {port} — reusing", + f"[{log_prefix}] Running container configuration changed — recreating {container_name}", flush=True, ) - if ready_file: - ready_file.touch() - signal.signal(signal.SIGINT, orig_int) - signal.signal(signal.SIGTERM, orig_term) - _lifecycle.idle_until_stopped(health_url, log_prefix) - return - - if container_exists(container_name) and not container_running(container_name): - # A container's command and entrypoint are immutable. Recreate failed - # containers so launcher fixes and changed service arguments take effect. + if not stop_container(container_name) or not remove_container(container_name): + log.error( + "Unable to replace stale running container %s", + container_name, + ) + sys.exit(1) + existing = False + running = False + + elif existing: + existing_fingerprint = container_label(container_name, _CONFIG_LABEL) + if existing_fingerprint != fingerprint: + print( + f"[{log_prefix}] Container configuration changed — recreating {container_name}", + flush=True, + ) + if running and not stop_container(container_name): + log.error( + "Unable to stop stale running container %s", + container_name, + ) + sys.exit(1) + if not remove_container(container_name): + log.error( + "Unable to remove stale container %s", + container_name, + ) + sys.exit(1) + existing = False + running = False + elif running: + print( + f"[{log_prefix}] Matching container still starting — waiting", + flush=True, + ) + proc = _wait_for_container(container_name) + else: + print( + f"[{log_prefix}] Restarting matching stopped container {container_name}", + flush=True, + ) + if not start_container(container_name): + log.error("Unable to restart container %s", container_name) + sys.exit(1) + proc = _wait_for_container(container_name) + + if proc is None: + _maybe_ngc_login(image) + argv = build_run_argv( + image=image, + container_name=container_name, + port=port, + model_cache=model_cache, + hf_token=hf_token, + cuda_visible_devices=cuda_visible_devices, + extra_env=extra_env, + extra_pip=extra_pip, + vllm_argv=vllm_argv, + ) print( - f"[{log_prefix}] Recreating stopped container {container_name}", + f"[{log_prefix}] Launching vLLM (docker) image={image} " + f"container={container_name} http://{host}:{port}/v1", flush=True, ) - if not remove_container(container_name): - log.error("Could not remove stopped container %s", container_name) - sys.exit(1) - - _maybe_ngc_login(image) - argv = build_run_argv( - image=image, - container_name=container_name, - port=port, - model_cache=model_cache, - hf_token=hf_token, - cuda_visible_devices=cuda_visible_devices, - extra_env=extra_env, - extra_pip=extra_pip, - vllm_argv=vllm_argv, - ) - print( - f"[{log_prefix}] Launching vLLM (docker) image={image} " - f"container={container_name} http://{host}:{port}/v1", - flush=True, - ) - proc = subprocess.Popen(argv, start_new_session=True) + proc = subprocess.Popen(argv, start_new_session=True) _state["proc"] = proc streamer = _LogStreamer(container_name)