xr_ai_vllm#

xr-ai-vllm — pluggable vLLM backend for xr-ai inference services.

Lets each vLLM-backed service host vllm via either:

pip — pip-installed vllm CLI in the wrapper’s venv (default; today’s behavior).
docker — docker run nvcr.io/nvidia/vllm:<tag> vllm serve … (NGC container).

The choice is per-server, set via vllm_backend: pip|docker in the service’s YAML. Both paths honor identical config keys (model, ports, vllm flags); only the runtime hosting vllm differs.

Stdlib-only by contract — no vllm or other heavy deps imported here, so the docker path stays light even when pip vllm is not installed.

Typical usage from a service wrapper:

from xr_ai_vllm import serve, DEFAULT_IMAGE

serve(
    backend=cfg.get("vllm_backend", "pip"),
    persistent=True,
    image=cfg.get("vllm_image", DEFAULT_IMAGE),
    container_name="xr-ai-vllm-vlm-server",
    log_prefix="vlm_server",
    model=model,
    extra_serve_args=[
        "--served-model-name", served_name,
        "--max-num-seqs", str(max_seqs),
        ...
    ],
    host=host, port=port,
    model_cache=model_cache,
    hf_token=os.environ.get("HF_TOKEN"),
    cuda_visible_devices=cfg.get("cuda_visible_devices"),
    ready_file=ns.ready_file,
)

Attributes#

DEFAULT_IMAGE

Default NGC vLLM image used when a service does not override vllm_image.

Functions#

`gpu_compute_major`(→ int)	Return the GPU's compute-capability major version, or 0 if unknown.
`load_config`(→ tuple[dict, pathlib.Path, ...)	Parse `--config`/`--ready-file` and load the YAML config.
`resolve_model_cache`(→ pathlib.Path)	Resolve `model_cache` (relative to the YAML dir) and ensure it exists.
`setup_hf_env`(→ str \| None)	Apply the shared HuggingFace / CUDA env block.
`serve_nim`(→ None)	Pull (if needed) and run a NIM container, blocking until stopped.
`serve`(→ None)	Launch vLLM via backend ("pip" or "docker").
`stop_persistent_servers`(→ bool)	Stop persisted servers and report whether every discovered server stopped.

Package Contents#

xr_ai_vllm.gpu_compute_major() → int#

Return the GPU’s compute-capability major version, or 0 if unknown.

Reads CUDA_VISIBLE_DEVICES from the env, so set the device filter before calling. Logs a warning on failure and falls back to 0.

xr_ai_vllm.load_config( argv: list[str] | None = None, ) → tuple[dict, pathlib.Path, pathlib.Path | None]#

Parse --config/--ready-file and load the YAML config.

Reconfigures stdout/stderr to line-buffered so logs flush under the launcher’s piped stdout. Returns (cfg, yaml_dir, ready_file); yaml_dir is the config’s directory (cwd when no config is given), used as the base for relative paths like model_cache.

xr_ai_vllm.resolve_model_cache( cfg: dict, yaml_dir: pathlib.Path, *, default: str, ) → pathlib.Path#: Resolve model_cache (relative to the YAML dir) and ensure it exists.

xr_ai_vllm.setup_hf_env(cfg: dict, model_cache: pathlib.Path) → str | None#

Apply the shared HuggingFace / CUDA env block.

Sets CUDA_VISIBLE_DEVICES (when configured), HF_TOKEN (when provided), HF_XET_HIGH_PERFORMANCE, HF_HOME, and TRANSFORMERS_CACHE.

A non-empty env value for HF_TOKEN, HF_HOME, or TRANSFORMERS_CACHE wins over the YAML value, per docs/source/getting_started/credentials.md. TRANSFORMERS_CACHE mirrors HF_HOME for Transformers <4.36 and libraries that have not adopted HF_HOME.

Returns the resolved cuda_visible_devices string (or None) so callers that run GPU detection can confirm the device filter is applied.

xr_ai_vllm.serve_nim( *, image: str, container_name: str, log_prefix: str, http_port: int, grpc_port: int | None = None, nim_cache: pathlib.Path, cuda_visible_devices: str | None = None, extra_env: dict[str, str] | None = None, ready_file: pathlib.Path | None = None, ) → None#

Pull (if needed) and run a NIM container, blocking until stopped.

Readiness is /v1/health/ready on http_port. First start includes the NGC engine download (multi-GB), so expect a long cold start; the mounted nim_cache makes subsequent starts fast.

xr_ai_vllm.DEFAULT_IMAGE = 'nvcr.io/nvidia/vllm:26.08-py3'#

Default NGC vLLM image used when a service does not override vllm_image.

Individual services may pin a newer image when required by their model.

xr_ai_vllm.serve( *, backend: str, persistent: bool, image: str = DEFAULT_IMAGE, container_name: str, log_prefix: str, model: str, extra_serve_args: list[str], host: str, port: int, model_cache: pathlib.Path, hf_token: str | None = None, cuda_visible_devices: str | None = None, extra_env: dict[str, str] | None = None, extra_pip: list[str] | None = None, ready_file: pathlib.Path | None = None, spark_uma: bool = False, ) → None#

Launch vLLM via backend (“pip” or “docker”).

extra_serve_args are the flags appended after vllm serve <model> — everything past the model id (e.g. --served-model-name, --max-num-seqs, --reasoning-parser, …). They are passed verbatim, so caller-side flag construction is unchanged from the per-service wrappers’ previous inline argv.

persistent controls the pip-mode lifecycle only:

True — vLLM pip process starts in a new session so it survives wrapper restarts. Cleanup is via stop_persistent_servers.
False — die with the wrapper.

For the docker backend persistent is ignored: the container always runs foreground with start_new_session=True, so it escapes the launcher’s process group regardless. Use Process(..., launch_mode="persist") in the orchestrator main.py to tell the launcher not to kill the wrapper on shutdown.

container_name is only consulted in docker mode. Use a stable, service-specific name (e.g. xr-ai-vllm-<entry-point>) so the stop helper can find it.

Docker mode verifies a Xet-capable Hub inside the container and repairs a missing or incompatible hf-xet wheel. extra_pip is an additional docker-only list of pip-installable package specs for model architectures whose wheels the image doesn’t bundle. Silently ignored in pip mode (deps belong in the wrapper’s pyproject.toml there).

spark_uma enables the DGX Spark cold-start safeguards in docker mode: prefetching and syncing the Hugging Face snapshot before CUDA starts, and one bounded restart of a container attempt launched by this wrapper after a non-container-OOM CUDA driver-allocation failure. Pip mode ignores this Docker lifecycle option, preserving the one-field backend switch used by the service YAML files.

xr_ai_vllm.stop_persistent_servers(services: list[tuple[str, int]]) → bool#

Stop persisted servers and report whether every discovered server stopped.

services is a list of (label, port) tuples. For each entry:

Look for a docker container labelled xr-ai-vllm.port=<port> (stamped at start time by the vLLM wrapper) and docker stop it.
Fall back to port → pid for pip-mode vLLM or in-process servers (e.g. STT). Pocket TTS is signalled as a complete process group; other local servers are signalled by PID.

A missing container and listener is already stopped. Output is print-style with [<label>] … prefixes. Discovery errors and listeners that are not identified as xr-ai services fail closed without sending a signal.