diff --git a/README.md b/README.md index d8e4fd478..b96276e8b 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ one, and both report live GPU and memory use throughout. | **Architectures** | x64 and arm64 on all three. Windows on ARM is experimental. | | **Installers** | Windows `.exe`; Linux `.deb`; macOS `.dmg`. On other Linux distributions, [build from source](docs/building.mdx). | | **Mixing nodes** | Windows, Linux, and macOS nodes can all be paired with each other | -| **Inference engines** | Ollama, LM Studio, and llama.cpp | +| **Inference engines** | Ollama, LM Studio, llama.cpp, and vLLM (Linux with an NVIDIA GPU, or Apple Silicon macOS) | **PAIR running on a machine does not mean an engine will.** PAIR itself runs on any supported Windows, Linux, or macOS machine. Each engine sets its own requirements @@ -100,7 +100,7 @@ you want by its full filename instead. **Settings → Service** and read the status there. - **Get an engine running.** On the node's card, open **Engine settings** and - select **Install** next to Ollama, LM Studio, or llama.cpp. PAIR downloads and + select **Install** next to Ollama, LM Studio, llama.cpp, or (on Linux or Apple Silicon macOS) vLLM. PAIR downloads and sets the engine up for you, so nothing needs to be in place beforehand. ![The Install engines dialog with Ollama downloading, reporting progress as it installs.](docs/assets/onboarding/engine-lifecycle/01-engine-installing.png) @@ -199,7 +199,7 @@ Your data means settings, logs, cluster identity and certificates, and any engin PAIR installed for you. An engine the uninstaller cannot remove — one still holding a file open, say — is left in place rather than stopping the uninstall; remove it with its own uninstaller. **Model weights are not touched** — they live in the -engine's own storage (`~/.ollama`, `~/.llamacpp`, `~/.lmstudio/models`), so +engine's own storage (`~/.ollama`, `~/.llamacpp`, `~/.lmstudio/models`, `~/.cache/huggingface`), so removing PAIR does not delete the models you downloaded. Delete those through the engine, or by removing its directory. An engine you installed yourself is also left alone. @@ -284,7 +284,7 @@ feedback and contributions will help shape priorities. ### Engines and integrations - [x] llama.cpp support. -- [ ] vLLM support. +- [x] vLLM support. - [ ] EXO support. - [ ] ComfyUI integration. - [ ] Unsloth support. diff --git a/desktop/docs/architecture.md b/desktop/docs/architecture.md index 44d48e57b..94928ec35 100644 --- a/desktop/docs/architecture.md +++ b/desktop/docs/architecture.md @@ -250,7 +250,7 @@ not select or pin proxy routes. An NVPAIR-launched engine binds to loopback and is never directly LAN-reachable. Peers reach it only through the node's proxy over a cluster-mTLS ingress, so discovery advertises the promoted proxy port for `ol`/`lm` rather than the -engine's private port; llama.cpp uses the additional `lc` key. Transport +engine's private port; llama.cpp uses the additional `lc` key and vLLM the `vl` key. Transport security is backend-owned; Electron only reflects the advertised proxy port and reads a remote engine's real port from `engine:remote-get-installed` facts. @@ -408,7 +408,7 @@ deletes finish. Unpackaged (`electron-vite dev`) builds omit relaunch and the UI tells the developer to run `npm start` again — quitting Electron tears down the Vite renderer server, so auto-relaunch would come up with nothing to load. -The engines' model stores (`~/.ollama`, `~/.llamacpp`, `~/.lmstudio/models`) and +The engines' model stores (`~/.ollama`, `~/.llamacpp`, `~/.lmstudio/models`, `~/.cache/huggingface`) and the application install tree are never removed. The reset does uninstall the engines PAIR installed, through `engine:uninstall-managed` while the broker is still running, so an engine a vendor installer placed in the user's home is diff --git a/desktop/docs/frontend-api.md b/desktop/docs/frontend-api.md index 6209ca269..ef9ef5cca 100644 --- a/desktop/docs/frontend-api.md +++ b/desktop/docs/frontend-api.md @@ -99,11 +99,19 @@ They do not imply a WebSocket connection. Browser clients are not supported. - `onSettingsDisconnected(callback)` reports that a node's settings authority became unreachable, so its cached snapshot is stale. -`EngineType` currently includes `ollama`, `lm-studio`, and `llama-cpp`. +`EngineType` currently includes `ollama`, `lm-studio`, `llama-cpp`, and `vllm`. +vLLM serves one model per process, so `unloadModel` (Eject) stops the engine, +which frees its GPU memory, and `loadModel` on a stopped engine starts it. Its model +list is the Hugging Face cache, which stays available while the engine is stopped, `deleteModel` removes a cached model (not the one +loaded in vLLM; deleting an ejected or stopped vLLM's model leaves it with no model +selected until `loadModel` picks one), `pullModel` downloads a Hugging Face repo into +that cache (no running engine needed), and `loadModel` on a cached model switches +the served model by restarting vLLM on it. `searchHub(engineType, query?)` serves Ollama's locked local catalog and cached -live Hugging Face catalogs for LM Studio and llama.cpp. A non-empty llama.cpp -query performs an explicit bounded search across public Hugging Face GGUF -repositories; other calls return the populated engine catalog. llama.cpp pull +live Hugging Face catalogs for LM Studio, llama.cpp and vLLM. A non-empty +llama.cpp or vLLM query performs an explicit bounded search across public +Hugging Face repositories (GGUF for llama.cpp, safetensors chat models for vLLM); +other calls return the populated engine catalog. llama.cpp pull IDs are exact `owner/repository:quantization` values; load, unload, and delete all use that same exact id. diff --git a/desktop/docs/services-api.md b/desktop/docs/services-api.md index 8137f326d..b60f8c1f1 100644 --- a/desktop/docs/services-api.md +++ b/desktop/docs/services-api.md @@ -43,6 +43,7 @@ - ⚠️ nvpair-proxy → node/set-local-backend - ⚠️ nvpair-ui-broker → discovery:unsubscribe - ⚠️ nvpair-ui-broker → engine:configure-launch +- ⚠️ nvpair-ui-broker → engine:restart - ⚠️ nvpair-ui-broker → engine:set-port - ⚠️ nvpair-ui-broker → engine:set-reserved-port - ⚠️ nvpair-ui-broker → engine:unsubscribe @@ -270,9 +271,11 @@ | `engine:get-settings` | request (we call) | ✅ yes | | `engine:install` | request (we call) | ✅ yes | | `engine:preview-settings` | request (we call) | ✅ yes | +| `engine:restart` | request (we call) | ⚠️ not called | | `engine:set-port` | request (we call) | ⚠️ not called | | `engine:set-reserved-port` | request (we call) | ⚠️ not called | | `engine:start` | request (we call) | ✅ yes | +| `engine:stop` | request (we call) | ✅ yes | | `engine:subscribe` | request (we call) | ✅ yes | | `engine:unsubscribe` | request (we call) | ⚠️ not called | | `errors:get-initial` | request (we call) | ✅ yes | diff --git a/desktop/docs/services-backend.md b/desktop/docs/services-backend.md index aea950318..d9c6b98df 100644 --- a/desktop/docs/services-backend.md +++ b/desktop/docs/services-backend.md @@ -116,7 +116,7 @@ reserved for inference clients. | ---------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------- | | `app:ready` | Complete broker startup and refresh snapshots | `state:request-refresh` | | `discovery:nodes-changed` | Replace discovery snapshot and diff nodes | `discovery:nodes-changed`, `nodes:upsert`, `nodes:remove` | -| `ollama-proxy:ready` / `lmstudio-proxy:ready` / `llamacpp-proxy:ready` | Record engine proxy port | `engines:state-changed` | +| `ollama-proxy:ready` / `lmstudio-proxy:ready` / `llamacpp-proxy:ready` / `vllm-proxy:ready` | Record engine proxy port | `engines:state-changed` | | proxy `node/*` | Update per-engine node presence; the advertised port is the peer's promoted proxy port (not the engine's private loopback port) | node and engine pushes | | `engine:ready` / `engine:state-changed` | Update engine facts and models | `engines:state-changed` | | `engine:settings-changed` | Validate and republish the owning node's settings snapshot | `engines:settings-changed` | diff --git a/desktop/docs/services-parity.md b/desktop/docs/services-parity.md index 2694a5d9c..b4dc07a17 100644 --- a/desktop/docs/services-parity.md +++ b/desktop/docs/services-parity.md @@ -16,25 +16,26 @@ history. ## Status summary -| Domain | Status | Current state | -| -------------------------- | ------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------- | -| Process supervision | Complete | Electron starts only `nvpair-ui-broker`; the broker supervises all workers | -| Discovery | Complete | Broker discovery snapshots drive available nodes and node state | -| Node telemetry | Integrated with direct poll | Electron polls advertised `/v1/node-info` (plain HTTP); remote OS and some remote telemetry are backend-limited | -| Manual nodes | Complete with local persistence | Broker owns probing, durable entries, restart replay, and proxy registration | -| Ollama routing | Complete | Broker relay and backend scheduler drive proxy routing | -| LM Studio routing | Complete | Parallel broker relay and scheduler path | -| llama.cpp backend | Integrated | Desktop and TUI expose install, lifecycle, catalog pull, inventory, load/unload, endpoints, routing, and demo traffic | -| Local engine lifecycle | Complete | Install, start, stop, uninstall, update, and port configuration | -| Remote engine lifecycle | Partial | Remote install, start, stop, status, and model pull are supported | -| Engine models | Partial | Core list, pull, load, unload, and supported delete actions are wired | -| Engine version metadata | Backend-limited | The backend status surface does not provide a complete version/install-owner/update contract | -| Errors | Complete | Broker-relayed registry with peer synchronization | -| Workloads | Complete after subscription | Local and peer workload events feed one Electron catalog | -| Cluster pairing | Complete | PIN pairing, identity, membership, leave, and removal | -| Cluster transport security | Backend-owned | Node-to-node transport security, including the proxies' cluster-mTLS inference ingress, is entirely backend; Personal AI Router implements none | -| Settings | Partial | Cluster identity plus per-engine ports and engine arguments, local and remote; inert backend settings are not surfaced | -| Model catalog search | Backend-owned | `engine:catalog` serves the locked Ollama list and the cached LM Studio and llama.cpp catalogs, and searches Hugging Face for llama.cpp | +| Domain | Status | Current state | +| -------------------------- | ------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Process supervision | Complete | Electron starts only `nvpair-ui-broker`; the broker supervises all workers | +| Discovery | Complete | Broker discovery snapshots drive available nodes and node state | +| Node telemetry | Integrated with direct poll | Electron polls advertised `/v1/node-info` (plain HTTP); remote OS and some remote telemetry are backend-limited | +| Manual nodes | Complete with local persistence | Broker owns probing, durable entries, restart replay, and proxy registration | +| Ollama routing | Complete | Broker relay and backend scheduler drive proxy routing | +| LM Studio routing | Complete | Parallel broker relay and scheduler path | +| llama.cpp backend | Integrated | Desktop and TUI expose install, lifecycle, catalog pull, inventory, load/unload, endpoints, routing, and demo traffic | +| vLLM engine | Integrated (Linux, macOS arm64) | Desktop and TUI expose install, lifecycle, the cached-model list, switching the served model (Load), Hugging Face catalog search and download, endpoints and routing; eject (stops the engine), delete of cached models; no demo traffic | +| Local engine lifecycle | Complete | Install, start, stop, uninstall, update, and port configuration | +| Remote engine lifecycle | Partial | Remote install, start, stop, status, and model pull are supported | +| Engine models | Partial | Core list, pull, load, unload, and supported delete actions are wired | +| Engine version metadata | Backend-limited | The backend status surface does not provide a complete version/install-owner/update contract | +| Errors | Complete | Broker-relayed registry with peer synchronization | +| Workloads | Complete after subscription | Local and peer workload events feed one Electron catalog | +| Cluster pairing | Complete | PIN pairing, identity, membership, leave, and removal | +| Cluster transport security | Backend-owned | Node-to-node transport security, including the proxies' cluster-mTLS inference ingress, is entirely backend; Personal AI Router implements none | +| Settings | Partial | Cluster identity plus per-engine ports and engine arguments, local and remote; inert backend settings are not surfaced | +| Model catalog search | Backend-owned | `engine:catalog` serves the locked Ollama list and the cached LM Studio and llama.cpp catalogs, and searches Hugging Face for llama.cpp | ## Supervision @@ -182,6 +183,49 @@ Current limits are explicit: the manual-node worker does not probe llama.cpp, the catalog offers no alternate quantizations, and package selection is fixed rather than driver-aware or offline-repacked. +### vLLM support + +The bundled `vllm` manifest covers Linux (NVIDIA GPU) and Apple Silicon macOS +(vLLM-Metal). It installs a pinned vLLM release into a private virtualenv +through a checksum-pinned `uv` bootstrap, or adopts a `vllm` already installed +at a detected path, then runs `vllm serve` on `8001` behind the proxy facade on +`8000`. vLLM serves one model per process and has no load or unload API, so +Eject (`unload_model`, a `stop_engine` action) stops the engine to free the GPU +and Load on a stopped engine starts it; its model list +(`list_models` on disk) is kept while stopped. It downloads and deletes +through the shared Hugging Face cache (`pull_model` is an `hf_pull` action, +`delete_model` an `hf_cache` delete that refuses the model an engine process +serves and, on an ejected or stopped engine, also clears it from the launch +arguments), +and has a Hugging Face catalog. Load switches the served model through `load_model` (a +`launch_model` action): it saves `--model ` into the launch arguments and +restarts the engine, restoring the previous model if the new one does not start. +Only the model being served is advertised to peers and the proxy, so a request +for another cached model is refused until it is loaded. A Load switch passes through a stopped engine for the length of the +restart. The pending Load in the renderer's pending-actions store (the one owner +of optimistic engine and model state) also covers the engine, because vLLM's +capabilities declare `loadRestartsEngine`: its card shows starting rather than +stopped, and its toggle stays locked, until the model is reported loaded, the +load fails, or the 20-minute switch budget runs out. Its model list is read from +disk while stopped, so it is kept throughout. Local and remote Loads are covered +the same way. + +A local Start, of any engine, shows starting until the engine-manager reports +the outcome. While it shows starting (a Start, or a vLLM Load restarting the +engine), its switch stays on and usable: turning it off sends an explicit +`stop` command (`engine:stop`, or `engine:remote-stop` for a peer), because a +toggle reads the engine's status and an engine that is still starting reports +not running. The engine-manager cancels the start in flight. Its safety net is the awaited lifecycle envelope +(`MODULAR_ENGINE_LIFECYCLE_CALL_TIMEOUT_MS`, 14 minutes), which outlasts every +engine's readiness window (vLLM 8 minutes, Ollama 10), so a slow start is never +shown as stopped while it is still loading. Weights live in the shared +Hugging Face cache (`~/.cache/huggingface`), which uninstall and app-data reset +never remove. + +Current limits are explicit: the manual-node worker does not probe vLLM, the +Inference Demo does not use it, and the icon is a placeholder lettermark rather +than the vLLM project logo. + ## Engine lifecycle `nvpair-engine-manager` is authoritative for installed, running, healthy, and @@ -284,16 +328,16 @@ restart-backed delete, the reply is withheld until the engine is ready again `MODULAR_MODEL_ACTION_TIMEOUT_MS` anchors the delete-specific observers. Ollama Load uses a separate 10/11-minute response-header budget: -| Budget | Value | Why | -| ------------------------------------------------------------- | ----- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `MODULAR_MODEL_ACTION_TIMEOUT_MS` (`modular-runtime.ts`) | 120s | The RPC must outlast a full stop + readiness-probed start. | -| `RESTART_DELETE_TIMEOUT_MS` (`pending-actions.store.ts`) | 180s | The optimistic-spinner safety net must outlast the RPC, or it expires mid-flight and drops the spinner while the delete is still running. | -| `engineResponseHeaderTimeout` (`executor.go`) | 30s | Downloads and ordinary local HTTP actions retain a prompt response-header bound; probe contexts stay shorter. | -| `pullProgressTimeout` (`executor.go`) | 30min | Streaming Ollama and llama.cpp pulls may run indefinitely while layer/file completed bytes advance; 30 minutes without advancing bytes cancels the stalled pull. | -| `ollamaLoadResponseHeaderTimeout` (`executor.go`) | 10min | Only Ollama's local `run_model` action gets the cold-load allowance. | -| `OLLAMA_LOAD_PENDING_TIMEOUT_MS` (`pending-actions.store.ts`) | 12min | Ollama's Load control stays locked beyond the remote path's 11-minute header budget, while backend success or failure still clears it immediately. | -| `PENDING_TIMEOUT_MS` | 60s | Unchanged for every other command, including an Ollama delete. | -| `remoteReadyResponseHeaderTimeout` (`remoteclient.go`) | 11min | Remote start/delete responses can wait on engine readiness, and a cold Ollama model load can withhold headers for minutes, so those calls use the readiness-sized client. Other model actions retain the ordinary 30s budget. | +| Budget | Value | Why | +| ------------------------------------------------------------- | ----- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `MODULAR_MODEL_ACTION_TIMEOUT_MS` (`modular-runtime.ts`) | 120s | The RPC must outlast a full stop + readiness-probed start. | +| `RESTART_DELETE_TIMEOUT_MS` (`pending-actions.store.ts`) | 180s | The optimistic-spinner safety net must outlast the RPC, or it expires mid-flight and drops the spinner while the delete is still running. | +| `engineResponseHeaderTimeout` (`executor.go`) | 30s | Downloads and ordinary local HTTP actions retain a prompt response-header bound; probe contexts stay shorter. | +| `pullProgressTimeout` (`executor.go`) | 30min | Streaming Ollama and llama.cpp pulls may run indefinitely while layer/file completed bytes advance; 30 minutes without advancing bytes cancels the stalled pull. | +| `ollamaLoadResponseHeaderTimeout` (`executor.go`) | 10min | Only Ollama's local `run_model` action gets the cold-load allowance. | +| `OLLAMA_LOAD_PENDING_TIMEOUT_MS` (`pending-actions.store.ts`) | 12min | Ollama's Load control stays locked beyond the remote path's 11-minute header budget, while backend success or failure still clears it immediately. | +| `PENDING_TIMEOUT_MS` | 60s | Unchanged for every other command, including an Ollama delete. | +| `remoteReadyResponseHeaderTimeout` (`remoteclient.go`) | 11min | Remote start/delete responses can wait on engine readiness, and a cold Ollama model load or a vLLM load (which restarts vLLM on the new model) can withhold headers for minutes, so those calls use the readiness-sized client. Other model actions retain the ordinary 30s budget. | The net is a backstop, not the mechanism: a delete now ends the spinner on real backend truth either way. On success the refreshed `list_models` drops the model; @@ -345,6 +389,15 @@ safety-net timeout (`pending-actions.store.ts`). Loaded state carries no `sizeVram`/`expiresAt` — the backend delivers the simpler `loadedByEngine` name-set, not structured details. +vLLM advertises only the model it is serving, so its `modelsByEngine` entry is +that one name, for routing. Its whole downloaded list arrives as +`downloadedByEngine`, read even while it is stopped, and `modelsForEngine` +lists a node's vLLM rows from it. Its `loadedByEngine` uses the same names, so +a model served under `--served-model-name` still marks its row loaded. When a +peer's vLLM list was not read, `EngineModels.inventoryKnown` is false and the +model manager says the downloaded models are unavailable instead of claiming +none are downloaded. + The model catalogue is owned by the backend: `engine:catalog` on `nvpair-engine-manager` serves the curated Ollama, LM Studio, and llama.cpp lists, and searches Hugging Face for llama.cpp, and both the desktop app and @@ -487,15 +540,15 @@ its effect is observable. The following responsibilities remain in Personal AI Router because the backend does not provide an equivalent client-facing contract: -| Responsibility | Location | -| ----------------------------------------------------------- | ------------------------------------------------ | -| Poll node telemetry over `/v1/node-info` | `node-info-poller.ts` | -| Bridge the local node into engine proxies | `modular-supervisor.ts` | -| Present optimistic engine transition state | `pending-actions.store.ts`, bridge state | -| Relay the backend model catalogue to the renderer | `service-bridge/model-catalog.ts` | -| Accumulate and reconcile receiver-side pending invites | `modular-state.ts`, `modular-supervisor.ts` | -| Mirror backend-coupled runtime defaults not yet reported | `modular-runtime.ts` | -| Collapse a superseded node row before the scanner proves it | `modular-state.ts`, `modular-runtime.ts` | +| Responsibility | Location | +| ----------------------------------------------------------- | ------------------------------------------- | +| Poll node telemetry over `/v1/node-info` | `node-info-poller.ts` | +| Bridge the local node into engine proxies | `modular-supervisor.ts` | +| Present optimistic engine transition state | `pending-actions.store.ts`, bridge state | +| Relay the backend model catalogue to the renderer | `service-bridge/model-catalog.ts` | +| Accumulate and reconcile receiver-side pending invites | `modular-state.ts`, `modular-supervisor.ts` | +| Mirror backend-coupled runtime defaults not yet reported | `modular-runtime.ts` | +| Collapse a superseded node row before the scanner proves it | `modular-state.ts`, `modular-runtime.ts` | These are current integration boundaries. Do not add compatibility shims or duplicate backend services around them. diff --git a/desktop/src/electron/service-bridge/empty-handlers.ts b/desktop/src/electron/service-bridge/empty-handlers.ts index e00893226..350cd8944 100644 --- a/desktop/src/electron/service-bridge/empty-handlers.ts +++ b/desktop/src/electron/service-bridge/empty-handlers.ts @@ -7,6 +7,7 @@ import type { ClusterNode, ClusterNodeIdentity, Invite } from '@/shared/types/cl import type { EngineType } from '@/shared/types/engines' import type { ServiceError } from '@/shared/types/errors' import { + MODEL_SWITCH_TIMEOUT_MS, MODULAR_CLUSTER_MANAGER_PORT, MODULAR_ENGINE_LIFECYCLE_CALL_TIMEOUT_MS } from '@/shared/constants/modular-runtime' @@ -57,6 +58,47 @@ function booleanValue(value: JsonValue | undefined): boolean { return typeof value === 'boolean' ? value : false } +/** + * When this app last asked each engine to stop, keyed by node and engine. A Stop + * or an Eject ends a vLLM Load that is still starting the engine, and the Load + * then replies `stopped`. A user who stopped it here needs no message about it; + * a stop from elsewhere (the terminal interface, or another node's app) is + * reported, or this app would go on showing the Load as in progress. + */ +const stopRequestedAt = new Map() + +function stopKey(nodeId: string, engineType: EngineType): string { + return `${nodeId}:${engineType}` +} + +function noteStopRequest(payload: WsInvokeRequest<'engine:command'>): void { + if (payload.command === 'stop' || payload.command === 'unloadModel') { + stopRequestedAt.set(stopKey(payload.nodeId, payload.engineType), Date.now()) + } +} + +/** Report a Load that a stop from outside this app ended. */ +function reportLoadStoppedElsewhere( + payload: WsInvokeRequest<'engine:command'>, + startedAt: number, + result: JsonValue | undefined +): void { + if (!payload.model || !booleanValue(objectValue(result)?.stopped)) return + const stoppedHere = stopRequestedAt.get(stopKey(payload.nodeId, payload.engineType)) ?? -1 + if (stoppedHere >= startedAt) return + getModularSupervisor().reportError( + `Loading ${payload.model} on ${engineManagerName(payload.engineType)} was stopped before it finished.`, + 'info', + `engine-cmd:load model:${engineManagerName(payload.engineType)}`, + { + nodeId: payload.nodeId, + engineType: payload.engineType, + operation: 'load', + modelName: payload.model + } + ) +} + /** Relay a `cluster:*` / `nodes:*` request to nvpair-cluster-manager via the broker. */ function callCluster( method: string, @@ -154,6 +196,31 @@ async function toggleLocalEngine(engine: string, engineType: EngineType): Promis } } +/** + * Stop a local engine, whatever it reports. A toggle reads the engine's status + * to decide, and an engine that is still starting reports not running, so the + * toggle would send another start; Stop is sent when the UI shows it starting. + * The engine-manager cancels the start in flight. + */ +function stopLocalEngine(engine: string, engineType: EngineType): void { + const supervisor = getModularSupervisor() + getModularBridgeState().beginLocalEngineOp(engineType, 'stopping') + supervisor.sendProcess( + 'broker', + 'engine:stop', + { engine }, + error => { + getModularBridgeState().clearPendingEngineOp(engineType) + supervisor.reportError( + `Failed to stop ${engine}: ${error}`, + 'error', + `engine-cmd:stop:${engine}` + ) + }, + true + ) +} + /** * Dispatch a UI engine command to the local `nvpair-engine-manager`. Lifecycle and * model operations are **fire-and-forget**: the engine-manager runs each in its @@ -241,6 +308,9 @@ function routeEngineManagerCommand(payload: WsInvokeRequest<'engine:command'>): case 'toggle': void toggleLocalEngine(engine, payload.engineType) break + case 'stop': + stopLocalEngine(engine, payload.engineType) + break case 'pullModel': // Pulls need UI feedback the backend doesn't emit (progress + failure), // so the supervisor owns the optimistic-progress + awaited-completion @@ -260,11 +330,38 @@ function routeEngineManagerCommand(payload: WsInvokeRequest<'engine:command'>): // (`/api/generate`) with no prompt: Ollama loads the model into VRAM // and returns immediately (`done_reason: "load"`) without generating. // LM Studio and llama.cpp declare real `load_model` manifest actions. - // All require the engine running (HTTP/CLI action), matching the - // disabled rule in ModelRow.tsx. + // vLLM's `load_model` serves the model by restarting a running engine + // on it, or starting a stopped one, which takes minutes and can fail + // (the model may not fit), so like Ollama's its rejection is observed + // and attributed to the pending model row. + // Ollama, LM Studio and llama.cpp require the engine running (HTTP/CLI + // action); vLLM declares modelOpsWhenStopped. ModelRow.tsx applies + // the same rule. // See docs/services-parity.md#models. if (payload.model) { - if (payload.engineType === 'ollama') { + if (payload.engineType === 'vllm') { + // The reply arrives once the new model is serving, or the + // previous one is restored. The renderer's pending Load shows + // the engine as starting while it restarts. A Stop or Eject + // ends the Load with `stopped` in the reply, not an error. + const startedAt = Date.now() + void supervisor + .callProcess( + 'broker', + 'engine:action', + { engine, action: 'load_model', params: { model: payload.model } }, + MODEL_SWITCH_TIMEOUT_MS + ) + .then(result => reportLoadStoppedElsewhere(payload, startedAt, result)) + .catch(err => + failAction('load model', { + nodeId: payload.nodeId, + engineType: payload.engineType, + operation: 'load', + modelName: payload.model + })(getErrorString(err)) + ) + } else if (payload.engineType === 'ollama') { supervisor.sendProcess( 'broker', 'engine:action', @@ -359,6 +456,9 @@ function routeRemoteEngineCommand(payload: WsInvokeRequest<'engine:command'>): v case 'toggle': void toggleRemoteEngine(nodeId, engine, payload.engineType) break + case 'stop': + void supervisor.toggleEngineRemote(nodeId, engine, payload.engineType, true) + break case 'pullModel': if (payload.model) { void supervisor.pullModelRemote(nodeId, engine, payload.engineType, payload.model) @@ -374,13 +474,20 @@ function routeRemoteEngineCommand(payload: WsInvokeRequest<'engine:command'>): v case 'loadModel': case 'unloadModel': if (payload.model) { - void supervisor.modelActionRemote( - nodeId, - engine, - payload.engineType, - payload.model, - payload.command - ) + const startedAt = Date.now() + void supervisor + .modelActionRemote( + nodeId, + engine, + payload.engineType, + payload.model, + payload.command + ) + .then(result => { + if (payload.command === 'loadModel') { + reportLoadStoppedElsewhere(payload, startedAt, result) + } + }) } break default: @@ -405,6 +512,7 @@ async function handleEngineCommand(payload?: WsInvokeRequest<'engine:command'>): // Remote peers: `nvpair-engine-manager` exposes `engine:remote-*` client methods // (cluster mTLS `ec` surface) for install/start/stop/pull and model ops. // Local-only ops are refused in {@link routeRemoteEngineCommand}. + noteStopRequest(payload) const selfId = getModularBridgeState().getSelfId() if (selfId && payload.nodeId && payload.nodeId !== selfId) { if (supervisor.hasProcess('broker')) { diff --git a/desktop/src/electron/service-bridge/model-catalog.ts b/desktop/src/electron/service-bridge/model-catalog.ts index efd24666b..60bd896bb 100644 --- a/desktop/src/electron/service-bridge/model-catalog.ts +++ b/desktop/src/electron/service-bridge/model-catalog.ts @@ -124,4 +124,5 @@ export async function getEngineHubModels( export function warmEngineHubs(): void { void getEngineHubModels('lm-studio') void getEngineHubModels('llama-cpp') + void getEngineHubModels('vllm') } diff --git a/desktop/src/electron/service-bridge/modular-state.ts b/desktop/src/electron/service-bridge/modular-state.ts index 0c7390bf7..e9fd0a31d 100644 --- a/desktop/src/electron/service-bridge/modular-state.ts +++ b/desktop/src/electron/service-bridge/modular-state.ts @@ -18,15 +18,19 @@ import type { Workload, WorkloadState } from '@/shared/types/workloads' import type { ClusterNode, Invite } from '@/shared/types/cluster' import type { WsInvokeResponse } from '@/shared/types/ws-channels' import { + MODULAR_ENGINE_LIFECYCLE_CALL_TIMEOUT_MS, MODULAR_NODE_INFO_SELF_HOST, MODULAR_SUPERSEDE_MIN_AGE_MS } from '@/shared/constants/modular-runtime' import { WorkloadStates } from '@/shared/constants/workloads' +import { EngineTypes, EnginesListingModelsWhenStopped } from '@/shared/constants/engines' import { + emptyEngineModels, emptyEngineStatus, engineManagerName, engineTypeFromManagerName, - isEngineType + isEngineType, + unsupportedEngineStatus } from '@/shared/utils/engines' import { engineProgressKey } from '@/shared/utils/engine-progress' import { workloadKey } from '@/shared/utils/workloads' @@ -77,10 +81,26 @@ interface EnginePresence { interface RemoteEngineFacts { installed: boolean running: boolean + starting: boolean healthy: boolean port: number } +/** + * The process status an engine-manager `EngineStatus` reports. `starting` covers + * a start in flight whoever began it (install-then-start, a settings restart, a + * restore at launch, the terminal interface), which for vLLM can take minutes. + */ +function factsProcessStatus(facts: { + installed: boolean + running: boolean + starting: boolean +}): EngineProcessStatus { + if (facts.running) return 'running' + if (facts.starting) return 'starting' + return facts.installed ? 'stopped' : 'not-installed' +} + interface ModularNode { // Canonical cross-domain node key = the backend's stable per-host UUID // (`shared/nodeid`). The broker carries it on `AvailableNode.hostUuid`, the @@ -156,6 +176,15 @@ interface ModularNode { // reported. Drives {@link ModelItem.status} `'loaded'`. Empty object when no // engine reports loaded state. Mirrors noderec.DirectoryNode.LoadedByEngine. loadedByEngine: Record + // Per-engine downloaded inventory of an engine that advertises only the model + // it is serving (vLLM), carried on `AvailableNode.downloadedByEngine`. Its + // {@link modelsByEngine} entry holds just the served name, for routing; this + // lists every model it can load, read even while it is stopped, so the model + // manager keeps its rows. Its {@link loadedByEngine} entry names models the + // same way. A present key with [] means the inventory was read and is empty; + // a missing key means it was not read. Mirrors + // noderec.DirectoryNode.DownloadedByEngine. + downloadedByEngine: Record engines: Record lastSeen: number } @@ -168,7 +197,8 @@ function emptyEngines(): Record { return { ollama: emptyPresence(), 'lm-studio': emptyPresence(), - 'llama-cpp': emptyPresence() + 'llama-cpp': emptyPresence(), + vllm: emptyPresence() } } @@ -404,8 +434,23 @@ const PENDING_OP_IDLE_TIMEOUT_MS = 90_000 // Vendor installers can be quiet for minutes; allow a bounded 30-minute window. const INSTALL_PENDING_OP_IDLE_TIMEOUT_MS = 30 * 60_000 +/** + * How long an optimistic status may stand without the backend resolving it. + * A start is resolved by `engine:state-changed` once the engine passes its + * readiness probe, and engines allow that probe minutes (vLLM 8, Ollama 10), + * so a start keeps its marker for the awaited lifecycle envelope that already + * outlasts every engine's readiness window. A shorter net shows a still-loading + * engine as stopped, and its toggle then sends another Start instead of Stop. + */ function pendingEngineOpIdleTimeoutMs(status: EngineProcessStatus): number { - return status === 'installing' ? INSTALL_PENDING_OP_IDLE_TIMEOUT_MS : PENDING_OP_IDLE_TIMEOUT_MS + switch (status) { + case 'installing': + return INSTALL_PENDING_OP_IDLE_TIMEOUT_MS + case 'starting': + return MODULAR_ENGINE_LIFECYCLE_CALL_TIMEOUT_MS + default: + return PENDING_OP_IDLE_TIMEOUT_MS + } } /** @@ -733,6 +778,7 @@ function parseProxyNode(params: JsonValue | undefined, engine: ProxyEngine): Mod // A proxy `node/*` event carries no loaded-model state either; the // broker-enriched value is preserved across the proxy merge (see mergeNode). loadedByEngine: {}, + downloadedByEngine: {}, nodeInfoPort: 0, nodeInfoUp: false, // A proxy `node/*` event carries no broker ranking; its `Addresses` are @@ -804,6 +850,9 @@ function parseBrokerNode(params: JsonValue | undefined): ModularNode | null { // `/v1/models`. `omitempty` on the wire — absent for a peer that reports // no loaded state, which stamps every model `idle`. loadedByEngine: parseModelsByEngine(obj.loadedByEngine), + // Broker `AvailableNode.downloadedByEngine`: what a vLLM peer has + // downloaded, reported even while it is stopped. `omitempty` on the wire. + downloadedByEngine: parseModelsByEngine(obj.downloadedByEngine), engines: emptyEngines(), lastSeen: normalizeLastSeen(numberValue(obj.lastSeen) || numberValue(obj.last_seen)) } @@ -855,6 +904,7 @@ function sameNode(left: ModularNode, right: ModularNode): boolean { // A residency-only change (a model loaded/evicted with the installed set // unchanged) must still count as a node change so the loaded dot updates. sameModelsByEngine(left.loadedByEngine, right.loadedByEngine) && + sameModelsByEngine(left.downloadedByEngine, right.downloadedByEngine) && sameEngines(left.engines, right.engines) && left.lastSeen === right.lastSeen ) @@ -890,7 +940,8 @@ class ModularBridgeState { private proxyPorts: Record = { ollama: 0, 'lm-studio': 0, - 'llama-cpp': 0 + 'llama-cpp': 0, + vllm: 0 } private selfId: string | null = null /** @@ -901,8 +952,14 @@ class ModularBridgeState { */ private engineManagerFacts = new Map< EngineType, - { installed: boolean; running: boolean; port: number } + { installed: boolean; running: boolean; starting: boolean; port: number } >() + /** + * The engines `nvpair-engine-manager` lists for this host. It leaves out an + * engine the host cannot run, such as vLLM on Windows, which is then + * reported as `unsupported`. Null until the first full listing. + */ + private localEngineList: ReadonlySet | null = null /** * Local model lists pulled from `nvpair-engine-manager`'s `list_models` action by * the supervisor. Cached so `engines:get-initial` can include them and survive @@ -958,6 +1015,8 @@ class ModularBridgeState { * evicted. */ private remoteEngineFacts = new Map() + /** The engines each peer's engine manager lists, as {@link localEngineList} is for this host. */ + private remoteEngineLists = new Map>() /** * The model of an in-flight remote pull keyed by {@link remoteOpKey}. The * backend's `engine:remote-progress` frames carry **no** `model` field @@ -1100,12 +1159,21 @@ class ModularBridgeState { models: cached.map(item => ({ ...item, status: loaded.has(item.name) ? 'loaded' : 'idle' - })) + })), + inventoryKnown: true }) } } } + if (selfId && this.localEngineList) { + for (const engineType of EngineTypes) { + if (!this.localEngineList.has(engineType)) { + statuses.push(unsupportedEngineStatus(selfId, engineType)) + } + } + } + // Remote nodes: per proxy-engine status + models from discovery. Status // prefers the authoritative `engine:remote-get-installed` facts (mTLS) // over the inferred presence, and is omitted entirely when neither is @@ -1327,6 +1395,27 @@ class ModularBridgeState { } } + /** + * Record a full `engine:get-installed` listing: each engine's status, and + * which engines the host can run at all. One left out is `unsupported`. + */ + applyEngineManagerList(result: JsonValue | undefined): void { + const list = objectValue(result)?.engines + if (!Array.isArray(list)) return + const listed = new Set() + for (const entry of list) { + this.applyEngineManagerStatus(entry) + const engineType = engineTypeFromManagerName(stringValue(objectValue(entry)?.engine)) + if (engineType) listed.add(engineType) + } + // An empty listing says nothing about what the host can run. + if (listed.size === 0) return + this.localEngineList = listed + for (const engineType of EngineTypes) { + if (!listed.has(engineType)) this.emitLocalEngineStatus(engineType) + } + } + /** * Record a `nvpair-engine-manager` `engine:state-changed` (or hydrated * `engine:get-installed` entry). Facts are kept regardless of whether the @@ -1341,6 +1430,7 @@ class ModularBridgeState { this.engineManagerFacts.set(engineType, { installed: booleanValue(obj.installed), running: booleanValue(obj.running), + starting: booleanValue(obj.starting), port: numberValue(obj.port) }) // A fresh authoritative state is the resolution of whatever op was in @@ -1687,6 +1777,8 @@ class ModularBridgeState { nodeId: string, engine: ProxyEngine ): EngineStatusData | null { + const listed = this.remoteEngineLists.get(nodeId) + if (listed && !listed.has(engine)) return unsupportedEngineStatus(nodeId, engine) const node = this.nodes.get(nodeId) const presence = node?.engines[engine] const facts = this.remoteEngineFacts.get(this.remoteOpKey(nodeId, engine)) @@ -1697,11 +1789,7 @@ class ModularBridgeState { base = { engineType: engine, nodeId, - processStatus: facts.running - ? 'running' - : facts.installed - ? 'stopped' - : 'not-installed', + processStatus: factsProcessStatus(facts), // Only the peer engine-manager's own port is a real engine port; // the discovery presence port is the proxy, so never fall back to // it for the engine port. @@ -1753,11 +1841,14 @@ class ModularBridgeState { this.remoteEngineFacts.set(this.remoteOpKey(nodeId, engineType), { installed: booleanValue(engineObj.installed), running: booleanValue(engineObj.running), + starting: booleanValue(engineObj.starting), healthy: booleanValue(engineObj.healthy), port: numberValue(engineObj.port) }) seen.add(engineType) } + // A peer leaves out an engine it cannot run. An empty listing says nothing. + if (seen.size > 0) this.remoteEngineLists.set(nodeId, seen) for (const engine of PROXY_ENGINES) { if (!seen.has(engine)) { this.remoteEngineFacts.delete(this.remoteOpKey(nodeId, engine)) @@ -1782,6 +1873,7 @@ class ModularBridgeState { this.remoteEngineFacts.set(this.remoteOpKey(nodeId, engineType), { installed: booleanValue(obj.installed), running: booleanValue(obj.running), + starting: booleanValue(obj.starting), healthy: booleanValue(obj.healthy), port: numberValue(obj.port) }) @@ -1790,6 +1882,7 @@ class ModularBridgeState { /** Forget a node's authoritative remote-engine facts (on eviction/leave). */ private dropRemoteEngineFacts(nodeId: string): void { + this.remoteEngineLists.delete(nodeId) for (const engine of PROXY_ENGINES) { this.remoteEngineFacts.delete(this.remoteOpKey(nodeId, engine)) this.remotePullModels.delete(this.remoteOpKey(nodeId, engine)) @@ -1839,7 +1932,7 @@ class ModularBridgeState { emitBridgePush('engines:state-changed', { nodeId, engineType, - models: { engineType, nodeId, models: items } + models: { engineType, nodeId, models: items, inventoryKnown: true } }) } @@ -1942,6 +2035,15 @@ class ModularBridgeState { * noderec.DirectoryNode.EngineModels. */ private modelsForEngine(node: ModularNode, engine: ProxyEngine): string[] { + // An engine that advertises only the model it is serving (vLLM) reports + // its downloaded inventory apart: that is the list to manage, plus + // anything loaded that is not in it, such as a model served from a local + // path. Its advertised entry is only the served name, for routing. + const downloaded = node.downloadedByEngine[engineManagerName(engine)] + if (downloaded !== undefined) { + const loaded = node.loadedByEngine[engineManagerName(engine)] ?? [] + return [...downloaded, ...loaded.filter(name => !downloaded.includes(name))] + } // Per-engine attribution is authoritative when the node carries any: // return exactly this engine's models so a dual-engine node attributes // each model to the engine that serves it, and a missing key means this @@ -1966,7 +2068,14 @@ class ModularBridgeState { nodeId: node.id, models: this.modelsForEngine(node, engine).map(name => modelItem(name, loaded.has(name)) - ) + ), + // An engine whose list is read from disk is listed while stopped, so + // an empty list reads as "nothing downloaded". That is true only when + // its inventory was actually read; a peer that could not be read, or + // whose list never arrived, is unknown. + inventoryKnown: + !EnginesListingModelsWhenStopped.includes(engine) || + engineManagerName(engine) in node.downloadedByEngine } } @@ -2003,11 +2112,12 @@ class ModularBridgeState { models: cached.map(item => ({ ...item, status: loaded.has(item.name) ? 'loaded' : 'idle' - })) + })), + inventoryKnown: true } } } - if (!node) return { engineType: engine, nodeId, models: [] } + if (!node) return emptyEngineModels(nodeId, engine) return this.toEngineModels(node, engine) } @@ -2028,6 +2138,9 @@ class ModularBridgeState { ): EngineStatusData | undefined { const nodeId = this.selfId if (!nodeId) return undefined + if (this.localEngineList && !this.localEngineList.has(engineType)) { + return unsupportedEngineStatus(nodeId, engineType) + } const facts = this.engineManagerFacts.get(engineType) // An in-flight lifecycle op outranks the fact-derived status so the UI @@ -2048,11 +2161,7 @@ class ModularBridgeState { return { engineType, nodeId, - processStatus: facts.running - ? 'running' - : facts.installed - ? 'stopped' - : 'not-installed', + processStatus: factsProcessStatus(facts), // The engine-manager reports the manifest's configured port for a // pinned-port engine (for example managed llama.cpp on 8081 behind // its 8080 proxy facade) whether it is running or merely installed, @@ -2587,6 +2696,7 @@ class ModularBridgeState { models: existing.models, modelsByEngine: existing.modelsByEngine, loadedByEngine: existing.loadedByEngine, + downloadedByEngine: existing.downloadedByEngine, engines: setEngine(existing.engines, engine, next.engines[engine]), lastSeen: Math.max(existing.lastSeen, next.lastSeen) } diff --git a/desktop/src/electron/service-bridge/modular-supervisor.ts b/desktop/src/electron/service-bridge/modular-supervisor.ts index 519f7b91a..77d5532a1 100644 --- a/desktop/src/electron/service-bridge/modular-supervisor.ts +++ b/desktop/src/electron/service-bridge/modular-supervisor.ts @@ -12,6 +12,7 @@ import { type JsonRpcNotification, type JsonValue } from './json-rpc-subprocess' +import { EnginesListingModelsWhenStopped } from '@/shared/constants/engines' import { createStructuredLogger } from '@/shared/utils/log' import getErrorString from '@/shared/utils/get-error-string' import { currentPlatform } from '@/shared/utils/platform' @@ -1080,8 +1081,8 @@ class ModularSupervisor { const result = await this.callProcess('broker', 'engine:get-installed') const obj = objectValue(result) if (!obj || !Array.isArray(obj.engines)) return + getModularBridgeState().applyEngineManagerList(result) for (const engine of obj.engines) { - getModularBridgeState().applyEngineManagerStatus(engine) this.refreshManagedEngineModels(engine) } } catch (err) { @@ -1908,7 +1909,8 @@ class ModularSupervisor { * message distinguishes the two cases exactly as the local path does. The * error is stamped with the *peer's* nodeId (not our selfId) plus * `engineType`/`modelName` so `pending-actions.store` clears the right row's - * spinner instead of leaving it to the safety-net timeout. + * spinner instead of leaving it to the safety-net timeout. Resolves with the + * peer's reply, or undefined once a failure has been reported. */ async modelActionRemote( nodeId: string, @@ -1916,7 +1918,7 @@ class ModularSupervisor { engineType: EngineType, model: string, command: 'loadModel' | 'unloadModel' | 'deleteModel' - ): Promise { + ): Promise { const method = command === 'loadModel' ? 'engine:remote-load-model' @@ -1924,13 +1926,14 @@ class ModularSupervisor { ? 'engine:remote-unload-model' : 'engine:remote-delete-model' try { - await this.callProcess( + const result = await this.callProcess( 'broker', method, { node: nodeId, engine, model }, PULL_TIMEOUT_MS ) await this.refreshRemoteEngineStatus(nodeId) + return result } catch (err) { const verb = command === 'loadModel' ? 'load' : command === 'unloadModel' ? 'unload' : 'delete' @@ -1945,6 +1948,7 @@ class ModularSupervisor { { nodeId, engineType, operation: verb, modelName: model } ) if (restartFailed) await this.refreshRemoteEngineStatus(nodeId) + return undefined } } @@ -1989,7 +1993,11 @@ class ModularSupervisor { const engineType = getModularBridgeState().modelPullTarget(engine) if (!engineType) return - if (!booleanValue(obj.running)) { + // An engine whose list is on disk keeps it while stopped: read it like a + // running engine's instead of clearing it. + // vLLM, the one engine whose Load restarts it, is among these, so a model + // switch passing through stopped never clears its list. + if (!booleanValue(obj.running) && !EnginesListingModelsWhenStopped.includes(engineType)) { this.beginModelRefresh(engineType) this.stoppedModelEngines.add(engineType) this.stoppedModelSentinels.add(engineType) diff --git a/desktop/src/electron/service-bridge/proxy-engines.ts b/desktop/src/electron/service-bridge/proxy-engines.ts index 447246af0..3be8fd0ec 100644 --- a/desktop/src/electron/service-bridge/proxy-engines.ts +++ b/desktop/src/electron/service-bridge/proxy-engines.ts @@ -1,62 +1,42 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -import type { EngineManagerName, EngineType } from '@/shared/types/engines' +import { EngineManagerNames } from '@/shared/constants/engines' +import type { EngineType } from '@/shared/types/engines' +import { engineTypeFromManagerName } from '@/shared/utils/engines' /** * Engines currently consumed through the broker's proxy plane. Keep this set * separate from the complete engine registry so an engine can be declared * before every proxy-state consumer is ready to handle it. */ -export type ProxyEngine = Extract -export type ProxyNodeSource = 'ollama-proxy' | 'lmstudio-proxy' | 'llamacpp-proxy' +export type ProxyEngine = Extract +/** A proxy's source name: the engine-manager's name for the engine, plus -proxy. */ +export type ProxyNodeSource = `${(typeof EngineManagerNames)[ProxyEngine]}-proxy` -interface ProxyEngineIdentity { - managerName: EngineManagerName - source: ProxyNodeSource -} - -const PROXY_ENGINE_IDENTITIES: Record = { - ollama: { - managerName: 'ollama', - source: 'ollama-proxy' - }, - 'lm-studio': { - managerName: 'lmstudio', - source: 'lmstudio-proxy' - }, - 'llama-cpp': { - managerName: 'llamacpp', - source: 'llamacpp-proxy' - } -} +const PROXY_SOURCE_SUFFIX = '-proxy' -export const PROXY_ENGINES: readonly ProxyEngine[] = ['ollama', 'lm-studio', 'llama-cpp'] +export const PROXY_ENGINES: readonly ProxyEngine[] = ['ollama', 'lm-studio', 'llama-cpp', 'vllm'] -export const PROXY_NODE_SOURCES: readonly ProxyNodeSource[] = PROXY_ENGINES.map( - engine => PROXY_ENGINE_IDENTITIES[engine].source -) +export const PROXY_NODE_SOURCES: readonly ProxyNodeSource[] = + PROXY_ENGINES.map(proxySourceForEngine) export function isProxyEngine(engine: EngineType): engine is ProxyEngine { return PROXY_ENGINES.some(candidate => candidate === engine) } export function proxySourceForEngine(engine: ProxyEngine): ProxyNodeSource { - return PROXY_ENGINE_IDENTITIES[engine].source + return `${EngineManagerNames[engine]}${PROXY_SOURCE_SUFFIX}` } export function proxyEngineFromSource(source: ProxyNodeSource): ProxyEngine export function proxyEngineFromSource(source: string): ProxyEngine | null export function proxyEngineFromSource(source: string): ProxyEngine | null { - for (const engine of PROXY_ENGINES) { - if (PROXY_ENGINE_IDENTITIES[engine].source === source) return engine - } - return null + if (!source.endsWith(PROXY_SOURCE_SUFFIX)) return null + return proxyEngineFromManagerId(source.slice(0, -PROXY_SOURCE_SUFFIX.length)) } export function proxyEngineFromManagerId(managerName: string): ProxyEngine | null { - for (const engine of PROXY_ENGINES) { - if (PROXY_ENGINE_IDENTITIES[engine].managerName === managerName) return engine - } - return null + const engine = engineTypeFromManagerName(managerName) + return engine !== null && isProxyEngine(engine) ? engine : null } diff --git a/desktop/src/shared/constants/engines.ts b/desktop/src/shared/constants/engines.ts index 15970f2ce..ddb1998dd 100644 --- a/desktop/src/shared/constants/engines.ts +++ b/desktop/src/shared/constants/engines.ts @@ -5,11 +5,18 @@ import { EngineType, ModelExpiry } from '@/shared/types/engines' // The engines `nvpair-engine-manager` ships a manifest for, and therefore the // only ones PAIR can install, run or route to. -export const EngineTypes = ['ollama', 'lm-studio', 'llama-cpp'] as const +export const EngineTypes = ['ollama', 'lm-studio', 'llama-cpp', 'vllm'] as const // Kept as a distinct export so a future engine can ship behind it rather than // appearing the moment its type exists. -export const EnabledEngineTypes: EngineType[] = ['ollama', 'lm-studio', 'llama-cpp'] +export const EnabledEngineTypes: EngineType[] = ['ollama', 'lm-studio', 'llama-cpp', 'vllm'] + +/** + * Engines whose model list is read from disk, so it is available while the engine + * is stopped. For every other engine the list comes from the running engine and is + * cleared when it stops. The `modelOpsWhenStopped` capability must agree. + */ +export const EnginesListingModelsWhenStopped: readonly EngineType[] = ['vllm'] /** * How `nvpair-engine-manager` spells each engine on the wire. This is the @@ -19,7 +26,8 @@ export const EnabledEngineTypes: EngineType[] = ['ollama', 'lm-studio', 'llama-c export const EngineManagerNames = { ollama: 'ollama', 'lm-studio': 'lmstudio', - 'llama-cpp': 'llamacpp' + 'llama-cpp': 'llamacpp', + vllm: 'vllm' } as const satisfies Record export const EngineSources = ['bundled', 'detected', 'installed'] as const @@ -27,7 +35,8 @@ export const EngineSources = ['bundled', 'detected', 'installed'] as const export const EngineDisplayNames: Record = { ollama: 'Ollama', 'lm-studio': 'LM Studio', - 'llama-cpp': 'llama.cpp' + 'llama-cpp': 'llama.cpp', + vllm: 'vLLM' } as const /** Default docs/install URLs for built-in backends. Single source of truth for UI and adapter buildInfo(). */ @@ -37,6 +46,10 @@ export const EngineDefaultLinks: Record fireCommand(transport, { command: 'toggle', engineType, nodeId }), + stop: (engineType, nodeId) => + fireCommand(transport, { command: 'stop', engineType, nodeId }), install: (engineType, nodeId) => fireCommand(transport, { command: 'install', engineType, nodeId }), uninstall: (engineType, nodeId) => diff --git a/desktop/src/ui/assets/engine-icons/vllm.svg b/desktop/src/ui/assets/engine-icons/vllm.svg new file mode 100644 index 000000000..47213b808 --- /dev/null +++ b/desktop/src/ui/assets/engine-icons/vllm.svg @@ -0,0 +1,10 @@ + + + + + + + + + + diff --git a/desktop/src/ui/components/BackendRow/BackendHeader.tsx b/desktop/src/ui/components/BackendRow/BackendHeader.tsx index 58abc1dee..dabaee576 100644 --- a/desktop/src/ui/components/BackendRow/BackendHeader.tsx +++ b/desktop/src/ui/components/BackendRow/BackendHeader.tsx @@ -25,6 +25,7 @@ export function BackendHeader({ isUnavailable, disabled, onToggle, + onStop, onExpand, expanded, isLocalNode, @@ -37,6 +38,8 @@ export function BackendHeader({ isUnavailable: boolean disabled: boolean onToggle: () => void + /** Stop an engine that is still starting; its switch stays on and usable. */ + onStop: () => void onExpand: () => void expanded: boolean isLocalNode: boolean @@ -198,6 +201,24 @@ export function BackendHeader({ aria-label={fullStatus} /> + {/* A start can take minutes (vLLM loads for up to + eight), so it can be stopped while it runs. */} + {backend.processStatus === 'starting' && ( + + e.stopPropagation()} + aria-label={`Stop ${backend.displayName}`} + title={`Stop ${backend.displayName}`} + /> + + )} ) diff --git a/desktop/src/ui/components/BackendRow/BackendRow.tsx b/desktop/src/ui/components/BackendRow/BackendRow.tsx index 8af69dce8..e7e2598cb 100644 --- a/desktop/src/ui/components/BackendRow/BackendRow.tsx +++ b/desktop/src/ui/components/BackendRow/BackendRow.tsx @@ -45,6 +45,12 @@ function syntheticPendingStatus( return 'uninstalling' case 'toggle': return current === 'running' ? 'stopping' : 'starting' + case 'stop': + return 'stopping' + case 'loadModel': + // An engine that restarts to serve the model reports stopped on the + // way; it is starting again, not turned off. + return current === 'running' ? null : 'starting' default: return null } @@ -151,6 +157,10 @@ export function BackendRow({ window.pairApi.engines.toggle(backend.type, nodeId) }, [backend.type, nodeId]) + const handleStop = useCallback(() => { + window.pairApi.engines.stop(backend.type, nodeId) + }, [backend.type, nodeId]) + const handleInstall = useCallback(() => { window.pairApi.engines.install(backend.type, nodeId) }, [backend.type, nodeId]) @@ -217,6 +227,7 @@ export function BackendRow({ isUnavailable={isUnavailable} disabled={controlsDisabled} onToggle={handleToggle} + onStop={handleStop} onExpand={handleExpand} expanded={expanded} isLocalNode={isLocalNode} diff --git a/desktop/src/ui/components/EngineIcon.tsx b/desktop/src/ui/components/EngineIcon.tsx index bb1618a0e..326d4dbae 100644 --- a/desktop/src/ui/components/EngineIcon.tsx +++ b/desktop/src/ui/components/EngineIcon.tsx @@ -6,6 +6,7 @@ import { type EngineType } from '@/shared/types/engines' import ollamaIcon from '@/ui/assets/engine-icons/ollama.png?inline' import lmStudioIcon from '@/ui/assets/engine-icons/lm-studio.png?inline' import llamaCppIcon from '@/ui/assets/engine-icons/llama-cpp.svg?inline' +import vllmIcon from '@/ui/assets/engine-icons/vllm.svg?inline' export default function EngineIcon({ type, size = 32 }: { type: EngineType; size?: number }) { const dimension = `${size}px` @@ -48,5 +49,13 @@ export default function EngineIcon({ type, size = 32 }: { type: EngineType; size ) } + if (type === 'vllm') { + return ( +
+ vLLM +
+ ) + } + return null } diff --git a/desktop/src/ui/components/ModelHub/EngineHubListRow.tsx b/desktop/src/ui/components/ModelHub/EngineHubListRow.tsx index 10c0ccdb2..fa4adf512 100644 --- a/desktop/src/ui/components/ModelHub/EngineHubListRow.tsx +++ b/desktop/src/ui/components/ModelHub/EngineHubListRow.tsx @@ -57,6 +57,14 @@ const EngineHubListRow = memo(function EngineHubListRow({ {displayName} + {/* The name drops its author, so two repositories that + share a name are told apart here. */} + {!!model.author.trim() && ( + + by {model.author.trim()} + + )} + {!!model.size && ( {formatBytes(model.size)} diff --git a/desktop/src/ui/components/ModelHub/ModelHubContent.tsx b/desktop/src/ui/components/ModelHub/ModelHubContent.tsx index 354709b75..8ee2c5be9 100644 --- a/desktop/src/ui/components/ModelHub/ModelHubContent.tsx +++ b/desktop/src/ui/components/ModelHub/ModelHubContent.tsx @@ -12,6 +12,7 @@ import { mergeModelHubResults, searchEngineHub } from '@/ui/utils/model-hub-sear import { resolveStoredSort, writeStoredSort } from '@/ui/utils/model-hub-content-storage' import { isHubEntryDownloaded } from '@/ui/utils/match-downloaded-model' import { EngineType } from '@/shared/types/engines' +import { EngineCapabilities } from '@/ui/constants/engine-capabilities' import type { ModelItem } from '@/ui/types/engine-info' import getErrorString from '@/shared/utils/get-error-string' @@ -184,7 +185,7 @@ export const ModelHubContent = ({ async (nextQuery: string) => { const normalized = nextQuery.trim() setQuery(nextQuery) - if (engine !== 'llama-cpp' || normalized.length === 0) return + if (!engine || !EngineCapabilities[engine].hubSearch || normalized.length === 0) return const generation = ++searchGenRef.current submittedQueryRef.current = normalized diff --git a/desktop/src/ui/components/ModelManager/ModelManager.tsx b/desktop/src/ui/components/ModelManager/ModelManager.tsx index e63dacf29..648448346 100644 --- a/desktop/src/ui/components/ModelManager/ModelManager.tsx +++ b/desktop/src/ui/components/ModelManager/ModelManager.tsx @@ -165,8 +165,8 @@ export function ModelManager({ backend, nodeId }: { backend: BackendInfo; nodeId ) const isRunning = useMemo(() => backend.processStatus === 'running', [backend.processStatus]) const supportsSearch = useMemo( - () => !hasModelSearchOnlyWhenRunning || isRunning, - [hasModelSearchOnlyWhenRunning, isRunning] + () => caps !== undefined && (!hasModelSearchOnlyWhenRunning || isRunning), + [caps, hasModelSearchOnlyWhenRunning, isRunning] ) const handleDownload = useCallback( @@ -202,11 +202,16 @@ export function ModelManager({ backend, nodeId }: { backend: BackendInfo; nodeId <> {(isRunning || modelOpsWhenStopped) && models.length === 0 && - incomingSyncs.length === 0 && ( + incomingSyncs.length === 0 && + (backend.modelsKnown ? ( No model/weight files downloaded - )} + ) : ( + + Downloaded models unavailable + + ))} {!modelOpsWhenStopped && models.length === 0 && incomingSyncs.length === 0 && diff --git a/desktop/src/ui/components/ModelManager/ModelRow.tsx b/desktop/src/ui/components/ModelManager/ModelRow.tsx index b571dffc3..80d77a7c9 100644 --- a/desktop/src/ui/components/ModelManager/ModelRow.tsx +++ b/desktop/src/ui/components/ModelManager/ModelRow.tsx @@ -49,7 +49,11 @@ function ModelRowInner({ items.push({ id: 'load', children: 'Load', - disabled: busy || isLoaded || !model.downloaded || !isRunning + disabled: + busy || + isLoaded || + !model.downloaded || + (!isRunning && !capabilities.modelOpsWhenStopped) }) if (capabilities.hasEject) { items.push({ @@ -63,12 +67,19 @@ function ModelRowInner({ }) } if (capabilities.hasDeleteModel) { - items.push({ id: 'delete', children: 'Delete', danger: true, disabled: busy }) + items.push({ + id: 'delete', + children: 'Delete', + danger: true, + disabled: busy || (isLoaded && !!capabilities.deleteBlockedWhileLoaded) + }) } return items }, [ capabilities.hasEject, capabilities.hasDeleteModel, + capabilities.deleteBlockedWhileLoaded, + capabilities.modelOpsWhenStopped, isLoaded, isRunning, model.downloaded, diff --git a/desktop/src/ui/components/NodeList/NodeEngineSettings.tsx b/desktop/src/ui/components/NodeList/NodeEngineSettings.tsx index 7581be681..37ef126ae 100644 --- a/desktop/src/ui/components/NodeList/NodeEngineSettings.tsx +++ b/desktop/src/ui/components/NodeList/NodeEngineSettings.tsx @@ -14,6 +14,7 @@ import { EngineDefaultLinks, EngineDisplayNames } from '@/shared/constants/engines' +import { getEnginesForNode } from '@/ui/utils/get-engines-for-node' import { BackendRow } from '@/ui/components/BackendRow/BackendRow' import { InlineErrorBanner } from '@/ui/components/InlineErrorBanner' @@ -75,35 +76,39 @@ export default function NodeEngineSettings({ nodeId }: { nodeId: string }) { void engineFingerprint void modelsFingerprint - const allBackends: { backend: BackendInfo; statusKnown: boolean }[] = EnabledEngineTypes.map( - type => { - const statusStore = useEngineStatusStore.getState() - // Distinguish "the service reported this engine" from "we synthesized - // a placeholder". The placeholder's processStatus is 'initializing'; - // without this flag a never-reported remote engine spins forever. - const statusKnown = statusStore.statusByNode.get(nodeId)?.has(type) ?? false - const status = statusStore.getStatus(nodeId, type) - const models = useEngineModelsStore.getState().getModels(nodeId, type) - const updateAvailable = useEngineUpdateAvailableStore.getState().getInfo(nodeId, type) + // One row per enabled engine, in order; an engine the node cannot run, + // such as vLLM on Windows, gets no row. + const statusStore = useEngineStatusStore.getState() + const allBackends: { backend: BackendInfo; statusKnown: boolean }[] = getEnginesForNode( + nodeId, + statusStore.statusByNode + ).map(status => { + const type = status.engineType + // Distinguish "the service reported this engine" from "we synthesized + // a placeholder". The placeholder's processStatus is 'initializing'; + // without this flag a never-reported remote engine spins forever. + const statusKnown = statusStore.statusByNode.get(nodeId)?.has(type) ?? false + const models = useEngineModelsStore.getState().getModels(nodeId, type) + const updateAvailable = useEngineUpdateAvailableStore.getState().getInfo(nodeId, type) - return { - backend: { - type, - displayName: EngineDisplayNames[type], - source: 'detected', - processStatus: status.processStatus, - port: status.enginePort, - proxyPort: status.proxyPort, - models: models.models, - docsUrl: EngineDefaultLinks[type]?.docsUrl, - installUrl: EngineDefaultLinks[type]?.installUrl, - installedVersion: status.installedVersion, - updateAvailable - }, - statusKnown - } + return { + backend: { + type, + displayName: EngineDisplayNames[type], + source: 'detected', + processStatus: status.processStatus, + port: status.enginePort, + proxyPort: status.proxyPort, + models: models.models, + modelsKnown: models.inventoryKnown, + docsUrl: EngineDefaultLinks[type]?.docsUrl, + installUrl: EngineDefaultLinks[type]?.installUrl, + installedVersion: status.installedVersion, + updateAvailable + }, + statusKnown } - ) + }) return ( diff --git a/desktop/src/ui/components/NodeList/NodeEnginesInline.tsx b/desktop/src/ui/components/NodeList/NodeEnginesInline.tsx index 8dfafe4ff..35d5473c2 100644 --- a/desktop/src/ui/components/NodeList/NodeEnginesInline.tsx +++ b/desktop/src/ui/components/NodeList/NodeEnginesInline.tsx @@ -18,15 +18,11 @@ export default function NodeEnginesInline({ nodeId }: { nodeId: string }) { const isDisconnected = useNodesStore(s => isRemote && s.nodes.get(nodeId)?.status === 'offline') // Re-render when a lifecycle command for this node begins/clears; the // per-engine pending state is read below via getState(). - const lifecyclePendingFingerprint = usePendingActionsStore(state => { - const parts: string[] = [] - for (const [key, p] of state.pending) { - if (p.nodeId === nodeId && key.endsWith(':lifecycle')) { - parts.push(`${p.engineType}:${p.action}`) - } - } - return parts.join('|') - }) + const lifecyclePendingFingerprint = usePendingActionsStore(state => + EnabledEngineTypes.map( + type => `${type}:${state.getLifecyclePending(nodeId, type) ?? ''}` + ).join('|') + ) void lifecyclePendingFingerprint const selfEngines = useMemo( () => getEnginesForNode(nodeId ?? '', statusByNode, EnabledEngineTypes), @@ -64,6 +60,13 @@ export default function NodeEnginesInline({ nodeId }: { nodeId: string }) { [nodeId] ) + const handleStop = useCallback( + (engineType: EngineType) => { + window.pairApi.engines.stop(engineType, nodeId) + }, + [nodeId] + ) + const handleInstall = useCallback( (engineType: EngineType) => { window.pairApi.engines.install(engineType, nodeId) @@ -74,9 +77,15 @@ export default function NodeEnginesInline({ nodeId }: { nodeId: string }) { return ( <> {allBackends.map(b => { - const pending = Boolean( - usePendingActionsStore.getState().getLifecyclePending(nodeId, b.type) - ) + const pendingAction = usePendingActionsStore + .getState() + .getLifecyclePending(nodeId, b.type) + const pending = Boolean(pendingAction) + // A start can take minutes, so it stays stoppable while it runs: + // a Start in progress, or a Load that restarts the engine. + const stoppable = + !(isRemote && isDisconnected) && + (b.status === 'starting' || pendingAction === 'loadModel') const isTransitioning = pending || b.status === 'installing' || @@ -115,7 +124,15 @@ export default function NodeEnginesInline({ nodeId }: { nodeId: string }) { className="shrink-0 no-drag-elements pair-engines-inline" > - {isTransitioning ? ( + {stoppable ? ( + handleStop(b.type)} + title={`Stop ${b.name}`} + aria-label={`Stop ${b.name}`} + /> + ) : isTransitioning ? ( Your downloaded models (for example ~/.ollama, ~/.llamacpp - , and ~/.lmstudio/models) are not deleted, and neither - are engines you installed yourself. + , ~/.lmstudio/models, and ~/.cache/huggingface) are{' '} + not deleted, and neither are engines you installed yourself. If an engine {APP_DISPLAY_NAME} installed cannot be removed, the reset stops before diff --git a/desktop/src/ui/components/Welcome/WelcomeModal.tsx b/desktop/src/ui/components/Welcome/WelcomeModal.tsx index 101da2471..4d394e6d7 100644 --- a/desktop/src/ui/components/Welcome/WelcomeModal.tsx +++ b/desktop/src/ui/components/Welcome/WelcomeModal.tsx @@ -51,7 +51,15 @@ export function WelcomeModal({ open, onOpenChange, selfId }: WelcomeModalProps) const failedTargets = useRef(new Set()) const knownErrors = useRef(new Map()) - const candidates = useMemo(() => getWelcomeEngineCandidates(window.windowApi.platform), []) + // The engine manager has the last word on what this machine can run, so an + // engine it leaves out is not offered here. + const candidates = useMemo( + () => + getWelcomeEngineCandidates(window.windowApi.platform).filter( + t => statusByNode.get(selfId)?.get(t)?.processStatus !== 'unsupported' + ), + [statusByNode, selfId] + ) const allEnginesInstalled = useMemo( () => areWelcomeEnginesInstalled(statusByNode, selfId, candidates), [candidates, selfId, statusByNode] diff --git a/desktop/src/ui/constants/engine-capabilities.ts b/desktop/src/ui/constants/engine-capabilities.ts index 2626a2e42..0816f37c2 100644 --- a/desktop/src/ui/constants/engine-capabilities.ts +++ b/desktop/src/ui/constants/engine-capabilities.ts @@ -57,9 +57,42 @@ export const EngineCapabilities: Record = { modelOpsWhenStopped: false, // The router removes cached models in-process via DELETE /models. hasDeleteModel: true, + hubSearch: true, engineHub: { label: 'llama.cpp', url: 'https://huggingface.co/models?library=gguf' } + }, + vllm: { + hasExpiry: false, + // vLLM serves one model per process and cannot drop it while running, so + // Eject stops the engine, which frees all of its GPU memory. Load starts + // it again on the chosen model. + hasEject: true, + // Linux (NVIDIA) and Apple Silicon macOS (vLLM-Metal). vLLM has no + // native Windows path. + hasInstall: ['linux', 'darwin'], + hasEnginePort: true, + hasInstallPath: false, + hasProxyWebUI: false, + hasPreferredNode: false, + hasCrashAlert: false, + // The cache is on disk, so the model list, Add model, Delete and Load all + // work while the engine is stopped (Load starts it). + hasModelSearchOnlyWhenRunning: false, + modelOpsWhenStopped: true, + // delete_model removes a model from the shared Hugging Face cache. The + // loaded model is refused, so Delete is disabled on it until Eject. + // Deleting the model an ejected or stopped vLLM is set to serve leaves it + // with no model selected until Load picks one. + hasDeleteModel: true, + deleteBlockedWhileLoaded: true, + // Load writes the model into the launch arguments and restarts vLLM on it. + loadRestartsEngine: true, + hubSearch: true, + engineHub: { + label: 'vLLM', + url: 'https://huggingface.co/models?pipeline_tag=text-generation&library=safetensors&sort=trending' + } } } diff --git a/desktop/src/ui/constants/welcome.ts b/desktop/src/ui/constants/welcome.ts index 4db425f71..d235f8d03 100644 --- a/desktop/src/ui/constants/welcome.ts +++ b/desktop/src/ui/constants/welcome.ts @@ -14,7 +14,8 @@ export const WELCOME_STEP_SUB_HEADINGS = ['', 'You can update later by clicking export const WELCOME_ENGINE_DEFAULT_SELECTED: Record = { ollama: true, 'lm-studio': true, - 'llama-cpp': true + 'llama-cpp': true, + vllm: true } export function getWelcomeEngineCandidates(os: PlatformDisplayName): EngineType[] { diff --git a/desktop/src/ui/stores/pending-actions.store.ts b/desktop/src/ui/stores/pending-actions.store.ts index 4b1e1a294..115eea22f 100644 --- a/desktop/src/ui/stores/pending-actions.store.ts +++ b/desktop/src/ui/stores/pending-actions.store.ts @@ -28,6 +28,7 @@ * state. */ import { create } from 'zustand' +import { MODEL_SWITCH_TIMEOUT_MS } from '@/shared/constants/modular-runtime' import { EngineCapabilities } from '@/ui/constants/engine-capabilities' import type { EngineCommandPayload, EngineCommandType } from '@/shared/types/engine-api' import type { EngineProcessStatus, EngineType } from '@/shared/types/engines' @@ -39,6 +40,8 @@ import { useEngineStatusStore } from '@/ui/stores/engine-status.store' const PENDING_TIMEOUT_MS = 60_000 /** Keep Ollama Load locked one safety-net interval beyond the longest (11m) backend path. */ const OLLAMA_LOAD_PENDING_TIMEOUT_MS = 12 * 60_000 +/** vLLM Load restarts the engine on the new model; keep the row locked as long as the app waits for it. */ +const VLLM_LOAD_PENDING_TIMEOUT_MS = MODEL_SWITCH_TIMEOUT_MS /** * Safety net for a delete on an engine whose backend restarts to pick the * deletion up (`EngineCaps.restartsOnModelDelete` / the manifest's @@ -56,6 +59,7 @@ const SWEEP_INTERVAL_MS = 1_000 /** Commands that flip an engine's lifecycle -- only one at a time per engine. */ const LIFECYCLE_COMMANDS: ReadonlySet = new Set([ 'toggle', + 'stop', 'install', 'uninstall', 'update' @@ -110,6 +114,9 @@ function timeoutFor(command: EngineCommandType, engineType: EngineType): number if (command === 'loadModel' && engineType === 'ollama') { return OLLAMA_LOAD_PENDING_TIMEOUT_MS } + if (command === 'loadModel' && engineType === 'vllm') { + return VLLM_LOAD_PENDING_TIMEOUT_MS + } if (command === 'deleteModel' && EngineCapabilities[engineType]?.restartsOnModelDelete) { return RESTART_DELETE_TIMEOUT_MS } @@ -124,6 +131,20 @@ function modelKey(nodeId: string, engineType: EngineType, model: string): string return `${nodeId}:${engineType}:model:${model}` } +/** The pending entries without the engine's model Loads on that node. */ +function withoutPendingLoads( + pending: Map, + nodeId: string, + engineType: EngineType +): Map { + const prefix = `${nodeId}:${engineType}:model:` + const next = new Map(pending) + for (const [key, p] of pending) { + if (key.startsWith(prefix) && p.action === 'loadModel') next.delete(key) + } + return next +} + function settingsKey(nodeId: string, engineType: EngineType): string { return `${nodeId}:${engineType}:settings` } @@ -132,7 +153,11 @@ interface PendingActionsState { pending: Map /** Record an optimistic pending entry for a fired command. No-op for non-lifecycle/non-model commands. */ begin(payload: EngineCommandPayload): void - /** The pending lifecycle action for an engine, or undefined when none is in flight. */ + /** + * The pending lifecycle action for an engine, or undefined when none is in + * flight. A Load on an engine whose Load restarts it (`loadRestartsEngine`) + * counts, as `loadModel`. + */ getLifecyclePending(nodeId: string, engineType: EngineType): EngineCommandType | undefined /** The pending action for a specific model, or undefined when none is in flight. */ getModelPending( @@ -177,11 +202,20 @@ export const usePendingActionsStore = create((set, get) => const { command, engineType, nodeId, model } = payload const expiresAt = Date.now() + timeoutFor(command, engineType) + // On an engine whose Load starts it, a Stop or Eject ends a Load still in + // flight. The engine-manager replies that it was stopped, which clears + // nothing here, so the stop itself ends the pending Load. + const current = + (command === 'stop' || command === 'unloadModel') && + EngineCapabilities[engineType]?.loadRestartsEngine + ? withoutPendingLoads(get().pending, nodeId, engineType) + : get().pending + if (LIFECYCLE_COMMANDS.has(command)) { const baselineStatus = useEngineStatusStore .getState() .getStatus(nodeId, engineType).processStatus - const next = new Map(get().pending) + const next = new Map(current) next.set(lifecycleKey(nodeId, engineType), { nodeId, engineType, @@ -194,7 +228,7 @@ export const usePendingActionsStore = create((set, get) => } if (MODEL_COMMANDS.has(command) && model) { - const next = new Map(get().pending) + const next = new Map(current) next.set(modelKey(nodeId, engineType, model), { nodeId, engineType, @@ -203,11 +237,23 @@ export const usePendingActionsStore = create((set, get) => expiresAt }) set({ pending: next }) + return } + if (current !== get().pending) set({ pending: current }) }, - getLifecyclePending: (nodeId, engineType) => - asCommand(get().pending.get(lifecycleKey(nodeId, engineType))?.action), + getLifecyclePending: (nodeId, engineType) => { + const lifecycle = asCommand(get().pending.get(lifecycleKey(nodeId, engineType))?.action) + if (lifecycle || !EngineCapabilities[engineType]?.loadRestartsEngine) return lifecycle + // A Load on an engine that restarts to serve the model is a lifecycle + // change too: one entry, the model row's, covers both the row and the + // engine, cleared by the same backend truth and timeout. + const prefix = `${nodeId}:${engineType}:model:` + for (const [key, p] of get().pending) { + if (key.startsWith(prefix) && p.action === 'loadModel') return 'loadModel' + } + return undefined + }, getModelPending: (nodeId, engineType, model) => asCommand(get().pending.get(modelKey(nodeId, engineType, model))?.action), diff --git a/desktop/src/ui/types/engine-info.ts b/desktop/src/ui/types/engine-info.ts index 3c4cb7807..b9d820c23 100644 --- a/desktop/src/ui/types/engine-info.ts +++ b/desktop/src/ui/types/engine-info.ts @@ -97,6 +97,11 @@ export interface BackendInfo { installedVersion?: string /** Models on this backend with per-model status */ models: ModelItem[] + /** + * Whether {@link models} was read. False when the engine's list could not + * be, so an empty list is not shown as "nothing downloaded". + */ + modelsKnown: boolean /** System-level dependencies required before install/run (local node only) */ prerequisites?: BackendPrerequisite[] /** Progress info during install/uninstall (local node only, ephemeral) */ diff --git a/desktop/src/ui/types/engine-manifest.ts b/desktop/src/ui/types/engine-manifest.ts index 1cd998bfe..fc0902253 100644 --- a/desktop/src/ui/types/engine-manifest.ts +++ b/desktop/src/ui/types/engine-manifest.ts @@ -36,12 +36,30 @@ export interface EngineCaps { modelOpsWhenStopped: boolean /** When true, show the Delete action in the model action menu. */ hasDeleteModel: boolean + /** + * When true, Load serves a model by restarting the engine on it (one model + * per process), so while a Load is in flight the engine passes through + * stopped and is shown as starting. + */ + loadRestartsEngine?: boolean + /** + * When true, the model hub can search the whole upstream catalog for a typed + * query (the backend's `engine:catalog` reports `searchable`), instead of only + * filtering the browse list. + */ + hubSearch?: boolean /** * When true, the engine restarts as part of deleting a model, so Delete asks * for confirmation first. The engine manager owns the restart; this flag only * tells the UI that deleting has that side effect. */ restartsOnModelDelete?: boolean + /** + * When true, the model an engine is serving cannot be deleted, so Delete is + * disabled for it. An engine that serves one model per process breaks if its + * model files go, so the backend refuses it and the UI does not offer it. + */ + deleteBlockedWhileLoaded?: boolean /** * If the engine has its own public model catalog (e.g. ollama.com), provide the * hub config here. Presence doubles as the truthy flag for the hub source selector. diff --git a/desktop/src/ui/utils/get-engines-for-node.ts b/desktop/src/ui/utils/get-engines-for-node.ts index 8d46d1894..7d7c9a748 100644 --- a/desktop/src/ui/utils/get-engines-for-node.ts +++ b/desktop/src/ui/utils/get-engines-for-node.ts @@ -19,8 +19,8 @@ interface EngineListRow { /** * One row per enabled engine type for this node. Types not yet reported by the * service appear as `initializing` placeholders ({@link emptyEngineStatus}) so - * compact rows can hide them until evidence arrives. When `engineFilter` is - * provided, only those types are included. + * compact rows can hide them until evidence arrives. Types the node cannot run + * are left out. When `engineFilter` is provided, only those types are included. */ export function getEnginesForNode( nodeId: string, @@ -29,11 +29,9 @@ export function getEnginesForNode( ): EngineStatusData[] { const inner = statusByNode.get(nodeId) const types = engineFilter ?? EnabledEngineTypes - return types.map(type => { - const s = inner?.get(type) - if (s) return s - return emptyEngineStatus(nodeId, type) - }) + return types + .map(type => inner?.get(type) ?? emptyEngineStatus(nodeId, type)) + .filter(s => s.processStatus !== 'unsupported') } function engineDisplayName(engineType: string): string { diff --git a/desktop/src/ui/utils/match-downloaded-model.ts b/desktop/src/ui/utils/match-downloaded-model.ts index cdce8576a..57610cf15 100644 --- a/desktop/src/ui/utils/match-downloaded-model.ts +++ b/desktop/src/ui/utils/match-downloaded-model.ts @@ -43,7 +43,8 @@ const matchExactName: DownloadedMatcher = (hubEntry, downloaded) => downloaded.n const MATCHERS: Partial> = { ollama: matchOllama, 'lm-studio': matchHfPullKeyOrName, - 'llama-cpp': matchExactName + 'llama-cpp': matchExactName, + vllm: matchExactName } export function isHubEntryDownloaded( diff --git a/desktop/src/ui/utils/status.ts b/desktop/src/ui/utils/status.ts index 0357e45b7..05ec43695 100644 --- a/desktop/src/ui/utils/status.ts +++ b/desktop/src/ui/utils/status.ts @@ -11,5 +11,6 @@ export const statusLabel: Record = { uninstalling: 'Uninstalling...', starting: 'Starting...', stopping: 'Stopping...', - initializing: 'Initializing...' + initializing: 'Initializing...', + unsupported: 'Not Supported' } diff --git a/desktop/tests/modular/engine-command-load.test.ts b/desktop/tests/modular/engine-command-load.test.ts index c26ae49db..ce6c19246 100644 --- a/desktop/tests/modular/engine-command-load.test.ts +++ b/desktop/tests/modular/engine-command-load.test.ts @@ -5,11 +5,13 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' const mocks = vi.hoisted(() => ({ state: { - getSelfId: vi.fn(() => 'local-node') + getSelfId: vi.fn(() => 'local-node'), + beginLocalEngineOp: vi.fn() }, supervisor: { hasProcess: vi.fn(() => true), sendProcess: vi.fn(), + callProcess: vi.fn(), reportError: vi.fn() } })) @@ -70,6 +72,93 @@ describe('local model load command', () => { ) }) + describe('vLLM model switch', () => { + const load = () => + handleServiceBridgeInvoke('engine:command', { + command: 'loadModel', + engineType: 'vllm', + nodeId: 'local-node', + model: 'Qwen/Qwen3-8B' + }) + const flush = () => new Promise(resolve => setTimeout(resolve, 0)) + + // The renderer's pending Load shows the engine as starting while vLLM + // restarts, and backend truth ends it; the main process only sends the + // request and attributes a failure to the model row. + it('awaits load_model for as long as a switch may take', async () => { + const reply = Promise.withResolvers() + mocks.supervisor.callProcess.mockReturnValue(reply.promise) + + await load() + + expect(mocks.supervisor.sendProcess).not.toHaveBeenCalled() + expect(mocks.supervisor.callProcess).toHaveBeenCalledWith( + 'broker', + 'engine:action', + { engine: 'vllm', action: 'load_model', params: { model: 'Qwen/Qwen3-8B' } }, + 20 * 60_000 + ) + reply.resolve(undefined) + await flush() + expect(mocks.supervisor.reportError).not.toHaveBeenCalled() + }) + + // A Stop or Eject ends the Load, and the reply says so rather than + // failing. The user who stopped it here needs no message; a stop made + // elsewhere is reported, so the Load does not stay shown in progress. + it('reports a stopped Load only when the stop came from elsewhere', async () => { + const stopped = { model: 'Qwen/Qwen3-8B', restarted: false, stopped: true } + mocks.supervisor.callProcess.mockResolvedValue(stopped) + + await load() + await flush() + expect(mocks.supervisor.reportError).toHaveBeenCalledWith( + 'Loading Qwen/Qwen3-8B on vllm was stopped before it finished.', + 'info', + 'engine-cmd:load model:vllm', + { + nodeId: 'local-node', + engineType: 'vllm', + operation: 'load', + modelName: 'Qwen/Qwen3-8B' + } + ) + + mocks.supervisor.reportError.mockClear() + const reply = Promise.withResolvers() + mocks.supervisor.callProcess.mockReturnValue(reply.promise) + await load() + await handleServiceBridgeInvoke('engine:command', { + command: 'stop', + engineType: 'vllm', + nodeId: 'local-node' + }) + reply.resolve(stopped) + await flush() + expect(mocks.supervisor.reportError).not.toHaveBeenCalled() + }) + + it('attributes a failure to the model row', async () => { + mocks.supervisor.callProcess.mockRejectedValue( + new Error('did not become ready; the previous model was restored') + ) + await load() + await flush() + + expect(mocks.supervisor.reportError).toHaveBeenCalledWith( + 'Failed to load model on vllm: did not become ready; the previous model was restored', + 'error', + 'engine-cmd:load model:vllm', + { + nodeId: 'local-node', + engineType: 'vllm', + operation: 'load', + modelName: 'Qwen/Qwen3-8B' + } + ) + }) + }) + it('leaves the LM Studio load path unobserved', async () => { await handleServiceBridgeInvoke('engine:command', { command: 'loadModel', diff --git a/desktop/tests/modular/engine-command-stop.test.ts b/desktop/tests/modular/engine-command-stop.test.ts new file mode 100644 index 000000000..9bc6155fb --- /dev/null +++ b/desktop/tests/modular/engine-command-stop.test.ts @@ -0,0 +1,85 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + state: { + getSelfId: vi.fn(() => 'local-node'), + beginLocalEngineOp: vi.fn(), + clearPendingEngineOp: vi.fn() + }, + supervisor: { + hasProcess: vi.fn(() => true), + sendProcess: vi.fn(), + callProcess: vi.fn(), + reportError: vi.fn(), + toggleEngineRemote: vi.fn() + } +})) + +vi.mock('@/electron/service-bridge/modular-supervisor', () => ({ + getModularSupervisor: () => mocks.supervisor +})) +vi.mock('@/electron/service-bridge/modular-state', () => ({ + getModularBridgeState: () => mocks.state, + isProxyEngine: () => false, + isUpstreamUnreachableError: () => false, + parseServiceErrors: () => [] +})) +vi.mock('@/electron/service-bridge/model-catalog', () => ({ getEngineHubModels: vi.fn() })) + +import { handleServiceBridgeInvoke } from '@/electron/service-bridge/empty-handlers' + +// A start can take minutes (vLLM loads for up to eight), and while it runs the +// engine reports not running, so a toggle would send a second start. Stop is +// explicit, and the engine-manager cancels the start in flight. +describe('stopping an engine that is still starting', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.state.getSelfId.mockReturnValue('local-node') + }) + + it('sends engine:stop locally without asking whether it is running', async () => { + await handleServiceBridgeInvoke('engine:command', { + command: 'stop', + engineType: 'vllm', + nodeId: 'local-node' + }) + + expect(mocks.state.beginLocalEngineOp).toHaveBeenCalledWith('vllm', 'stopping') + expect(mocks.supervisor.callProcess).not.toHaveBeenCalled() + expect(mocks.supervisor.sendProcess).toHaveBeenCalledOnce() + const [name, method, params, onFailure, observeResponse] = + mocks.supervisor.sendProcess.mock.calls[0] + expect([name, method, params, observeResponse]).toEqual([ + 'broker', + 'engine:stop', + { engine: 'vllm' }, + true + ]) + + onFailure('cannot stop engine "vllm": it is running under external management') + expect(mocks.state.clearPendingEngineOp).toHaveBeenCalledWith('vllm') + expect(mocks.supervisor.reportError).toHaveBeenCalledWith( + 'Failed to stop vllm: cannot stop engine "vllm": it is running under external management', + 'error', + 'engine-cmd:stop:vllm' + ) + }) + + it('sends a stop to a remote peer whatever its last reported state', async () => { + await handleServiceBridgeInvoke('engine:command', { + command: 'stop', + engineType: 'vllm', + nodeId: 'remote-node' + }) + + expect(mocks.supervisor.toggleEngineRemote).toHaveBeenCalledWith( + 'remote-node', + 'vllm', + 'vllm', + true + ) + }) +}) diff --git a/desktop/tests/modular/engine-identity.test.ts b/desktop/tests/modular/engine-identity.test.ts index aaefae02e..ac3633169 100644 --- a/desktop/tests/modular/engine-identity.test.ts +++ b/desktop/tests/modular/engine-identity.test.ts @@ -7,6 +7,7 @@ import { EngineDefaultLinks, EngineDisplayNames, EngineManagerNames, + EnginesListingModelsWhenStopped, EngineTypes } from '@/shared/constants/engines' import { engineManagerName, engineTypeFromManagerName, isEngineType } from '@/shared/utils/engines' @@ -16,8 +17,8 @@ import { gatewayEndpointDisplayUrl } from '@/ui/utils/gateway-inference-paths' describe('engine identity', () => { it('enables llama.cpp workflows on every platform', () => { - expect(EngineTypes).toEqual(['ollama', 'lm-studio', 'llama-cpp']) - expect(EnabledEngineTypes).toEqual(['ollama', 'lm-studio', 'llama-cpp']) + expect(EngineTypes).toEqual(['ollama', 'lm-studio', 'llama-cpp', 'vllm']) + expect(EnabledEngineTypes).toEqual(['ollama', 'lm-studio', 'llama-cpp', 'vllm']) expect(getWelcomeEngineCandidates('Windows')).toContain('llama-cpp') expect(getWelcomeEngineCandidates('MacOS')).toContain('llama-cpp') expect(getWelcomeEngineCandidates('Linux')).toContain('llama-cpp') @@ -27,10 +28,47 @@ describe('engine identity', () => { expect(WELCOME_ENGINE_DEFAULT_SELECTED).toEqual({ ollama: true, 'lm-studio': true, - 'llama-cpp': true + 'llama-cpp': true, + vllm: true }) }) + it('offers vLLM on Linux and macOS, not Windows', () => { + expect(getWelcomeEngineCandidates('Linux')).toContain('vllm') + expect(getWelcomeEngineCandidates('Windows')).not.toContain('vllm') + expect(getWelcomeEngineCandidates('MacOS')).toContain('vllm') + }) + + it('maps vLLM to its engine-manager wire id and one-model-per-process capabilities', () => { + expect(EngineManagerNames.vllm).toBe('vllm') + expect(engineManagerName('vllm')).toBe('vllm') + expect(engineTypeFromManagerName('vllm')).toBe('vllm') + expect(isEngineType('vllm')).toBe(true) + expect(EngineDisplayNames.vllm).toBe('vLLM') + // vLLM cannot unload a model while running, so Eject stops the engine; it + // downloads from and deletes in the Hugging Face cache, except the model it + // is serving, and its model list is on disk, so it works while stopped. + expect(EngineCapabilities.vllm).toMatchObject({ + hasEject: true, + modelOpsWhenStopped: true, + hasModelSearchOnlyWhenRunning: false, + hasDeleteModel: true, + deleteBlockedWhileLoaded: true, + hubSearch: true, + hasInstall: ['linux', 'darwin'], + hasEnginePort: true, + engineHub: { label: 'vLLM' } + }) + }) + + it('keeps the capability and the supervisor agreeing on which model lists survive a stop', () => { + for (const engineType of EngineTypes) { + expect(EngineCapabilities[engineType].modelOpsWhenStopped).toBe( + EnginesListingModelsWhenStopped.includes(engineType) + ) + } + }) + it('maps the desktop id to the sole engine-manager wire id', () => { expect(EngineManagerNames['llama-cpp']).toBe('llamacpp') expect(engineManagerName('llama-cpp')).toBe('llamacpp') diff --git a/desktop/tests/modular/engine-start-pending.test.ts b/desktop/tests/modular/engine-start-pending.test.ts new file mode 100644 index 000000000..28ac3840c --- /dev/null +++ b/desktop/tests/modular/engine-start-pending.test.ts @@ -0,0 +1,134 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +vi.mock('electron', () => ({ BrowserWindow: { getAllWindows: () => [] } })) +vi.mock('@/electron/window', () => ({ createOverviewWindow: vi.fn() })) + +import { getModularBridgeState } from '@/electron/service-bridge/modular-state' +import type { EngineType } from '@/shared/types/engines' +import { engineManagerName } from '@/shared/utils/engines' + +const NODE = 'start-pending-self' + +function processStatus(engineType: EngineType) { + return getModularBridgeState() + .getEngineInitialState() + .statuses.find(status => status.nodeId === NODE && status.engineType === engineType) + ?.processStatus +} + +function reportStopped(engineType: EngineType): void { + getModularBridgeState().applyEngineManagerStatus({ + engine: engineManagerName(engineType), + installed: true, + running: false, + healthy: false, + port: 8001 + }) +} + +// vLLM and Ollama allow their readiness probe minutes. Until the engine-manager +// reports the engine up, or the start failed, the engine is starting: shown as +// stopped, its toggle would send another Start instead of Stop. +describe('a local engine start', () => { + beforeEach(() => { + vi.useFakeTimers() + const state = getModularBridgeState() + state.setSelfId(NODE) + state.handleNotification({ + source: 'broker', + method: 'discovery:nodes-changed', + params: { + nodes: [{ hostUuid: NODE, name: 'start-host', ipAddress: '127.0.0.1', port: 14318 }] + } + }) + }) + + afterEach(() => { + for (const engineType of ['vllm', 'ollama'] as const) { + getModularBridgeState().clearPendingEngineOp(engineType) + } + vi.useRealTimers() + }) + + it('stays starting through the readiness window of a slow engine', () => { + const test = (engineType: EngineType, readinessMs: number) => { + reportStopped(engineType) + getModularBridgeState().beginLocalEngineOp(engineType, 'starting') + + vi.advanceTimersByTime(readinessMs) + + expect(processStatus(engineType)).toBe('starting') + } + test('vllm', 480_000) + test('ollama', 600_000) + }) + + it('ends when the engine-manager reports the outcome', () => { + reportStopped('vllm') + getModularBridgeState().beginLocalEngineOp('vllm', 'starting') + + getModularBridgeState().applyEngineManagerStatus({ + engine: 'vllm', + installed: true, + running: true, + healthy: true, + port: 8001 + }) + + expect(processStatus('vllm')).toBe('running') + }) + + it('gives up after the lifecycle envelope so a lost notification cannot leave it spinning', () => { + reportStopped('vllm') + getModularBridgeState().beginLocalEngineOp('vllm', 'starting') + + vi.advanceTimersByTime(14 * 60_000 + 1_000) + + expect(processStatus('vllm')).toBe('stopped') + }) + + it('keeps the short net for a stop', () => { + getModularBridgeState().applyEngineManagerStatus({ + engine: 'vllm', + installed: true, + running: true, + healthy: true, + port: 8001 + }) + getModularBridgeState().beginLocalEngineOp('vllm', 'stopping') + + vi.advanceTimersByTime(91_000) + + expect(processStatus('vllm')).toBe('running') + }) + + // A start this app did not send (install-then-start, a settings restart, + // a restore at launch, the terminal interface) has no pending op here. The + // engine-manager reports it as starting, so it is not shown as off. + it('shows a start the engine-manager reports, whoever began it', () => { + const report = (running: boolean, starting: boolean) => + getModularBridgeState().applyEngineManagerStatus({ + engine: 'vllm', + installed: true, + running, + starting, + healthy: running, + port: 8001 + }) + + report(false, true) + expect(processStatus('vllm')).toBe('starting') + + vi.advanceTimersByTime(20 * 60_000) + expect(processStatus('vllm')).toBe('starting') + + report(true, false) + expect(processStatus('vllm')).toBe('running') + + report(false, false) + expect(processStatus('vllm')).toBe('stopped') + }) +}) diff --git a/desktop/tests/modular/inference-demo-lifecycle.test.ts b/desktop/tests/modular/inference-demo-lifecycle.test.ts index 54cb20fb0..f355a7663 100644 --- a/desktop/tests/modular/inference-demo-lifecycle.test.ts +++ b/desktop/tests/modular/inference-demo-lifecycle.test.ts @@ -50,7 +50,8 @@ const POISONED_ENV: Record = { const proxyPorts: Record = { ollama: 11434, 'lm-studio': 1234, - 'llama-cpp': 8080 + 'llama-cpp': 8080, + vllm: 8000 } /** Model inventory each probe returns. Mutable so a test can return none. */ diff --git a/desktop/tests/modular/install-state.test.ts b/desktop/tests/modular/install-state.test.ts index 0f4c1189d..a65e7fd1b 100644 --- a/desktop/tests/modular/install-state.test.ts +++ b/desktop/tests/modular/install-state.test.ts @@ -77,8 +77,10 @@ describe('LM Studio install state', () => { vi.advanceTimersByTime(1) expect(statuses.at(-1)).toBe('not-installed') + // A start has its own, shorter net: the awaited lifecycle envelope that + // outlasts every engine's readiness window, not the installer's 30 minutes. state.beginLocalEngineOp('lm-studio', 'starting') - vi.advanceTimersByTime(90_000) + vi.advanceTimersByTime(14 * 60_000) expect(statuses.at(-1)).toBe('not-installed') }) }) diff --git a/desktop/tests/modular/lmstudio-stale-model.test.ts b/desktop/tests/modular/lmstudio-stale-model.test.ts index 23a79fa9a..a07032e67 100644 --- a/desktop/tests/modular/lmstudio-stale-model.test.ts +++ b/desktop/tests/modular/lmstudio-stale-model.test.ts @@ -506,3 +506,67 @@ describe('LM Studio model reconciliation', () => { hasProcess.mockRestore() }) }) + +describe('model list during a vLLM model switch', () => { + const nodeId = 'vllm-switch-list-node' + + function vllmModels() { + return getModularBridgeState() + .getEngineInitialState() + .models.find(models => models.nodeId === nodeId && models.engineType === 'vllm') + ?.models.map(model => model.name) + } + + function startVllmNode() { + const state = getModularBridgeState() + state.setSelfId(nodeId) + state.handleNotification({ + source: 'broker', + method: 'discovery:nodes-changed', + params: { + nodes: [ + { + hostUuid: nodeId, + name: 'vllm-switch-host', + ipAddress: '127.0.0.1', + port: 14318 + } + ] + } + }) + state.applyEngineManagerStatus({ + engine: 'vllm', + installed: true, + running: true, + healthy: true, + port: 8001 + }) + state.setLocalEngineModels('vllm', ['Qwen/Qwen3-8B', 'Qwen/Qwen3-32B']) + return state + } + + // vLLM's list is the Hugging Face cache on disk, so a stopped (or ejected) + // engine still has it: the list is read, not cleared, unlike every other engine. + it('reads the cached-model list from disk while the engine is stopped', async () => { + const state = startVllmNode() + const supervisor = getModularSupervisor() + const hasProcess = vi.spyOn(supervisor, 'hasProcess').mockReturnValue(true) + const call = vi + .spyOn(supervisor, 'callProcess') + .mockResolvedValue({ data: [{ id: 'Qwen/Qwen3-8B' }, { id: 'org/new' }] }) + + supervisor.refreshManagedEngineModels({ engine: 'vllm', running: false }) + expect(vllmModels()).toEqual(['Qwen/Qwen3-8B', 'Qwen/Qwen3-32B']) + await Promise.resolve() + await Promise.resolve() + + expect(call).toHaveBeenCalledWith('broker', 'engine:action', { + engine: 'vllm', + action: 'list_models' + }) + expect(vllmModels()).toEqual(['Qwen/Qwen3-8B', 'org/new']) + expect(state.getEngineInitialState().statuses.length).toBeGreaterThan(0) + call.mockRestore() + hasProcess.mockRestore() + }) +}) diff --git a/desktop/tests/modular/model-catalog.test.ts b/desktop/tests/modular/model-catalog.test.ts index e65d0953f..d2931eb96 100644 --- a/desktop/tests/modular/model-catalog.test.ts +++ b/desktop/tests/modular/model-catalog.test.ts @@ -127,8 +127,8 @@ describe('model catalog relay', () => { // cost a multi-megabyte reply for nothing. it('warms only the live-fetched catalogs', () => { warmEngineHubs() - expect(mocks.supervisor.callProcess).toHaveBeenCalledTimes(2) - for (const engine of ['lmstudio', 'llamacpp']) { + expect(mocks.supervisor.callProcess).toHaveBeenCalledTimes(3) + for (const engine of ['lmstudio', 'llamacpp', 'vllm']) { expect(mocks.supervisor.callProcess).toHaveBeenCalledWith( 'broker', 'engine:catalog', diff --git a/desktop/tests/modular/model-switch-status.test.ts b/desktop/tests/modular/model-switch-status.test.ts new file mode 100644 index 000000000..bb0f5c806 --- /dev/null +++ b/desktop/tests/modular/model-switch-status.test.ts @@ -0,0 +1,197 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import type { ServiceError } from '@/shared/types/errors' +import type { EngineStatePatch } from '@/shared/types/engine-api' +import type { EngineType, ModelItem } from '@/shared/types/engines' +import type { ServiceStatus } from '@/shared/types/ipc-channels' +import { usePendingActionsStore } from '@/ui/stores/pending-actions.store' + +/** + * vLLM serves one model per process, so Load restarts the engine on the model + * and it reports stopped on the way. The pending Load, one entry in the + * pending-actions store, is what shows the engine as busy for that whole time: + * there is no second tracker with its own timer and clearing rules. + */ + +let errorsCb: ((errors: ServiceError[]) => void) | null = null +let stateCb: ((patch: EngineStatePatch) => void) | null = null + +function loadedModel(name: string): ModelItem { + return { + name, + size: 1, + downloaded: true, + status: 'loaded', + parameterSize: '', + quantization: '', + family: '', + digest: '', + sizeVram: null, + expiresAt: null, + expiry: '-1', + capabilities: [] + } +} + +function makeFakeWindow() { + const noopUnsub = (): void => {} + return { + pairApi: { + engines: { + onStateChanged: (cb: (patch: EngineStatePatch) => void) => { + stateCb = cb + return noopUnsub + }, + onProgress: () => noopUnsub, + onSettingsChanged: () => noopUnsub + }, + errors: { + onUpdate: (cb: (errors: ServiceError[]) => void) => { + errorsCb = cb + return noopUnsub + } + } + }, + windowApi: { + service: { + getStatus: vi + .fn() + .mockResolvedValue({ connectorStatus: 'connected', weSpawned: true }), + // The store's cached status already defaults to 'connected', so + // these tests never need to push one. + onStatusChanged: (_cb: (status: ServiceStatus) => void) => noopUnsub + } + } + } +} + +function load(nodeId: string, engineType: EngineType, model: string): void { + usePendingActionsStore.getState().begin({ command: 'loadModel', engineType, nodeId, model }) +} + +function lifecyclePending(nodeId: string, engineType: EngineType) { + return usePendingActionsStore.getState().getLifecyclePending(nodeId, engineType) +} + +describe('a vLLM Load is one pending action for the row and the engine', () => { + beforeEach(() => { + vi.useFakeTimers() + errorsCb = null + stateCb = null + vi.stubGlobal('window', makeFakeWindow()) + usePendingActionsStore.getState().initialize() + }) + + afterEach(() => { + usePendingActionsStore.getState().cleanup() + vi.unstubAllGlobals() + vi.useRealTimers() + }) + + it('holds the engine busy, local or remote, until the model is reported loaded', () => { + for (const nodeId of ['local-node', 'remote-node']) load(nodeId, 'vllm', 'Qwen/Qwen3-8B') + + expect(lifecyclePending('local-node', 'vllm')).toBe('loadModel') + expect(lifecyclePending('remote-node', 'vllm')).toBe('loadModel') + + stateCb?.({ + nodeId: 'local-node', + engineType: 'vllm', + models: { + nodeId: 'local-node', + engineType: 'vllm', + models: [loadedModel('Qwen/Qwen3-8B')], + inventoryKnown: true + } + }) + + expect(lifecyclePending('local-node', 'vllm')).toBeUndefined() + expect(lifecyclePending('remote-node', 'vllm')).toBe('loadModel') + }) + + it('releases the engine when the Load fails', () => { + load('local-node', 'vllm', 'org/too-big') + + errorsCb?.([ + { + id: 'pair-ui:engine-load:vllm:org/too-big', + message: 'Failed to load org/too-big', + timestamp: Date.now(), + severity: 'error', + nodeId: 'local-node', + engineType: 'vllm', + operation: 'load', + modelName: 'org/too-big' + } + ]) + + expect(lifecyclePending('local-node', 'vllm')).toBeUndefined() + }) + + it('gives up after the switch budget so a lost notification cannot leave it busy', () => { + load('local-node', 'vllm', 'Qwen/Qwen3-8B') + + vi.advanceTimersByTime(19 * 60_000) + expect(lifecyclePending('local-node', 'vllm')).toBe('loadModel') + + vi.advanceTimersByTime(2 * 60_000) + expect(lifecyclePending('local-node', 'vllm')).toBeUndefined() + }) + + it('leaves engines whose Load does not restart them alone', () => { + load('local-node', 'llama-cpp', 'org/model:Q4_K_M') + load('local-node', 'ollama', 'llama3.2') + + expect(lifecyclePending('local-node', 'llama-cpp')).toBeUndefined() + expect(lifecyclePending('local-node', 'ollama')).toBeUndefined() + }) + + it('shows a Stop sent during the Load as the lifecycle action', () => { + load('local-node', 'vllm', 'Qwen/Qwen3-8B') + usePendingActionsStore + .getState() + .begin({ command: 'stop', engineType: 'vllm', nodeId: 'local-node' }) + + expect(lifecyclePending('local-node', 'vllm')).toBe('stop') + }) + + it('ends the Load when the user stops the engine or ejects', () => { + const modelPending = (nodeId: string, engineType: EngineType, model: string) => + usePendingActionsStore.getState().getModelPending(nodeId, engineType, model) + const begin = usePendingActionsStore.getState().begin + + // The engine-manager replies that a stopped Load was stopped, not that + // it failed, so no error clears the entry: the stop has to. + load('local-node', 'vllm', 'Qwen/Qwen3-8B') + begin({ command: 'stop', engineType: 'vllm', nodeId: 'local-node' }) + expect(modelPending('local-node', 'vllm', 'Qwen/Qwen3-8B')).toBeUndefined() + + load('remote-node', 'vllm', 'Qwen/Qwen3-8B') + begin({ + command: 'unloadModel', + engineType: 'vllm', + nodeId: 'remote-node', + model: 'Qwen/Qwen2.5-0.5B-Instruct' + }) + expect(modelPending('remote-node', 'vllm', 'Qwen/Qwen3-8B')).toBeUndefined() + expect(modelPending('remote-node', 'vllm', 'Qwen/Qwen2.5-0.5B-Instruct')).toBe( + 'unloadModel' + ) + + // An engine whose Load does not start it keeps its Load. + load('local-node', 'ollama', 'llama3.2') + begin({ command: 'stop', engineType: 'ollama', nodeId: 'local-node' }) + expect(modelPending('local-node', 'ollama', 'llama3.2')).toBe('loadModel') + }) + + it('lets a lifecycle command already in flight win', () => { + usePendingActionsStore + .getState() + .begin({ command: 'toggle', engineType: 'vllm', nodeId: 'local-node' }) + load('local-node', 'vllm', 'Qwen/Qwen3-8B') + + expect(lifecyclePending('local-node', 'vllm')).toBe('toggle') + }) +}) diff --git a/desktop/tests/modular/pending-actions-delete.test.ts b/desktop/tests/modular/pending-actions-delete.test.ts index fb1640d8f..4d8e42a65 100644 --- a/desktop/tests/modular/pending-actions-delete.test.ts +++ b/desktop/tests/modular/pending-actions-delete.test.ts @@ -209,7 +209,8 @@ describe('pending-actions store: model actions', () => { models: { nodeId: 'local-node', engineType: 'ollama', - models: [loadedModel('llama3.2')] + models: [loadedModel('llama3.2')], + inventoryKnown: true } }) errorsCb?.([ @@ -244,6 +245,26 @@ describe('pending-actions store: model actions', () => { ).toBeUndefined() }) + it('keeps a vLLM load locked while the engine restarts on the new model', () => { + usePendingActionsStore.getState().begin({ + command: 'loadModel', + engineType: 'vllm', + nodeId: 'local-node', + model: 'Qwen/Qwen3-8B' + }) + + // Well past the default net: readiness alone can take 8 minutes. + vi.advanceTimersByTime(12 * 60_000) + expect( + usePendingActionsStore.getState().getModelPending('local-node', 'vllm', 'Qwen/Qwen3-8B') + ).toBe('loadModel') + + vi.advanceTimersByTime(8 * 60_000 + 1_000) + expect( + usePendingActionsStore.getState().getModelPending('local-node', 'vllm', 'Qwen/Qwen3-8B') + ).toBeUndefined() + }) + it('leaves the default net in place for LM Studio loads', () => { usePendingActionsStore.getState().begin({ command: 'loadModel', diff --git a/desktop/tests/modular/proxy-engine-identity.test.ts b/desktop/tests/modular/proxy-engine-identity.test.ts index 69031984f..b52cbbec0 100644 --- a/desktop/tests/modular/proxy-engine-identity.test.ts +++ b/desktop/tests/modular/proxy-engine-identity.test.ts @@ -19,11 +19,17 @@ import { engineManagerName } from '@/shared/utils/engines' describe('proxy engine identity', () => { it('declares the complete proxy set in stable order', () => { - expect(PROXY_ENGINES).toEqual(['ollama', 'lm-studio', 'llama-cpp']) - expect(PROXY_NODE_SOURCES).toEqual(['ollama-proxy', 'lmstudio-proxy', 'llamacpp-proxy']) + expect(PROXY_ENGINES).toEqual(['ollama', 'lm-studio', 'llama-cpp', 'vllm']) + expect(PROXY_NODE_SOURCES).toEqual([ + 'ollama-proxy', + 'lmstudio-proxy', + 'llamacpp-proxy', + 'vllm-proxy' + ]) expect(isProxyEngine('ollama')).toBe(true) expect(isProxyEngine('lm-studio')).toBe(true) expect(isProxyEngine('llama-cpp')).toBe(true) + expect(isProxyEngine('vllm')).toBe(true) }) it('round-trips engine, manager, and relay source identities', () => { @@ -35,6 +41,7 @@ describe('proxy engine identity', () => { expect(proxyEngineFromSource('other-proxy')).toBeNull() expect(proxyEngineFromManagerId('other')).toBeNull() expect(proxyEngineFromManagerId('llamacpp')).toBe('llama-cpp') + expect(proxyEngineFromManagerId('vllm')).toBe('vllm') }) it('routes each proxy notification to the mapped engine', () => { @@ -99,15 +106,19 @@ describe('proxy engine identity', () => { ] } }) + // A peer lists every engine it can run; one it leaves out is unsupported. state.applyRemoteEngineFacts(nodeId, { engines: [ + { engine: 'ollama', installed: true, running: true, healthy: true, port: 11435 }, + { engine: 'lmstudio', installed: true, running: true, healthy: true, port: 1235 }, { engine: 'llamacpp', installed: true, running: true, healthy: true, port: 8081 - } + }, + { engine: 'vllm', installed: false, running: false, healthy: false, port: 0 } ] }) @@ -130,7 +141,8 @@ describe('proxy engine identity', () => { processStatus: 'running', enginePort: 8081, proxyPort: 8080 - }) + }), + expect.objectContaining({ engineType: 'vllm', processStatus: 'not-installed' }) ]) const llamaModels = state @@ -192,4 +204,53 @@ describe('proxy engine identity', () => { )?.models ).toEqual([expect.objectContaining({ name: model, status: 'loaded' })]) }) + it('projects local vLLM lifecycle and its single served model into the initial snapshot', () => { + const state = getModularBridgeState() + const nodeId = 'proxy-identity-vllm-local' + const model = 'Qwen/Qwen2.5-0.5B-Instruct' + state.setSelfId(nodeId) + state.handleNotification({ + source: 'broker', + method: 'discovery:nodes-changed', + params: { + nodes: [ + { + hostUuid: nodeId, + name: 'proxy-identity-vllm-host', + ipAddress: '127.0.0.1', + port: 14318 + } + ] + } + }) + state.handleNotification({ + source: 'vllm-proxy', + method: 'ready', + params: { port: 8000 } + }) + state.applyEngineManagerStatus({ + engine: 'vllm', + installed: true, + running: true, + healthy: true, + port: 8001 + }) + state.setLocalEngineModels('vllm', [model]) + state.applyLocalLoadedModels({ vllm: [model] }) + + const initial = state.getEngineInitialState() + expect( + initial.statuses.find( + status => status.nodeId === nodeId && status.engineType === 'vllm' + ) + ).toMatchObject({ + processStatus: 'running', + enginePort: 8001, + proxyPort: 8000 + }) + expect( + initial.models.find(models => models.nodeId === nodeId && models.engineType === 'vllm') + ?.models + ).toEqual([expect.objectContaining({ name: model, status: 'loaded' })]) + }) }) diff --git a/desktop/tests/modular/remote-vllm-models.test.ts b/desktop/tests/modular/remote-vllm-models.test.ts new file mode 100644 index 000000000..7519034fe --- /dev/null +++ b/desktop/tests/modular/remote-vllm-models.test.ts @@ -0,0 +1,101 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, expect, it, vi } from 'vitest' + +vi.mock('electron', () => ({ BrowserWindow: { getAllWindows: () => [] } })) +vi.mock('@/electron/window', () => ({ createOverviewWindow: vi.fn() })) + +import { getModularBridgeState } from '@/electron/service-bridge/modular-state' +import type { EngineModels, EngineType } from '@/shared/types/engines' +import type { JsonObject } from '@/electron/service-bridge/json-rpc-subprocess' + +// A vLLM peer advertises only the model it is serving, for routing, and reports +// what it has downloaded apart, so the model manager can list it. +function remoteModels( + nodeId: string, + node: JsonObject, + engine: EngineType +): EngineModels | undefined { + const state = getModularBridgeState() + state.setSelfId('remote-vllm-self') + state.handleNotification({ + source: 'broker', + method: 'discovery:nodes-changed', + params: { + nodes: [ + { + hostUuid: nodeId, + name: `${nodeId}-host`, + ipAddress: '192.0.2.50', + port: 14318, + ...node + } + ] + } + }) + return state + .getEngineInitialState() + .models.find(models => models.nodeId === nodeId && models.engineType === engine) +} + +describe('remote vLLM model rows', () => { + it('keeps the downloaded models of a stopped peer as idle rows', () => { + const models = remoteModels( + 'remote-vllm-stopped', + { + loadedByEngine: { vllm: [] }, + downloadedByEngine: { + vllm: ['Qwen/Qwen3.5-4B', 'mistralai/Ministral-3-3B-Instruct-2512'] + } + }, + 'vllm' + ) + + expect(models?.inventoryKnown).toBe(true) + expect(models?.models).toEqual([ + expect.objectContaining({ name: 'Qwen/Qwen3.5-4B', status: 'idle' }), + expect.objectContaining({ + name: 'mistralai/Ministral-3-3B-Instruct-2512', + status: 'idle' + }) + ]) + }) + + it('marks the cached model loaded when it is served under an alias', () => { + const models = remoteModels( + 'remote-vllm-alias', + { + models: ['qwen-api'], + modelsByEngine: { vllm: ['qwen-api'] }, + loadedByEngine: { vllm: ['Qwen/Qwen3.5-4B'] }, + downloadedByEngine: { + vllm: ['Qwen/Qwen3.5-4B', 'mistralai/Ministral-3-3B-Instruct-2512'] + } + }, + 'vllm' + ) + + expect(models?.models).toEqual([ + expect.objectContaining({ name: 'Qwen/Qwen3.5-4B', status: 'loaded' }), + expect.objectContaining({ + name: 'mistralai/Ministral-3-3B-Instruct-2512', + status: 'idle' + }) + ]) + }) + + it('reports an unread inventory as unknown rather than empty', () => { + const nodeId = 'remote-vllm-unread' + const vllm = remoteModels(nodeId, {}, 'vllm') + expect(vllm?.models).toEqual([]) + expect(vllm?.inventoryKnown).toBe(false) + + // An engine whose list comes from the running server is unchanged: it + // shows "Start engine to see models" while stopped. + const ollama = getModularBridgeState() + .getEngineInitialState() + .models.find(models => models.nodeId === nodeId && models.engineType === 'ollama') + expect(ollama?.inventoryKnown).toBe(true) + }) +}) diff --git a/desktop/tests/modular/unsupported-engine.test.ts b/desktop/tests/modular/unsupported-engine.test.ts new file mode 100644 index 000000000..b60ea23fe --- /dev/null +++ b/desktop/tests/modular/unsupported-engine.test.ts @@ -0,0 +1,132 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, expect, it, vi } from 'vitest' + +vi.mock('electron', () => ({ BrowserWindow: { getAllWindows: () => [] } })) +vi.mock('@/electron/window', () => ({ createOverviewWindow: vi.fn() })) + +import type { EngineStatusByNode, EngineStatusData } from '@/shared/types/engines' +import { getModularBridgeState } from '@/electron/service-bridge/modular-state' +import { getEnginesForNode } from '@/ui/utils/get-engines-for-node' + +// An engine manager lists only the engines its host can run. On Windows the +// listing has no vLLM, the way these do. +const windowsListing = { + engines: [ + { engine: 'ollama', installed: true, running: true, healthy: true, port: 11434 }, + { engine: 'lmstudio', installed: false, running: false, healthy: false, port: 0 }, + { engine: 'llamacpp', installed: false, running: false, healthy: false, port: 0 } + ] +} + +function statusOf(nodeId: string, engineType: string): EngineStatusData | undefined { + return getModularBridgeState() + .getEngineInitialState() + .statuses.find(status => status.nodeId === nodeId && status.engineType === engineType) +} + +function byNode(statuses: EngineStatusData[]): EngineStatusByNode { + const out: EngineStatusByNode = new Map() + for (const status of statuses) { + const inner = out.get(status.nodeId) ?? new Map() + inner.set(status.engineType, status) + out.set(status.nodeId, inner) + } + return out +} + +describe('engines a node cannot run', () => { + it('marks an engine the local engine manager leaves out as unsupported', () => { + const state = getModularBridgeState() + state.setSelfId('windows-self') + + state.applyEngineManagerList(windowsListing) + + expect(statusOf('windows-self', 'vllm')).toMatchObject({ processStatus: 'unsupported' }) + expect(statusOf('windows-self', 'lm-studio')).toMatchObject({ + processStatus: 'not-installed' + }) + }) + + it('marks an engine a peer leaves out as unsupported', () => { + const state = getModularBridgeState() + const peerId = 'windows-peer' + state.setSelfId('linux-self') + state.handleNotification({ + source: 'broker', + method: 'discovery:nodes-changed', + params: { + nodes: [ + { + hostUuid: peerId, + name: 'windows-peer', + ipAddress: '192.0.2.70', + port: 14318, + trusted: true, + clustered: true + } + ] + } + }) + + state.applyRemoteEngineFacts(peerId, windowsListing) + + expect(statusOf(peerId, 'vllm')).toMatchObject({ processStatus: 'unsupported' }) + expect(statusOf(peerId, 'lm-studio')).toMatchObject({ processStatus: 'not-installed' }) + }) + + it('keeps an engine a peer has not listed yet as unknown, not unsupported', () => { + const state = getModularBridgeState() + const peerId = 'unpolled-peer' + state.setSelfId('linux-self') + state.handleNotification({ + source: 'broker', + method: 'discovery:nodes-changed', + params: { + nodes: [ + { + hostUuid: peerId, + name: 'unpolled-peer', + ipAddress: '192.0.2.71', + port: 14318, + trusted: true, + clustered: true + } + ] + } + }) + + state.applyRemoteEngineFacts(peerId, { engines: [] }) + + expect(statusOf(peerId, 'vllm')).toBeUndefined() + }) + + it('gives an unsupported engine no row, and keeps one for an unreported engine', () => { + const nodeId = 'windows-node' + const statuses = byNode([ + { + engineType: 'ollama', + nodeId, + processStatus: 'running', + enginePort: 11434, + proxyPort: null + }, + { + engineType: 'vllm', + nodeId, + processStatus: 'unsupported', + enginePort: null, + proxyPort: null + } + ]) + + const rows = getEnginesForNode(nodeId, statuses).map(s => [s.engineType, s.processStatus]) + + expect(rows).toEqual([ + ['ollama', 'running'], + ['lm-studio', 'initializing'], + ['llama-cpp', 'initializing'] + ]) + }) +}) diff --git a/docs/engine-lifecycle.mdx b/docs/engine-lifecycle.mdx index 63eb159e5..b705c7aa6 100644 --- a/docs/engine-lifecycle.mdx +++ b/docs/engine-lifecycle.mdx @@ -9,7 +9,7 @@ SPDX-License-Identifier: Apache-2.0 */} An **engine** is the local inference runtime NVIDIA PAIR uses to -run models. Today that means Ollama, LM Studio, or llama.cpp on a given machine. +run models. Today that means Ollama, LM Studio, llama.cpp, or vLLM on a given machine. PAIR can install and run those engines for you, or work with a supported existing Ollama or LM Studio installation. This page explains what you can expect when you install, start, stop, or remove an engine, and how to move one @@ -89,6 +89,47 @@ When you install an engine, consider the following: demand and is not part of the PAIR installer. Package choice is not driver-aware. llama.cpp retains CPU fallback when GPU acceleration is unavailable. +- vLLM is available on Linux with an NVIDIA GPU and on Apple Silicon Macs + running macOS 15 or newer (through vLLM-Metal). If a `vllm` is already + installed at a known location (`~/.local/bin/vllm` on Linux, + `~/.venv-vllm-metal/bin/vllm` on macOS), PAIR uses it and downloads nothing. + Otherwise PAIR downloads a checksum-pinned `uv`, creates a private Python + environment and installs a pinned vLLM release into it. On Linux this is + several GiB because it includes PyTorch and the CUDA libraries, and PAIR first + checks for an NVIDIA driver that supports CUDA 13 or newer and about 25 GB of + free disk, enough for the engine and its default model. On macOS it installs + checksum-pinned vLLM and vLLM-Metal wheels and needs about 20 GB of free disk. Either way PAIR stops with a clear message if a + requirement is missing. vLLM-Metal cannot run FP8 models (FP8 weights or an + FP8 KV cache, such as `nvidia/Llama-3.1-8B-Instruct-FP8`), apart from Qwen3.5 + and Qwen3.6, so on a Mac they are left out of **Add model**, and loading or + downloading one is refused with a message to choose the BF16 version. Linux, + DGX Spark included, runs them. + Nothing is compiled. vLLM serves a single model per process. Until you choose + one it serves Qwen3.5-4B, downloaded on the first start (about 9.3 GB, and a GPU + with at least 16 GB of memory); on a smaller GPU, add and load a smaller model + first. Its model list + shows the complete text models already in your Hugging Face cache + (`~/.cache/huggingface`), and **Load** on one of them switches vLLM to it by + restarting the engine, which takes as long as the model needs to load; if the + new model does not start, the previous one is restored. **Add model** browses + and searches Hugging Face (open safetensors chat models) and downloads the one + you pick into the same cache, with a progress bar, even while vLLM is stopped; + it then appears in the list. A repository whose files would not make a model + vLLM can load, such as several unnumbered weight files with no index, is + refused before anything is downloaded. While vLLM starts, which can take + minutes, it is shown as starting, however the start began. **Delete** removes a cached model from disk, + except the one vLLM has loaded (eject it first); if vLLM cannot say which model + it is serving, nothing is deleted. Deleting the model a stopped + or ejected vLLM is set to serve leaves it with no model selected, and + **Start** refuses with an error until you **Load** another model. **Eject** + stops vLLM, which frees all of its GPU memory (vLLM cannot drop its model and + keep running). Eject, or turning the engine's switch off, also works while + vLLM is still loading a model: it stops at once, without an error, and a Load cut short keeps + the model that was selected before; the model list stays visible and **Load** starts it again on + the model you pick. vLLM does not unload a model on its own when it sits idle, + so eject it when you want the GPU back. Other nodes and the proxy see only the + model vLLM is serving, so a request for a different cached model is refused + until you load it. PAIR-managed llama.cpp models sleep after five minutes without inference work. This releases model and KV-cache memory, and the next request wakes the model @@ -131,8 +172,9 @@ an engine PAIR installed. Use it when you no longer want PAIR to own that engine on the machine. Uninstalling an engine does **not** remove its models. Downloaded model files stay -on disk in the engine's own storage (`~/.ollama`, `~/.llamacpp`, or -`~/.lmstudio/models`), so uninstalling and reinstalling does not cost you +on disk in the engine's own storage (`~/.ollama`, `~/.llamacpp`, +`~/.cache/huggingface`, or `~/.lmstudio/models`), so uninstalling and +reinstalling does not cost you re-downloading them. The same is true of removing PAIR itself, or of resetting its data. Both remove the engines PAIR installed as far as they can, and never your models. A reset stops and tells you if an engine cannot be removed, while diff --git a/docs/engine-settings.mdx b/docs/engine-settings.mdx index 2e93c6f8b..ae9991bfd 100644 --- a/docs/engine-settings.mdx +++ b/docs/engine-settings.mdx @@ -8,7 +8,7 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 */} -Open a device's Ollama, LM Studio, or llama.cpp row and expand **Settings**. +Open a device's Ollama, LM Studio, llama.cpp, or vLLM row and expand **Settings**. Edit the server port, proxy port, and engine arguments, then select the green **Apply** button. A paired device that supports settings can be edited from another device in the cluster. @@ -85,6 +85,7 @@ configuration consistent: | Ollama | `OLLAMA_HOST=127.0.0.1:` | The host in `OLLAMA_HOST` | `OLLAMA_ORIGINS` | | LM Studio | `--port`, `-p` | `--bind`, `LMS_SERVER_HOST` | `--cors` | | llama.cpp | `--port` | `--host` | `--cors-origins` | +| vLLM | `--port` | `--host` | Not managed | For LM Studio, `--port 1235`, `--port=1235`, `-p 1235`, `-p1235`, and `-p=1235` all synchronize the same server-port field. Conflicting repetitions and invalid @@ -94,6 +95,31 @@ Add or remove the switch locally. Origin lists are normalized before validation, and conflicting repeated lists are rejected. Other vendor options remain opaque. llama.cpp must remain in router mode. Model-selection arguments such as `-m` or `-hf` switch it to an incompatible single-model server and fail readiness. +vLLM always serves one model: its launch arguments name that model with +`--model `, and changing it restarts the engine. At every start +PAIR sizes the rest for that model: a 64K context (`--max-model-len`, shrunk if +the free GPU memory cannot hold it), YaRN rope scaling when the model's own window +is smaller (for example Qwen3-8B's 40,960 tokens), a GPU memory share just large +enough for the weights and that context, so other engines can run alongside it, +and a limit of 32 requests at a time (`--max-num-seqs`). A model that is not +downloaded yet, such as the default on a fresh install, has its configuration +and chat template fetched first and is then set up the same way, tool calling +included, before vLLM downloads its weights. Hugging Face settings in vLLM's +launch settings apply to PAIR's own requests, downloads and deletes too: +`HF_HUB_OFFLINE` or `TRANSFORMERS_OFFLINE`, `HF_TOKEN` (not sent when +`HF_HUB_DISABLE_IMPLICIT_TOKEN` is set), a mirror in `HF_ENDPOINT`, and a cache folder in +`HF_HUB_CACHE` or `HF_HOME`, or `--download-dir` in the launch text, which the model list +then shows. Arguments you add to the launch text, +such as your own `--max-model-len`, `--gpu-memory-utilization` or +`--max-num-seqs`, take precedence, and a launch with +`--tensor-parallel-size` or `--pipeline-parallel-size` is left as written. PAIR +also turns on tool calling, which agent clients require: it checks how the +model's chat template writes a tool call and picks the vLLM parser that reads +that format (for example `hermes` for Qwen3, `qwen3_coder` for Qwen3.5 and +Granite 4.2). For a model no vLLM parser can read, such as Nemotron Nano v2, +tool calling stays off and vLLM answers a request that asks for tool calls with an +error. To choose a parser yourself, add `--tool-call-parser ` to the +launch text. Changing the server-port field updates the port in the command. Editing the command's port updates the numeric field. If both are edited before validation diff --git a/docs/getting-started.mdx b/docs/getting-started.mdx index 721638d6f..90ea86cff 100644 --- a/docs/getting-started.mdx +++ b/docs/getting-started.mdx @@ -35,7 +35,7 @@ inference. Two or more on the same local network let you try pairing and routing Nothing else has to be in place first: -- **Engines.** PAIR installs and starts Ollama, LM Studio, or llama.cpp for you +- **Engines.** PAIR installs and starts Ollama, LM Studio, llama.cpp, or (on Linux or Apple Silicon macOS) vLLM for you in step 4. Existing Ollama and LM Studio installations can also be detected and used. - **Models.** PAIR downloads models for you in step 4. A request needs only one @@ -210,6 +210,7 @@ engine would normally listen on: - `11434` for Ollama. - `1234` for LM Studio. - `8080` for llama.cpp. +- `8000` for vLLM. Tools already pointed at those ports keep working without reconfiguration, and the engine moves to the next free port. @@ -239,6 +240,7 @@ The engine therefore has to understand the style you send. | Ollama | Works | Works | | LM Studio | Works | Not available | | llama.cpp | Works | Not available | +| vLLM | Works | Not available | **If you are unsure, send the OpenAI-style request.** It works with every supported engine, and most tools and SDKs use it. Use the Ollama style diff --git a/docs/overview.mdx b/docs/overview.mdx index 8d61d9928..b2f1ed6e8 100644 --- a/docs/overview.mdx +++ b/docs/overview.mdx @@ -42,7 +42,7 @@ These terms have specific meanings in PAIR: - **Cluster** — the set of nodes you have paired together. A node belongs to at most one cluster, and it must leave before it can join another. - **Engine** — the local inference server that runs models: Ollama, LM Studio, - or llama.cpp. PAIR can install, start, stop, and uninstall these engines. It can + llama.cpp, or vLLM. PAIR can install, start, stop, and uninstall these engines. It can also adopt supported existing Ollama and LM Studio installations. - **Model** — what you prepare on each node. Nodes do not share models, so a node can serve a request only for a model it already holds. Preparing the same diff --git a/services/installer/linux/INSTALL.md b/services/installer/linux/INSTALL.md index fad347ca5..171fa49fe 100644 --- a/services/installer/linux/INSTALL.md +++ b/services/installer/linux/INSTALL.md @@ -66,7 +66,7 @@ The workers store configuration in `~/.config/Nvidia Corporation/Personal AI Rou that directory too if you want a fully clean slate. Downloaded models are not in either directory. Each engine keeps them in its own -store — `~/.ollama`, `~/.lmstudio/models`, and `~/.llamacpp` — so deleting the +store — `~/.ollama`, `~/.lmstudio/models`, `~/.llamacpp`, and `~/.cache/huggingface` (vLLM) — so deleting the bundle or the configuration directory leaves them in place. ## Requirements @@ -84,9 +84,9 @@ bundle or the configuration directory leaves them in place. that other machines on the LAN are actually advertising. Run the broker with `--log-level debug` to see live discovery logs on stderr. - **"Bind: address already in use"**: the proxies listen on `11434` (Ollama), - `1234` (LM Studio), and `8080` (llama.cpp); another instance of the proxy or + `1234` (LM Studio), `8080` (llama.cpp), and `8000` (vLLM); another instance of the proxy or another process is holding one of them. Find the offender with - `ss -tlnp | grep -E ':(11434|1234|8080)\b'` and stop or kill it. + `ss -tlnp | grep -E ':(11434|1234|8080|8000)\b'` and stop or kill it. Report issues at the project's tracker. Include the broker's `--log-level debug` output and your distro / kernel version. diff --git a/services/installer/macos/INSTALL.md b/services/installer/macos/INSTALL.md index 27357ab98..83d12c06a 100644 --- a/services/installer/macos/INSTALL.md +++ b/services/installer/macos/INSTALL.md @@ -52,7 +52,7 @@ rm -rf "$HOME/Library/Application Support/Nvidia Corporation/Personal AI Router" ``` Downloaded models are not in either directory. Each engine keeps them in its own -store — `~/.ollama`, `~/.lmstudio/models`, and `~/.llamacpp` — so deleting the +store — `~/.ollama`, `~/.lmstudio/models`, `~/.llamacpp`, and `~/.cache/huggingface` (vLLM) — so deleting the bundle or the configuration directory leaves them in place. ## Requirements @@ -81,9 +81,9 @@ xattr -dr com.apple.quarantine /path/to/NVIDIA-Personal-AI-Router- advertising. Run the broker with `--log-level debug` to see live discovery logs on stderr. - **"Bind: address already in use"**: the proxies listen on `11434` (Ollama), - `1234` (LM Studio), and `8080` (llama.cpp); another instance of the proxy or + `1234` (LM Studio), `8080` (llama.cpp), and `8000` (vLLM); another instance of the proxy or another process is holding one of them. Find the offender with - `lsof -nP -iTCP:11434 -iTCP:1234 -iTCP:8080 -sTCP:LISTEN` and stop it. + `lsof -nP -iTCP:11434 -iTCP:1234 -iTCP:8080 -iTCP:8000 -sTCP:LISTEN` and stop it. Report issues at the project's tracker. Include the broker's `--log-level debug` output and your macOS version (`sw_vers`) and architecture (`uname -m`). diff --git a/services/nvpair-engine-manager/MANIFEST.md b/services/nvpair-engine-manager/MANIFEST.md index 5c9b2d3b0..20bbc5638 100644 --- a/services/nvpair-engine-manager/MANIFEST.md +++ b/services/nvpair-engine-manager/MANIFEST.md @@ -238,7 +238,7 @@ probes. ### Actions Each action is a config-declared operation exposed over `engine:action`. -Exactly one of `http`, `cmd`, or `remove_path`: +Exactly one of `http`, `cmd`, `remove_path`, `hf_cache`, or `launch_model`: - **`http`** — call the engine's loopback control API. The caller's `params` are sent as the JSON request body by default. Set @@ -254,6 +254,45 @@ Exactly one of `http`, `cmd`, or `remove_path`: action `params` as placeholders); the resolved target must stay under the declared root. Symlink escapes and missing targets are rejected. Does **not** require the engine to be running. +- **`hf_cache`** — list the models in a Hugging Face hub cache directory. + `hf_cache.dir` is templated, normally `"{hf_hub_cache}"`. The reply is + `{"data":[{"id","size"}]}`, so `result: {"array": "data", "field": "id"}` + turns it into a model list. A repository is listed when any of its snapshots + is a complete download holding a config, a tokenizer and the weights vLLM + loads, judged on that snapshot's own files (huggingface_hub links a file into + a snapshot only once it is whole, so a download still running for another + revision or an unused file does not hide a complete one): vision-only and + GGUF-only repositories, unfinished downloads (weights not yet linked, shards + missing from the index or from a `-0000N-of-0000M` set, a lone shard of a + larger set, several unnumbered weight files with no index) and links whose + blob is gone are left out. `size` is what the snapshot takes on disk. This is + a check on file layout, not on whether the engine supports the architecture. + A missing directory is an empty cache. Does **not** require the engine to be + running. With `"delete": true` and `"tool": "hf"` (only as `delete_model`), + the action instead removes the model named by `{"model": ""}` by + running ` cache rm model/ --cache-dir --yes`, the tool + being looked up beside the engine binary as for `hf_pull`. The Hugging Face + tool keeps shared-store files another repository still links to and takes the + store's locks; the manager adds the served-model checks around it, asking a + running engine which model it serves through its `loaded_models` action and + refusing when it cannot tell, and holds the engine's operation lock until the + removal and any rollback of the launch settings are done. +- **`launch_model`** — switch the model an engine serves, for an engine with no + load call because the model is a launch argument (vLLM). `launch_model.flag` + names the flag (`"--model"`). The action takes `{"model": ""}` (a + Hugging Face id such as `org/name`; anything that could be read as a flag, a + path or a shell fragment is refused), writes it into the saved launch + settings exactly as the settings editor would, and restarts a running engine + on it. If the new model does not come up, the previous one is restored and the + error says so; a stopped engine is started on it, and asking a running engine + for the model already in the launch arguments restarts nothing. The engine's + operation lock is held from reading the launch to the last restore. An engine + with a `launch_model` action will not start while its launch arguments lack + the flag: it has no model selected, and refusing beats letting the engine + pick and download a default of its own. The refusal is reported as a failed + start, like the other refusals before launch, so a client that does not wait + for the reply still learns of it. Requires `runtime.editable_launch` + and `runtime.ready` on every platform. An optional **`result`** declares how `engine:models` extracts a normalized name list from the action's JSON response: `array` is the top-level array @@ -313,6 +352,86 @@ placeholder before the action runs: declared `remove_path.root` via `lms ls --json`. When no indexed file matches, the runner falls back to `{root}/{model}` if that path exists. +A manifest-level **`advertised_models`** (`"list"`, the default, or `"loaded"`) +chooses what the node tells its peers and its own proxy it can serve. The proxy +sends a request to any node advertising the model, so an engine that serves one +model per process sets `"loaded"` and advertises only `loaded_models`, even when +its `list_models` offers many models it could be switched to. With `"loaded"`, +a `loaded_models` action is required, and `engine:models` reports the engine in +two parts: + +- `modelsByEngine` (and the flat `models`) holds the `loaded_models` names: what + a request may be routed to. It is absent while the engine is stopped. +- `downloadedByEngine` holds the `list_models` names: every model the engine can + load, so a model manager on this node or a peer keeps its rows for models that + are on disk but not being served. It is read whenever the engine is installed, + running or not, so its `list_models` must not need a running server (vLLM's + `hf_cache` list reads the cache). An empty list means the inventory was read + and is empty; a missing key means it was not read. + +Its `loaded_models` `result` may also set **`inventory_field`**, the element field +that names each loaded model the way `list_models` does. vLLM's `/v1/models` +gives the name a request uses as `id` (a `--served-model-name` alias when one is +set) and the model it loaded as `root` (its `--model`), so its +`{"array":"data","field":"id","inventory_field":"root"}` advertises the `id` and +reports `loadedByEngine` as the `root`, which marks the cached model's row as +loaded. `inventory_field` is valid only there. + +An **`hf_pull`** action (`{"tool": "hf", "dir": "{models_dir}/hub"}`, only as +`pull_model`) downloads a Hugging Face repository with the named tool, looked up +in the directory that holds the engine binary, and reports byte progress taken +from the cache directory. It needs no running engine. + +A **`stop_engine`** action (only as `unload_model`) stops the engine, for one that +serves a single model per process and so cannot unload it while running. A +`launch_model` action on a stopped engine saves the choice and starts it. + +A runtime **`auto_launch`** (`{"kind": "hf_context", "context_tokens": 65536, "max_num_seqs": 32}`) +lets the manager add start-time arguments that depend on the model and the +machine. For `hf_context` it reads the served model's `config.json` from the +Hugging Face cache, at the revision the engine loads (`--revision` or `main`, +resolved through `refs/` as `huggingface_hub` does), and the free GPU memory, +then appends `--max-model-len` +(the target, the model's own limit when it cannot be extended, or what fits), +`--hf-overrides` with YaRN rope scaling when the target exceeds the model's +native window, and `--gpu-memory-utilization` sized to the weights the engine +loads plus KV cache, never below 15% of the GPU and never above what is free +less a 1 GiB margin (the free-memory cap wins). +A model that is not in the cache yet, which vLLM downloads during the start, +first has its files other than weights fetched into the cache with the engine's +Hugging Face tool (the one an `hf_pull` action names) at that revision, and is +then sized and given a tool-call parser the same way; when Hugging Face cannot +answer, the share is the free GPU memory less the margin. That fetch, and the +`hf_pull` and `hf_cache` work, run with the engine's launch environment, read +the way `huggingface_hub` reads it: `HF_HUB_OFFLINE` (else +`TRANSFORMERS_OFFLINE`), the token (`HF_TOKEN`, `HUGGING_FACE_HUB_TOKEN`, or +the stored login under `HF_TOKEN_PATH` or `HF_HOME`, none of them sent when +`HF_HUB_DISABLE_IMPLICIT_TOKEN` is set), and a mirror in `HF_ENDPOINT` all +apply. The readiness `timeout_s` counts from when the start begins, so this +work before launch is inside it. +The optional `max_num_seqs` adds `--max-num-seqs`, the most requests the engine +runs at once (0 or absent leaves the engine's default). +Arguments the user saved always win and nothing added is written back to saved +settings, so a model switch is re-sized for the new model. It also adds +`--enable-auto-tool-choice` and a `--tool-call-parser` found from the model's own +chat template with the installed vLLM (see the README's vLLM section), or nothing +when no parser reads the model's tool calls. Tensor- or +pipeline-parallel launches are left alone. It requires a process-mode runtime +with `editable_launch`. + +A runtime **`unsupported_weights`** (`{"formats": ["fp8"], "except_model_types": +[...], "reason": "..."}`) names checkpoint formats the platform's engine cannot +load. `fp8` is the only format: FP8 weights or an FP8 KV cache, read the way +vLLM reads them (`quant_method` `fp8` or `fbgemm_fp8`, a compressed-tensors +float 8-bit scheme, or ModelOpt's FP8 `quant_algo` / `kv_cache_quant_algo` in +`config.json` or `hf_quant_config.json`). A model whose `config.json` +`model_type` is listed in `except_model_types` is allowed. A start of such a +model is refused before the engine loads it, a `pull_model` before it +downloads, and `engine:catalog` leaves FP8 repositories out for a target +platform that declares the rule. `reason` is shown to the user. Validation +requires a known format and a reason. The bundled vLLM manifest declares it for +`darwin/arm64` only. + An optional **`restart_after`** restarts the engine once the action succeeds. Declare it on an action that changes state the engine only reads at startup: LM Studio's `delete_model` removes the files, but its server answers @@ -376,6 +495,7 @@ validation at load: | `{download}` | Path of the verified download | `install.run` | | `{download_}` | Path of a verified member of `install.artifacts` | `install.run` | | `{install_dir}` | Per-engine user-scoped install dir | `detect`, `install`, runtime | +| `{hf_hub_cache}` | The Hugging Face hub cache the engine itself uses, from its launch environment as `huggingface_hub` resolves it: `HF_HUB_CACHE`, else `HUGGINGFACE_HUB_CACHE`, else `HF_HOME/hub`, else `$XDG_CACHE_HOME/huggingface/hub`, else `~/.cache/huggingface/hub` | the `dir` of an `hf_cache` or `hf_pull` action only | | `{models_dir}` | The platform's `models_dir`, expanded | `detect`, runtime args/env, `uninstall.run`, actions. **Not** `install`, which resolves only `{install_dir}` and the download paths — and not `uninstall.remove`, where naming the store is rejected at load. Unset when the platform declares no store, so a template referencing it then fails to resolve rather than silently becoming a path under `/` | A `cmd` action additionally templates the action's own `params` as @@ -393,12 +513,18 @@ field is missing, `manifest_version` is unsupported, a platform key isn't nor nonempty `artifacts`, `fetch` or nonempty `artifacts` is present without nonempty `install.run`, `install.script` is combined with `fetch`/`artifacts`/`run`, `install.mode` or -`runtime.mode` is invalid, `models_dir` is a template rather than a literal +`runtime.mode` is invalid, `runtime.unsupported_weights` lacks a reason or +names a format other than `fp8`, `models_dir` is a template rather than a literal path, an `uninstall` block sets neither `run` nor `remove`, `uninstall.remove` is present without a `models_dir` to preserve, an `uninstall.remove` target is the model store or inside it (whether spelled literally or as `{models_dir}`), an action sets none or more than one of -`http`/`cmd`/`remove_path`, a `remove_path` action omits `root` or +`http`/`cmd`/`remove_path`/`hf_cache`/`launch_model`, an `hf_cache` action omits +`dir` or sets only one of `delete` and `tool`, a `launch_model` flag is not a flag or a platform lacks +`editable_launch` or `ready`, `advertised_models` is unknown or `"loaded"` +without a `loaded_models` action, `result.inventory_field` is set anywhere but +the `loaded_models` action of an engine with `advertised_models` `"loaded"` or is +not a single field name, a `remove_path` action omits `root` or `path`, `http.params_in` is not `body` or `query`, a `result` is set without both `array` and `field` (or a `result.match` without both `match.field` and a non-empty `match.in`), or @@ -518,7 +644,7 @@ GPU selection, and auth. | **Ollama** | strong (env-first) | `ollama serve` (foreground) | env: `OLLAMA_HOST`, `OLLAMA_MODELS`, `OLLAMA_KEEP_ALIVE`, `OLLAMA_NUM_PARALLEL`, `OLLAMA_MAX_LOADED_MODELS`, `OLLAMA_MAX_QUEUE`, `OLLAMA_CONTEXT_LENGTH`, `OLLAMA_FLASH_ATTENTION` | HTTP `/api/tags`, `/api/pull`; CLI `ollama pull/ls/ps/stop` | | **llama.cpp** | strong (env+flags) — **reference design** | `llama-server --host 127.0.0.1 --port {port}` (foreground) | flags + `LLAMA_ARG_*` (host/port, ctx-size, n-parallel, cont-batching, flash-attn, device, n-gpu-layers, tensor-split, main-gpu, api-key, models-dir/max) | HTTP `/v1/models`, `/models` download/delete, `/models/load`, `/models/unload`, `/health`, `/slots` | | **LM Studio** | command / daemon | `lms daemon up` → `lms server start --port {port}` | small env (`LMS_SERVER_HOST`, `LM_API_TOKEN`); most config is flags/API/settings (`lms load --context-length/--gpu/--ttl`) | HTTP `/api/v1/models[/download\|load\|unload]`; CLI `lms ls/get/load/unload/ps` | -| **vLLM** | flags-first — **Linux/WSL only** | `vllm serve --host 127.0.0.1 --port {port}` | flags (host/port/api-key); `HF_HOME` for cache. `VLLM_PORT`/`VLLM_HOST_IP` are **not** the API bind | OpenAI `/v1/models`; one model per process (unload = restart) | +| **vLLM** | flags-first — **Linux/WSL and Apple Silicon macOS** | `vllm serve --host 127.0.0.1 --port {port}` | flags (host/port/api-key); `HF_HOME` for cache. `VLLM_PORT`/`VLLM_HOST_IP` are **not** the API bind | OpenAI `/v1/models`; one model per process (unload = restart) | | **Jan** | hybrid (on llama.cpp router) | `jan serve --port {port}` (CLI) | forwards `LLAMA_ARG_*`; perf settings are router-preset-driven | `jan serve` auto-downloads HF repos | | **GPT4All** | weak (GUI app) | none (Local API Server toggled in the GUI, `:4891`) | desktop settings / Python SDK — not env | OpenAI `/v1` once enabled in-app | @@ -526,5 +652,5 @@ Caveats worth encoding when authoring these: - **Model store location**: leave each engine's model cache at the engine's own default — `~/.ollama`, `~/.llamacpp`, `~/.lmstudio/models` — and declare that path as `models_dir`. Never point one inside `{install_dir}` or anywhere else under the app data root: `uninstall` skips `models_dir`, but the app-level "remove all data" uninstall deletes the entire app data root and nothing can exempt a subdirectory of it. See Model stores. - **llama.cpp** ships prebuilt archives with **published SHA-256s**, so it's an ideal checksum-pinned `fetch`+`run` target (no placeholder SHA needed). - **LM Studio** model dir is settings-controlled (no documented relocation env); surface its config as flags/API, not env. -- **vLLM**: native Windows is unsupported (WSL only); the bind is a flag, not env. +- **vLLM**: native Windows is unsupported (WSL only), and macOS is Apple Silicon only through the vLLM-Metal plugin (`darwin/arm64`); the bind is a flag, not env. The bundled `vllm.json` is the reference: it passes the model as `--model`, since manifest `runtime.args` have no `{model}` token, and installs through a sha256-pinned `uv` bootstrap with a detect-first path. - **GPT4All**: no first-class headless CLI today — out of scope for a config-only manager unless targeting its Python SDK. diff --git a/services/nvpair-engine-manager/README.md b/services/nvpair-engine-manager/README.md index 8ca398cc2..33e04c8ca 100644 --- a/services/nvpair-engine-manager/README.md +++ b/services/nvpair-engine-manager/README.md @@ -6,7 +6,7 @@ SPDX-License-Identifier: Apache-2.0 # nvpair-engine-manager A config-driven control plane for local inference engines. The bundled -manifests support Ollama, LM Studio, and llama.cpp. It manages everything about +manifests support Ollama, LM Studio, llama.cpp, and vLLM. It manages everything about an engine **except serving inference**: detect, user-mode install, start/stop/restart, health, and config-declared actions. Adding an engine is a JSON manifest, not code. @@ -39,6 +39,173 @@ supported NVIDIA/Metal devices can accelerate, while the dynamic CPU backend remains the fallback. Hardware acceptance, not `/health` alone, is required to claim GPU activation. +### vLLM + +The `vllm` manifest covers Linux (`linux/amd64` and `linux/arm64`, NVIDIA GPU +with a CUDA 13 driver) and Apple Silicon macOS (`darwin/arm64`, macOS 15 or newer, +through the vLLM-Metal plugin). There is no native Windows path, so on Windows +`engine:get-installed` leaves vLLM out and PAIR does not offer it. The +macOS install builds a managed Python 3.12 venv with a sha256-pinned `uv` and +installs the pinned vLLM core and vLLM-Metal wheels, each checked against a +digest, so no compiler is needed. Every platform also pins `mistral_common==1.12.0`: +vLLM 0.32 asks only for 1.11.6 or newer, but 1.13 made an argument of the +tool-call grammar builder required that vLLM 0.32 does not pass, so with the +Mistral tool-call parser every request to a Mistral model failed. vLLM 0.33 +makes the same call, so the pin changes only with an upgrade to a vLLM release +that itself requires `mistral_common` 1.13. `vllm serve` listens on loopback `8001`, and +the proxy facade owns `8000`. vLLM serves **one model per process** and has no +load, unload, pull or delete API of its own, so the manifest provides them as +manager actions: `load_model`, `unload_model`, `pull_model` and `delete_model`, +described below. PAIR passes requests straight to vLLM and does not load or +unload models on its own: a model is loaded with Load and released with Eject. + +The default launch serves `Qwen/Qwen3.5-4B` through `--model `. +The model thinks before it answers, and with no `--reasoning-parser` in the launch +that reasoning arrives in the reply ahead of the answer; adding +`--reasoning-parser qwen3` in the launch settings moves it to +`reasoning_content`. The model is about 9.3 GB of weights, so the default needs a GPU with at least 16 GB of +memory; on a smaller GPU, download and load a smaller model before starting. +Edit the `--model` argument, or use `load_model`, to serve another model. Only +`--host` and `--port` are managed controls, and there is no CORS control. +The context window and GPU share are chosen automatically at every start +(`auto_launch`, kind `hf_context`), not pinned in the saved launch. The manager +reads the model's `config.json` from the Hugging Face cache, at the revision +vLLM loads (`--revision`, or `main` through `refs/main`, the way +`huggingface_hub` resolves it; another cached revision is never read in its +place), and the free GPU +memory (`nvidia-smi`; on a unified-memory machine such as DGX Spark, where it +reports `[N/A]`, the total and available system memory from `/proc/meminfo`), +then adds: `--max-model-len` of 64K (PAIR's agent-ready default; a +model with a larger window is capped there, and a smaller one is shrunk to what +fits), a YaRN `--hf-overrides` when the model's own window is below 64K (for +example Qwen3-8B, which stops at 40,960), and a `--gpu-memory-utilization` just +large enough for the weights, that KV cache and overhead, instead of vLLM's +92% default. The weights are the files vLLM loads, not everything in the +snapshot: an unused PyTorch copy beside the safetensors is not counted. The +share is at least 15% of the GPU but never more than is free less 1 GiB, which +wins when the two disagree. A model with its own rope scaling is never overridden. Anything +the user saves in the launch arguments wins, and a tensor- or pipeline-parallel +launch is left untouched. A model that is not yet downloaded, such as the +default on a fresh install, first has its small files (config, tokenizer, chat +template; everything but the weights) fetched into the cache with the `hf` tool +beside `vllm`, at the revision vLLM loads, and its weight sizes read from +Hugging Face. It is then sized, and its tool-call parser chosen, exactly as a +cached model is, and vLLM finds those files already there. The fetch and every +other Hugging Face request PAIR makes for vLLM (pull, delete, size lookup) run +with the engine's own environment, read as `huggingface_hub` reads it, so its +launch settings apply to them too: `HF_HUB_OFFLINE`, else `TRANSFORMERS_OFFLINE`; +the token (`HF_TOKEN`, `HUGGING_FACE_HUB_TOKEN`, or the one `hf auth login` +stored under `HF_HOME` or `HF_TOKEN_PATH`), none of which is sent when +`HF_HUB_DISABLE_IMPLICIT_TOKEN` is set; a mirror in `HF_ENDPOINT`; and the cache folder (`HF_HUB_CACHE`, +else `HF_HOME/hub`), or `--download-dir` when the launch arguments set one, since vLLM +keeps a model's weights, config and tokenizer there. (vLLM's Transformers config lookup +also caches a few MB of config and tokenizer files in the default cache; Delete leaves +that shared cache alone, since other Hugging Face tools may use it, and no weights go +there.) The model list, Pull, Delete and +start-time sizing all use the cache vLLM itself will use. If Hugging Face cannot answer (offline, +gated without a token, unreachable), the model gets the free GPU memory less +1 GiB and no parser. Every start also adds `--max-num-seqs 32`: vLLM's default of 256 +concurrent requests does not fit a hybrid model such as Qwen3.5-4B on a 16 GB +GPU, and a desktop does not need it. The added arguments appear in the engine-manager log +(`engine launch adjusted for the model and GPU`) and not in the saved settings, +so switching models re-sizes for the new one. +It also enables tool calling, which agent clients need on every request +(vLLM answers HTTP 400 to `tool_choice: "auto"` otherwise). The parser is found +from the model, not from a list of families: `hf` when the tokenizer ships +`response_template` metadata, otherwise the parser of the installed vLLM that +reads the model's own tool-call format. To find it, `toolparser_probe.py` runs +with vLLM's Python, renders the model's chat template with one example tool call, +and offers that text to every parser vLLM provides; among parsers that read it +equally, `hermes`, `qwen3_coder`, `llama3_json`, `mistral` and `pythonic` are +preferred in that order. The probe downloads nothing, runs no code from the model +repository, takes a few seconds, and is done once per model and vLLM while the +service runs. It runs under the start, so cancelling the start (Stop, a caller +giving up, shutdown) stops it. It reads the chat format vLLM will use: the +tokenizer `--tokenizer` and `--tokenizer-revision` name, if set, and a +`--chat-template` file or inline template. When that tokenizer is not cached, +or the template is a relative path vLLM would resolve against its own working +directory, no parser is chosen rather than one read from the wrong template. The chosen parser appears in the log +(`vLLM tool-call parser chosen from the model's chat template`). +When no parser reads the model's format (Nemotron Nano v2's ``, or a +template with no tool support), nothing is added, and a request that asks for +tool calls gets vLLM's own error (HTTP 400). A parser named in the launch +arguments is always used as given. +The environment sets `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's FlashInfer sampler +compiles GPU kernels on first start, and that build cannot handle the spaces in +PAIR's app-data folder (`Nvidia Corporation/Personal AI Router`), so the +built-in sampler is used instead. As a side effect no CUDA toolkit or C++ +compiler is needed. The same limit means anything that forces a FlashInfer +kernel build (for example choosing a FlashInfer attention backend) fails until +the install can live on a path without spaces. +Readiness polls `/health` for up to eight minutes while weights load. + +The model list is the Hugging Face cache (`~/.cache/huggingface/hub`, or the +one `HF_HUB_CACHE` or `HF_HOME` in vLLM's launch settings names), not what +vLLM happens to be serving: `list_models` is an `hf_cache` action listing the +complete, text-model-shaped downloads (see `MANIFEST.md`), and it works while +vLLM is stopped. Peers and the proxy, though, are only told about the model being +served (`advertised_models: loaded`, read from vLLM's `GET /v1/models`), because +the proxy routes a request to any node advertising its model and vLLM cannot +answer for the others. A stopped vLLM advertises nothing. `load_model` switches +the served model: it writes `--model ` into the saved launch arguments and +restarts vLLM on it, restoring the previous model if the new one does not start +(for example it does not fit in GPU memory). + +`unload_model` is a `stop_engine` action: vLLM cannot drop its model while running, so ejecting stops the engine, which frees all of its GPU memory and keeps the saved model choice; a model the engine reports it is not serving is refused. Eject and `engine:stop` cancel a start that is still in flight, as shutdown does, instead of waiting out its readiness window: a Load cut short this way puts the previous model back in the launch arguments without starting it and replies `{model, restarted: false, stopped: true}` rather than an error, since the user asked for it, and a start stopped by the user is not reported as a failed start. Likewise `engine:start` (including `engine:install` with `start`), `engine:restart` and a peer's remote start that a Stop cut short reply with the engine's status, stopped, rather than an error. The broker relays an `engine:stop` without waiting for the lock a Start or Restart holds, so it reaches the engine-manager while that start is still waiting for readiness. `load_model` on a stopped engine saves the choice and starts it, restoring the previous choice if it does not start. + +`delete_model` removes a model from the hub cache with the `hf cache rm` command of the `hf` tool beside `vllm` (an `hf_cache` action with `delete` and `tool`). Hugging Face's own tool decides which files go: it keeps any file in the cache's shared `blobs/` store that another model, dataset or Space still links to, and takes that store's locks, so a delete cannot remove a file another program is using or linking at the same moment. The cache is the user's standard Hugging Face cache, shared with every other Hugging Face tool on the machine. The model a running vLLM is serving is refused; which model that is is asked of vLLM itself (`loaded_models`), because an adopted vLLM may serve a model other than the one its saved arguments name, and the delete is refused when vLLM cannot say. Deleting the model an ejected or stopped vLLM is set to serve also takes `--model` out of its launch arguments, and a start with no `--model` is refused with an error rather than letting vLLM download a default model, until `load_model` selects one. The engine's operation lock is held from that check through the removal and any rollback, so a settings save or start made meanwhile waits and is not undone by a failed delete. `load_model`, `unload_model` and `delete_model` share that one launch-change path. + +`pull_model` is an `hf_pull` action: it asks Hugging Face for the repo's file list +and sizes, picks the files vLLM reads (safetensors, config and tokenizer; not +GGUF, ONNX, pickle copies of weights that also exist as safetensors, or +`original/` checkpoints), and runs the `hf` tool +that ships beside `vllm` in the engine's virtualenv to download them into the hub +cache. Like the other engines' pulls (one Ollama tag, one llama.cpp +quantization), it downloads what the engine loads, not the whole repository: the +same set vLLM fetches when it downloads a model itself. A Mistral repository +ships each model in two formats of the same size; vLLM's default +`load_format auto` loads Mistral's own `consolidated*.safetensors` with +`params.json` and `tekken.json` whenever a repository has them, so only those +weights are downloaded and the Hugging Face-format `model-*.safetensors` are +skipped. A repository that ships its weights both as indexed shards and as one +file (`model.safetensors` beside `model-0000N-of-…` and +`model.safetensors.index.json`) has its index downloaded first, as vLLM's own +download does, and only the shards it lists are fetched; vLLM loads only those +(`filter_duplicate_safetensors_files`). Pull, the first-start size lookup and +the cache listing choose weights with one rule, so a first start and a cached +start size a model alike, and launch sizing counts only the indexed shards. +The cache listing judges such a model complete on those files. Other formats in the +repository are not fetched; a full `hf download` adds them to the same cache +without downloading the rest again. Progress is the +share of those bytes that has landed in the cache, a download that makes no +progress for the pull timeout is abandoned, and a cancelled one is stopped. It +needs no running engine. A gated repo needs its terms accepted and `HF_TOKEN` +(or a prior `hf auth login`), in the service's environment or the engine's launch +settings. + +On Apple Silicon, vLLM-Metal stores the KV cache in the model's dtype and +refuses an FP8 KV cache, and its loader has no FP8 path for most model families. +The `darwin/arm64` platform therefore declares `unsupported_weights` (see +`MANIFEST.md`): a start or download of an FP8 checkpoint is refused with that +reason before anything is loaded or downloaded, a failed Load keeps the model +that was serving, and the Mac catalog leaves FP8 repositories out. Qwen3.5 and +Qwen3.6, which vLLM-Metal dequantizes itself, are allowed. Linux, DGX Spark +included, loads FP8 as before. + +Install is detect-first. `{install_dir}/venv/bin/vllm` (the managed install) and +`~/.local/bin/vllm` are checked in that order, and the first one found is +launched, so an existing install is adopted and nothing is downloaded. Otherwise +the install downloads a sha256-pinned `uv` release, lets `uv` fetch a managed +Python 3.12, creates `{install_dir}/venv` and installs a pinned `vllm` wheel from +PyPI. No source build happens. The pinned wheel is a CUDA 13 build, so the install +first checks that `nvidia-smi` exists, that the driver supports CUDA 13 or newer +(it reads either the `CUDA UMD Version` or the older `CUDA Version` field), and +that at least 25 GB of disk is free, enough for the engine and its default model. It fails with a clear message otherwise, +before downloading anything. The +download is several GB, because `torch` and the CUDA libraries come with it. +Model weights download from Hugging Face on first start into `~/.cache/huggingface` +(`HF_HOME`), which no removal this service performs will delete. + ## Communication Bidirectional newline-delimited JSON-RPC 2.0 — the same conventions as @@ -55,9 +222,9 @@ Requests (caller → service): | Method | Params | Result | |---|---|---| -| `engine:get-installed` | — | `{ engines: [EngineStatus] }` | +| `engine:get-installed` | — | `{ engines: [EngineStatus] }` for every engine this host can run; an engine whose manifest has no block for the host's OS and CPU is left out | | `engine:describe` | `{ engine }` | the engine's manifest | -| `engine:status` | `{ engine }` | `EngineStatus` | +| `engine:status` | `{ engine }` | `EngineStatus`; an error for an engine this host cannot run | | `engine:install` | `{ engine, start?, port?, bind? }` | `EngineStatus` (after install; also starts it if `start:true`) | | `engine:uninstall` | `{ engine }` | `EngineStatus` (after removal) | | `engine:uninstall-managed` | — | `{ engines: [{ engine, removed, error? }] }` — removes every engine this service installed, selected by its install marker, each through the `engine:uninstall` path. Backs "reset all data" in both clients; downloaded models are kept. Neither `removed` nor `error` means the engine was left alone as not ours | @@ -68,8 +235,8 @@ Requests (caller → service): | `engine:action` | `{ engine, action, params }` | the engine's raw response. `action:"pull_model"` is streamed: it emits live `engine:pull-progress` notifications and returns the pull's terminal result (see below). An action whose manifest declares `restart_after` (LM Studio's `delete_model`) restarts a running engine before replying, so the response also means the engine is back and healthy | | `engine:logs` | `{ engine }` | `{ lines: [LogLine] }` | | `engine:errors` | — | `{ errors: [ServiceError] }` | -| `engine:catalog` | `{ engine, platform?, arch?, query? }` | `{ models: [{ id, name, author, url, size?, downloads, likes, updatedAt, tags, family?, parameterSize?, appleOnly? }], source, platform?, arch?, fetchedAt?, searchable?, query? }` — the models an engine can **download**, as opposed to `engine:models`, which reports what is already installed. Each engine has one curated source. Ollama's has no public API, so it is a locked list compiled into this binary from `catalog/ollama-models.json`, regenerated on demand by a developer (`npm run scrape:ollama-models` in `desktop/`) who reviews the diff; serving it needs no network. LM Studio's catalog *is* the `lmstudio-community` Hugging Face org, whose repo ids are exactly what `lms get` accepts, so it is fetched live and cached for six hours, with concurrent callers coalesced onto one request, a failed attempt — with or without a list to fall back on — followed by a backoff before the next. llama.cpp downloads any GGUF repo as `repo:quantization`; its browse list is the GGUF repos of `ggml-org`, `bartowski`, and `unsloth`, fetched and cached the same way, keeping only open, text-generating repos that ship a `Q4_K_M` file, each offered as `repo:Q4_K_M`. It is the one `searchable` source: a `query` searches all of Hugging Face under the same rules, each query cached on its own (the 20 most recently used are kept), and the reply echoes it. A query longer than 100 characters is refused rather than shortened, since a shortened one searches for something the operator did not type. Other sources ignore a query and return their whole list for the client to filter. `platform` and `arch` are the GOOS and GOARCH the models will be **installed on**, which is not always this host — a client driving a peer should name that peer. Omitting both means this host; naming a platform without an arch leaves the arch unknown rather than borrowing this host's. MLX quantizations only install on Apple Silicon, so they are marked `appleOnly` and kept only for a `darwin`/`arm64` target — an Intel Mac does not get them; the reply echoes the `platform` and `arch` it filtered for. `name` is pull-ready: it can be handed to `pull_model` verbatim. An engine with no curated source is an error, not an empty list. The Ollama reply is a single multi-megabyte frame, so every hop on its path must allow `jsonrpc.WorkerFrameBytes`. | -| `engine:models` | — | `{ models: [string], modelsByEngine: { : [string] }, loadedByEngine: { : [string] } }` — the flat de-duplicated union of every running engine's models, the per-engine breakdown keyed by engine name, and the per-engine set of models currently **loaded in memory** (all normalized from each engine's `list_models` / `loaded_models` action `result` spec). `modelsByEngine` carries a key for every running engine whose inventory was successfully queried, including an empty list = "running, no models available"; a missing key means not running / not queryable / invalid response. `loadedByEngine` uses the same known-empty distinction for residency and also omits engines with no loaded endpoint. The `/v1/models` HTTP surface returns the same shape. | +| `engine:catalog` | `{ engine, platform?, arch?, query? }` | `{ models: [{ id, name, author, url, size?, downloads, likes, updatedAt, tags, family?, parameterSize?, appleOnly? }], source, platform?, arch?, fetchedAt?, searchable?, query? }` — the models an engine can **download**, as opposed to `engine:models`, which reports what is already installed. Each engine has one curated source. Ollama's has no public API, so it is a locked list compiled into this binary from `catalog/ollama-models.json`, regenerated on demand by a developer (`npm run scrape:ollama-models` in `desktop/`) who reviews the diff; serving it needs no network. LM Studio's catalog *is* the `lmstudio-community` Hugging Face org, whose repo ids are exactly what `lms get` accepts, so it is fetched live and cached for six hours, with concurrent callers coalesced onto one request, a failed attempt — with or without a list to fall back on — followed by a backoff before the next. vLLM serves Hugging Face repos by their repo id: its browse list is the open, text-generating safetensors repos of `Qwen`, `nvidia`, `mistralai`, `microsoft` and `ibm-granite` (both `text-generation` and `image-text-to-text`), with parameter count and size, and, like llama.cpp, a `query` searches all of Hugging Face under the same rules. llama.cpp downloads any GGUF repo as `repo:quantization`; its browse list is the GGUF repos of `ggml-org`, `bartowski`, and `unsloth`, fetched and cached the same way, keeping only open, text-generating repos that ship a `Q4_K_M` file, each offered as `repo:Q4_K_M`. It is the one `searchable` source: a `query` searches all of Hugging Face under the same rules, each query cached on its own (the 20 most recently used are kept), and the reply echoes it. A query longer than 100 characters is refused rather than shortened, since a shortened one searches for something the operator did not type. Other sources ignore a query and return their whole list for the client to filter. `platform` and `arch` are the GOOS and GOARCH the models will be **installed on**, which is not always this host — a client driving a peer should name that peer. Omitting both means this host; naming a platform without an arch leaves the arch unknown rather than borrowing this host's. MLX quantizations only install on Apple Silicon, so they are marked `appleOnly` and kept only for a `darwin`/`arm64` target — an Intel Mac does not get them; the reply echoes the `platform` and `arch` it filtered for. `name` is pull-ready: it can be handed to `pull_model` verbatim. An engine with no curated source is an error, not an empty list. The Ollama reply is a single multi-megabyte frame, so every hop on its path must allow `jsonrpc.WorkerFrameBytes`. | +| `engine:models` | — | `{ models: [string], modelsByEngine: { : [string] }, loadedByEngine: { : [string] }, downloadedByEngine: { : [string] } }` — the flat de-duplicated union of every running engine's models, the per-engine breakdown keyed by engine name, and the per-engine set of models currently **loaded in memory** (all normalized from each engine's `list_models` / `loaded_models` action `result` spec). `modelsByEngine` carries a key for every running engine whose inventory was successfully queried, including an empty list = "running, no models available"; a missing key means not running / not queryable / invalid response. `loadedByEngine` uses the same known-empty distinction for residency and also omits engines with no loaded endpoint. For an engine with `advertised_models: "loaded"` (vLLM), `modelsByEngine` holds only the served model's name, which is what a request may be routed to, and `downloadedByEngine` holds its `list_models` inventory, read whenever it is installed, running or not, so a model manager keeps its rows after an Eject. Its `loadedByEngine` names the served model by its `--model` (the `root` of vLLM's `/v1/models`), not a `--served-model-name` alias, so it matches the inventory, and is `[]` while it is stopped. A missing `downloadedByEngine` key means the inventory was not read. The `/v1/models` HTTP surface returns the same shape. | | `engine:remote-get-installed` | `{ node }` | `{ engines: [EngineStatus] }` fetched from the remote node over `ec` mTLS | | `engine:remote-install` | `{ node, engine, start? }` | `{ opId, status: EngineStatus }` after the remote install (live progress via `engine:remote-progress`) | | `engine:remote-pull-model` | `{ node, engine, model?, params? }` | `{ opId, result }` after the remote pull (live progress via `engine:remote-progress`) | @@ -78,14 +245,20 @@ Requests (caller → service): | `shutdown` | — | `null` | | `log/set-level` | `{ level }` | `{ level }` | -`EngineStatus` = `{ engine, display_name, installed, running, healthy, port }`. +`EngineStatus` = `{ engine, display_name, installed, running, starting?, healthy, port }`. +`starting` is true while a start launches the engine and waits for it to be +ready, whoever began it (a Start, install-then-start, a settings restart, a +Load, or the restore at launch); `engine:state-changed` is pushed when it +begins and ends. A start refused before launch never reports it. Notifications (service → caller): `engine:ready{version}`, `engine:state-changed{EngineStatus}`, -`engine:models-changed{engine, models}` — pushed when an engine's set of -loaded (in-memory) models changes (explicit load/unload, JIT auto-load, or -TTL/idle eviction); `models` is the full `engine:models` shape (incl. -`loadedByEngine`) so a consumer swaps its whole snapshot, +`engine:models-changed{engine, models}` — pushed when one of an engine's +model lists changes: what it serves (`modelsByEngine`), what is loaded in +memory (`loadedByEngine`: explicit load/unload, JIT auto-load, or TTL/idle +eviction) or what it has downloaded (`downloadedByEngine`), so a pull or delete +is pushed too; `models` is the full `engine:models` shape so a consumer swaps +its whole snapshot, `engine:install-progress{engine, stage, percent?, error?}` (`error` on the terminal `failed` frame), `engine:pull-progress{engine, op, stage, percent?, message}` (live progress for a @@ -186,8 +359,14 @@ allowance on supported Windows systems. The deadline remains finite: if Ollama never serves its readiness endpoint, engine-manager stops the owned process and reports the failed start. On Unix, stopping an owned process sends SIGTERM to its process group, waits the manifest's `stop.grace_s` (five seconds by -default), then escalates to SIGKILL; `signal:"kill"` skips the grace. Failed -startup cleanup uses the same policy. On Windows, the windowless engines we +default), then escalates to SIGKILL, which also takes any helper processes the +engine forked (vLLM's `EngineCore` holds the GPU) that outlived it; `signal:"kill"` +skips the grace. Failed startup cleanup uses the same policy. While the whole +service is shutting down the grace is capped at three seconds, whatever the +manifest says: the parent allows engine-manager only a few seconds to stop its +engines before it kills it, and engine-manager cannot escalate once killed, so a +longer grace would leave the engine orphaned (an engine still loading its model +does not answer a polite signal at all). On Windows, the windowless engines we spawn cannot receive a graceful (non-`/F`) close, so stopping uses immediate `taskkill /T /F`. @@ -204,7 +383,9 @@ would only collide on the port. Consequences worth knowing: - **An adopted engine has no child process the service owns**, so `engine:stop` (and `engine:restart`'s stop phase, and shutdown's cleanup) resolves the PID bound to the engine's port and terminates it **only when that process is - running the very binary NVPAIR manages for the engine** — reclaiming an orphan + running the very binary NVPAIR manages for the engine** (for an engine + installed as a Python package, such as vLLM, that is the interpreter running the + engine's script, read from the process's command line on Linux) — reclaiming an orphan a prior run left on our own managed port (e.g. an `ollama serve` on `11435` whose handle was lost after a crash). This is precise to the port, so a genuine third-party listener on a *different* port (Ollama's own desktop app @@ -352,7 +533,9 @@ termination, console hiding) are the only build-tagged Go Shuts down on stdin EOF (parent closed the pipe), `SIGINT`/`SIGTERM`, or a `shutdown` JSON-RPC request — stopping any running engines first so none -are orphaned. +are orphaned. A signal shuts down at once even while stdin is idle and an engine +is still starting; the start is cancelled and its process group force-killed +within the three-second shutdown grace (see Process-mode stops above). That engine sweep (`StopAll`) runs **exactly once per process**, and every caller returns only once it has finished. An ordinary quit asks for it twice: @@ -367,3 +550,8 @@ the sweep, which would re-pay the readiness probe for every adopted engine whose stop is declined. They do retry any engine whose stop failed through a process handle or a stop command, such as an `lms server stop` that errored, because another attempt can succeed where a declined stop cannot. + +`engine:status` and `engine:get-installed` never wait for a lifecycle operation +(a start, stop, install or model switch can hold an engine for minutes): while +one runs they report what is already known and the operation reports its own +result. diff --git a/services/nvpair-engine-manager/actions.go b/services/nvpair-engine-manager/actions.go index d01d8e522..e95bddb00 100644 --- a/services/nvpair-engine-manager/actions.go +++ b/services/nvpair-engine-manager/actions.go @@ -31,6 +31,14 @@ func (e *Executor) Action(ctx context.Context, engine, action string, params jso if !ok { return nil, fmt.Errorf("engine %q has no action %q", engine, action) } + if act.LaunchModel != nil { + // Like restart_after, the restart gets the caller's context: it is a stop + // plus a readiness-probed start that the action timeout does not budget for. + return e.switchLaunchModel(ctx, engine, act.LaunchModel, params) + } + if act.StopEngine { + return e.ejectByStopping(ctx, engine, params) + } res, err := e.dispatchAction(ctx, st, engine, action, act, params) if err != nil { return nil, err @@ -101,6 +109,9 @@ func (e *Executor) dispatchAction(ctx context.Context, st *engineState, engine, if act.RemovePath != nil { return e.runRemovePathAction(ctx, st, act, params) } + if act.HFCache != nil { + return e.runHFCacheAction(ctx, st, engine, act, params) + } if len(act.Cmd) > 0 { return e.runCmdAction(ctx, st, act, port, params) } diff --git a/services/nvpair-engine-manager/catalog.go b/services/nvpair-engine-manager/catalog.go index ed52b1199..6433b1b03 100644 --- a/services/nvpair-engine-manager/catalog.go +++ b/services/nvpair-engine-manager/catalog.go @@ -159,10 +159,10 @@ const ( llamaCPPQuantization = "Q4_K_M" // llamaCPPLimit is the page size for each publisher and for each search. llamaCPPLimit = 50 - // llamaCPPSearchCacheLimit bounds how many distinct queries are kept, oldest - // use evicted first. Each is a short list, but the set of possible queries - // is not bounded at all. - llamaCPPSearchCacheLimit = 20 + // searchCacheLimit bounds how many distinct queries are kept across every + // engine's catalog, oldest use evicted first. Each is a short list, but the + // set of possible queries is not bounded at all. + searchCacheLimit = 20 // maxCatalogQuery bounds a query passed upstream, in characters. A longer // one is refused rather than shortened, which would search for something // the operator did not type. @@ -215,11 +215,19 @@ type hfModelRow struct { Gated json.RawMessage `json:"gated"` PipelineTag string `json:"pipeline_tag"` Siblings []hfSibling `json:"siblings"` + // Safetensors is present when the listing asks to expand it: the parameter + // count per tensor type, which gives a model's size without a second request. + Safetensors struct { + Parameters map[string]int64 `json:"parameters"` + Total int64 `json:"total"` + } `json:"safetensors"` } // hfSibling is one file in a Hugging Face repo. type hfSibling struct { RFilename string `json:"rfilename"` + // Size is reported only when the repo is asked for with blobs=true. + Size int64 `json:"size"` } // catalogService answers engine:catalog. The embedded list is parsed once; each @@ -234,12 +242,13 @@ type catalogService struct { lmStudio *listingCache llamaCPP *listingCache + vllm *listingCache // searchMu guards the llama.cpp searches, kept by lowercased query with // searchOrder from least to most recently used. - searchMu sync.Mutex - llamaSearches map[string]*listingCache - searchOrder []string + searchMu sync.Mutex + searches map[string]*listingCache + searchOrder []string // baseURL overrides the listing endpoint in tests. Empty means the real one; // the caching and coalescing around the fetch is the part worth testing, and @@ -281,11 +290,12 @@ func newCatalogService() *catalogService { // endpoint has no reason to redirect, and following one blindly would let an // upstream (or a captive portal) move the request somewhere else. c := &catalogService{ - client: newEngineHTTPClient(catalogHTTPTimeout), - llamaSearches: map[string]*listingCache{}, + client: newEngineHTTPClient(catalogHTTPTimeout), + searches: map[string]*listingCache{}, } c.lmStudio = &listingCache{label: "lm studio catalog", fetch: c.fetchLmStudio} c.llamaCPP = &listingCache{label: "llama.cpp catalog", fetch: c.fetchLlamaCPP} + c.vllm = &listingCache{label: "vllm catalog", fetch: c.fetchVLLM} return c } @@ -345,6 +355,28 @@ func (c *catalogService) Catalog(ctx context.Context, engine, platform, arch, qu Searchable: true, Query: query, }, nil + case "vllm": + query, err := catalogQuery(query) + if err != nil { + return CatalogResult{}, err + } + models, fetchedAt, err := c.vllmCatalog(ctx, query) + if err != nil { + return CatalogResult{}, err + } + source := hfModelsAPI + "?filter=safetensors&author=" + strings.Join(vllmPublishers, ",") + if query != "" { + source = hfModelsAPI + "?filter=safetensors&search=" + url.QueryEscape(query) + } + return CatalogResult{ + Models: models, + Source: source, + Platform: platform, + Arch: arch, + FetchedAt: fetchedAt.UTC().Format(time.RFC3339), + Searchable: true, + Query: query, + }, nil default: return CatalogResult{}, fmt.Errorf("no model catalog for engine %q", engine) } @@ -402,6 +434,8 @@ func normalizeCatalogEngine(engine string) string { return "lmstudio" case "llama-cpp", "llama.cpp", "llamacpp": return "llamacpp" + case "vllm": + return "vllm" default: return e } @@ -579,23 +613,25 @@ func (c *catalogService) llamaCPPCatalog(ctx context.Context, query string) ([]C // llamaCPPSearch is the cache for one query, created on first use. Queries // differing only in case share one. func (c *catalogService) llamaCPPSearch(query string) *listingCache { - key := strings.ToLower(query) + return c.searchCache("llamacpp:"+strings.ToLower(query), "llama.cpp search", func(ctx context.Context) ([]CatalogModel, error) { + rows, err := c.fetchHF(ctx, llamaCPPListing("search", query)) + if err != nil { + return nil, err + } + return normalizeListing(rows, llamaCPPModel), nil + }) +} + +// searchCache returns the cache for one search, creating it on first use and +// marking it most recently used. The set of possible queries is unbounded, so +// the least recently used are evicted past searchCacheLimit. +func (c *catalogService) searchCache(key, label string, fetch func(context.Context) ([]CatalogModel, error)) *listingCache { c.searchMu.Lock() defer c.searchMu.Unlock() - search, ok := c.llamaSearches[key] + search, ok := c.searches[key] if !ok { - search = &listingCache{ - label: "llama.cpp search", - emptyOK: true, - fetch: func(ctx context.Context) ([]CatalogModel, error) { - rows, err := c.fetchHF(ctx, llamaCPPListing("search", query)) - if err != nil { - return nil, err - } - return normalizeLlamaCPPRows(rows), nil - }, - } - c.llamaSearches[key] = search + search = &listingCache{label: label, emptyOK: true, fetch: fetch} + c.searches[key] = search } for i, k := range c.searchOrder { if k == key { @@ -604,26 +640,38 @@ func (c *catalogService) llamaCPPSearch(query string) *listingCache { } } c.searchOrder = append(c.searchOrder, key) - for len(c.searchOrder) > llamaCPPSearchCacheLimit { - delete(c.llamaSearches, c.searchOrder[0]) + for len(c.searchOrder) > searchCacheLimit { + delete(c.searches, c.searchOrder[0]) c.searchOrder = c.searchOrder[1:] } return search } -// fetchLlamaCPP reads every approved publisher's GGUF repos. Any publisher -// failing fails the whole fetch, so a refresh never replaces the cached list -// with one that is silently missing a publisher. +// fetchLlamaCPP reads every approved publisher's GGUF repos. func (c *catalogService) fetchLlamaCPP(ctx context.Context) ([]CatalogModel, error) { - batches := make([][]hfModelRow, len(llamaCPPPublishers)) - errs := make([]error, len(llamaCPPPublishers)) + queries := make([]url.Values, 0, len(llamaCPPPublishers)) + for _, publisher := range llamaCPPPublishers { + queries = append(queries, llamaCPPListing("author", publisher)) + } + rows, err := c.fetchListings(ctx, queries) + if err != nil { + return nil, err + } + return normalizeListing(rows, llamaCPPModel), nil +} + +// fetchListings reads several listing requests at once and returns their +// entries in request order. Any request failing fails the whole fetch, so a +// refresh never replaces the cached list with one that is silently missing a +// publisher. +func (c *catalogService) fetchListings(ctx context.Context, queries []url.Values) ([]hfModelRow, error) { + batches := make([][]hfModelRow, len(queries)) + errs := make([]error, len(queries)) var wg sync.WaitGroup - for i, publisher := range llamaCPPPublishers { - wg.Add(1) - go func() { - defer wg.Done() - batches[i], errs[i] = c.fetchHF(ctx, llamaCPPListing("author", publisher)) - }() + for i, q := range queries { + wg.Go(func() { + batches[i], errs[i] = c.fetchHF(ctx, q) + }) } wg.Wait() if err := errors.Join(errs...); err != nil { @@ -633,7 +681,7 @@ func (c *catalogService) fetchLlamaCPP(ctx context.Context) ([]CatalogModel, err for _, batch := range batches { rows = append(rows, batch...) } - return normalizeLlamaCPPRows(rows), nil + return rows, nil } // withoutURL drops the request URL a net/http error carries, keeping the cause. @@ -750,12 +798,12 @@ func isMLX(id string, tags []string) bool { return mlxPattern.MatchString(id) } -// The shapes a llama.cpp download depends on. +// The shapes a catalog download depends on. var ( - // llamaCPPRepoPattern is a repo id safe to pass to llama.cpp: owner/name, - // each starting with a letter or digit, nothing that could read as a path - // or a flag. - llamaCPPRepoPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$`) + // hfRepoPattern is a Hugging Face repo id safe to pass to an engine's + // download: owner/name, each starting with a letter or digit, nothing that + // could read as a path or a flag. + hfRepoPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$`) // llamaCPPQuantFile matches the offered quantization in a GGUF file name. llamaCPPQuantFile = regexp.MustCompile(`(?i)` + llamaCPPQuantization + `[.-]`) // splitGGUFPattern matches a shard of a model split across files. @@ -770,14 +818,16 @@ var generativePipelines = map[string]bool{ "any-to-any": true, } -// normalizeLlamaCPPRows converts listing entries into llama.cpp downloads, -// dropping the ones llama.cpp cannot fetch or serve, de-duplicating across -// publishers, and putting the most-downloaded first. -func normalizeLlamaCPPRows(rows []hfModelRow) []CatalogModel { +// normalizeListing converts listing entries into one engine's downloads with +// toModel, which drops the entries that engine cannot fetch or serve. Entries +// are de-duplicated across publishers and pipelines, the later one winning, and +// the most-downloaded come first. llama.cpp and vLLM read the same Hugging Face +// listing and differ only in toModel. +func normalizeListing(rows []hfModelRow, toModel func(hfModelRow) (CatalogModel, bool)) []CatalogModel { seen := make(map[string]int, len(rows)) out := make([]CatalogModel, 0, len(rows)) for _, r := range rows { - model, ok := llamaCPPModel(r) + model, ok := toModel(r) if !ok { continue } @@ -803,7 +853,7 @@ func llamaCPPModel(r hfModelRow) (CatalogModel, bool) { if repo == "" { repo = strings.TrimSpace(r.ModelID) } - if !llamaCPPRepoPattern.MatchString(repo) { + if !hfRepoPattern.MatchString(repo) { return CatalogModel{}, false } // Only an explicit false is open. A gated repo reports how it is gated, and diff --git a/services/nvpair-engine-manager/catalog_test.go b/services/nvpair-engine-manager/catalog_test.go index 98127021f..a99198c16 100644 --- a/services/nvpair-engine-manager/catalog_test.go +++ b/services/nvpair-engine-manager/catalog_test.go @@ -214,14 +214,17 @@ func TestNormalizeCatalogEngine_ResolvesSupportedEngineNames(t *testing.T) { test("hyphenated llama.cpp", "llama-cpp", "llamacpp") test("dotted llama.cpp", "llama.cpp", "llamacpp") test("other engine preserved", "vllm", "vllm") + test("unknown engine preserved", "nosuchengine", "nosuchengine") } // TestCatalogRejectsUnknownEngine checks an engine with no curated source errors // rather than returning an empty list that reads as "nothing available". func TestCatalog_RejectsUnknownEngine(t *testing.T) { c := newCatalogService() - _, err := c.Catalog(context.Background(), "vllm", "linux", "amd64", "") - assert.Error(t, err, "unknown engine returned a catalog") + for _, engine := range []string{"nosuchengine"} { + _, err := c.Catalog(context.Background(), engine, "linux", "amd64", "") + assert.Error(t, err, "%s returned a catalog", engine) + } } // TestNormalizeHFRowsMarksMLX checks Apple-only quantizations are labelled @@ -394,7 +397,7 @@ func TestNormalizeLlamaCPPRows_KeepsOnlyDownloadableModels(t *testing.T) { ggufRow("../escape", 1), ggufRow("a/has space", 1), } - got := normalizeLlamaCPPRows(rows) + got := normalizeListing(rows, llamaCPPModel) ids := make([]string, len(got)) for i, m := range got { ids[i] = m.ID @@ -411,7 +414,7 @@ func TestNormalizeLlamaCPPRows_KeepsOnlyDownloadableModels(t *testing.T) { // TestNormalizeLlamaCPPRowsDedupesAcrossPublishers checks a repo listed twice, // as it is when a search and a publisher page overlap, yields one row. func TestNormalizeLlamaCPPRows_DedupesAcrossPublishers(t *testing.T) { - got := normalizeLlamaCPPRows([]hfModelRow{ggufRow("a/model", 1), ggufRow("a/model", 7)}) + got := normalizeListing([]hfModelRow{ggufRow("a/model", 1), ggufRow("a/model", 7)}, llamaCPPModel) if assert.Len(t, got, 1, "one row must carry the later entry") { assert.Equal(t, 7, got[0].Downloads) } @@ -558,10 +561,10 @@ func TestSearchAtTheLengthLimit_IsAnswered(t *testing.T) { func TestLlamaCPPSearch_BoundsConcurrentQueries(t *testing.T) { c := newCatalogService() first := c.llamaCPPSearch("first") - for i := range llamaCPPSearchCacheLimit { + for i := range searchCacheLimit { c.llamaCPPSearch(fmt.Sprint("query ", i)) } - assert.Len(t, c.llamaSearches, llamaCPPSearchCacheLimit) + assert.Len(t, c.searches, searchCacheLimit) assert.NotSame(t, first, c.llamaCPPSearch("first"), "the least recently used search was not evicted") } diff --git a/services/nvpair-engine-manager/catalog_vllm.go b/services/nvpair-engine-manager/catalog_vllm.go new file mode 100644 index 000000000..6d25b457c --- /dev/null +++ b/services/nvpair-engine-manager/catalog_vllm.go @@ -0,0 +1,147 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "context" + "fmt" + "net/url" + "strings" + "time" +) + +// The vLLM catalogue: Hugging Face repos that ship safetensors weights and +// answer chat or completion requests. The download name is the repo id itself, +// which is what vLLM serves and what the Hugging Face cache stores. +const vllmListingLimit = 50 + +// vllmPublishers are the Hugging Face accounts whose repos make up the browse +// list: publishers of widely used open-weight models that vLLM serves directly. +// Gated repos are dropped, so a publisher that gates most of its models +// contributes only its open ones. +var vllmPublishers = []string{"Qwen", "nvidia", "mistralai", "microsoft", "ibm-granite"} + +// vllmPipelines are the pipeline tags listed. Vision-language chat models carry +// image-text-to-text and are served by vLLM like any other chat model. +var vllmPipelines = []string{"text-generation", "image-text-to-text"} + +func (c *catalogService) vllmCatalog(ctx context.Context, query string) ([]CatalogModel, time.Time, error) { + if query == "" { + return c.vllm.get(ctx) + } + return c.searchCache("vllm:"+strings.ToLower(query), "vllm search", func(ctx context.Context) ([]CatalogModel, error) { + rows, err := c.fetchListings(ctx, vllmListings("search", []string{query})) + if err != nil { + return nil, err + } + return normalizeListing(rows, vllmModel), nil + }).get(ctx) +} + +func (c *catalogService) fetchVLLM(ctx context.Context) ([]CatalogModel, error) { + rows, err := c.fetchListings(ctx, vllmListings("author", vllmPublishers)) + if err != nil { + return nil, err + } + return normalizeListing(rows, vllmModel), nil +} + +// vllmListings is one listing request per value and pipeline. +func vllmListings(field string, values []string) []url.Values { + var queries []url.Values + for _, v := range values { + for _, p := range vllmPipelines { + queries = append(queries, vllmListing(field, v, p)) + } + } + return queries +} + +// vllmListing is a listing request narrowed by one field, author or search, and +// one pipeline. expand asks for exactly the fields used, including the tensor +// parameter counts that give a model's size. +func vllmListing(field, value, pipeline string) url.Values { + q := url.Values{} + q.Set(field, value) + q.Set("pipeline_tag", pipeline) + q.Set("filter", "safetensors") + q.Set("sort", "downloads") + q.Set("direction", "-1") + q.Set("limit", fmt.Sprint(vllmListingLimit)) + for _, f := range []string{"safetensors", "downloads", "likes", "lastModified", "createdAt", "tags", "gated", "private", "pipeline_tag"} { + q.Add("expand[]", f) + } + return q +} + +// vllmModel turns one listing entry into a download, or reports it is not one. +// Gated and private repos fail to download without credentials, MLX is Apple's +// framework, and a listing can still include embedding models the chat endpoint +// cannot serve. +func vllmModel(r hfModelRow) (CatalogModel, bool) { + repo := strings.TrimSpace(r.ID) + if repo == "" { + repo = strings.TrimSpace(r.ModelID) + } + if !hfRepoPattern.MatchString(repo) { + return CatalogModel{}, false + } + if r.Private || string(bytes.TrimSpace(r.Gated)) != "false" { + return CatalogModel{}, false + } + if !hasTag(r.Tags, "safetensors") || !isGenerative(r) || isMLX(repo, r.Tags) { + return CatalogModel{}, false + } + updated := r.LastModified + if updated == "" { + updated = r.CreatedAt + } + if _, err := time.Parse(time.RFC3339, updated); err != nil { + return CatalogModel{}, false + } + return CatalogModel{ + ID: repo, + Name: repo, + Author: repo[:strings.Index(repo, "/")], + URL: "https://huggingface.co/" + repo, + Size: safetensorsBytes(r.Safetensors.Parameters), + Downloads: max(r.Downloads, 0), + Likes: max(r.Likes, 0), + UpdatedAt: updated, + Tags: r.Tags, + ParameterSize: parameterLabel(r.Safetensors.Total), + }, true +} + +// tensorBytes is the storage size of a safetensors element type. A type not +// listed makes the total unknown, which is reported as no size at all. +var tensorBytes = map[string]uint64{ + "F64": 8, "F32": 4, "BF16": 2, "F16": 2, "I64": 8, "I32": 4, "I16": 2, + "I8": 1, "U8": 1, "F8_E4M3": 1, "F8_E5M2": 1, "BOOL": 1, +} + +func safetensorsBytes(parameters map[string]int64) uint64 { + var total uint64 + for kind, count := range parameters { + width, ok := tensorBytes[kind] + if !ok || count < 0 { + return 0 + } + total += uint64(count) * width + } + return total +} + +// parameterLabel renders a parameter count the way model names do: 600M, 8B. +func parameterLabel(n int64) string { + switch { + case n >= 1_000_000_000: + return strings.TrimSuffix(fmt.Sprintf("%.1f", float64(n)/1e9), ".0") + "B" + case n >= 1_000_000: + return fmt.Sprintf("%dM", (n+500_000)/1_000_000) + default: + return "" + } +} diff --git a/services/nvpair-engine-manager/catalog_vllm_test.go b/services/nvpair-engine-manager/catalog_vllm_test.go new file mode 100644 index 000000000..fb126131d --- /dev/null +++ b/services/nvpair-engine-manager/catalog_vllm_test.go @@ -0,0 +1,127 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func vllmRow(id string, mutate func(*map[string]any)) map[string]any { + row := map[string]any{ + "id": id, "gated": false, "private": false, "downloads": 100, "likes": 5, + "lastModified": "2026-01-02T03:04:05.000Z", "pipeline_tag": "text-generation", + "tags": []string{"transformers", "safetensors", "text-generation"}, + "safetensors": map[string]any{"parameters": map[string]int64{"BF16": 8_000_000_000}, "total": 8_000_000_000}, + } + if mutate != nil { + mutate(&row) + } + return row +} + +func TestVLLMCatalog_KeepsOpenSafetensorsChatModels(t *testing.T) { + var mu sync.Mutex + var queries []map[string][]string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + queries = append(queries, r.URL.Query()) + mu.Unlock() + if r.URL.Query().Get("pipeline_tag") == "image-text-to-text" { + assert.NoError(t, json.NewEncoder(w).Encode([]any{vllmRow("Qwen/Vision-8B", func(m *map[string]any) { (*m)["pipeline_tag"] = "image-text-to-text"; (*m)["downloads"] = 900 })})) + return + } + assert.NoError(t, json.NewEncoder(w).Encode([]any{ + vllmRow("Qwen/Chat-8B", nil), + vllmRow("Qwen/Gated", func(m *map[string]any) { (*m)["gated"] = "auto" }), + vllmRow("Qwen/Private", func(m *map[string]any) { (*m)["private"] = true }), + vllmRow("Qwen/Embed", func(m *map[string]any) { (*m)["pipeline_tag"] = "feature-extraction" }), + vllmRow("Qwen/NoSafetensors", func(m *map[string]any) { (*m)["tags"] = []string{"gguf", "text-generation"} }), + vllmRow("mlx-community/Chat-4bit", func(m *map[string]any) { (*m)["tags"] = []string{"safetensors", "mlx", "text-generation"} }), + vllmRow("../bad", nil), + })) + })) + defer srv.Close() + c := newCatalogService() + c.baseURL = srv.URL + + res, err := c.Catalog(t.Context(), "vllm", "", "", "") + + require.NoError(t, err) + var ids []string + for _, m := range res.Models { + ids = append(ids, m.ID) + } + require.Equal(t, []string{"Qwen/Vision-8B", "Qwen/Chat-8B"}, ids, "want the open chat models, most downloaded first") + assert.True(t, res.Searchable, "the vLLM catalog can be searched") + m := res.Models[1] + // A pull-ready name, 8B parameters and 16 GB of BF16 weights. + assert.Equal(t, "Qwen/Chat-8B", m.Name) + assert.Equal(t, "8B", m.ParameterSize) + assert.Equal(t, uint64(16_000_000_000), m.Size) + assert.Equal(t, "Qwen", m.Author) + pipelines := map[string]bool{} + for _, q := range queries { + require.NotEmpty(t, q["pipeline_tag"]) + require.NotEmpty(t, q["filter"]) + pipelines[q["pipeline_tag"][0]] = true + assert.Equal(t, "safetensors", q["filter"][0], "a listing must filter on safetensors") + assert.Contains(t, q["expand[]"], "safetensors", "a listing must expand safetensors") + } + assert.True(t, pipelines["text-generation"], "both chat pipelines must be listed") + assert.True(t, pipelines["image-text-to-text"], "both chat pipelines must be listed") +} + +func TestVLLMCatalog_SearchesWithTheTrimmedQuery(t *testing.T) { + var mu sync.Mutex + var searches []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + searches = append(searches, r.URL.Query().Get("search")) + mu.Unlock() + assert.NoError(t, json.NewEncoder(w).Encode([]any{vllmRow("someone/Llama-Chat", nil)})) + })) + defer srv.Close() + c := newCatalogService() + c.baseURL = srv.URL + + res, err := c.Catalog(t.Context(), "vllm", "", "", " llama ") + + require.NoError(t, err) + assert.Equal(t, "llama", res.Query) + require.Len(t, res.Models, 1) + assert.Equal(t, "someone/Llama-Chat", res.Models[0].ID) + require.NotEmpty(t, searches) + assert.Equal(t, "llama", searches[0], "the upstream must be searched with the trimmed query") + _, err = c.Catalog(t.Context(), "vllm", "", "", string(make([]byte, maxCatalogQuery+1))) + assert.Error(t, err, "an over-long query is refused") +} + +func TestParameterLabel_RoundsToMillionsOrBillions(t *testing.T) { + test := func(name string, parameters int64, want string) { + t.Run(name, func(t *testing.T) { + assert.Equal(t, want, parameterLabel(parameters)) + }) + } + test("millions round to the nearest million", 751_632_384, "752M") + test("billions keep one decimal", 8_190_000_000, "8.2B") + test("whole billions drop the decimal", 32_000_000_000, "32B") + test("too few to label", 5_000, "") +} + +func TestSafetensorsBytes_SumsKnownTensorTypesOnly(t *testing.T) { + test := func(name string, parameters map[string]int64, want uint64) { + t.Run(name, func(t *testing.T) { + assert.Equal(t, want, safetensorsBytes(parameters)) + }) + } + test("two bytes per BF16 and four per F32", map[string]int64{"BF16": 10, "F32": 1}, 24) + test("unknown when a type is not recognised", map[string]int64{"MYSTERY": 1}, 0) +} diff --git a/services/nvpair-engine-manager/controllifecycle.go b/services/nvpair-engine-manager/controllifecycle.go index d12235d51..be9ba7281 100644 --- a/services/nvpair-engine-manager/controllifecycle.go +++ b/services/nvpair-engine-manager/controllifecycle.go @@ -12,6 +12,7 @@ package main import ( "encoding/json" + "errors" "io" "net/http" ) @@ -56,7 +57,8 @@ func (s *controlServer) handleStart(w http.ResponseWriter, r *http.Request) { http.Error(w, "port must be between 0 and 65535", http.StatusBadRequest) return } - if err := s.exec.StartWith(r.Context(), req.Engine, startOpts{Port: req.Port, Bind: "127.0.0.1"}); err != nil { + // A start the peer's own Stop cut short answers with the status, stopped. + if err := s.exec.StartWith(r.Context(), req.Engine, startOpts{Port: req.Port, Bind: "127.0.0.1"}); err != nil && !errors.Is(err, errStartStopped) { http.Error(w, err.Error(), http.StatusInternalServerError) return } diff --git a/services/nvpair-engine-manager/e2e_test.go b/services/nvpair-engine-manager/e2e_test.go index 3ecbf36a4..39070bfa9 100644 --- a/services/nvpair-engine-manager/e2e_test.go +++ b/services/nvpair-engine-manager/e2e_test.go @@ -131,7 +131,9 @@ func TestEngineManager_PreservesDesiredStateAcrossShutdown(t *testing.T) { second := startE2EManager(t, cfg, home) notify(t, second.stdin, restoreEnabledMethod, nil) - waitNotify(t, second.frames, "engine:state-changed", 20*time.Second) + // The restore pushes a state when the start begins (starting) and when the + // engine is up (running). + waitRunning(t, second.frames, 20*time.Second) send(t, second.stdin, 1, "engine:status", map[string]any{"engine": "fake"}) require.Contains(t, string(waitResult(t, second.frames, "1", 5*time.Second)), `"running":true`, "saved ON state was not restored") send(t, second.stdin, 2, "engine:stop", map[string]any{"engine": "fake"}) @@ -362,6 +364,23 @@ func waitNotify(t *testing.T, frames <-chan frame, method string, timeout time.D } } +// waitRunning waits for an engine:state-changed that reports the engine running. +func waitRunning(t *testing.T, frames <-chan frame, timeout time.Duration) { + t.Helper() + deadline := time.After(timeout) + for { + select { + case f := <-frames: + var status EngineStatus + if f.Method == "engine:state-changed" && json.Unmarshal(f.Params, &status) == nil && status.Running { + return + } + case <-deadline: + require.FailNow(t, "timed out waiting for the engine to report running") + } + } +} + // Configure can request a parent rebind while both stdio readers continue // servicing unrelated messages; literal settings survive a fresh worker. func TestEngineManager_RelaysSettingsRebindAndReloadsWorker(t *testing.T) { diff --git a/services/nvpair-engine-manager/executor.go b/services/nvpair-engine-manager/executor.go index 2b5c03f85..4e633c5ed 100644 --- a/services/nvpair-engine-manager/executor.go +++ b/services/nvpair-engine-manager/executor.go @@ -38,8 +38,11 @@ type EngineStatus struct { DisplayName string `json:"display_name"` Installed bool `json:"installed"` Running bool `json:"running"` - Healthy bool `json:"healthy"` - Port int `json:"port,omitempty"` + // Starting is true while a start is launching the engine and waiting for + // it to become ready. Running is false until it is. + Starting bool `json:"starting,omitempty"` + Healthy bool `json:"healthy"` + Port int `json:"port,omitempty"` } // engineState is the per-engine runtime state. @@ -60,19 +63,27 @@ type engineState struct { // never takes opMu. opMu sync.Mutex - mu sync.Mutex - installed bool - running bool - healthy bool - stopping bool + mu sync.Mutex + installed bool + running bool + healthy bool + stopping bool + // starting is true while a start launches the engine and waits for it to + // become ready, after the checks that can refuse it. + starting bool adopted bool // started externally; we can't stop it gen int64 // bumped each start; gates stale health loops port int binPath string proc *managedProc healthStop context.CancelFunc - // startCancel lets StopAll unblock doStart before waiting on opMu. + // startCancel lets StopAll, Stop and Eject unblock doStart before waiting + // on opMu. startCancel context.CancelFunc + // stopsWaiting counts Stop and Eject calls waiting for opMu. While it is + // above zero no start begins, and a model switch that a stop cancelled + // puts its previous launch back without starting it. + stopsWaiting int } // Executor owns engine lifecycle for every engine known on this host. @@ -233,7 +244,7 @@ func (e *Executor) state(engine string) (*engineState, error) { } plat, ok := m.HostPlatform() if !ok { - return nil, fmt.Errorf("engine %q has no platform block for %s/%s", engine, runtime.GOOS, runtime.GOARCH) + return nil, fmt.Errorf("%s is not supported on %s/%s", m.DisplayName, runtime.GOOS, runtime.GOARCH) } st := &engineState{ manifest: m, diff --git a/services/nvpair-engine-manager/executor_test.go b/services/nvpair-engine-manager/executor_test.go index 678d392cb..d36aa047f 100644 --- a/services/nvpair-engine-manager/executor_test.go +++ b/services/nvpair-engine-manager/executor_test.go @@ -63,6 +63,20 @@ func newTestExecutor(t *testing.T, m *Manifest) *Executor { return NewExecutor(reg, NewReporter(nil), func(string, any) {}, t.TempDir()) } +// The readiness allowance counts from when the start began, so time spent +// before the engine launched is not granted a second time. +func TestWaitReady_CountsTheAllowanceFromWhenTheStartBegan(t *testing.T) { + e := newTestExecutor(t, &Manifest{Engine: "test"}) + probe := &Probe{TCP: "127.0.0.1:{port}", TimeoutS: 60} + began := time.Now().Add(-61 * time.Second) + + finished := time.Now() + err := e.waitReady(t.Context(), probe, 1, began) + + require.Error(t, err) + assert.Less(t, time.Since(finished), 10*time.Second, "the allowance had already run out") +} + func responseHeaderTimeout(t *testing.T, client *http.Client) time.Duration { t.Helper() transport, ok := client.Transport.(*http.Transport) @@ -408,6 +422,29 @@ func TestGetInstalled_AdoptsExternallyRunningEngineWithoutDetectPath(t *testing. require.True(t, statuses[0].Healthy, "expected get-installed to adopt the external engine (%v)", statuses) } +// An engine whose manifest has no block for this OS and CPU, the way vLLM has +// none for Windows, can never be installed here, so it is not offered. +func TestGetInstalled_OmitsEnginesThisHostCannotRun(t *testing.T) { + runsHere := testEngineManifest(fakeEngineBin) + elsewhere := &Manifest{ + Engine: "elsewhere", DisplayName: "Elsewhere", ManifestVersion: 1, + Platforms: map[string]Platform{"plan9/mips": {Runtime: Runtime{Bin: "elsewhere", Port: 1}}}, + } + reg := NewRegistry() + reg.engines[runsHere.Engine] = runsHere + reg.engines[elsewhere.Engine] = elsewhere + ex := NewExecutor(reg, NewReporter(nil), func(string, any) {}, t.TempDir()) + + var listed []string + for _, st := range ex.GetInstalled() { + listed = append(listed, st.Engine) + } + assert.Equal(t, []string{runsHere.Engine}, listed) + + _, err := ex.Status("elsewhere") + require.ErrorContains(t, err, "Elsewhere is not supported on "+runtime.GOOS+"/"+runtime.GOARCH) +} + func TestInstall_AdoptsExternalServiceWithoutDownloading(t *testing.T) { engineSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/api/version" { @@ -1044,3 +1081,38 @@ func rmCmd(path string) []string { } return []string{"rm", "-f", path} } + +// A start that takes a minute (vLLM loading weights) holds the engine's +// lifecycle lock for that minute. Status and the installed-engines listing are +// what the UI polls and what a toggle awaits with a short timeout, so they must +// answer from what is already known instead of queueing behind the operation. +func TestStatus_AnswersWithoutWaitingForARunningLifecycleOperation(t *testing.T) { + ex := newTestExecutor(t, testEngineManifest(fakeEngineBin)) + st, err := ex.state("fake") + require.NoError(t, err) + st.opMu.Lock() // an operation is in flight + defer st.opMu.Unlock() + + answered := make(chan EngineStatus, 2) + go func() { + status, err := ex.Status("fake") + assert.NoError(t, err) + answered <- status + }() + go func() { + for _, status := range ex.GetInstalled() { + if status.Engine == "fake" { + answered <- status + } + } + }() + for i := 0; i < 2; i++ { + select { + case status := <-answered: + assert.Equal(t, "fake", status.Engine) + assert.NotEmpty(t, status.DisplayName) + case <-time.After(2 * time.Second): + require.FailNow(t, "status waited for a running lifecycle operation") + } + } +} diff --git a/services/nvpair-engine-manager/hfcache.go b/services/nvpair-engine-manager/hfcache.go new file mode 100644 index 000000000..9f476287c --- /dev/null +++ b/services/nvpair-engine-manager/hfcache.go @@ -0,0 +1,465 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io/fs" + "maps" + "os" + "os/exec" + "path/filepath" + "regexp" + "slices" + "strconv" + "strings" +) + +// hfCacheModel is one usable model found in a Hugging Face hub cache. Size is +// the bytes the snapshot occupies on disk. +type hfCacheModel struct { + ID string `json:"id"` + Size int64 `json:"size"` +} + +// hfCacheResult is the wire shape of an hf_cache action, chosen so the ordinary +// `result: {array: "data", field: "id"}` extractor turns it into a model list. +type hfCacheResult struct { + Data []hfCacheModel `json:"data"` +} + +var hfTokenizerFiles = []string{"tokenizer.json", "tokenizer_config.json", "tokenizer.model", "vocab.json"} + +// hfShardRe matches the shard names Hugging Face repos use when they carry no +// index file, such as model-00003-of-00004.safetensors. +var hfShardRe = regexp.MustCompile(`-(\d+)-of-(\d+)\.(?:safetensors|bin)$`) + +// runHFCacheAction lists the models in act.HFCache.Dir. It needs no running +// engine, so the list is available while the engine is stopped. +func (e *Executor) runHFCacheAction(ctx context.Context, st *engineState, engine string, act Action, params json.RawMessage) (json.RawMessage, error) { + dir, err := resolvePlaceholders(act.HFCache.Dir, hfActionVars(st)) + if err != nil { + return nil, err + } + if act.HFCache.Delete { + return e.deleteHFModel(ctx, st, engine, expandPath(dir), act.HFCache.Tool, params) + } + models, err := scanHFCache(expandPath(dir)) + if err != nil { + return nil, err + } + return json.Marshal(hfCacheResult{Data: models}) +} + +// scanHFCache returns the models in a Hugging Face hub cache directory that +// look servable: a model repository whose latest download is complete and holds +// a config, a tokenizer and weights. Vision-only and GGUF-only repositories +// carry no tokenizer or no safetensors/pytorch weights, and interrupted +// downloads leave shards or .incomplete blobs behind, so all of those are left +// out. This is a heuristic on file layout, not a check that the engine +// supports the architecture. +// +// A missing directory is an empty cache, not an error. +func scanHFCache(dir string) ([]hfCacheModel, error) { + entries, err := os.ReadDir(dir) + if errors.Is(err, fs.ErrNotExist) { + return []hfCacheModel{}, nil + } + if err != nil { + return nil, fmt.Errorf("read model cache: %w", err) + } + models := []hfCacheModel{} + for _, entry := range entries { + name := entry.Name() + if !strings.HasPrefix(name, "models--") || !entry.IsDir() { + continue + } + // Hugging Face stores "org/name" as "org--name" and forbids "--" inside + // either part, so the mapping reverses cleanly. + id := strings.ReplaceAll(strings.TrimPrefix(name, "models--"), "--", "/") + if id == "" { + continue + } + if size, ok := completeHFModel(filepath.Join(dir, name)); ok { + models = append(models, hfCacheModel{ID: id, Size: size}) + } + } + slices.SortFunc(models, func(a, b hfCacheModel) int { + return strings.Compare(strings.ToLower(a.ID), strings.ToLower(b.ID)) + }) + return models, nil +} + +// completeHFModel reports whether any snapshot of the repository at repoPath is +// a complete, servable model, and the size of the first such snapshot. Each +// snapshot is judged on its own files: huggingface_hub links a file into a +// snapshot only once its blob is whole, so a download still running, for +// another revision or a file vLLM does not load, cannot make a complete +// snapshot look incomplete, or an incomplete one look complete. +func completeHFModel(repoPath string) (int64, bool) { + snapshots, err := os.ReadDir(filepath.Join(repoPath, "snapshots")) + if err != nil { + return 0, false + } + for _, snapshot := range snapshots { + if !snapshot.IsDir() { + continue + } + if disk, _, ok := completeHFSnapshot(filepath.Join(repoPath, "snapshots", snapshot.Name())); ok { + return disk, true + } + } + return 0, false +} + +// completeHFSnapshot reports whether the snapshot at path is a complete, +// servable model, the bytes it occupies on disk, and the bytes of the weights +// vLLM will load from it. The two differ when a repository ships the same +// weights in another format as well: the cache list shows what the snapshot +// takes on disk, and launch sizing uses what is loaded. +func completeHFSnapshot(path string) (disk, weightBytes int64, ok bool) { + entries, err := os.ReadDir(path) + if err != nil { + return 0, 0, false + } + // Files are symlinks into blobs/. Stat follows them, so a link whose blob is + // missing is not a file we have. + files := map[string]int64{} + for _, entry := range entries { + info, err := os.Stat(filepath.Join(path, entry.Name())) + if err != nil || info.IsDir() { + continue + } + files[entry.Name()] = info.Size() + } + weights, index, ok := hfSnapshotWeights(files) + if ok { + weights, ok = hfCompleteWeights(path, files, weights, index) + } + if !ok { + return 0, 0, false + } + + for _, bytes := range files { + disk += bytes + } + for _, name := range weights { + weightBytes += files[name] + } + return disk, weightBytes, true +} + +// hfSnapshotWeights finds, in a snapshot's top-level files, the weights vLLM +// will load, before an index narrows them, and the index that lists them. ok +// is false when the snapshot lacks the config or tokenizer vLLM needs, or has +// no weights. +func hfSnapshotWeights(files map[string]int64) (weights []string, index string, ok bool) { + names := slices.Sorted(maps.Keys(files)) + weights, index = hfChooseWeights(names) + // A Mistral-format model is loaded with params.json and the Mistral + // tokenizer (tekken.json or tokenizer.model.v*), and may lack config.json. + if index == hfMistralIndex { + if !hasAny(files, []string{"config.json", "params.json"}) || (!hasMistralTokenizer(files) && !hasAny(files, hfTokenizerFiles)) { + return nil, "", false + } + } else if _, ok := files["config.json"]; !ok || !hasAny(files, hfTokenizerFiles) { + return nil, "", false + } + return weights, index, len(weights) > 0 +} + +const ( + hfIndex = "model.safetensors.index.json" + hfMistralIndex = "consolidated.safetensors.index.json" +) + +// hfTrainingFiles are files a training run leaves beside the weights, which +// vLLM never loads (its filter_files_not_needed_for_inference). +var hfTrainingFiles = []string{"training_args.bin", "optimizer.bin", "optimizer.pt", "scheduler.pt", "scaler.pt"} + +// hfChooseWeights picks, from a model's file names, the weight files vLLM's +// default loader reads, before an index narrows them, and the name of the index +// that would list them. It is the one rule Pull, the first-start metadata fetch +// and the cache all use, and follows vLLM's DefaultModelLoader._prepare_weights +// with load_format auto: Mistral-format consolidated*.safetensors whenever there +// are any, else *.safetensors, else *.bin, else *.pt, never training state, and +// only files at the top of the repository, the only place vLLM looks. +func hfChooseWeights(names []string) (weights []string, index string) { + top := slices.DeleteFunc(slices.Clone(names), func(name string) bool { return strings.Contains(name, "/") }) + if slices.ContainsFunc(top, isConsolidatedWeights) { + return slices.DeleteFunc(top, func(name string) bool { return !isConsolidatedWeights(name) }), hfMistralIndex + } + for _, suffix := range []string{".safetensors", ".bin", ".pt"} { + for _, name := range top { + if strings.HasSuffix(name, suffix) && !slices.Contains(hfTrainingFiles, name) { + weights = append(weights, name) + } + } + if len(weights) > 0 { + break + } + } + return weights, hfIndex +} + +// hfIndexedWeights narrows weights to the shards an index lists, as vLLM's +// filter_duplicate_safetensors_files does. A repository can ship the same +// weights twice, sharded and as one file, and vLLM loads only the indexed copy, +// so only that copy is downloaded and counted. ok is false when the index lists +// a shard that is not among weights. +func hfIndexedWeights(weights, shards []string) ([]string, bool) { + for _, shard := range shards { + if !slices.Contains(weights, shard) { + return nil, false + } + } + return slices.DeleteFunc(slices.Clone(weights), func(name string) bool { return !slices.Contains(shards, name) }), true +} + +// hfReadIndex returns the shard files the safetensors index at path lists, +// sorted. ok is false when it cannot be read or lists none. +func hfReadIndex(path string) ([]string, bool) { + raw, err := os.ReadFile(path) + if err != nil { + return nil, false + } + var index struct { + WeightMap map[string]string `json:"weight_map"` + } + if json.Unmarshal(raw, &index) != nil || len(index.WeightMap) == 0 { + return nil, false + } + shards := slices.Collect(maps.Values(index.WeightMap)) + slices.Sort(shards) + return slices.Compact(shards), true +} + +// checkHFLayout refuses, before anything is downloaded, a repository whose +// files would not make a complete model once they are: the rule +// completeHFSnapshot applies to the cache, applied to the file names. Without +// it, a repository with two unnumbered weight files and no index downloaded in +// full and was then reported incomplete, and stayed in the cache unlisted. A +// repository with an index is accepted here; the index is checked after the +// download. +func checkHFLayout(names []string) error { + files := map[string]int64{} + for _, name := range names { + if !strings.Contains(name, "/") { + files[name] = 0 + } + } + weights, index, ok := hfSnapshotWeights(files) + if !ok { + return errors.New("the repository lacks the config, tokenizer or weights vLLM needs to load it") + } + if _, hasIndex := files[index]; hasIndex { + return nil + } + if _, complete := hfCompleteWeights("", files, weights, index); !complete { + return fmt.Errorf("the repository's %d weight files are not one model vLLM can load: there is no index listing them, and they are not numbered shards of one set", len(weights)) + } + return nil +} + +// hfCompleteWeights verifies the weight files are all present, and returns +// those vLLM loads. An index lists every shard, and only the indexed files are +// loaded; without one, only shard names like model-00003-of-00004.safetensors +// say how many there should be, and several unnumbered weight files cannot be +// told apart from an interrupted download. This is the rule vLLM applies: it +// loads the index's shards when there is an index, and otherwise every weight +// file it finds. +func hfCompleteWeights(snapshot string, files map[string]int64, weights []string, index string) ([]string, bool) { + if _, ok := files[index]; ok { + shards, ok := hfReadIndex(filepath.Join(snapshot, index)) + if !ok { + return nil, false + } + return hfIndexedWeights(weights, shards) + } + // One unnumbered weight file is the whole model. A file named as one shard + // of several is not, however many files there are. + if len(weights) == 1 && !hfShardRe.MatchString(weights[0]) { + return weights, true + } + seen := map[int]bool{} + total := 0 + for _, name := range weights { + match := hfShardRe.FindStringSubmatch(name) + if match == nil { + return nil, false + } + n, _ := strconv.Atoi(match[1]) + count, _ := strconv.Atoi(match[2]) + if total != 0 && count != total { + return nil, false + } + total = count + seen[n] = true + } + if total == 0 { + return nil, false + } + // Shard numbers come from file names, so count alone is not enough: + // -00000-of-00002 and -00001-of-00002 are two files with shard 2 missing. + for n := 1; n <= total; n++ { + if !seen[n] { + return nil, false + } + } + return weights, true +} + +// hasMistralTokenizer reports whether a snapshot has the tokenizer vLLM uses +// for a Mistral-format model: tekken.json or a tokenizer.model.v* file. +func hasMistralTokenizer(files map[string]int64) bool { + for name := range files { + if name == "tekken.json" || strings.HasPrefix(name, "tokenizer.model.v") { + return true + } + } + return false +} + +func hasAny(files map[string]int64, names []string) bool { + for _, name := range names { + if _, ok := files[name]; ok { + return true + } + } + return false +} + +// deleteHFModel removes one model from the hub cache with the engine's own +// Hugging Face tool (`hf cache rm`), so the cache's owner decides which files go: +// it keeps any file in the cache-wide shared store that another model, dataset or +// Space still links to, and takes the store's locks while it does. +// +// The model a running engine is serving is refused. That is asked of the engine +// itself, because an adopted engine may serve a model other than the one its +// saved arguments name, and the delete is refused when the engine cannot say. A +// running engine's selected model is refused too, since its next start would +// download it again. The model a stopped engine is set to serve is deleted and +// taken out of its launch arguments, so the engine has no model selected and +// will not start until one is loaded; if the removal fails, that selection is +// put back. The engine's operation lock is held from the first check to the +// rollback, so a settings save or start cannot land in between and be undone. +func (e *Executor) deleteHFModel(ctx context.Context, st *engineState, engine, hubDir, toolName string, params json.RawMessage) (json.RawMessage, error) { + var p struct { + Model string `json:"model"` + } + if err := json.Unmarshal(params, &p); err != nil || !hfModelIDRe.MatchString(strings.TrimSpace(p.Model)) { + return nil, fmt.Errorf("delete: params.model must be a Hugging Face model id such as org/name") + } + id := strings.TrimSpace(p.Model) + tool, err := hfToolPath(st, toolName) + if err != nil { + return nil, fmt.Errorf("delete %s: %w", id, err) + } + name := "models--" + strings.ReplaceAll(id, "/", "--") + repo := filepath.Join(hubDir, name) + if rel, err := filepath.Rel(hubDir, repo); err != nil || rel != name { + return nil, fmt.Errorf("delete: refusing a path outside the model cache") + } + if _, err := os.Stat(repo); err != nil { + return nil, fmt.Errorf("delete: %s is not in the model cache", id) + } + err = e.editLaunch(engine, func(edit *launchEdit) error { + refused := fmt.Errorf("delete: %s is the model %s is serving; eject it or load another model first", id, engine) + if edit.before.Running { + serving, err := e.servingModels(ctx, st, engine) + if err != nil { + return fmt.Errorf("delete: cannot tell which model %s is serving, so nothing was deleted: %w", engine, err) + } + if slices.Contains(serving, id) { + return refused + } + } + selected, err := e.selectedModelLocked(edit) + if err != nil { + return fmt.Errorf("delete: cannot tell which model %s is set to serve, so nothing was deleted: %w", engine, err) + } + unselected := false + if selected == id { + flag := modelLaunchFlag(st.manifest) + if edit.before.Running || flag == "" || !edit.before.Editable { + return refused + } + text, _, err := withoutLaunchFlag(edit.before.LaunchText, flag) + if err != nil { + return err + } + if err := edit.save(ctx, text, false); err != nil { + return fmt.Errorf("delete: could not clear %s from the launch settings of %s: %w", id, engine, err) + } + unselected = true + } + failed := func(err error) error { + if !unselected { + return err + } + if restoreErr := edit.restore(ctx); restoreErr != nil { + return fmt.Errorf("%w; putting %s back in the launch settings of %s also failed: %v", err, id, engine, restoreErr) + } + return err + } + cmd := exec.CommandContext(ctx, tool, "cache", "rm", "model/"+id, "--cache-dir", hubDir, "--yes") + cmd.Env = os.Environ() + for key, value := range engineHFEnv(st).toolEnv() { + cmd.Env = append(cmd.Env, key+"="+value) + } + if out, err := cmd.CombinedOutput(); err != nil { + return failed(fmt.Errorf("delete %s: %s failed: %w: %s", id, toolName, err, lastLines(string(out), hfErrorTail))) + } + if _, err := os.Stat(repo); err == nil { + return failed(fmt.Errorf("delete %s: %s finished but the model is still in the cache", id, toolName)) + } + return nil + }) + if err != nil { + return nil, err + } + return json.Marshal(map[string]any{"deleted": id}) +} + +// lastLines returns up to n trailing non-empty lines of out, joined with "; ". +func lastLines(out string, n int) string { + var lines []string + for _, line := range strings.Split(out, "\n") { + if line = strings.TrimSpace(line); line != "" { + lines = append(lines, line) + } + } + if len(lines) > n { + lines = lines[len(lines)-n:] + } + return strings.Join(lines, "; ") +} + +// selectedModelLocked returns the model the engine's launch arguments select, +// or "" when they name none. It fails when that cannot be determined, so a +// caller that must not touch the selected model fails closed. The engine's +// binary is detected first when the launch cannot be resolved without it. +func (e *Executor) selectedModelLocked(edit *launchEdit) (string, error) { + state := edit.before + if state.LaunchText == "" { + if _, err := e.Detect(edit.engine); err != nil { + return "", err + } + state = e.launchStateLocked(edit.engine, edit.st) + edit.before = state + } + tokens, err := parseLaunchText(state.LaunchText) + if err != nil { + return "", err + } + flag := modelLaunchFlag(edit.st.manifest) + if flag == "" { + return "", nil + } + model, _ := launchFlagValue(tokens, flag) + return model, nil +} diff --git a/services/nvpair-engine-manager/hfcache_test.go b/services/nvpair-engine-manager/hfcache_test.go new file mode 100644 index 000000000..5e5c641dc --- /dev/null +++ b/services/nvpair-engine-manager/hfcache_test.go @@ -0,0 +1,598 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "runtime" + "slices" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + settings "nvpair-shared/enginesettings" +) + +// writeHFRepo builds a repository the way the Hugging Face hub cache lays it +// out: content in blobs/, and snapshots// symlinks into it. A +// file content of "" is written empty; names ending in ".incomplete" go to +// blobs/ only, as an interrupted download leaves them. +func writeHFRepo(t *testing.T, cache, id string, files map[string]string) string { + t.Helper() + repo := filepath.Join(cache, "models--"+strings.ReplaceAll(id, "/", "--")) + blobs := filepath.Join(repo, "blobs") + snapshot := filepath.Join(repo, "snapshots", "abc123") + for _, dir := range []string{blobs, snapshot, filepath.Join(repo, "refs")} { + require.NoError(t, os.MkdirAll(dir, 0o755)) + } + // main points at the snapshot, as after an hf download of the main branch. + require.NoError(t, os.WriteFile(filepath.Join(repo, "refs", "main"), []byte("abc123"), 0o644)) + for name, content := range files { + blob := filepath.Join(blobs, "blob-"+name) + require.NoError(t, os.WriteFile(blob, []byte(content), 0o644)) + if strings.HasSuffix(name, ".incomplete") { + continue + } + require.NoError(t, os.Symlink(blob, filepath.Join(snapshot, name))) + } + return repo +} + +func llmFiles(extra map[string]string) map[string]string { + files := map[string]string{"config.json": "{}", "tokenizer.json": "{}"} + for name, content := range extra { + files[name] = content + } + return files +} + +func scannedIDs(t *testing.T, cache string) []string { + t.Helper() + models, err := scanHFCache(cache) + require.NoError(t, err) + ids := []string{} + for _, model := range models { + ids = append(ids, model.ID) + } + return ids +} + +func TestScanHFCache_ListsCompleteModelsOnly(t *testing.T) { + cache := t.TempDir() + + writeHFRepo(t, cache, "Qwen/Qwen2.5-0.5B-Instruct", llmFiles(map[string]string{"model.safetensors": "weights"})) + writeHFRepo(t, cache, "plain/no-org-model", llmFiles(map[string]string{"pytorch_model.bin": "weights"})) + // Both formats of the same weights count once, not as an unverifiable shard set. + writeHFRepo(t, cache, "sentence-transformers/all-MiniLM-L6-v2", llmFiles(map[string]string{ + "model.safetensors": "a", "pytorch_model.bin": "b", "tf_model.h5": "c", + })) + // Sharded, with an index naming every shard. + writeHFRepo(t, cache, "Qwen/Qwen3-8B", llmFiles(map[string]string{ + "model-00001-of-00002.safetensors": "a", + "model-00002-of-00002.safetensors": "b", + "model.safetensors.index.json": `{"weight_map":{"x":"model-00001-of-00002.safetensors","y":"model-00002-of-00002.safetensors"}}`, + })) + // Sharded, no index, but the numbered names prove all four are present. + writeHFRepo(t, cache, "Qwen/Qwen3-VL-8B-Instruct", llmFiles(map[string]string{ + "model-00001-of-00004.safetensors": "a", "model-00002-of-00004.safetensors": "b", + "model-00003-of-00004.safetensors": "c", "model-00004-of-00004.safetensors": "d", + })) + + // Left out: each is missing something a servable model has. + writeHFRepo(t, cache, "Qwen/partial-index", llmFiles(map[string]string{ + "model-00001-of-00002.safetensors": "a", + "model.safetensors.index.json": `{"weight_map":{"x":"model-00001-of-00002.safetensors","y":"model-00002-of-00002.safetensors"}}`, + })) + writeHFRepo(t, cache, "Qwen/partial-numbered", llmFiles(map[string]string{ + "model-00001-of-00004.safetensors": "a", "model-00002-of-00004.safetensors": "b", "model-00004-of-00004.safetensors": "d", + })) + // Two files for a two-shard model, but numbered from zero, so shard 2 is missing. + writeHFRepo(t, cache, "Qwen/misnumbered-shards", llmFiles(map[string]string{ + "model-00000-of-00002.safetensors": "a", "model-00001-of-00002.safetensors": "b", + })) + writeHFRepo(t, cache, "Qwen/unnumbered-shards-no-index", llmFiles(map[string]string{ + "layers-0.safetensors": "a", "layers-1.safetensors": "b", + })) + writeHFRepo(t, cache, "vision/detr", map[string]string{"config.json": "{}", "model.safetensors": "w", "preprocessor_config.json": "{}"}) + writeHFRepo(t, cache, "bartowski/model-GGUF", map[string]string{"model-Q4_K_M.gguf": "w"}) + writeHFRepo(t, cache, "no/config", map[string]string{"tokenizer.json": "{}", "model.safetensors": "w"}) + writeHFRepo(t, cache, "no/weights", llmFiles(nil)) + // Weights still downloading are not linked into the snapshot yet. + writeHFRepo(t, cache, "downloading/now", llmFiles(map[string]string{ + "model.safetensors.incomplete": "partial", + })) + + // A file whose blob is gone is not a file we have. + broken := writeHFRepo(t, cache, "broken/link", llmFiles(map[string]string{"model.safetensors": "w"})) + require.NoError(t, os.Remove(filepath.Join(broken, "blobs", "blob-model.safetensors"))) + + // Not a model repository, and not a directory at all. + require.NoError(t, os.MkdirAll(filepath.Join(cache, "datasets--org--data", "snapshots", "x"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(cache, "models--stray-file"), nil, 0o644)) + + // plain/no-org-model is a two-part id here; the order is case-insensitive. + assert.Equal(t, []string{ + "plain/no-org-model", + "Qwen/Qwen2.5-0.5B-Instruct", + "Qwen/Qwen3-8B", + "Qwen/Qwen3-VL-8B-Instruct", + "sentence-transformers/all-MiniLM-L6-v2", + }, scannedIDs(t, cache)) +} + +func TestScanHFCache_ReportsSnapshotSize(t *testing.T) { + cache := t.TempDir() + writeHFRepo(t, cache, "org/model", llmFiles(map[string]string{"model.safetensors": strings.Repeat("x", 1000)})) + + models, err := scanHFCache(cache) + + require.NoError(t, err) + require.Len(t, models, 1) + // The weights plus config.json and tokenizer.json, "{}" each. + assert.Equal(t, int64(1000+2+2), models[0].Size) +} + +func TestScanHFCache_TreatsAMissingDirectoryAsAnEmptyCache(t *testing.T) { + models, err := scanHFCache(filepath.Join(t.TempDir(), "never-created")) + + require.NoError(t, err) + // An empty non-nil list encodes as [] rather than null. + assert.NotNil(t, models) + assert.Empty(t, models) +} + +// The wire shape has to be readable by the ordinary list extractor, so that +// list_models for such an engine needs no per-engine code. +func TestHFCacheAction_ResultFeedsTheModelListExtractor(t *testing.T) { + cache := t.TempDir() + writeHFRepo(t, cache, "Qwen/Qwen3-8B", llmFiles(map[string]string{"model.safetensors": "w"})) + e := &Executor{} + + raw, err := e.runHFCacheAction(t.Context(), &engineState{}, "vllm", Action{HFCache: &ActionHFCache{Dir: cache}}, nil) + + require.NoError(t, err) + ids, ok := extractStringsResult(raw, &ActionResult{Array: "data", Field: "id"}) + assert.True(t, ok, "the list extractor must read %s", raw) + assert.Equal(t, []string{"Qwen/Qwen3-8B"}, ids) + var decoded hfCacheResult + require.NoError(t, json.Unmarshal(raw, &decoded)) + require.Len(t, decoded.Data, 1) + assert.NotZero(t, decoded.Data[0].Size, "size missing from %s", raw) +} + +func TestHFCacheAction_ValidatesItsDirAndKind(t *testing.T) { + accept := func(name string, action Action) { + t.Run(name, func(t *testing.T) { + assert.NoError(t, action.validate("list_models")) + }) + } + reject := func(name string, action Action) { + t.Run(name, func(t *testing.T) { + assert.Error(t, action.validate("list_models")) + }) + } + accept("a valid hf_cache action", Action{HFCache: &ActionHFCache{Dir: "{models_dir}/hub"}}) + reject("an hf_cache action with no dir", Action{HFCache: &ActionHFCache{Dir: " "}}) + reject("an action with two kinds", Action{HFCache: &ActionHFCache{Dir: "{models_dir}/hub"}, Cmd: []string{"x"}}) +} + +func TestHFCacheDeleteAction_RequiresDeleteModelAndTheTool(t *testing.T) { + accept := func(name, actionName string, cache ActionHFCache) { + t.Run(name, func(t *testing.T) { + assert.NoError(t, (&Action{HFCache: &cache}).validate(actionName)) + }) + } + reject := func(name, actionName string, cache ActionHFCache) { + t.Run(name, func(t *testing.T) { + assert.Error(t, (&Action{HFCache: &cache}).validate(actionName)) + }) + } + accept("delete_model with the tool", "delete_model", ActionHFCache{Dir: "{models_dir}/hub", Delete: true, Tool: "hf"}) + reject("delete outside delete_model", "list_models", ActionHFCache{Dir: "{models_dir}/hub", Delete: true, Tool: "hf"}) + reject("delete without the tool that performs it", "delete_model", ActionHFCache{Dir: "{models_dir}/hub", Delete: true}) + reject("a tool on a listing", "list_models", ActionHFCache{Dir: "{models_dir}/hub", Tool: "hf"}) +} + +// fakeHFCacheRm stands in for `hf cache rm model/ --cache-dir --yes`. +// It records its arguments in /hf-calls and its HF_TOKEN in /hf-token and removes the model's folder, or +// fails the way the real tool does when the model's files include one named +// hf-rm-fails. A model with a file named hf-rm-waits makes it create +// /hf-running and wait for /hf-release first. Which shared files survive is the real tool's business, and is +// checked against the real tool rather than here. +const fakeHFCacheRm = `#!/bin/sh +[ "$1 $2" = "cache rm" ] || { echo "unexpected: $*" >&2; exit 2; } +dir=$5 +echo "$*" >> "$dir/hf-calls" +printf '%s' "${HF_TOKEN:-}" > "$dir/hf-token" +repo="$dir/models--$(echo "${3#model/}" | sed 's#/#--#g')" +if [ -n "$(find "$repo" -name hf-rm-waits)" ]; then touch "$dir/hf-running"; while [ -d "$dir" ] && [ ! -e "$dir/hf-release" ]; do sleep 0.05; done; fi +if [ -n "$(find "$repo" -name hf-rm-fails)" ]; then echo "PermissionError: cannot remove $repo" >&2; exit 1; fi +rm -rf "$repo" +` + +// skipWithoutFakeHF skips a test that runs the fake hf tool, a shell script. +func skipWithoutFakeHF(t *testing.T) { + t.Helper() + if runtime.GOOS == "windows" { + t.Skip("the fake hf tool is a shell script; vLLM has no Windows install") + } +} + +func hfCalls(t *testing.T, cache string) []string { + t.Helper() + raw, err := os.ReadFile(filepath.Join(cache, "hf-calls")) + if err != nil { + return nil + } + return strings.Split(strings.TrimSpace(string(raw)), "\n") +} + +func TestDeleteHFModel_RemovesTheModelWithTheHuggingFaceTool(t *testing.T) { + skipWithoutFakeHF(t) + cache := t.TempDir() + e := switchModelExecutor(t, "--model") + _, err := loadModel(t, e, "org/served") + require.NoError(t, err) + m, _ := e.reg.Get("fake") + m.Actions["delete_model"] = Action{HFCache: &ActionHFCache{Dir: cache, Delete: true, Tool: "hf"}} + writeHFRepo(t, cache, "org/served", llmFiles(map[string]string{"model.safetensors": "w"})) + writeHFRepo(t, cache, "org/other", llmFiles(map[string]string{"model.safetensors": "wwww"})) + del := func(model string) (json.RawMessage, error) { + return e.Action(t.Context(), "fake", "delete_model", json.RawMessage(`{"model":`+quoteJSON(model)+`}`)) + } + + _, err = del("org/served") + require.ErrorContains(t, err, "or load another model first", "the model being served must be refused") + require.DirExists(t, filepath.Join(cache, "models--org--served"), "a refused delete must leave the model in place") + + res, err := del("org/other") + require.NoError(t, err) + var result struct { + Deleted string `json:"deleted"` + } + require.NoError(t, json.Unmarshal(res, &result)) + assert.Equal(t, "org/other", result.Deleted) + assert.Equal(t, []string{"org/served"}, scannedIDs(t, cache)) + + for _, bad := range []string{"org/missing", "../escape", "--help", ""} { + _, err := del(bad) + assert.Error(t, err, "deleting %q should fail", bad) + } + // Only the one allowed delete reached the tool, as a single model target. + assert.Equal(t, []string{"cache rm model/org/other --cache-dir " + cache + " --yes"}, hfCalls(t, cache)) +} + +// When the tool fails, the model stays and a stopped engine keeps it selected, +// so the failed delete changes nothing. +func TestDeleteHFModel_KeepsTheSelectionWhenTheToolFails(t *testing.T) { + skipWithoutFakeHF(t) + cache := t.TempDir() + e := switchModelExecutor(t, "--model") + _, err := loadModel(t, e, "org/served") + require.NoError(t, err) + m, _ := e.reg.Get("fake") + m.Actions["delete_model"] = Action{HFCache: &ActionHFCache{Dir: cache, Delete: true, Tool: "hf"}} + writeHFRepo(t, cache, "org/served", llmFiles(map[string]string{"model.safetensors": "w", "hf-rm-fails": "x"})) + require.NoError(t, e.Stop("fake")) + before := launchTextOf(t, e) + + _, err = e.Action(t.Context(), "fake", "delete_model", json.RawMessage(`{"model":"org/served"}`)) + + assert.ErrorContains(t, err, "PermissionError", "the tool's own error is reported") + assert.Equal(t, []string{"org/served"}, scannedIDs(t, cache), "the model is kept") + assert.Equal(t, before, launchTextOf(t, e), "the launch text is restored") +} + +// Without the Hugging Face tool nothing is deleted and the selection is untouched. +func TestDeleteHFModel_RefusesWithoutTheHuggingFaceTool(t *testing.T) { + cache := t.TempDir() + e := switchModelExecutor(t, "--model") + _, err := loadModel(t, e, "org/served") + require.NoError(t, err) + m, _ := e.reg.Get("fake") + m.Actions["delete_model"] = Action{HFCache: &ActionHFCache{Dir: cache, Delete: true, Tool: "no-such-hf"}} + writeHFRepo(t, cache, "org/served", llmFiles(map[string]string{"model.safetensors": "w"})) + require.NoError(t, e.Stop("fake")) + before := launchTextOf(t, e) + + _, err = e.Action(t.Context(), "fake", "delete_model", json.RawMessage(`{"model":"org/served"}`)) + + assert.ErrorContains(t, err, "was not found", "a missing tool is reported") + assert.Equal(t, []string{"org/served"}, scannedIDs(t, cache)) + assert.Equal(t, before, launchTextOf(t, e), "the launch text is unchanged") +} + +// A stopped engine's model can be deleted. The engine is then left with no model +// selected, so it refuses to start rather than download one, until a model is +// loaded. +func TestDeleteHFModel_LeavesAStoppedEngineWithNoModelSelected(t *testing.T) { + skipWithoutFakeHF(t) + cache := t.TempDir() + e := switchModelExecutor(t, "--model") + _, err := loadModel(t, e, "org/served") + require.NoError(t, err) + m, _ := e.reg.Get("fake") + m.Actions["delete_model"] = Action{HFCache: &ActionHFCache{Dir: cache, Delete: true, Tool: "hf"}} + writeHFRepo(t, cache, "org/served", llmFiles(map[string]string{"model.safetensors": "w"})) + writeHFRepo(t, cache, "org/other", llmFiles(map[string]string{"model.safetensors": "w"})) + require.NoError(t, e.Stop("fake")) + + _, err = e.Action(t.Context(), "fake", "delete_model", json.RawMessage(`{"model":"org/served"}`)) + + require.NoError(t, err, "deleting the model of a stopped engine") + assert.Equal(t, []string{"org/other"}, scannedIDs(t, cache)) + assert.NotContains(t, launchTextOf(t, e), "--model", "the launch text must no longer name a model") + assert.ErrorContains(t, e.Start(t.Context(), "fake"), "no model selected") + _, err = loadModel(t, e, "org/other") + require.NoError(t, err, "loading another model") + state, err := e.LaunchSettings("fake") + require.NoError(t, err) + assert.True(t, state.Running, "the engine runs after a model is loaded") +} + +// With the bundled manifest and no saved settings the served model is the +// default launch's, even in a manager that has not been asked about the engine +// yet. Deleting also fails closed, rather than assuming nothing is served, when +// the served model cannot be told. +func TestDeleteHFModel_RefusesTheBundledDefaultModel(t *testing.T) { + if runtime.GOOS != "linux" { + t.Skip("the bundled vllm manifest is Linux only") + } + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, "config")) + cache := filepath.Join(home, ".cache", "huggingface", "hub") + writeHFRepo(t, cache, "Qwen/Qwen3.5-4B", llmFiles(map[string]string{"model.safetensors": "w"})) + writeHFRepo(t, cache, "org/other", llmFiles(map[string]string{"model.safetensors": "w"})) + + reg := loadWithOverrides(t, t.TempDir()) + base := t.TempDir() + // The bundled manifest runs `hf` from the vLLM virtualenv. + venvBin := filepath.Join(base, "vllm", "venv", "bin") + require.NoError(t, os.MkdirAll(venvBin, 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(venvBin, "hf"), []byte(fakeHFCacheRm), 0o755)) + e := NewExecutor(reg, NewReporter(nil), func(string, any) {}, base) + act := func(model string) error { + _, err := e.Action(t.Context(), "vllm", "delete_model", json.RawMessage(`{"model":`+quoteJSON(model)+`}`)) + return err + } + + // The bundled default launch serves Qwen/Qwen3.5-4B. + require.Error(t, act("Qwen/Qwen3.5-4B"), "the default served model was deleted by a manager that had not detected the engine") + require.DirExists(t, filepath.Join(cache, "models--Qwen--Qwen3.5-4B"), "the served model must stay in the cache") + assert.NoError(t, act("org/other"), "another model should be deletable") +} + +func TestScanHFCache_JudgesAMistralFormatModelOnTheWeightsVLLMLoads(t *testing.T) { + index := func(shards ...string) string { + m := map[string]string{} + for i, shard := range shards { + m[fmt.Sprintf("layer.%d", i)] = shard + } + raw, err := json.Marshal(map[string]any{"weight_map": m}) + require.NoError(t, err) + return string(raw) + } + test := func(name string, files map[string]string, wantListed bool) { + t.Run(name, func(t *testing.T) { + cache := t.TempDir() + writeHFRepo(t, cache, "mistralai/M", files) + + listed := slices.Contains(scannedIDs(t, cache), "mistralai/M") + + assert.Equal(t, wantListed, listed) + }) + } + test("consolidated weights with params.json and tekken.json", + map[string]string{"params.json": "{}", "tekken.json": "{}", "consolidated.safetensors": "w"}, true) + test("Mistral 7B v0.3 style tokenizer.model.v3", + map[string]string{"config.json": "{}", "tokenizer.model.v3": "t", "consolidated.safetensors": "w"}, true) + test("sharded consolidated weights all present", + map[string]string{"params.json": "{}", "tekken.json": "{}", + "consolidated.safetensors.index.json": index("consolidated-00001-of-00002.safetensors", "consolidated-00002-of-00002.safetensors"), + "consolidated-00001-of-00002.safetensors": "w", "consolidated-00002-of-00002.safetensors": "w"}, true) + test("sharded consolidated weights with one missing", + map[string]string{"params.json": "{}", "tekken.json": "{}", + "consolidated.safetensors.index.json": index("consolidated-00001-of-00002.safetensors", "consolidated-00002-of-00002.safetensors"), + "consolidated-00001-of-00002.safetensors": "w"}, false) + test("Hugging Face shards complete, consolidated copy still downloading", + map[string]string{"config.json": "{}", "tokenizer.json": "{}", + "model.safetensors.index.json": index("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"), + "model-00001-of-00002.safetensors": "w", "model-00002-of-00002.safetensors": "w", + "consolidated.safetensors.index.json": index("consolidated-00001-of-00002.safetensors", "consolidated-00002-of-00002.safetensors"), + "consolidated-00001-of-00002.safetensors": "w"}, false) + test("no tokenizer at all", + map[string]string{"params.json": "{}", "consolidated.safetensors": "w"}, false) +} + +// A repository can ship the same weights twice, as indexed shards and as one +// file. vLLM loads only the shards the index lists, so launch sizing counts +// only those, while the cache list still reports everything on disk. +func TestCompleteHFSnapshot_CountsOnlyTheWeightsTheIndexLists(t *testing.T) { + cache := t.TempDir() + repo := writeHFRepo(t, cache, "org/both-copies", llmFiles(map[string]string{ + "model.safetensors.index.json": `{"weight_map": {"a": "model-00001-of-00002.safetensors", "b": "model-00002-of-00002.safetensors"}}`, + "model-00001-of-00002.safetensors": "11111", + "model-00002-of-00002.safetensors": "22222", + "model.safetensors": "the same ten", + })) + + disk, weights, ok := completeHFSnapshot(filepath.Join(repo, "snapshots", "abc123")) + + require.True(t, ok) + assert.Equal(t, int64(5+5), weights, "the indexed shards only") + assert.Greater(t, disk, int64(5+5+12), "everything on disk") +} + +// A settings save made while hf cache rm runs waits for the delete, so a failed +// delete puts back only its own change and the save is kept. +func TestDeleteHFModel_KeepsASettingsSaveMadeWhileTheToolRuns(t *testing.T) { + skipWithoutFakeHF(t) + cache := t.TempDir() + e := switchModelExecutor(t, "--model") + _, err := loadModel(t, e, "org/served") + require.NoError(t, err) + m, _ := e.reg.Get("fake") + m.Actions["delete_model"] = Action{HFCache: &ActionHFCache{Dir: cache, Delete: true, Tool: "hf"}} + writeHFRepo(t, cache, "org/served", llmFiles(map[string]string{"model.safetensors": "w", "hf-rm-waits": "x", "hf-rm-fails": "x"})) + require.NoError(t, e.Stop("fake")) + state, err := e.LaunchSettings("fake") + require.NoError(t, err) + edited := state.LaunchText + " --max-model-len 4096" + + // However the test ends, release the fake hf so it does not wait forever. + t.Cleanup(func() { _ = os.WriteFile(filepath.Join(cache, "hf-release"), nil, 0o600) }) + deleted := make(chan error, 1) + go func() { + _, err := e.Action(t.Context(), "fake", "delete_model", json.RawMessage(`{"model":"org/served"}`)) + deleted <- err + }() + require.Eventually(t, func() bool { + _, err := os.Stat(filepath.Join(cache, "hf-running")) + return err == nil + }, 10*time.Second, 20*time.Millisecond, "hf cache rm never started") + + saved := make(chan error, 1) + go func() { + _, err := e.ConfigureLaunch(t.Context(), settings.Configure{ + Engine: "fake", + Settings: settings.Config{ServerPort: state.ServerPort, LaunchText: edited}, + }, func() error { return nil }) + saved <- err + }() + select { + case err := <-saved: + require.FailNow(t, "the settings save finished while hf cache rm was running", "save error: %v", err) + case <-time.After(300 * time.Millisecond): + } + + require.NoError(t, os.WriteFile(filepath.Join(cache, "hf-release"), nil, 0o600)) + require.ErrorContains(t, <-deleted, "PermissionError") + require.NoError(t, <-saved) + assert.Equal(t, edited, launchTextOf(t, e), "the settings save must survive the delete's rollback") +} + +// An adopted engine can serve another model than its saved arguments name. That +// model is still refused, because the check asks the engine what it serves. +func TestDeleteHFModel_RefusesTheModelAnAdoptedEngineServes(t *testing.T) { + skipWithoutFakeHF(t) + t.Setenv("FAKE_SERVE_MODEL", "org/adopted") + cache := t.TempDir() + e := switchModelExecutor(t, "--model") + m, _ := e.reg.Get("fake") + m.Actions["delete_model"] = Action{HFCache: &ActionHFCache{Dir: cache, Delete: true, Tool: "hf"}} + writeHFRepo(t, cache, "org/adopted", llmFiles(map[string]string{"model.safetensors": "w"})) + require.Contains(t, launchTextOf(t, e), "--model org/initial") + + _, err := e.Action(t.Context(), "fake", "delete_model", json.RawMessage(`{"model":"org/adopted"}`)) + require.ErrorContains(t, err, "is the model fake is serving") + assert.Equal(t, []string{"org/adopted"}, scannedIDs(t, cache), "a refused delete must leave the model in place") + assert.Empty(t, hfCalls(t, cache), "a refused delete must not reach hf cache rm") +} + +// A running engine that cannot say what it serves gets no delete at all, rather +// than a guess from its saved arguments. +func TestDeleteHFModel_RefusesWhenTheRunningEngineCannotSayWhatItServes(t *testing.T) { + skipWithoutFakeHF(t) + cache := t.TempDir() + e := switchModelExecutor(t, "--model") + m, _ := e.reg.Get("fake") + m.Actions["delete_model"] = Action{HFCache: &ActionHFCache{Dir: cache, Delete: true, Tool: "hf"}} + delete(m.Actions, "loaded_models") + writeHFRepo(t, cache, "org/other", llmFiles(map[string]string{"model.safetensors": "w"})) + + _, err := e.Action(t.Context(), "fake", "delete_model", json.RawMessage(`{"model":"org/other"}`)) + require.ErrorContains(t, err, "cannot tell which model fake is serving") + assert.Equal(t, []string{"org/other"}, scannedIDs(t, cache)) + assert.Empty(t, hfCalls(t, cache)) +} + +// One shard of several is an interrupted download, not a model: without an +// index, the shard's own name says how many files there should be. +func TestScanHFCache_LeavesOutALoneShardOfASplitModel(t *testing.T) { + test := func(name string, weights map[string]string, wantListed bool) { + t.Run(name, func(t *testing.T) { + cache := t.TempDir() + writeHFRepo(t, cache, "org/m", llmFiles(weights)) + + listed := slices.Contains(scannedIDs(t, cache), "org/m") + + assert.Equal(t, wantListed, listed) + }) + } + test("the first of four shards alone", map[string]string{"model-00001-of-00004.safetensors": "w"}, false) + test("the only shard of one", map[string]string{"model-00001-of-00001.safetensors": "w"}, true) + test("one unnumbered weight file", map[string]string{"model.safetensors": "w"}, true) + test("all four shards", map[string]string{ + "model-00001-of-00004.safetensors": "w", "model-00002-of-00004.safetensors": "w", + "model-00003-of-00004.safetensors": "w", "model-00004-of-00004.safetensors": "w", + }, true) +} + +// HF_TOKEN set in the engine's launch settings reaches hf cache rm, as it does +// the engine itself. +func TestDeleteHFModel_RunsTheToolWithTheEnginesHuggingFaceSettings(t *testing.T) { + skipWithoutFakeHF(t) + t.Setenv("HF_TOKEN", "") + cache := t.TempDir() + e := switchModelExecutor(t, "--model") + m, _ := e.reg.Get("fake") + m.Actions["delete_model"] = Action{HFCache: &ActionHFCache{Dir: cache, Delete: true, Tool: "hf"}} + writeHFRepo(t, cache, "org/other", llmFiles(map[string]string{"model.safetensors": "w"})) + st, err := e.state("fake") + require.NoError(t, err) + st.mu.Lock() + st.plat.Runtime.LaunchEnv = []string{"HF_TOKEN=engine-token"} + st.mu.Unlock() + + _, err = e.Action(t.Context(), "fake", "delete_model", json.RawMessage(`{"model":"org/other"}`)) + + require.NoError(t, err) + token, err := os.ReadFile(filepath.Join(cache, "hf-token")) + require.NoError(t, err) + assert.Equal(t, "engine-token", string(token)) +} + +// A complete snapshot stays listed whatever else in the repository is still +// downloading or was interrupted: huggingface_hub links a file into a snapshot +// only once it is whole, so each snapshot is judged on its own files. A +// leftover or running download used to hide the whole repository, and so a +// usable model could not be loaded or deleted from the model list. +func TestScanHFCache_ListsACompleteSnapshotBesideAnUnfinishedDownload(t *testing.T) { + test := func(name string, unfinished func(t *testing.T, repo string), want []string) { + t.Run(name, func(t *testing.T) { + cache := t.TempDir() + repo := writeHFRepo(t, cache, "org/model", llmFiles(map[string]string{"model.safetensors": "w"})) + unfinished(t, repo) + + assert.Equal(t, want, scannedIDs(t, cache)) + }) + } + partial := func(t *testing.T, repo, name string) { + require.NoError(t, os.WriteFile(filepath.Join(repo, "blobs", name), []byte("part"), 0o644)) + } + test("a leftover from an interrupted download the retry finished", func(t *testing.T, repo string) { + partial(t, repo, "5af571cbf074e6d21a03528d2330792e532ca608f24ac70a143f6b369968ab8c.86b892a3.incomplete") + }, []string{"org/model"}) + test("a file vLLM does not load, still downloading", func(t *testing.T, repo string) { + partial(t, repo, "0123abcd.incomplete") + }, []string{"org/model"}) + test("another revision, still downloading", func(t *testing.T, repo string) { + newer := filepath.Join(repo, "snapshots", "def456") + require.NoError(t, os.MkdirAll(newer, 0o755)) + require.NoError(t, os.Symlink(filepath.Join(repo, "blobs", "blob-config.json"), filepath.Join(newer, "config.json"))) + require.NoError(t, os.Symlink(filepath.Join(repo, "blobs", "blob-tokenizer.json"), filepath.Join(newer, "tokenizer.json"))) + partial(t, repo, "fedcba98.incomplete") + }, []string{"org/model"}) + test("only an unfinished revision", func(t *testing.T, repo string) { + require.NoError(t, os.Remove(filepath.Join(repo, "snapshots", "abc123", "model.safetensors"))) + partial(t, repo, "fedcba98.incomplete") + }, []string{}) +} diff --git a/services/nvpair-engine-manager/hfpull.go b/services/nvpair-engine-manager/hfpull.go new file mode 100644 index 000000000..4dfe52b17 --- /dev/null +++ b/services/nvpair-engine-manager/hfpull.go @@ -0,0 +1,631 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "net/url" + "os" + "os/exec" + "path/filepath" + "slices" + "strings" + "sync" + "time" +) + +// The download itself is done by Hugging Face's own `hf` CLI, which ships with +// the huggingface_hub package vLLM depends on, so it writes the cache exactly as +// vLLM reads it. This file only chooses the files to fetch, reports progress and +// stops a download that stalls. + +const ( + hfPollInterval = time.Second + // hfStopGrace is how long a cancelled download gets to exit before it is killed. + hfStopGrace = 3 * time.Second + // hfErrorTail is how many trailing output lines a failed download reports. + hfErrorTail = 4 +) + +// hfDroppedSuffixes are file types a vLLM-served model never reads. Skipping +// them keeps a repo that ships several formats from downloading each of them. +var hfDroppedSuffixes = []string{".gguf", ".onnx", ".h5", ".msgpack", ".ot", ".tflite"} + +// isConsolidatedWeights reports whether a file holds weights in Mistral's own +// format, consolidated*.safetensors. vLLM matches that pattern on the file name. +func isConsolidatedWeights(name string) bool { + base := strings.ToLower(filepath.Base(name)) + return strings.HasPrefix(base, "consolidated") && strings.HasSuffix(base, ".safetensors") +} + +// isWeightFile reports whether a repository file holds model weights in a +// format vLLM or PyTorch reads. +func isWeightFile(name string) bool { + return hasSuffixAny(strings.ToLower(name), []string{".safetensors", ".bin", ".pt", ".pth"}) +} + +// selectHFFiles chooses the repository files to download, and the size of the +// weights among them: the weights vLLM will load (hfChooseWeights, narrowed to +// shards when the repository's index lists them), plus config and tokenizer +// files. Any other weight file is a copy vLLM would not load, such as pickle +// weights beside safetensors, Hugging Face-format shards beside Mistral-format +// consolidated*.safetensors (with params.json and the Mistral tokenizer), or a +// single-file copy beside indexed shards, and is not downloaded. Launch sizing +// uses the weight size, the bytes vLLM loads into memory. File names that could +// be read as an option or escape the cache are refused. +func selectHFFiles(siblings []hfSibling, shards []string) (files []string, weights int64, err error) { + names := make([]string, 0, len(siblings)) + for _, f := range siblings { + names = append(names, f.RFilename) + } + chosen, _ := hfChooseWeights(names) + if shards != nil { + var ok bool + if chosen, ok = hfIndexedWeights(chosen, shards); !ok { + return nil, 0, errors.New("the repository's weight index lists files the repository does not have") + } + } + hasWeights := false + for _, f := range siblings { + name := f.RFilename + lower := strings.ToLower(name) + if name == "" || strings.HasPrefix(name, "-") || strings.HasPrefix(name, "/") || strings.Contains(name, "..") { + continue + } + if strings.HasPrefix(lower, "original/") || hasSuffixAny(lower, hfDroppedSuffixes) { + continue + } + switch { + case slices.Contains(chosen, name): + hasWeights = true + weights += f.Size + case isWeightFile(name): + continue + } + files = append(files, name) + } + if !hasWeights { + return nil, 0, errors.New("the repository has no safetensors or PyTorch weights this engine can load") + } + return files, weights, nil +} + +// hfSiblingIndex returns the index that decides which of a repository's weight +// files vLLM loads, when it has one and there is more than one candidate. vLLM +// reads the index in exactly that case before downloading. +func hfSiblingIndex(siblings []hfSibling) (string, bool) { + names := make([]string, 0, len(siblings)) + for _, f := range siblings { + names = append(names, f.RFilename) + } + weights, index := hfChooseWeights(names) + return index, len(weights) > 1 && slices.Contains(names, index) +} + +func hasSuffixAny(s string, suffixes []string) bool { + for _, suffix := range suffixes { + if strings.HasSuffix(s, suffix) { + return true + } + } + return false +} + +// hfRepoFiles lists a repository's files with sizes at a revision. +func (e *Executor) hfRepoFiles(ctx context.Context, id, revision string, env hfEnv) ([]hfSibling, error) { + ctx, cancel := context.WithTimeout(ctx, catalogHTTPTimeout) + defer cancel() + listing := env.endpoint() + "/api/models/" + id + if revision != "" && revision != "main" { + listing += "/revision/" + url.PathEscape(revision) + } + body, err := e.hfGet(ctx, listing+"?blobs=true", maxCatalogBody, env) + if err != nil { + return nil, err + } + var repo struct { + Siblings []hfSibling `json:"siblings"` + } + if err := json.Unmarshal(body, &repo); err != nil { + return nil, fmt.Errorf("decode the Hugging Face reply: %w", err) + } + return repo.Siblings, nil +} + +// hfGet fetches a Hugging Face URL with the engine's HF_TOKEN, when set, and +// returns at most limit bytes of the body. It makes no request when the engine +// is set to work offline. Errors say what the user can do about them. +func (e *Executor) hfGet(ctx context.Context, rawURL string, limit int64, env hfEnv) ([]byte, error) { + if env.offline() { + return nil, errors.New("HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE is set for this engine, so Hugging Face is not contacted") + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) + if err != nil { + return nil, err + } + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "PAIR/1.0") + if token := env.token(); token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + resp, err := e.client.Do(req) + if err != nil { + return nil, fmt.Errorf("look up the model on Hugging Face: %w", withoutURL(err)) + } + defer resp.Body.Close() + switch resp.StatusCode { + case http.StatusOK: + case http.StatusNotFound: + return nil, errors.New("no such model on Hugging Face") + case http.StatusUnauthorized, http.StatusForbidden: + return nil, errors.New("this model needs a Hugging Face login: accept its terms on huggingface.co and set HF_TOKEN") + default: + return nil, fmt.Errorf("Hugging Face returned %s", resp.Status) + } + return io.ReadAll(io.LimitReader(resp.Body, limit)) +} + +// hfEnv is the environment an engine's Hugging Face work runs in: the +// engine's launch environment (the manifest's env and the user's assignments +// in its launch settings) over the service's own, which is what the engine +// process itself sees. PAIR's own requests and hf tool runs use it, so +// HF_HUB_OFFLINE or an HF_TOKEN set for the engine applies to them too. +type hfEnv map[string]string + +// engineHFEnv returns the environment an engine's next start would get. A +// manifest value that cannot be resolved without a running launch (a port, for +// instance) is left out; Hugging Face settings never need one. +func engineHFEnv(st *engineState) hfEnv { + env := hfEnv{} + if st.plat == nil { + return env + } + st.mu.Lock() + rt := st.plat.Runtime + st.mu.Unlock() + vars := st.pathVars() + for key, template := range rt.Env { + if value, err := resolvePlaceholders(template, vars); err == nil { + env[key] = value + } + } + if rt.LaunchEnv != nil { + if literal, err := literalEnvironment(rt.LaunchEnv); err == nil { + for key, value := range literal { + for existing := range env { + if environmentKey(existing) == environmentKey(key) { + delete(env, existing) + } + } + env[key] = value + } + } + } + return env +} + +// hfActionVars are the placeholders an hf_cache or hf_pull dir may use: the +// engine's paths and the hub cache its environment selects. +func hfActionVars(st *engineState) map[string]string { + vars := st.pathVars() + var args []string + if st.plat != nil { + st.mu.Lock() + args = slices.Concat(st.plat.Runtime.Args, st.plat.Runtime.LaunchArgs) + st.mu.Unlock() + } + vars[hfHubCachePlaceholder] = hfLaunchHub(args, engineHFEnv(st)) + return vars +} + +// hfLaunchHub is the hub cache vLLM uses for a launch: its --download-dir when +// the launch sets one, else the cache the engine's environment names. vLLM +// passes --download-dir to huggingface_hub as cache_dir for the weights and +// for the tokenizer and config files alike, in the usual hub layout, so the +// model list, Pull, Delete and launch sizing all read and write there too. The +// directory is read as huggingface_hub reads cache_dir: ~ is the home folder, +// and a relative path is relative to the working directory vLLM shares with +// this service. +func hfLaunchHub(args []string, env hfEnv) string { + dir, ok := launchFlagValue(args, "--download-dir") + if !ok || dir == "" { + return env.hubCache() + } + if rest, ok := strings.CutPrefix(dir, "~"); ok && (rest == "" || strings.HasPrefix(rest, "/")) { + if home, err := os.UserHomeDir(); err == nil { + dir = home + rest + } + } + if abs, err := filepath.Abs(dir); err == nil { + dir = abs + } + return dir +} + +// hfDefaultEndpoint is the Hugging Face Hub, when HF_ENDPOINT names no mirror. +// A variable so tests can point it at a local server. +var hfDefaultEndpoint = "https://huggingface.co" + +// home is the Hugging Face home: HF_HOME, else $XDG_CACHE_HOME/huggingface, +// else ~/.cache/huggingface, as huggingface_hub's constants resolve it. +func (env hfEnv) home() string { + if home := env.get("HF_HOME"); home != "" { + return env.expand(home) + } + if cache := env.get("XDG_CACHE_HOME"); cache != "" { + return filepath.Join(env.expand(cache), "huggingface") + } + return filepath.Join(expandPath("~"), ".cache", "huggingface") +} + +// hubCache is the hub cache huggingface_hub reads and writes for the engine: +// HF_HUB_CACHE, else the legacy HUGGINGFACE_HUB_CACHE, else /hub. +func (env hfEnv) hubCache() string { + for _, key := range []string{"HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE"} { + if dir := env.get(key); dir != "" { + return env.expand(dir) + } + } + return filepath.Join(env.home(), "hub") +} + +// endpoint is the Hub the engine downloads from: HF_ENDPOINT, such as a +// mirror, else huggingface.co. +func (env hfEnv) endpoint() string { + if endpoint := strings.TrimRight(strings.TrimSpace(env.get("HF_ENDPOINT")), "/"); endpoint != "" { + return endpoint + } + return hfDefaultEndpoint +} + +// token is the token huggingface_hub would send: none when +// HF_HUB_DISABLE_IMPLICIT_TOKEN is set, else HF_TOKEN, else +// HUGGING_FACE_HUB_TOKEN, else the one a `hf auth login` stored in +// HF_TOKEN_PATH or /token. +func (env hfEnv) token() string { + if hfTrue(env.get("HF_HUB_DISABLE_IMPLICIT_TOKEN")) { + return "" + } + for _, key := range []string{"HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"} { + if token := strings.TrimSpace(env.get(key)); token != "" { + return token + } + } + path := env.get("HF_TOKEN_PATH") + if path == "" { + path = filepath.Join(env.home(), "token") + } + raw, err := os.ReadFile(env.expand(path)) + if err != nil { + return "" + } + return strings.TrimSpace(string(raw)) +} + +// expand expands $VARS, from the engine's environment, and ~ in a Hugging Face +// path setting, as huggingface_hub does with os.path.expandvars and +// os.path.expanduser. +func (env hfEnv) expand(path string) string { + return expandPath(os.Expand(path, env.get)) +} + +// get returns the engine's value for key, else the service's. +func (env hfEnv) get(key string) string { + for name, value := range env { + if environmentKey(name) == environmentKey(key) { + return value + } + } + return os.Getenv(key) +} + +// offline reports whether HF_HUB_OFFLINE, else TRANSFORMERS_OFFLINE, asks +// Hugging Face tools not to use the network, read the way huggingface_hub +// reads them. +func (env hfEnv) offline() bool { + value := env.get("HF_HUB_OFFLINE") + if value == "" { + value = env.get("TRANSFORMERS_OFFLINE") + } + return hfTrue(value) +} + +// hfTrue reports whether a Hugging Face setting is on, as huggingface_hub's +// _is_true decides it. +func hfTrue(value string) bool { + switch strings.ToUpper(strings.TrimSpace(value)) { + case "1", "ON", "YES", "TRUE": + return true + } + return false +} + +// toolEnv is the environment for one run of the engine's hf tool: the engine's +// own, without progress bars or telemetry. +func (env hfEnv) toolEnv() map[string]string { + out := map[string]string{ + "HF_HUB_DISABLE_PROGRESS_BARS": "1", + "HF_HUB_DISABLE_TELEMETRY": "1", + } + for key, value := range env { + out[key] = value + } + return out +} + +// maxHFConfigBody is far above any model's config.json. +const maxHFConfigBody = 1 << 20 + +// hfMetadataTimeout bounds the metadata fetch a first start makes, so an +// unreachable Hugging Face delays the start by at most this much. +const hfMetadataTimeout = 2 * time.Minute + +// hfPendingModel is a model a start is about to download: the snapshot that +// already holds its config, tokenizer and chat template, fetched ahead of the +// weights, and the bytes of the weights vLLM will load. +type hfPendingModel struct { + snapshot string + weights int64 +} + +// hfModelFetcher prepares the model a start serves, at the revision it loads, +// when that model is not in the cache yet, and reports whether it could. A start +// builds one with sync.OnceValues, so sizing, the tool-call parser and the +// weight check share a single fetch. +type hfModelFetcher func() (hfPendingModel, bool) + +// fetchHFMetadata downloads a model's small files (config, tokenizer, chat +// template and the like, everything a download fetches except the weights) +// into the hub cache with the engine's hf tool, at the revision vLLM will load, +// and sizes its weights from the repository listing. A start that is about to +// download the model can then size its launch and choose its tool-call parser +// from the same files as for a cached model, and vLLM finds those files +// already there. It reports false, and the start goes ahead without them, when +// the engine is offline, the model is gated without a token, or Hugging Face +// cannot be reached. +func (e *Executor) fetchHFMetadata(ctx context.Context, st *engineState, hubDir, id, revision string, env hfEnv) (hfPendingModel, bool) { + if !hfModelIDRe.MatchString(id) || env.offline() { + return hfPendingModel{}, false + } + toolName := "" + for _, act := range st.manifest.Actions { + if act.HFPull != nil { + toolName = act.HFPull.Tool + } + } + tool, err := hfToolPath(st, toolName) + if toolName == "" || err != nil { + slog.Info("no Hugging Face tool to fetch the model's metadata; the first start is sized from free GPU memory", "model", id, "err", err) + return hfPendingModel{}, false + } + ctx, cancel := context.WithTimeout(ctx, hfMetadataTimeout) + defer cancel() + siblings, err := e.hfRepoFiles(ctx, id, revision, env) + if err != nil { + slog.Info("model lookup failed; the first start is sized from free GPU memory", "model", id, "err", err) + return hfPendingModel{}, false + } + files, weights, err := selectHFFiles(siblings, nil) + if err != nil { + slog.Info("model lookup found no loadable weights", "model", id, "err", err) + return hfPendingModel{}, false + } + metadata := slices.DeleteFunc(files, isWeightFile) + args := append([]string{"download", id}, metadata...) + args = append(args, "--cache-dir", hubDir, "--revision", revision) + cmd := exec.CommandContext(ctx, tool, args...) + cmd.Env = os.Environ() + for key, value := range env.toolEnv() { + cmd.Env = append(cmd.Env, key+"="+value) + } + if out, err := cmd.CombinedOutput(); err != nil { + slog.Info("fetching the model's metadata failed; the first start is sized from free GPU memory", "model", id, "err", err, "output", lastLines(string(out), hfErrorTail)) + return hfPendingModel{}, false + } + snapshot, ok := hfSnapshotDir(hubDir, id, revision) + if !ok { + return hfPendingModel{}, false + } + // The index came down with the metadata. Size only the shards it lists, the + // copy vLLM loads. + if index, ok := hfSiblingIndex(siblings); ok { + shards, read := hfReadIndex(filepath.Join(snapshot, index)) + if read { + _, weights, err = selectHFFiles(siblings, shards) + } + if !read || err != nil { + slog.Info("the model's weight index could not be read; the first start is sized from free GPU memory", "model", id, "err", err) + return hfPendingModel{}, false + } + } + return hfPendingModel{snapshot: snapshot, weights: weights}, true +} + +// fetchHFIndex downloads a repository's weight index ahead of its weights, as +// vLLM's own download does, and returns the shards it lists, so a Pull fetches +// only the copy of the weights vLLM loads. +func (e *Executor) fetchHFIndex(ctx context.Context, tool, hubDir, id, index string, env hfEnv) ([]string, error) { + cmd := exec.CommandContext(ctx, tool, "download", id, index, "--cache-dir", hubDir) + cmd.Env = os.Environ() + for key, value := range env.toolEnv() { + cmd.Env = append(cmd.Env, key+"="+value) + } + if out, err := cmd.CombinedOutput(); err != nil { + return nil, fmt.Errorf("download the weight index: %s", lastLines(string(out), hfErrorTail)) + } + snapshot, ok := hfSnapshotDir(hubDir, id, "main") + if !ok { + return nil, errors.New("the weight index did not reach the cache") + } + shards, ok := hfReadIndex(filepath.Join(snapshot, index)) + if !ok { + return nil, errors.New("the repository's weight index could not be read") + } + return shards, nil +} + +// hfBlobBytes totals what is in a repository's blobs directory, which holds +// finished files and the .incomplete ones still being written. +func hfBlobBytes(repoDir string) int64 { + entries, err := os.ReadDir(filepath.Join(repoDir, "blobs")) + if err != nil { + return 0 + } + var total int64 + for _, entry := range entries { + // os.Stat, not entry.Info: huggingface_hub replaces a finished blob + // with a link to its shared store, and the link's own size is tiny. + if info, err := os.Stat(filepath.Join(repoDir, "blobs", entry.Name())); err == nil && !info.IsDir() { + total += info.Size() + } + } + return total +} + +// hfToolPath finds the engine's Hugging Face CLI, named by a manifest's +// hf_pull or hf_cache action, in the directory that holds the engine binary, +// after following any link to it (a vLLM virtualenv ships `hf` beside `vllm`). +func hfToolPath(st *engineState, name string) (string, error) { + st.mu.Lock() + binPath := st.binPath + st.mu.Unlock() + if binPath == "" { + bin, err := resolvePlaceholders(st.plat.Runtime.Bin, st.pathVars()) + if err != nil { + return "", err + } + binPath = expandPath(bin) + } + // A vllm on PATH is often a link into the environment it was installed in + // (pipx, uv tool), and hf sits beside the real one, as Python does for + // vllmPython. + if resolved, err := filepath.EvalSymlinks(binPath); err == nil { + binPath = resolved + } + tool := filepath.Join(filepath.Dir(binPath), name) + if _, err := os.Stat(tool); err != nil { + return "", fmt.Errorf("the Hugging Face tool %q was not found next to %s; install it with `pip install huggingface_hub` or reinstall the engine", name, binPath) + } + return tool, nil +} + +// pullHFModel downloads a Hugging Face repository into the hub cache with the +// engine's `hf` tool. Progress is the share of the repository's bytes that has +// landed in the cache, which stays honest across resumed and parallel files. +// A download that makes no progress for the pull timeout is abandoned. +func (e *Executor) pullHFModel(ctx context.Context, st *engineState, engine, model string, act Action) (json.RawMessage, error) { + id := strings.TrimSpace(model) + if !hfModelIDRe.MatchString(id) { + return nil, fmt.Errorf("pull %q: not a valid Hugging Face model id (use org/name)", model) + } + hubDir, err := resolvePlaceholders(act.HFPull.Dir, hfActionVars(st)) + if err != nil { + return nil, err + } + hubDir = expandPath(hubDir) + + tool, err := hfToolPath(st, act.HFPull.Tool) + if err != nil { + return nil, fmt.Errorf("pull %q: %w", id, err) + } + + e.emitPullProgress(ProgressEvent{Engine: engine, Op: "pull", Stage: "pulling", Message: id}) + env := engineHFEnv(st) + siblings, err := e.hfRepoFiles(ctx, id, "main", env) + if err != nil { + return nil, fmt.Errorf("pull %q: %w", id, err) + } + files, _, err := selectHFFiles(siblings, nil) + if index, ok := hfSiblingIndex(siblings); err == nil && ok { + var shards []string + if shards, err = e.fetchHFIndex(ctx, tool, hubDir, id, index, env); err == nil { + files, _, err = selectHFFiles(siblings, shards) + } + } + if err == nil { + err = checkHFLayout(files) + } + if err != nil { + return nil, fmt.Errorf("pull %q: %w", id, err) + } + st.mu.Lock() + rule := st.plat.Runtime.UnsupportedWeights + st.mu.Unlock() + if err := e.checkRepoWeights(ctx, rule, id, files, env); err != nil { + return nil, fmt.Errorf("pull %q: %w", id, err) + } + var total int64 + for _, f := range siblings { + if slices.Contains(files, f.RFilename) { + total += f.Size + } + } + + repoDir := filepath.Join(hubDir, "models--"+strings.ReplaceAll(id, "/", "--")) + args := append([]string{"download", id}, files...) + args = append(args, "--cache-dir", hubDir) + var tailMu sync.Mutex + var tail []string + proc, err := startManagedProc(tool, args, env.toolEnv(), func(_, line string) { + if line = strings.TrimSpace(line); line != "" { + tailMu.Lock() + defer tailMu.Unlock() + tail = append(tail, line) + if len(tail) > hfErrorTail { + tail = tail[1:] + } + } + }) + if err != nil { + return nil, fmt.Errorf("pull %q: start the download: %w", id, err) + } + + pullCtx, watchdog := newPullProgressWatchdog(ctx, e.pullProgressTimeout) + defer watchdog.stop() + ticker := time.NewTicker(hfPollInterval) + defer ticker.Stop() + lastPct := -1 + for { + select { + case <-proc.done: + if proc.cmd.ProcessState == nil || !proc.cmd.ProcessState.Success() { + tailMu.Lock() + reason := strings.Join(tail, "; ") + tailMu.Unlock() + return nil, fmt.Errorf("pull %q: download failed: %s", id, reason) + } + // The revision just downloaded must be complete; another, older + // complete revision in the cache would not make it so. + snapshot, ok := hfSnapshotDir(hubDir, id, "main") + if _, _, complete := completeHFSnapshot(snapshot); !ok || !complete { + return nil, fmt.Errorf("pull %q: the download finished but the model in the cache is incomplete", id) + } + e.emitPullProgress(ProgressEvent{Engine: engine, Op: "pull", Stage: "success", Percent: 100, Message: id}) + return json.Marshal(map[string]any{"model": id, "bytes": total}) + case <-pullCtx.Done(): + proc.stop(hfStopGrace) + return nil, fmt.Errorf("pull %q: %w", id, watchdog.resolveError(pullCtx, context.Cause(pullCtx))) + case <-ticker.C: + done := hfBlobBytes(repoDir) + watchdog.recordProgress("bytes", done) + if total <= 0 { + continue + } + pct := int(done * 100 / total) + if pct > 99 { + pct = 99 + } + if pct != lastPct && pct > 0 { + lastPct = pct + e.emitPullProgress(ProgressEvent{Engine: engine, Op: "pull", Stage: "downloading", Percent: pct, Message: id}) + } + } + } +} diff --git a/services/nvpair-engine-manager/hfpull_test.go b/services/nvpair-engine-manager/hfpull_test.go new file mode 100644 index 000000000..21d8aae71 --- /dev/null +++ b/services/nvpair-engine-manager/hfpull_test.go @@ -0,0 +1,519 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "runtime" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSelectHFFiles_KeepsOnlyTheWeightsVLLMCanLoad(t *testing.T) { + file := func(name string, size int64) hfSibling { return hfSibling{RFilename: name, Size: size} } + t.Run("safetensors repo skips pickle weights and other formats", func(t *testing.T) { + files, total, err := selectHFFiles([]hfSibling{ + file("config.json", 10), file("tokenizer.json", 20), + file("model-00001-of-00002.safetensors", 1000), file("model-00002-of-00002.safetensors", 1000), + file("model.safetensors.index.json", 5), + file("pytorch_model.bin", 2000), + file("original/params.json", 5), file("model.gguf", 3000), file("-rf", 1), file("../escape", 1), + }, nil) + + require.NoError(t, err) + assert.Equal(t, []string{"config.json", "tokenizer.json", "model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors", "model.safetensors.index.json"}, files) + assert.Equal(t, int64(1000+1000), total) + }) + t.Run("a repo with only pickle weights keeps them", func(t *testing.T) { + files, _, err := selectHFFiles([]hfSibling{file("config.json", 1), file("pytorch_model.bin", 9)}, nil) + + require.NoError(t, err) + assert.Contains(t, files, "pytorch_model.bin") + }) + t.Run("only a consolidated file is kept when nothing else exists", func(t *testing.T) { + files, _, err := selectHFFiles([]hfSibling{file("params.json", 1), file("consolidated.safetensors", 9)}, nil) + + require.NoError(t, err) + assert.Contains(t, files, "consolidated.safetensors") + }) + t.Run("a GGUF-only repo is refused", func(t *testing.T) { + _, _, err := selectHFFiles([]hfSibling{file("config.json", 1), file("m.gguf", 9)}, nil) + + assert.Error(t, err, "nothing in the repo is something vLLM can load") + }) +} + +// A repository can ship the same weights twice, as indexed shards and as one +// file. vLLM loads only what the index lists, so only that copy is downloaded +// and counted. +func TestSelectHFFiles_KeepsOnlyTheShardsTheIndexLists(t *testing.T) { + file := func(name string, size int64) hfSibling { return hfSibling{RFilename: name, Size: size} } + siblings := []hfSibling{ + file("config.json", 1), file("tokenizer.json", 1), + file("model-00001-of-00002.safetensors", 24), file("model-00002-of-00002.safetensors", 24), + file("model.safetensors.index.json", 1), file("model.safetensors", 48), + } + + files, total, err := selectHFFiles(siblings, []string{"model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"}) + + require.NoError(t, err) + assert.Equal(t, []string{"config.json", "tokenizer.json", "model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors", "model.safetensors.index.json"}, files) + assert.Equal(t, int64(48), total) + + _, _, err = selectHFFiles(siblings, []string{"model-00001-of-00002.safetensors", "missing.safetensors"}) + assert.ErrorContains(t, err, "index lists files the repository does not have") +} + +func TestSelectHFFiles_DownloadsTheMistralFormatWeightsVLLMLoads(t *testing.T) { + // Mistral's repositories carry each model twice, in Mistral's own format and + // in Hugging Face's. vLLM's load_format auto loads consolidated*.safetensors + // whenever a repository has any, so those are the weights to download. + file := func(name string, size int64) hfSibling { return hfSibling{RFilename: name, Size: size} } + test := func(name string, siblings []hfSibling, wantFiles []string, wantWeights int64) { + t.Run(name, func(t *testing.T) { + files, total, err := selectHFFiles(siblings, nil) + + require.NoError(t, err) + assert.Equal(t, wantFiles, files) + assert.Equal(t, wantWeights, total) + }) + } + test("consolidated beside sharded Hugging Face weights", + []hfSibling{ // as mistralai/Mistral-Small-3.2-24B-Instruct-2506 + file("config.json", 1), file("params.json", 1), file("tekken.json", 2), file("tokenizer_config.json", 1), + file("consolidated.safetensors", 48), + file("model-00001-of-00002.safetensors", 24), file("model-00002-of-00002.safetensors", 24), + file("model.safetensors.index.json", 1), + }, + []string{"config.json", "params.json", "tekken.json", "tokenizer_config.json", "consolidated.safetensors", "model.safetensors.index.json"}, 48) + test("consolidated beside a single Hugging Face file", + []hfSibling{ // as mistralai/Ministral-3-3B-Instruct-2512 + file("config.json", 1), file("params.json", 1), file("tekken.json", 2), file("tokenizer.json", 1), + file("consolidated.safetensors", 5), file("model.safetensors", 5), + }, + []string{"config.json", "params.json", "tekken.json", "tokenizer.json", "consolidated.safetensors"}, 5) + test("sharded consolidated weights with their index", + []hfSibling{ + file("params.json", 1), file("tekken.json", 2), + file("consolidated-00001-of-00002.safetensors", 24), file("consolidated-00002-of-00002.safetensors", 24), + file("consolidated.safetensors.index.json", 1), + file("model-00001-of-00002.safetensors", 24), file("model-00002-of-00002.safetensors", 24), + }, + []string{"params.json", "tekken.json", "consolidated-00001-of-00002.safetensors", "consolidated-00002-of-00002.safetensors", "consolidated.safetensors.index.json"}, 48) +} + +// hfPullFixture builds an executor whose engine binary sits beside a fake `hf` +// tool, plus a fake Hugging Face API. script is the body of the fake tool. +func hfPullFixture(t *testing.T, siblings []hfSibling, apiStatus int, script string) (e *Executor, hub string, events *[]ProgressEvent) { + t.Helper() + if runtime.GOOS == "windows" { + t.Skip("uses a shell script as the hf tool") + } + dir := t.TempDir() + hub = filepath.Join(dir, "hub") + bin := filepath.Join(dir, "vllm") + require.NoError(t, os.WriteFile(bin, []byte("#!/bin/sh\nexit 0\n"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "hf"), []byte("#!/bin/sh\n"+script), 0o755)) + api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if apiStatus != http.StatusOK { + http.Error(w, "x", apiStatus) + return + } + assert.NoError(t, json.NewEncoder(w).Encode(map[string]any{"siblings": siblings})) + })) + t.Cleanup(api.Close) + old := hfDefaultEndpoint + hfDefaultEndpoint = api.URL + t.Cleanup(func() { hfDefaultEndpoint = old }) + + m := testEngineManifest(bin) + m.Engine = "vllm" + m.Actions = map[string]Action{"pull_model": {HFPull: &ActionHFPull{Tool: "hf", Dir: hub}}} + reg := NewRegistry() + reg.engines[m.Engine] = m + var mu sync.Mutex + collected := []ProgressEvent{} + e = NewExecutor(reg, NewReporter(nil), func(method string, params any) { + fields, ok := params.(map[string]any) + if !ok || method != "engine:pull-progress" { + return + } + ev := ProgressEvent{Stage: fields["stage"].(string)} + if pct, ok := fields["percent"].(int); ok { + ev.Percent = pct + } + mu.Lock() + collected = append(collected, ev) + mu.Unlock() + }, t.TempDir()) + return e, hub, &collected +} + +// The fake tool lays out a Hugging Face cache the way `hf download` does. +const fakeHFDownload = ` +repo="$(echo "$@" | sed 's/.*--cache-dir //')/models--Qwen--Test" +mkdir -p "$repo/blobs" "$repo/snapshots/abc" "$repo/refs" +printf abc > "$repo/refs/main" +printf '{}' > "$repo/blobs/c"; printf '{}' > "$repo/blobs/t"; head -c 1000 /dev/zero > "$repo/blobs/w" +ln -s ../../blobs/c "$repo/snapshots/abc/config.json" +ln -s ../../blobs/t "$repo/snapshots/abc/tokenizer.json" +ln -s ../../blobs/w "$repo/snapshots/abc/model.safetensors" +` + +func TestPullHFModel_DownloadsAndReportsSuccess(t *testing.T) { + e, hub, events := hfPullFixture(t, []hfSibling{ + {RFilename: "config.json", Size: 2}, {RFilename: "tokenizer.json", Size: 2}, {RFilename: "model.safetensors", Size: 1000}, + }, http.StatusOK, fakeHFDownload) + + res, err := e.PullModelStream(t.Context(), "vllm", "Qwen/Test", nil) + + require.NoError(t, err) + var result struct { + Model string `json:"model"` + } + require.NoError(t, json.Unmarshal(res, &result)) + assert.Equal(t, "Qwen/Test", result.Model) + models, err := scanHFCache(hub) + require.NoError(t, err) + require.Len(t, models, 1, "the downloaded model must show in the cache list") + assert.Equal(t, "Qwen/Test", models[0].ID) + require.NotEmpty(t, *events) + assert.Equal(t, ProgressEvent{Stage: "success", Percent: 100}, (*events)[len(*events)-1]) +} + +// fakeHFDownloadWithIndex records each run's files and lays them out in the +// cache, writing FAKE_HF_INDEX as the index. +const fakeHFDownloadWithIndex = ` +id=$2; shift 2 +files="" +while [ $# -gt 0 ]; do + case "$1" in + --cache-dir) dir=$2; shift 2 ;; + *) files="$files $1"; shift ;; + esac +done +echo "$files" >> "$(dirname "$dir")/hf-calls" +repo="$dir/models--$(echo "$id" | sed 's#/#--#g')" +mkdir -p "$repo/refs" "$repo/blobs" "$repo/snapshots/abc" +printf abc > "$repo/refs/main" +for f in $files; do + case "$f" in + model.safetensors.index.json) printf '%s' "$FAKE_HF_INDEX" > "$repo/blobs/$f" ;; + *) printf '{}' > "$repo/blobs/$f" ;; + esac + ln -sf "../../blobs/$f" "$repo/snapshots/abc/$f" +done +` + +// A Pull reads the index first, as vLLM's own download does, and then fetches +// only the shards it lists, not a second copy of the same weights. +func TestPullHFModel_DownloadsOnlyTheShardsTheIndexLists(t *testing.T) { + file := func(name string) hfSibling { return hfSibling{RFilename: name, Size: 10} } + e, hub, _ := hfPullFixture(t, []hfSibling{ + file("config.json"), file("tokenizer.json"), + file("model-00001-of-00002.safetensors"), file("model-00002-of-00002.safetensors"), + file("model.safetensors.index.json"), file("model.safetensors"), + }, http.StatusOK, fakeHFDownloadWithIndex) + t.Setenv("FAKE_HF_INDEX", `{"weight_map": {"a": "model-00001-of-00002.safetensors", "b": "model-00002-of-00002.safetensors"}}`) + + res, err := e.PullModelStream(t.Context(), "vllm", "Qwen/Test", nil) + + require.NoError(t, err) + var result struct { + Bytes int64 `json:"bytes"` + } + require.NoError(t, json.Unmarshal(res, &result)) + assert.Equal(t, int64(50), result.Bytes, "config, tokenizer, two shards and the index") + calls, err := os.ReadFile(filepath.Join(filepath.Dir(hub), "hf-calls")) + require.NoError(t, err) + assert.Equal(t, " model.safetensors.index.json\n config.json tokenizer.json model-00001-of-00002.safetensors model-00002-of-00002.safetensors model.safetensors.index.json\n", string(calls)) +} + +// Pull checks the revision it downloaded. An older complete revision of the +// same model in the cache does not make a broken download a success. +func TestPullHFModel_ReportsAnIncompleteDownloadBesideAnOlderCompleteRevision(t *testing.T) { + e, hub, _ := hfPullFixture(t, oneFileRepo, http.StatusOK, ` +dir="$(echo "$@" | sed 's/.*--cache-dir //')" +repo="$dir/models--Qwen--Test" +mkdir -p "$repo/snapshots/new" "$repo/refs" +printf new > "$repo/refs/main" +ln -sf ../../blobs/blob-config.json "$repo/snapshots/new/config.json" +ln -sf ../../blobs/blob-tokenizer.json "$repo/snapshots/new/tokenizer.json" +`) + writeHFRepo(t, hub, "Qwen/Test", llmFiles(map[string]string{"model.safetensors": "w"})) + + _, err := e.PullModelStream(t.Context(), "vllm", "Qwen/Test", nil) + + assert.ErrorContains(t, err, "the model in the cache is incomplete") +} + +func TestPullHFModel_ReportsTheToolsOwnMessageWhenTheDownloadFails(t *testing.T) { + e, _, _ := hfPullFixture(t, oneFileRepo, http.StatusOK, + `echo "connection reset by peer" >&2; exit 1`) + + _, err := e.PullModelStream(t.Context(), "vllm", "Qwen/Test", nil) + + assert.ErrorContains(t, err, "connection reset by peer") +} + +// oneFileRepo is the smallest repository vLLM can load: config, tokenizer and +// one weights file. +var oneFileRepo = []hfSibling{{RFilename: "config.json", Size: 2}, {RFilename: "tokenizer.json", Size: 2}, {RFilename: "model.safetensors", Size: 10}} + +// A repository whose files would not make a model vLLM can load is refused +// before anything is downloaded. It used to download in full, be reported +// incomplete, and stay in the cache where the model list does not show it. +func TestPullHFModel_RefusesALayoutVLLMCannotLoadBeforeDownloading(t *testing.T) { + file := func(name string) hfSibling { return hfSibling{RFilename: name, Size: 10} } + test := func(name string, siblings []hfSibling, wantErr string) { + t.Run(name, func(t *testing.T) { + e, hub, _ := hfPullFixture(t, siblings, http.StatusOK, `touch "$(dirname "$0")/ran-anyway"; exit 0`) + + _, err := e.PullModelStream(t.Context(), "vllm", "Qwen/Test", nil) + + if wantErr == "" { + assert.NotContains(t, fmt.Sprint(err), "not one model") + assert.FileExists(t, filepath.Join(filepath.Dir(hub), "ran-anyway"), "the download ran") + return + } + assert.ErrorContains(t, err, wantErr) + assert.NoFileExists(t, filepath.Join(filepath.Dir(hub), "ran-anyway"), "nothing was downloaded") + }) + } + test("two unnumbered weight files and no index", + []hfSibling{file("config.json"), file("tokenizer.json"), file("model.safetensors"), file("adapter_model.safetensors")}, + "2 weight files are not one model") + test("a shard missing from a numbered set", + []hfSibling{file("config.json"), file("tokenizer.json"), file("model-00001-of-00003.safetensors"), file("model-00002-of-00003.safetensors")}, + "not one model") + test("no tokenizer", + []hfSibling{file("config.json"), file("model.safetensors")}, + "lacks the config, tokenizer or weights") + test("several files with an index", // the index is checked after the download + []hfSibling{file("config.json"), file("tokenizer.json"), file("a.safetensors"), file("b.safetensors"), file("model.safetensors.index.json")}, + "") + test("a complete numbered set", + []hfSibling{file("config.json"), file("tokenizer.json"), file("model-00001-of-00002.safetensors"), file("model-00002-of-00002.safetensors")}, + "") +} + +func TestPullHFModel_RefusesBadInputsBeforeDownloading(t *testing.T) { + reject := func(name, model string, apiStatus int, wantMessage string) { + t.Run(name, func(t *testing.T) { + e, _, _ := hfPullFixture(t, nil, apiStatus, "touch ran-anyway; exit 0") + + _, err := e.PullModelStream(t.Context(), "vllm", model, nil) + + assert.ErrorContains(t, err, wantMessage) + }) + } + reject("not a repo id", "../etc/passwd", http.StatusOK, "valid Hugging Face model id") + reject("option-looking id", "--help", http.StatusOK, "valid Hugging Face model id") + reject("unknown repo", "Qwen/Nope", http.StatusNotFound, "no such model") + reject("gated repo", "meta/Gated", http.StatusForbidden, "HF_TOKEN") +} + +func TestPullHFModel_StopsPromptlyWhenCancelled(t *testing.T) { + e, _, _ := hfPullFixture(t, oneFileRepo, http.StatusOK, "sleep 30") + ctx, cancel := context.WithTimeout(t.Context(), 1500*time.Millisecond) + defer cancel() + start := time.Now() + + _, err := e.PullModelStream(ctx, "vllm", "Qwen/Test", nil) + + assert.Error(t, err, "a cancelled download must fail") + assert.LessOrEqual(t, time.Since(start), 10*time.Second, "a cancelled download must stop promptly") +} + +// A vllm on PATH is often a link into the environment it was installed in, as +// pipx and uv tool install it. hf sits beside the real binary, not the link. +func TestHFToolPath_FindsHFBesideTheEngineALinkPointsTo(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("vLLM has no Windows install") + } + env := filepath.Join(t.TempDir(), "venv", "bin") + bin := filepath.Join(t.TempDir(), "bin") + require.NoError(t, os.MkdirAll(env, 0o755)) + require.NoError(t, os.MkdirAll(bin, 0o755)) + for _, name := range []string{"vllm", "hf"} { + require.NoError(t, os.WriteFile(filepath.Join(env, name), []byte("#!/bin/sh\n"), 0o755)) + } + link := filepath.Join(bin, "vllm") + require.NoError(t, os.Symlink(filepath.Join(env, "vllm"), link)) + st := &engineState{binPath: link} + + tool, err := hfToolPath(st, "hf") + + require.NoError(t, err) + assert.Equal(t, filepath.Join(env, "hf"), tool) +} + +// PAIR reads the same Hugging Face settings vLLM's huggingface_hub reads, in +// the same order, from the engine's environment over the service's. +func TestHFEnv_ResolvesTheCacheEndpointAndTokenAsHuggingFaceHubDoes(t *testing.T) { + for _, key := range []string{"HF_HOME", "HF_HUB_CACHE", "HUGGINGFACE_HUB_CACHE", "XDG_CACHE_HOME", "HF_ENDPOINT", "HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HF_TOKEN_PATH", "HF_HUB_DISABLE_IMPLICIT_TOKEN", "HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"} { + t.Setenv(key, "") + } + home, err := os.UserHomeDir() + require.NoError(t, err) + cache := func(name string, env hfEnv, want string) { + t.Run("cache: "+name, func(t *testing.T) { + assert.Equal(t, want, env.hubCache()) + }) + } + cache("HF_HUB_CACHE first", hfEnv{"HF_HUB_CACHE": "/hub", "HUGGINGFACE_HUB_CACHE": "/legacy", "HF_HOME": "/home"}, "/hub") + cache("the legacy HUGGINGFACE_HUB_CACHE next", hfEnv{"HUGGINGFACE_HUB_CACHE": "/legacy", "HF_HOME": "/home"}, "/legacy") + cache("then HF_HOME/hub", hfEnv{"HF_HOME": "/home"}, filepath.Join("/home", "hub")) + cache("then XDG_CACHE_HOME/huggingface/hub", hfEnv{"XDG_CACHE_HOME": "/xdg"}, filepath.Join("/xdg", "huggingface", "hub")) + cache("else ~/.cache/huggingface/hub", hfEnv{}, filepath.Join(home, ".cache", "huggingface", "hub")) + cache("~ expanded", hfEnv{"HF_HOME": "~/hf"}, filepath.Join(home, "hf", "hub")) + cache("$VARS expanded from the engine's environment", hfEnv{"MODELS": "/models", "HF_HOME": "$MODELS/hf"}, filepath.Join("/models", "hf", "hub")) + + endpoint := func(name string, env hfEnv, want string) { + t.Run("endpoint: "+name, func(t *testing.T) { + assert.Equal(t, want, env.endpoint()) + }) + } + endpoint("the Hub by default", hfEnv{}, "https://huggingface.co") + endpoint("a mirror, without its trailing slash", hfEnv{"HF_ENDPOINT": "https://mirror.example/"}, "https://mirror.example") + + tokenFile := filepath.Join(t.TempDir(), "token") + require.NoError(t, os.WriteFile(tokenFile, []byte("from-file\n"), 0o600)) + hfHome := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(hfHome, "token"), []byte("from-login\n"), 0o600)) + token := func(name string, env hfEnv, want string) { + t.Run("token: "+name, func(t *testing.T) { + assert.Equal(t, want, env.token()) + }) + } + token("HF_TOKEN first", hfEnv{"HF_TOKEN": "env", "HUGGING_FACE_HUB_TOKEN": "legacy", "HF_HOME": hfHome}, "env") + token("the legacy HUGGING_FACE_HUB_TOKEN next", hfEnv{"HUGGING_FACE_HUB_TOKEN": "legacy", "HF_HOME": hfHome}, "legacy") + token("then HF_TOKEN_PATH", hfEnv{"HF_TOKEN_PATH": tokenFile, "HF_HOME": hfHome}, "from-file") + token("then the token hf auth login stored in HF_HOME", hfEnv{"HF_HOME": hfHome}, "from-login") + token("none when HF_HUB_DISABLE_IMPLICIT_TOKEN is set", hfEnv{"HF_HUB_DISABLE_IMPLICIT_TOKEN": "true", "HF_TOKEN": "env", "HF_HOME": hfHome}, "") + + offline := func(name string, env hfEnv, want bool) { + t.Run("offline: "+name, func(t *testing.T) { + assert.Equal(t, want, env.offline()) + }) + } + offline("online by default", hfEnv{}, false) + offline("HF_HUB_OFFLINE", hfEnv{"HF_HUB_OFFLINE": "1"}, true) + offline("TRANSFORMERS_OFFLINE when HF_HUB_OFFLINE is unset", hfEnv{"TRANSFORMERS_OFFLINE": "YES"}, true) + offline("HF_HUB_OFFLINE wins over TRANSFORMERS_OFFLINE", hfEnv{"HF_HUB_OFFLINE": "0", "TRANSFORMERS_OFFLINE": "1"}, false) + token("none", hfEnv{"HF_HOME": t.TempDir()}, "") +} + +// The model list, Pull and Delete use the cache vLLM itself uses. With +// HF_HUB_CACHE set in its launch settings, that is not {models_dir}/hub. +// A finished blob that huggingface_hub moved to its shared store is a link; +// download progress counts the file it points to, so it does not fall back. +func TestHFBlobBytes_CountsAFinishedBlobLinkedToTheSharedStore(t *testing.T) { + repo := filepath.Join(t.TempDir(), "models--org--model") + blobs := filepath.Join(repo, "blobs") + require.NoError(t, os.MkdirAll(blobs, 0o755)) + shared := filepath.Join(t.TempDir(), "shared-blob") + require.NoError(t, os.WriteFile(shared, make([]byte, 4096), 0o644)) + if err := os.Symlink(shared, filepath.Join(blobs, "finished")); err != nil { + t.Skipf("cannot create a symlink here: %v", err) + } + require.NoError(t, os.WriteFile(filepath.Join(blobs, "other.incomplete"), make([]byte, 100), 0o644)) + + assert.Equal(t, int64(4096+100), hfBlobBytes(repo)) +} + +func TestHFCacheAction_ListsTheCacheTheEngineUses(t *testing.T) { + t.Setenv("HF_HUB_CACHE", "") + modelsDir := t.TempDir() + elsewhere := t.TempDir() + writeHFRepo(t, filepath.Join(modelsDir, "hub"), "org/default-place", llmFiles(map[string]string{"model.safetensors": "w"})) + writeHFRepo(t, elsewhere, "org/engine-place", llmFiles(map[string]string{"model.safetensors": "w"})) + m := testEngineManifest(fakeEngineBin) + key := runtime.GOOS + "/" + runtime.GOARCH + p := m.Platforms[key] + p.Runtime.Env = map[string]string{"HF_HOME": "{models_dir}"} + p.ModelsDir = modelsDir + m.Platforms[key] = p + m.Actions["list_models"] = Action{HFCache: &ActionHFCache{Dir: "{" + hfHubCachePlaceholder + "}"}, Result: &ActionResult{Array: "data", Field: "id"}} + require.NoError(t, m.Validate()) + e := newTestExecutor(t, m) + st, err := e.state("fake") + require.NoError(t, err) + listed := func() []string { + raw, err := e.Action(t.Context(), "fake", "list_models", nil) + require.NoError(t, err) + return extractStrings(raw, &ActionResult{Array: "data", Field: "id"}) + } + + assert.Equal(t, []string{"org/default-place"}, listed(), "by default the manifest's HF_HOME, {models_dir}/hub") + + st.mu.Lock() + st.plat.Runtime.LaunchEnv = []string{"HF_HUB_CACHE=" + elsewhere} + st.mu.Unlock() + assert.Equal(t, []string{"org/engine-place"}, listed(), "HF_HUB_CACHE in the engine's launch settings") + + downloadDir := t.TempDir() + writeHFRepo(t, downloadDir, "org/download-dir-place", llmFiles(map[string]string{"model.safetensors": "w"})) + st.mu.Lock() + st.plat.Runtime.LaunchArgs = []string{"--download-dir", downloadDir} + st.mu.Unlock() + assert.Equal(t, []string{"org/download-dir-place"}, listed(), "--download-dir in the launch arguments, where vLLM keeps its models") +} + +// vLLM hands --download-dir to huggingface_hub as cache_dir, which expands ~ +// and resolves a relative path against the working directory. +func TestHFLaunchHub_UsesTheDownloadDirVLLMUses(t *testing.T) { + home, err := os.UserHomeDir() + require.NoError(t, err) + cwd, err := os.Getwd() + require.NoError(t, err) + env := hfEnv{"HF_HUB_CACHE": filepath.Join(home, "env-hub")} + test := func(name string, args []string, want string) { + t.Run(name, func(t *testing.T) { + assert.Equal(t, want, hfLaunchHub(args, env)) + }) + } + test("no --download-dir: the environment's cache", []string{"serve", "--model", "org/m"}, filepath.Join(home, "env-hub")) + test("an absolute directory", []string{"--download-dir", filepath.Join(home, "models")}, filepath.Join(home, "models")) + test("the = form", []string{"--download-dir=" + filepath.Join(home, "models")}, filepath.Join(home, "models")) + test("~ is the home folder", []string{"--download-dir", "~/models"}, filepath.Join(home, "models")) + test("a relative directory", []string{"--download-dir", "models"}, filepath.Join(cwd, "models")) + test("an empty value: the environment's cache", []string{"--download-dir="}, filepath.Join(home, "env-hub")) +} + +func TestHFHubCachePlaceholder_IsOnlyForHuggingFaceActionDirs(t *testing.T) { + accept := func(name string, edit func(*Manifest)) { + t.Run(name, func(t *testing.T) { + m := testEngineManifest(fakeEngineBin) + edit(m) + assert.NoError(t, m.Validate()) + }) + } + reject := func(name string, edit func(*Manifest)) { + t.Run(name, func(t *testing.T) { + m := testEngineManifest(fakeEngineBin) + edit(m) + assert.ErrorContains(t, m.Validate(), "hf_hub_cache") + }) + } + accept("an hf_cache dir", func(m *Manifest) { + m.Actions["list_models"] = Action{HFCache: &ActionHFCache{Dir: "{hf_hub_cache}"}} + }) + accept("an hf_pull dir", func(m *Manifest) { + m.Actions["pull_model"] = Action{HFPull: &ActionHFPull{Tool: "hf", Dir: "{hf_hub_cache}"}} + }) + reject("a remove_path root", func(m *Manifest) { + m.Actions["delete_model"] = Action{RemovePath: &ActionRemovePath{Root: "{hf_hub_cache}", Path: "{hf_hub_cache}/{model}"}} + }) +} diff --git a/services/nvpair-engine-manager/launch_controls_test.go b/services/nvpair-engine-manager/launch_controls_test.go index 7427a3ab1..3ea815e8f 100644 --- a/services/nvpair-engine-manager/launch_controls_test.go +++ b/services/nvpair-engine-manager/launch_controls_test.go @@ -112,7 +112,7 @@ func TestBundledLlamaCPP_DisablesCORSByDefault(t *testing.T) { func TestBundledManifests_DeclareReviewedNetworkingControls(t *testing.T) { reg := loadWithOverrides(t, t.TempDir()) // Adding a bundled engine requires an explicit networking review and cases. - wantEngines := []string{"llamacpp", "lmstudio", "ollama"} + wantEngines := []string{"llamacpp", "lmstudio", "ollama", "vllm"} names := reg.Names() slices.Sort(names) assert.Equal(t, wantEngines, names, "review networking controls for every bundled engine") @@ -135,6 +135,11 @@ func TestBundledManifests_DeclareReviewedNetworkingControls(t *testing.T) { require.Equal(t, []LaunchControl{{Value: "{server.host}", Flags: []string{"--host"}}, {Value: "{server.port}", Flags: []string{"--port"}}, {Value: "{cors.origins}", Flags: []string{"--cors-origins"}}}, policy.Controls, "incomplete llama.cpp controls") valid = []string{"--host 127.0.0.1 --port 23456", "--host=127.0.0.1 --port=23456", "--port 23456 --cors-origins https://example.test"} invalid = []string{"--host 0.0.0.0", "--host=::", "--port 0", "--port 65536", "--port", "--cors-origins=*", "--port 23456 --port 23457", "-- --host 0.0.0.0"} + case "vllm": + require.Equal(t, []string{"serve"}, policy.FixedArgs, "vLLM startup subcommand is not fixed") + require.Equal(t, []LaunchControl{{Value: "{server.host}", Flags: []string{"--host"}}, {Value: "{server.port}", Flags: []string{"--port"}}}, policy.Controls, "incomplete vLLM controls") + valid = []string{"--host 127.0.0.1 --port 23456", "--host=127.0.0.1 --port=23456", "--port 23456 --model Qwen/Qwen2.5-0.5B-Instruct"} + invalid = []string{"--host 0.0.0.0", "--host=::", "--port 0", "--port 65536", "--port", "--port 23456 --port 23457", "-- --host 0.0.0.0"} default: require.Equal(t, []LaunchControl{{Value: "{server.host}:{server.port}", Env: []string{"OLLAMA_HOST"}}, {Value: "{cors.origins}", Env: []string{"OLLAMA_ORIGINS"}}}, policy.Controls, "incomplete Ollama controls") valid = []string{`OLLAMA_HOST="127.0.0.1:23456"`, `OLLAMA_HOST='127.0.0.1:23456'`, "OLLAMA_HOST=127.0.0.1:23456"} diff --git a/services/nvpair-engine-manager/launchauto.go b/services/nvpair-engine-manager/launchauto.go new file mode 100644 index 000000000..24497dacf --- /dev/null +++ b/services/nvpair-engine-manager/launchauto.go @@ -0,0 +1,576 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "math" + "os" + "os/exec" + "path/filepath" + "slices" + "strconv" + "strings" + "time" +) + +const ( + gib = int64(1) << 30 + // vllmOverhead covers activations, CUDA graphs and the CUDA context on top of + // weights and KV cache. Measured on Qwen3-8B it is well under this. + vllmOverhead = 3 * gib + // weightsSlack is headroom over the snapshot size for runtime copies. + weightsSlack = 1.08 + // freeMargin is left unclaimed on the GPU for the desktop and other engines. + freeMargin = gib + // minAutoContext is the smallest window worth shrinking to; below it the + // model is left to vLLM, whose own error then names the real shortfall. + minAutoContext = 4096 + maxAutoUtil = 0.95 + minAutoUtil = 0.15 +) + +// gpuMemory reports total and free bytes on GPU 0. It is a variable so tests +// can supply a machine. +var gpuMemory = nvidiaSMIMemory + +func nvidiaSMIMemory() (total, free int64, ok bool) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + out, err := exec.CommandContext(ctx, "nvidia-smi", "--query-gpu=memory.total,memory.free", "--format=csv,noheader,nounits").Output() + if err != nil { + return 0, 0, false + } + return parseGPUMemory(string(out), systemMemory) +} + +// parseGPUMemory reads the first row of nvidia-smi's memory.total,memory.free +// query, in MiB. On a unified-memory machine such as DGX Spark the GPU shares +// system RAM and nvidia-smi reports [N/A]; system memory then stands in, as it +// does in nvpair-node-info, and as vLLM itself sees it there. +func parseGPUMemory(out string, system func() (total, available int64, ok bool)) (total, free int64, ok bool) { + line, _, _ := strings.Cut(strings.TrimSpace(out), "\n") + parts := strings.Split(line, ",") + if len(parts) != 2 { + return 0, 0, false + } + if nvidiaSMINotApplicable(parts[0]) { + return system() + } + t, errT := strconv.ParseInt(strings.TrimSpace(parts[0]), 10, 64) + f, errF := strconv.ParseInt(strings.TrimSpace(parts[1]), 10, 64) + if errT != nil || errF != nil || t <= 0 { + return 0, 0, false + } + return t * 1024 * 1024, f * 1024 * 1024, true +} + +// nvidiaSMINotApplicable reports an nvidia-smi field that holds no number, such +// as the [N/A] or [Not Supported] a unified-memory GPU gives for its memory. +func nvidiaSMINotApplicable(field string) bool { + value := strings.ToLower(strings.Trim(strings.TrimSpace(field), "[]")) + return value == "n/a" || value == "not supported" +} + +// systemMemory reports total and available system memory from /proc/meminfo. +func systemMemory() (total, available int64, ok bool) { + raw, err := os.ReadFile("/proc/meminfo") + if err != nil { + return 0, 0, false + } + return parseMeminfo(string(raw)) +} + +// parseMeminfo reads MemTotal and MemAvailable, which /proc/meminfo gives in kB. +// Available rather than free counts the page cache the kernel can reclaim. +func parseMeminfo(text string) (total, available int64, ok bool) { + for _, line := range strings.Split(text, "\n") { + key, rest, found := strings.Cut(line, ":") + if !found { + continue + } + fields := strings.Fields(rest) + if len(fields) == 0 { + continue + } + kb, err := strconv.ParseInt(fields[0], 10, 64) + if err != nil { + continue + } + switch key { + case "MemTotal": + total = kb * 1024 + case "MemAvailable": + available = kb * 1024 + } + } + return total, available, total > 0 && available > 0 +} + +// applyAutoLaunch appends the arguments the manifest's auto_launch asks for. +// It runs after the user's launch has been validated and never changes an +// argument the user supplied, so anything set in the engine settings wins. +// +// hub is the Hugging Face hub cache the engine will use. fetch, when not nil, +// prepares the model the launch serves when it is not in the cache yet, so a +// first start that downloads it is configured for it too. Sizing and the +// tool-call parser both read it, so it must return the same result on every +// call; a start builds it with sync.OnceValues. +func applyAutoLaunch(ctx context.Context, engine string, rt Runtime, command launchCommand, hub string, fetch hfModelFetcher) launchCommand { + if rt.AutoLaunch == nil || rt.AutoLaunch.Kind != autoLaunchHFContext { + return command + } + extra, note := planHFContext(command.Args, hub, int64(rt.AutoLaunch.ContextTokens), fetch) + if tools, parser := planToolCalling(ctx, command.Args, hub, command.Bin, fetch); len(tools) > 0 { + extra = append(extra, tools...) + note = strings.TrimPrefix(note+"; tool calls parsed with "+parser, "; ") + } + if seqs := planMaxNumSeqs(command.Args, rt.AutoLaunch.MaxNumSeqs); len(seqs) > 0 { + extra = append(extra, seqs...) + note = strings.TrimPrefix(note+"; at most "+seqs[1]+" requests at a time", "; ") + } + if len(extra) == 0 { + return command + } + slog.Info("engine launch adjusted for the model and GPU", "engine", engine, "added_args", extra, "reason", note) + command.Args = slices.Concat(command.Args, extra) + return command +} + +// planMaxNumSeqs returns --max-num-seqs with the manifest's limit, unless the +// limit is 0 or the user already set the flag. It does not depend on the model, +// so it applies to every launch, sharded ones included. +func planMaxNumSeqs(args []string, limit int) []string { + if limit <= 0 { + return nil + } + if _, set := launchFlagValue(args, "--max-num-seqs"); set { + return nil + } + return []string{"--max-num-seqs", strconv.Itoa(limit)} +} + +// launchFlagValue returns the value of a flag in either "--flag value" or +// "--flag=value" form, and whether the flag is present at all. When the flag +// is repeated, the last value wins, as it does for vLLM's argument parser. +// Arguments after a "--" terminator are not options, so they are not read. +func launchFlagValue(args []string, flag string) (value string, found bool) { + for i := 0; i < len(args); i++ { + arg := args[i] + switch { + case arg == "--": + return value, found + case arg == flag: + value, found = "", true + if i+1 < len(args) { + value = args[i+1] + i++ // its value + } + default: + if v, ok := strings.CutPrefix(arg, flag+"="); ok { + value, found = v, true + } + } + } + return value, found +} + +type hfModelConfig struct { + MaxPositionEmbeddings int64 `json:"max_position_embeddings"` + NumHiddenLayers int64 `json:"num_hidden_layers"` + NumAttentionHeads int64 `json:"num_attention_heads"` + NumKeyValueHeads int64 `json:"num_key_value_heads"` + HiddenSize int64 `json:"hidden_size"` + HeadDim int64 `json:"head_dim"` + ModelType string `json:"model_type"` + TorchDtype string `json:"torch_dtype"` + Dtype string `json:"dtype"` + RopeScaling json.RawMessage `json:"rope_scaling"` + RopeParameters json.RawMessage `json:"rope_parameters"` + TextConfig *hfModelConfig `json:"text_config"` +} + +// ropeExtended reports whether the model already carries a rope-scaling +// definition other than plain defaults. Overriding one would change the model's +// behaviour in ways this code cannot judge, so such models are only capped. +func (c hfModelConfig) ropeExtended() bool { + for _, raw := range []json.RawMessage{c.RopeScaling, c.RopeParameters} { + var fields map[string]any + if json.Unmarshal(raw, &fields) != nil { + continue + } + if kind, ok := fields["rope_type"].(string); ok && kind != "" && kind != "default" { + return true + } + if kind, ok := fields["type"].(string); ok && kind != "" && kind != "default" { + return true + } + if _, ok := fields["factor"]; ok { + return true + } + } + return false +} + +// text is the language model's part of the config: text_config for a +// vision-language model, the config itself otherwise. +func (c hfModelConfig) text() hfModelConfig { + if c.TextConfig != nil { + return *c.TextConfig + } + return c +} + +// kvShape returns what KV cache sizing needs from a text config, and whether +// all of it is there. +func (c hfModelConfig) kvShape() (native, layers, kvHeads, headDim int64, ok bool) { + native, layers = c.MaxPositionEmbeddings, c.NumHiddenLayers + kvHeads = c.NumKeyValueHeads + if kvHeads == 0 { + kvHeads = c.NumAttentionHeads + } + headDim = c.HeadDim + if headDim == 0 && c.NumAttentionHeads > 0 { + headDim = c.HiddenSize / c.NumAttentionHeads + } + return native, layers, kvHeads, headDim, native > 0 && layers > 0 && kvHeads > 0 && headDim > 0 +} + +func (c hfModelConfig) dtypeBytes() int64 { + d := c.Dtype + if d == "" { + d = c.TorchDtype + } + if d == "float32" { + return 4 + } + return 2 +} + +// hfRevision is the revision vLLM loads for the model in args: its --revision, +// or main. +func hfRevision(args []string) string { + if revision, ok := launchFlagValue(args, "--revision"); ok && revision != "" { + return revision + } + return "main" +} + +// hfSnapshotDir returns the snapshot directory the hub cache holds for a model +// at a revision, resolved the way huggingface_hub resolves it: a branch or tag +// is read from refs/, a commit hash names its snapshot directly, and no other +// snapshot is taken in its place. +func hfSnapshotDir(hubDir, id, revision string) (string, bool) { + repo := filepath.Join(hubDir, "models--"+strings.ReplaceAll(id, "/", "--")) + commit := revision + if ref, err := os.ReadFile(filepath.Join(repo, "refs", revision)); err == nil { + commit = strings.TrimSpace(string(ref)) + } + snapshots := filepath.Join(repo, "snapshots") + dir := filepath.Join(snapshots, commit) + if rel, err := filepath.Rel(snapshots, dir); err != nil || rel != commit || strings.ContainsAny(commit, `/\`) { + return "", false + } + if info, err := os.Stat(dir); err != nil || !info.IsDir() { + return "", false + } + return dir, true +} + +// hfLaunchModel is the model a launch serves, as auto-configuration reads it. +type hfLaunchModel struct { + snapshot string + cfg hfModelConfig + // weights is the bytes of the weights vLLM loads. + weights int64 + // cached is false for a model this start downloads. + cached bool +} + +// findLaunchModel reads the model args name at the revision vLLM will load: +// from its complete snapshot in the cache, or, when that revision is not fully +// cached, from fetch, which prepares a model this start is about to download. +// It reports false when neither has the model's config, and auto-configuration +// then leaves the model alone rather than read another revision's. +func findLaunchModel(args []string, hubDir string, fetch hfModelFetcher) (hfLaunchModel, bool) { + model, ok := launchFlagValue(args, "--model") + if !ok || model == "" { + return hfLaunchModel{}, false + } + revision := hfRevision(args) + if dir, ok := hfSnapshotDir(hubDir, model, revision); ok { + if _, weights, complete := completeHFSnapshot(dir); complete { + cfg, ok := readHFConfig(dir) + return hfLaunchModel{snapshot: dir, cfg: cfg, weights: weights, cached: true}, ok + } + } + if fetch == nil { + return hfLaunchModel{}, false + } + pending, ok := fetch() + if !ok { + return hfLaunchModel{}, false + } + cfg, ok := readHFConfig(pending.snapshot) + return hfLaunchModel{snapshot: pending.snapshot, cfg: cfg, weights: pending.weights}, ok +} + +func readHFConfig(snapshot string) (hfModelConfig, bool) { + var cfg hfModelConfig + raw, err := os.ReadFile(filepath.Join(snapshot, "config.json")) + if err != nil || json.Unmarshal(raw, &cfg) != nil { + return cfg, false + } + return cfg, true +} + +// autoGPUShare returns the --gpu-memory-utilization for a launch that needs +// need bytes on a GPU with total and free bytes, and whether to set one. The +// share is never below minAutoUtil, so a small model still gets room to work, +// but what is free, less freeMargin for the desktop and other engines, caps it +// before that floor: claiming more than is free fails the launch. The cap +// rounds down so the request never exceeds what is free. Nothing is set when +// nothing is free beyond the margin; vLLM then reports the shortfall itself. +func autoGPUShare(need, total, free int64) (float64, bool) { + if total <= 0 { + return 0, false + } + capShare := math.Floor(math.Min(maxAutoUtil, float64(free-freeMargin)/float64(total))*100) / 100 + if capShare <= 0 { + return 0, false + } + share := math.Max(minAutoUtil, math.Ceil(float64(need)/float64(total)*100)/100) + return math.Min(share, capShare), true +} + +// planHFContext decides the extra vLLM arguments for the model named in args: +// +// - --max-model-len: the target window, or the model's own limit when that is +// lower and cannot be extended, or what fits in GPU memory when less. +// - --hf-overrides: YaRN rope scaling when the target exceeds the model's +// native window, so a 40K model can serve 64K. +// - --gpu-memory-utilization: just enough for weights, KV cache and overhead, +// not vLLM's 92% default that would starve other engines. +// +// A flag the user already set is never touched. Tensor or pipeline parallel +// launches are left alone entirely: sizing them needs more than GPU 0. +// +// The model is read at the revision vLLM loads (see findLaunchModel). One that +// is not in the cache yet is read through fetch, when given, since vLLM +// downloads it on this start. Without its size and config, only the memory +// share is bounded. +func planHFContext(args []string, hubDir string, target int64, fetch hfModelFetcher) (extra []string, note string) { + if _, ok := launchFlagValue(args, "--tensor-parallel-size"); ok { + return nil, "" + } + if _, ok := launchFlagValue(args, "--pipeline-parallel-size"); ok { + return nil, "" + } + if model, ok := launchFlagValue(args, "--model"); !ok || model == "" { + return nil, "" + } + _, userLen := launchFlagValue(args, "--max-model-len") + _, userOverrides := launchFlagValue(args, "--hf-overrides") + userUtilText, userUtil := launchFlagValue(args, "--gpu-memory-utilization") + _, userKV := launchFlagValue(args, "--kv-cache-memory-bytes") + if userLen && userUtil { + return nil, "" + } + + total, free, haveGPU := gpuMemory() + found, known := findLaunchModel(args, hubDir, fetch) + if !known { + // Not downloaded yet and Hugging Face could not answer, so its size and + // limits are unknown. Give it what is free, less the margin left for the + // desktop and other engines, under the same ceiling as a sized launch. + // vLLM's own default would claim memory that is already in use. + if haveGPU && !userUtil && !userKV { + if util, ok := autoGPUShare(total, total, free); ok { + return []string{"--gpu-memory-utilization", fmt.Sprintf("%.2f", util)}, "model not in cache and not looked up" + } + } + return nil, "" + } + source := "" + if !found.cached { + source = "model not in cache, sized from Hugging Face" + } + cfg, weights := found.cfg, found.weights + text := cfg.text() + native, layers, kvHeads, headDim, ok := text.kvShape() + if !ok { + return nil, "" + } + perToken := 2 * layers * kvHeads * headDim * text.dtypeBytes() + fixed := int64(float64(weights)*weightsSlack) + vllmOverhead + + ctx := target + var overrides string + extend := native < target && !text.ropeExtended() + if native < target && !extend { + ctx = native + } + + // The GPU share and window must fit together. + var budget int64 + switch { + case userUtil: + util, err := strconv.ParseFloat(userUtilText, 64) + if err != nil || !haveGPU { + return nil, "" + } + budget = int64(util * float64(total)) + case haveGPU: + budget = free - freeMargin + } + if haveGPU && budget > 0 && fixed+perToken*ctx > budget { + fit := (budget - fixed) / perToken + fit -= fit % 1024 + if fit >= minAutoContext { + ctx = fit + } + } + if !userLen { + if ctx > native { + if !extend { + ctx = native + } else { + factor := math.Ceil(float64(ctx)/float64(native)*100) / 100 + scaling := map[string]any{ + "rope_type": "yarn", "factor": factor, + "original_max_position_embeddings": native, + } + fields := map[string]any{"rope_scaling": scaling, "max_position_embeddings": int64(math.Ceil(float64(native) * factor))} + if cfg.TextConfig != nil { + fields = map[string]any{"text_config": fields} + } + if raw, err := json.Marshal(fields); err == nil { + overrides = string(raw) + } + } + } + extra = append(extra, "--max-model-len", strconv.FormatInt(ctx, 10)) + if overrides != "" && !userOverrides { + extra = append(extra, "--hf-overrides", overrides) + } + } + if !userUtil && !userKV && haveGPU { + utilCtx := ctx + if userLen { + if v, _ := launchFlagValue(args, "--max-model-len"); v != "" { + if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 { + utilCtx = n + } + } + } + if util, ok := autoGPUShare(fixed+perToken*utilCtx, total, free); ok { + extra = append(extra, "--gpu-memory-utilization", fmt.Sprintf("%.2f", util)) + } + } + if len(extra) == 0 { + return nil, "" + } + note = fmt.Sprintf("native window %d, target %d, weights %.1f GiB, GPU free %.1f GiB", native, target, float64(weights)/float64(gib), float64(free)/float64(gib)) + if source != "" { + note = source + "; " + note + } + return extra, note +} + +// launchChatFormat returns where vLLM will read this launch's chat format: the +// directory of the tokenizer it loads, and the --chat-template it is given, if +// any. By default that is the model's own snapshot. --tokenizer names another +// repository or an absolute directory, and --tokenizer-revision a revision of +// it, which defaults to the model's revision as in vLLM. ok is false when what +// the launch names cannot be read here: a tokenizer repository or revision that +// is not cached, or a chat template given as a relative path, which vLLM would +// resolve against its own working directory. A wrong guess picks a parser that +// cannot read the model's tool calls, so none is chosen instead. +func launchChatFormat(args []string, hubDir, modelSnapshot string) (tokenizer, template string, ok bool) { + tokenizer = modelSnapshot + name, setName := launchFlagValue(args, "--tokenizer") + revision, setRevision := launchFlagValue(args, "--tokenizer-revision") + if (setName && name != "") || (setRevision && revision != "") { + model, _ := launchFlagValue(args, "--model") + if name == "" { + name = model + } + if revision == "" { + revision = hfRevision(args) + } + switch { + case filepath.IsAbs(name): + tokenizer = name + case name == model && revision == hfRevision(args): + tokenizer = modelSnapshot + case hfModelIDRe.MatchString(name): + if tokenizer, ok = hfSnapshotDir(hubDir, name, revision); !ok { + return "", "", false + } + default: + return "", "", false + } + } + if value, set := launchFlagValue(args, "--chat-template"); set && value != "" { + literal := strings.ContainsAny(value, "{}\n") + if !literal && !filepath.IsAbs(value) { + return "", "", false + } + template = value + } + return tokenizer, template, true +} + +// planToolCalling returns the arguments that let clients send tool_choice +// "auto", which agent clients do on every request. vLLM answers HTTP 400 +// unless both --enable-auto-tool-choice and --tool-call-parser are set. The +// parser is the one the installed vLLM provides for the model's own tool-call +// format (see detectToolParser), so no list of model families is kept here. A +// flag the user already supplied is left as given. When no parser fits, nothing +// is added, and vLLM answers requests that ask for tool calls with its own error. +// +// The model's tokenizer and chat template are read at the revision vLLM loads. +// For a model this start downloads, fetch puts them in the cache first, so the +// first start is configured the same way as every later one. +func planToolCalling(ctx context.Context, args []string, hubDir, bin string, fetch hfModelFetcher) (extra []string, parser string) { + model, ok := launchFlagValue(args, "--model") + if !ok || model == "" { + return nil, "" + } + userParser, hasParser := launchFlagValue(args, "--tool-call-parser") + _, hasEnable := launchFlagValue(args, "--enable-auto-tool-choice") + if hasParser && hasEnable { + return nil, userParser + } + if hasParser { + return []string{"--enable-auto-tool-choice"}, userParser + } + found, ok := findLaunchModel(args, hubDir, fetch) + if !ok { + return nil, "" + } + tokenizer, template, ok := launchChatFormat(args, hubDir, found.snapshot) + if !ok { + slog.Info("the tokenizer or chat template this launch names cannot be read here, so no tool-call parser is chosen; vLLM will refuse requests that ask for tool calls", "model", model) + return nil, "" + } + parser, matched, err := detectToolParser(ctx, bin, tokenizer, template) + if err != nil { + slog.Warn("could not work out the model's tool-call format; vLLM will refuse requests that ask for tool calls", "model", model, "err", err) + return nil, "" + } + if parser == "" { + slog.Info("no vLLM tool-call parser reads this model's tool calls; vLLM will refuse requests that ask for tool calls", "model", model) + return nil, "" + } + slog.Info("vLLM tool-call parser chosen from the model's chat template", "model", model, "parser", parser, "matched", matched) + extra = []string{"--tool-call-parser", parser} + if !hasEnable { + extra = append(extra, "--enable-auto-tool-choice") + } + return extra, parser +} diff --git a/services/nvpair-engine-manager/launchauto_test.go b/services/nvpair-engine-manager/launchauto_test.go new file mode 100644 index 000000000..2b3161c88 --- /dev/null +++ b/services/nvpair-engine-manager/launchauto_test.go @@ -0,0 +1,881 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const qwen8BConfig = `{"max_position_embeddings":40960,"num_hidden_layers":36,"num_attention_heads":32,"num_key_value_heads":8,"head_dim":128,"torch_dtype":"bfloat16"}` + +// cachedModel writes a repository whose weights file reports the given size +// without using disk (a sparse file). +func cachedModel(t *testing.T, id, config string, weightBytes int64) string { + t.Helper() + cache := t.TempDir() + cachedModelAt(t, cache, id, config, weightBytes) + return cache +} + +func cachedModelAt(t *testing.T, cache, id, config string, weightBytes int64) { + t.Helper() + repo := writeHFRepo(t, cache, id, llmFiles(map[string]string{"config.json": config, "model.safetensors": "x"})) + require.NoError(t, os.Truncate(filepath.Join(repo, "blobs", "blob-model.safetensors"), weightBytes)) +} + +// pendingModel stands in for a model this start downloads: fetch puts its +// config and tokenizer in cache, as the metadata fetch does, and reports its +// weight bytes. Like the fetch a start builds, it runs once. +func pendingModel(t *testing.T, cache, id, config string, weightBytes int64) hfModelFetcher { + t.Helper() + return sync.OnceValues(func() (hfPendingModel, bool) { + repo := writeHFRepo(t, cache, id, map[string]string{"config.json": config, "tokenizer.json": "{}"}) + return hfPendingModel{snapshot: filepath.Join(repo, "snapshots", "abc123"), weights: weightBytes}, true + }) +} + +func withGPU(t *testing.T, totalGiB, freeGiB int64) { + t.Helper() + old := gpuMemory + gpuMemory = func() (int64, int64, bool) { return totalGiB * gib, freeGiB * gib, true } + t.Cleanup(func() { gpuMemory = old }) +} + +func withoutGPU(t *testing.T) { + t.Helper() + old := gpuMemory + gpuMemory = func() (int64, int64, bool) { return 0, 0, false } + t.Cleanup(func() { gpuMemory = old }) +} + +func flagOf(t *testing.T, args []string, flag string) string { + t.Helper() + v, _ := launchFlagValue(args, flag) + return v +} + +func utilOf(t *testing.T, args []string) float64 { + t.Helper() + v, err := strconv.ParseFloat(flagOf(t, args, "--gpu-memory-utilization"), 64) + require.NoError(t, err, "no usable --gpu-memory-utilization in %v", args) + return v +} + +func TestLaunchFlagValue_ReadsTheValueVLLMWillUse(t *testing.T) { + test := func(name string, args []string, wantValue string, wantFound bool) { + t.Run(name, func(t *testing.T) { + value, found := launchFlagValue(args, "--model") + + assert.Equal(t, wantValue, value) + assert.Equal(t, wantFound, found) + }) + } + test("separate value", []string{"--model", "org/a", "--port", "8001"}, "org/a", true) + test("equals form", []string{"--model=org/a"}, "org/a", true) + test("repeated: the last one wins", []string{"--model", "org/a", "--port", "8001", "--model", "org/b"}, "org/b", true) + test("repeated in both forms", []string{"--model=org/a", "--model", "org/b", "--model=org/c"}, "org/c", true) + test("after the terminator is not an option", []string{"--model", "org/a", "--", "--model", "org/b"}, "org/a", true) + test("no value at the end", []string{"--port", "8001", "--model"}, "", true) + test("a longer flag with the same prefix", []string{"--model-loader-extra-config", "{}"}, "", false) + test("absent", []string{"--port", "8001"}, "", false) +} + +// The GPU share holds what the launch needs, at least minAutoUtil, and never +// more than is free: a share above free memory fails the launch, so the cap on +// free memory wins over the floor. +func TestAutoGPUShare_NeverClaimsMoreThanIsFree(t *testing.T) { + share := func(name string, needGiB, totalGiB, freeGiB float64, want float64) { + t.Run(name, func(t *testing.T) { + got, ok := autoGPUShare(int64(needGiB*float64(gib)), int64(totalGiB*float64(gib)), int64(freeGiB*float64(gib))) + + require.True(t, ok) + assert.InDelta(t, want, got, 0.0001) + }) + } + none := func(name string, needGiB, totalGiB, freeGiB float64) { + t.Run(name, func(t *testing.T) { + _, ok := autoGPUShare(int64(needGiB*float64(gib)), int64(totalGiB*float64(gib)), int64(freeGiB*float64(gib))) + + assert.False(t, ok) + }) + } + share("what the launch needs, rounded up", 26, 48, 48, 0.55) + share("a small launch still gets the floor", 2, 48, 48, 0.15) + // 7 GiB free less the 1 GiB margin is 0.125 of 48 GiB, below the 0.15 floor + // (7.2 GiB), which would claim more than is free. + share("free memory wins over the floor", 6, 48, 7, 0.12) + share("free memory wins over the need", 30, 48, 20, 0.39) + share("never above the ceiling", 60, 48, 48, 0.95) + none("nothing free beyond the margin", 6, 48, 1) + none("no total reading", 6, 0, 0) +} + +func TestAutoContext_SizesTheModelNamedLast(t *testing.T) { + withGPU(t, 48, 48) + // Only the second --model is cached. vLLM serves the second, so sizing must + // use it rather than fall back to the uncached first one. + cache := cachedModel(t, "Qwen/Qwen3-8B", qwen8BConfig, 16_400_000_000) + + extra, _ := planHFContext([]string{"serve", "--model", "org/not-cached", "--model", "Qwen/Qwen3-8B"}, cache, 65536, nil) + + assert.Equal(t, "65536", flagOf(t, extra, "--max-model-len"), "sized from Qwen3-8B's config: %v", extra) +} + +func TestAutoContext_ExtendsAShortModelToTheTarget(t *testing.T) { + withGPU(t, 48, 48) + cache := cachedModel(t, "Qwen/Qwen3-8B", qwen8BConfig, 16_400_000_000) + + extra, _ := planHFContext([]string{"serve", "--model", "Qwen/Qwen3-8B"}, cache, 65536, nil) + + assert.Equal(t, "65536", flagOf(t, extra, "--max-model-len"), "the model alone stops at 40960") + var overrides struct { + RopeScaling struct { + Type string `json:"rope_type"` + Factor float64 `json:"factor"` + Original int64 `json:"original_max_position_embeddings"` + } `json:"rope_scaling"` + MaxPositions int64 `json:"max_position_embeddings"` + } + require.NoError(t, json.Unmarshal([]byte(flagOf(t, extra, "--hf-overrides")), &overrides)) + assert.Equal(t, "yarn", overrides.RopeScaling.Type) + assert.Equal(t, int64(40960), overrides.RopeScaling.Original) + assert.InDelta(t, 1.6, overrides.RopeScaling.Factor, 0.01) + assert.GreaterOrEqual(t, overrides.MaxPositions, int64(65536)) + // 16.4 GB of weights plus about 9.7 GB of KV cache for 64K plus overhead is + // about 60% of a 48 GiB card; vLLM's own default would take 92%. + assert.InDelta(t, 0.62, utilOf(t, extra), 0.07) +} + +func TestAutoContext_CapsALongModelAtTheTarget(t *testing.T) { + withGPU(t, 48, 48) + cfg := strings.Replace(qwen8BConfig, "40960", "262144", 1) + cache := cachedModel(t, "Qwen/Long", cfg, 16_400_000_000) + + extra, _ := planHFContext([]string{"--model", "Qwen/Long"}, cache, 65536, nil) + + assert.Equal(t, "65536", flagOf(t, extra, "--max-model-len"), "the target, not the 262144 native window") + assert.NotContains(t, extra, "--hf-overrides", "a model that already supports the target needs no rope override") +} + +func TestAutoContext_LeavesModelsWithRopeScalingAlone(t *testing.T) { + withGPU(t, 48, 48) + cfg := strings.Replace(qwen8BConfig, `"torch_dtype"`, `"rope_scaling":{"rope_type":"yarn","factor":4.0},"torch_dtype"`, 1) + cache := cachedModel(t, "Qwen/Scaled", cfg, 16_400_000_000) + + extra, _ := planHFContext([]string{"--model", "Qwen/Scaled"}, cache, 131072, nil) + + assert.NotContains(t, extra, "--hf-overrides", "must not override a model's own rope scaling") + assert.Equal(t, "40960", flagOf(t, extra, "--max-model-len")) +} + +func TestAutoContext_ShrinksToWhatFitsInFreeMemory(t *testing.T) { + withGPU(t, 48, 22) + cache := cachedModel(t, "Qwen/Qwen3-8B", qwen8BConfig, 16_400_000_000) + + extra, _ := planHFContext([]string{"--model", "Qwen/Qwen3-8B"}, cache, 65536, nil) + + n, err := strconv.ParseInt(flagOf(t, extra, "--max-model-len"), 10, 64) + require.NoError(t, err) + assert.GreaterOrEqual(t, n, int64(4096)) + assert.Less(t, n, int64(65536)) + assert.Zero(t, n%1024, "a multiple of 1024") + assert.NotContains(t, extra, "--hf-overrides", "a window inside the native limit needs no rope override") + assert.LessOrEqual(t, utilOf(t, extra), 21.0/48.0, "no more than is free less the margin") +} + +// The free-memory cap wins over the minimum share: on a 48 GiB GPU with 7 GiB +// free, the context is fitted into 6 GiB and the share must not then be raised +// to 0.15 (7.2 GiB), more than is free. +func TestAutoContext_KeepsTheShareWithinFreeMemoryOnALargeBusyGPU(t *testing.T) { + withGPU(t, 48, 7) + cache := cachedModel(t, "Qwen/Qwen3-0.6B", qwen8BConfig, 1_200_000_000) + + extra, _ := planHFContext([]string{"--model", "Qwen/Qwen3-0.6B"}, cache, 65536, nil) + + assert.InDelta(t, 0.12, utilOf(t, extra), 0.0001) + assert.LessOrEqual(t, utilOf(t, extra)*48, 6.0, "the share must fit the 6 GiB free beyond the margin") +} + +func TestAutoContext_NeverOverridesWhatTheUserSet(t *testing.T) { + withGPU(t, 48, 48) + cache := cachedModel(t, "Qwen/Qwen3-8B", qwen8BConfig, 16_400_000_000) + + t.Run("the user's context", func(t *testing.T) { + extra, _ := planHFContext([]string{"--model", "Qwen/Qwen3-8B", "--max-model-len=8192"}, cache, 65536, nil) + + assert.NotContains(t, extra, "--max-model-len") + assert.NotContains(t, extra, "--hf-overrides") + assert.LessOrEqual(t, utilOf(t, extra), 0.45, "an 8K window needs far less of the GPU") + }) + t.Run("the user's GPU share", func(t *testing.T) { + extra, _ := planHFContext([]string{"--model", "Qwen/Qwen3-8B", "--gpu-memory-utilization", "0.5"}, cache, 65536, nil) + + assert.NotContains(t, extra, "--gpu-memory-utilization") + n, err := strconv.ParseInt(flagOf(t, extra, "--max-model-len"), 10, 64) + require.NoError(t, err) + assert.Less(t, n, int64(65536), "a 0.5 share (24 GiB) cannot hold 64K of Qwen3-8B") + assert.GreaterOrEqual(t, n, int64(4096)) + }) + t.Run("both", func(t *testing.T) { + extra, _ := planHFContext([]string{"--model", "Qwen/Qwen3-8B", "--max-model-len", "4096", "--gpu-memory-utilization", "0.3"}, cache, 65536, nil) + + assert.Empty(t, extra) + }) +} + +func TestAutoContext_SkipsShardedLaunches(t *testing.T) { + withGPU(t, 48, 48) + cache := cachedModel(t, "Qwen/Qwen3-8B", qwen8BConfig, 16_400_000_000) + test := func(flag string) { + t.Run(flag, func(t *testing.T) { + extra, _ := planHFContext([]string{"--model", "Qwen/Qwen3-8B", flag, "2"}, cache, 65536, nil) + + assert.Empty(t, extra, "sizing a sharded launch needs more than GPU 0") + }) + } + test("--tensor-parallel-size") + test("--pipeline-parallel-size") +} + +func TestAutoContext_GivesAModelNotInTheCacheWhatIsFree(t *testing.T) { + withGPU(t, 48, 20) + + extra, _ := planHFContext([]string{"--model", "org/not-downloaded"}, t.TempDir(), 65536, nil) + + assert.NotContains(t, extra, "--max-model-len", "an unknown model's limits are unknown") + assert.LessOrEqual(t, utilOf(t, extra), 19.0/48.0+0.001, "must leave a margin of the 20 GiB free") +} + +func TestAutoContext_BoundsTheWindowWithoutAGPUReading(t *testing.T) { + withoutGPU(t) + cache := cachedModel(t, "Qwen/Qwen3-8B", qwen8BConfig, 16_400_000_000) + + extra, _ := planHFContext([]string{"--model", "Qwen/Qwen3-8B"}, cache, 65536, nil) + + assert.Equal(t, "65536", flagOf(t, extra, "--max-model-len"), "the window does not depend on the GPU reading") + assert.NotContains(t, extra, "--gpu-memory-utilization", "without a reading the GPU share is left to vLLM") +} + +func TestAutoContext_WrapsOverridesForAMultimodalConfig(t *testing.T) { + withGPU(t, 48, 48) + cfg := `{"architectures":["Qwen3VLForConditionalGeneration"],"text_config":` + qwen8BConfig + `}` + cache := cachedModel(t, "Qwen/VL", cfg, 16_400_000_000) + + extra, _ := planHFContext([]string{"--model", "Qwen/VL"}, cache, 65536, nil) + + var overrides map[string]json.RawMessage + require.NoError(t, json.Unmarshal([]byte(flagOf(t, extra, "--hf-overrides")), &overrides)) + assert.Contains(t, overrides, "text_config", "a multimodal model's rope settings live under text_config") +} + +// A snapshot that carries the same weights twice, in safetensors and as an +// unused PyTorch copy, is sized by the safetensors vLLM loads, not by what the +// snapshot takes on disk. +func TestAutoContext_SizesByTheWeightsVLLMLoadsNotTheDiskTheSnapshotUses(t *testing.T) { + withGPU(t, 48, 48) + single := cachedModel(t, "Qwen/Qwen3-8B", qwen8BConfig, 16_400_000_000) + both := cachedModel(t, "Qwen/Qwen3-8B", qwen8BConfig, 16_400_000_000) + repo := writeHFRepo(t, both, "Qwen/Qwen3-8B", map[string]string{"pytorch_model.bin": "x"}) + require.NoError(t, os.Truncate(filepath.Join(repo, "blobs", "blob-pytorch_model.bin"), 16_400_000_000)) + + alone, _ := planHFContext([]string{"--model", "Qwen/Qwen3-8B"}, single, 65536, nil) + withCopy, _ := planHFContext([]string{"--model", "Qwen/Qwen3-8B"}, both, 65536, nil) + + assert.Equal(t, alone, withCopy) + models, err := scanHFCache(both) + require.NoError(t, err) + require.Len(t, models, 1) + assert.Greater(t, models[0].Size, int64(32_000_000_000), "the cache list still shows what the snapshot takes on disk") +} + +// hfRevisions writes the cache Kaylee's example describes: two snapshots of one +// model, with refs/main pointing at the newer. +func hfRevisions(t *testing.T) string { + t.Helper() + cache := t.TempDir() + short := strings.Replace(qwen8BConfig, "40960", "8192", 1) + long := strings.Replace(qwen8BConfig, "40960", "131072", 1) + repo := writeHFRepo(t, cache, "org/m", llmFiles(map[string]string{"model.safetensors": "w"})) + require.NoError(t, os.Rename(filepath.Join(repo, "snapshots", "abc123"), filepath.Join(repo, "snapshots", "aaa111"))) + writeHFRepo(t, cache, "org/m", llmFiles(map[string]string{"model.safetensors": "w"})) + require.NoError(t, os.Rename(filepath.Join(repo, "snapshots", "abc123"), filepath.Join(repo, "snapshots", "bbb222"))) + // Each revision has its own config.json; the other files are shared blobs. + for revision, config := range map[string]string{"aaa111": short, "bbb222": long} { + path := filepath.Join(repo, "snapshots", revision, "config.json") + require.NoError(t, os.Remove(path)) + require.NoError(t, os.WriteFile(path, []byte(config), 0o644)) + } + require.NoError(t, os.WriteFile(filepath.Join(repo, "refs", "main"), []byte("bbb222"), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(repo, "refs", "v1"), []byte("aaa111"), 0o644)) + return cache +} + +func TestHFSnapshotDir_ResolvesARevisionTheWayHuggingFaceDoes(t *testing.T) { + cache := hfRevisions(t) + resolves := func(name, revision, wantSnapshot string) { + t.Run(name, func(t *testing.T) { + dir, ok := hfSnapshotDir(cache, "org/m", revision) + + require.True(t, ok) + assert.Equal(t, filepath.Join(cache, "models--org--m", "snapshots", wantSnapshot), dir) + }) + } + missing := func(name, revision string) { + t.Run(name, func(t *testing.T) { + _, ok := hfSnapshotDir(cache, "org/m", revision) + + assert.False(t, ok) + }) + } + resolves("main through refs, not the first snapshot by name", "main", "bbb222") + resolves("a tag through refs", "v1", "aaa111") + resolves("a commit names its snapshot", "aaa111", "aaa111") + missing("a revision that is not cached", "ccc333") + missing("a revision that is not a name in the cache", "../../models--org--m") +} + +// Automatic configuration reads the revision vLLM loads: refs/main by default, +// or --revision. A revision that is not cached is not configured from another. +func TestAutoContext_ReadsTheRevisionVLLMLoads(t *testing.T) { + withGPU(t, 48, 48) + cache := hfRevisions(t) + + t.Run("main", func(t *testing.T) { + extra, _ := planHFContext([]string{"--model", "org/m"}, cache, 65536, nil) + + assert.Equal(t, "65536", flagOf(t, extra, "--max-model-len")) + assert.NotContains(t, extra, "--hf-overrides", "main's 131072-token window needs no rope override") + }) + t.Run("an explicit older revision", func(t *testing.T) { + extra, _ := planHFContext([]string{"--model", "org/m", "--revision", "aaa111"}, cache, 65536, nil) + + assert.Contains(t, extra, "--hf-overrides", "aaa111's 8192-token window is extended to the target") + }) + t.Run("a revision that is not cached", func(t *testing.T) { + extra, _ := planHFContext([]string{"--model", "org/m", "--revision", "ccc333"}, cache, 65536, nil) + + assert.NotContains(t, extra, "--max-model-len", "no other revision's config is read") + }) +} + +func TestApplyAutoLaunch_AppendsAfterTheUsersArguments(t *testing.T) { + withGPU(t, 48, 48) + // The engine's hub cache. + parent := t.TempDir() + cachedModelAt(t, filepath.Join(parent, "hub"), "Qwen/Qwen3-8B", qwen8BConfig, 16_400_000_000) + rt := Runtime{AutoLaunch: &AutoLaunch{Kind: autoLaunchHFContext, ContextTokens: 65536}} + in := launchCommand{Bin: "vllm", Args: []string{"serve", "--model", "Qwen/Qwen3-8B"}} + + out := applyAutoLaunch(t.Context(), "vllm", rt, in, filepath.Join(parent, "hub"), nil) + + assert.Equal(t, in.Args, out.Args[:3]) + assert.Equal(t, "65536", flagOf(t, out.Args, "--max-model-len")) + assert.Len(t, in.Args, 3, "the caller's arguments were modified") + assert.Equal(t, in.Args, applyAutoLaunch(t.Context(), "vllm", Runtime{}, in, filepath.Join(parent, "hub"), nil).Args, "an engine without auto_launch must be untouched") +} + +func TestPlanToolCalling_UsesTheParserThatReadsTheModelsToolCalls(t *testing.T) { + cache := cachedModel(t, "org/m", qwen8BConfig, 1_000_000) + var probedSnapshot string + withToolParserProbe(t, func(bin, snapshot string) (toolParserProbe, error) { + probedSnapshot = snapshot + return toolParserProbe{Matches: []string{"ernie45", "granite4", "hermes", "qwen3_json"}}, nil + }) + + extra, parser := planToolCalling(t.Context(), []string{"--model", "org/m"}, cache, "/opt/vllm/venv/bin/vllm", nil) + + assert.Equal(t, "hermes", parser) + assert.Equal(t, "hermes", flagOf(t, extra, "--tool-call-parser")) + assert.Contains(t, extra, "--enable-auto-tool-choice") + assert.Equal(t, filepath.Join(cache, "models--org--m", "snapshots", "abc123"), probedSnapshot) +} + +func TestPlanToolCalling_PrefersTheModelsOwnResponseTemplate(t *testing.T) { + cache := cachedModel(t, "org/m", qwen8BConfig, 1_000_000) + withToolParserProbe(t, func(bin, snapshot string) (toolParserProbe, error) { + return toolParserProbe{HF: true, Matches: []string{"hermes"}}, nil + }) + + extra, _ := planToolCalling(t.Context(), []string{"--model", "org/m"}, cache, "vllm", nil) + + assert.Equal(t, "hf", flagOf(t, extra, "--tool-call-parser"), "the hf parser for a model that ships response_template") +} + +func TestPlanToolCalling_AddsNothingWhenNoParserCanBeChosen(t *testing.T) { + cache := cachedModel(t, "org/m", qwen8BConfig, 1_000_000) + test := func(name, model string, probe func(bin, snapshot string) (toolParserProbe, error)) { + t.Run(name, func(t *testing.T) { + withToolParserProbe(t, probe) + + extra, parser := planToolCalling(t.Context(), []string{"--model", model}, cache, "vllm", nil) + + assert.Empty(t, extra) + assert.Empty(t, parser) + }) + } + test("no parser reads the format", "org/m", func(bin, snapshot string) (toolParserProbe, error) { return toolParserProbe{}, nil }) + test("the probe fails", "org/m", func(bin, snapshot string) (toolParserProbe, error) { return toolParserProbe{}, errors.New("no python") }) + test("the model is not cached and cannot be fetched", "org/not-cached", func(bin, snapshot string) (toolParserProbe, error) { + t.Fatal("a model with no template has nothing to probe") + return toolParserProbe{}, nil + }) +} + +// A model this start downloads has its tokenizer and chat template fetched +// first, so its first start gets the same tool-call parser as every later one. +// Without that, vLLM would serve the model and refuse agent tool requests until +// a restart. +func TestPlanToolCalling_ChoosesTheParserOnTheFirstStartOfAnUncachedModel(t *testing.T) { + cache := t.TempDir() + var probedSnapshot string + withToolParserProbe(t, func(bin, snapshot string) (toolParserProbe, error) { + probedSnapshot = snapshot + return toolParserProbe{Matches: []string{"hermes"}}, nil + }) + + extra, parser := planToolCalling(t.Context(), []string{"--model", "org/new"}, cache, "vllm", pendingModel(t, cache, "org/new", qwen8BConfig, 1_000_000)) + + assert.Equal(t, "hermes", parser) + assert.Contains(t, extra, "--enable-auto-tool-choice") + assert.Equal(t, filepath.Join(cache, "models--org--new", "snapshots", "abc123"), probedSnapshot) +} + +func TestPlanToolCalling_KeepsTheParserTheUserChose(t *testing.T) { + cache := cachedModel(t, "org/q", qwen8BConfig, 1_000_000) + withToolParserProbe(t, func(bin, snapshot string) (toolParserProbe, error) { + t.Fatal("a parser the user named needs no probe") + return toolParserProbe{}, nil + }) + + extra, _ := planToolCalling(t.Context(), []string{"--model", "org/q", "--tool-call-parser", "qwen3_xml", "--enable-auto-tool-choice"}, cache, "vllm", nil) + assert.Empty(t, extra, "the user set both") + + extra, parser := planToolCalling(t.Context(), []string{"--model", "org/q", "--tool-call-parser", "qwen3_xml"}, cache, "vllm", nil) + assert.Equal(t, "qwen3_xml", parser) + assert.Equal(t, []string{"--enable-auto-tool-choice"}, extra, "only the enable flag beside the user's parser") +} + +func TestParseGPUMemory_ReadsADiscreteGPUAndFallsBackOnUnifiedMemory(t *testing.T) { + system := func() (int64, int64, bool) { return 128 * gib, 100 * gib, true } + reads := func(name, out string, wantTotal, wantFree int64) { + t.Run(name, func(t *testing.T) { + total, free, ok := parseGPUMemory(out, system) + + require.True(t, ok) + assert.Equal(t, wantTotal, total) + assert.Equal(t, wantFree, free) + }) + } + unreadable := func(name, out string) { + t.Run(name, func(t *testing.T) { + _, _, ok := parseGPUMemory(out, system) + + assert.False(t, ok) + }) + } + reads("discrete GPU, in MiB", "49140, 48600\n", 49140<<20, 48600<<20) + reads("first GPU of several", "24576, 20000\n49140, 48600\n", 24576<<20, 20000<<20) + reads("DGX Spark reports N/A", "[N/A], [N/A]\n", 128*gib, 100*gib) + reads("not supported", "[Not Supported], [Not Supported]\n", 128*gib, 100*gib) + unreadable("malformed", "error\n") + unreadable("empty", "") +} + +func TestParseMeminfo_ReadsTotalAndAvailable(t *testing.T) { + text := "MemTotal: 130567000 kB\nMemFree: 2000000 kB\nMemAvailable: 104000000 kB\nBuffers: 100 kB\n" + + total, available, ok := parseMeminfo(text) + + require.True(t, ok) + assert.Equal(t, int64(130567000*1024), total) + assert.Equal(t, int64(104000000*1024), available) + _, _, ok = parseMeminfo("MemTotal: 100 kB\n") + assert.False(t, ok, "a reading without MemAvailable must not count as known") +} + +// qwen35_4BConfig is the part of Qwen/Qwen3.5-4B's config.json that sizing +// reads: a vision-language model whose language model sits in text_config. +const qwen35_4BConfig = `{"model_type":"qwen3_5","text_config":{"max_position_embeddings":262144,"num_hidden_layers":32,"num_attention_heads":16,"num_key_value_heads":4,"head_dim":256,"hidden_size":2560,"dtype":"bfloat16"}}` + +// qwen35_4BBytes is about the weights a download of Qwen/Qwen3.5-4B fetches. +const qwen35_4BBytes = 9_440_000_000 + +func TestAutoContext_FirstStartOfAnUncachedModelGetsRoomForItsWeights(t *testing.T) { + // The default model on a 16 GiB GPU, before vLLM has downloaded it. A blind + // half share (8 GiB) cannot even hold the weights. + withGPU(t, 16, 15) + cache := t.TempDir() + + extra, note := planHFContext([]string{"--model", "Qwen/Qwen3.5-4B"}, cache, 65536, pendingModel(t, cache, "Qwen/Qwen3.5-4B", qwen35_4BConfig, qwen35_4BBytes)) + + weightsAndOverhead := float64(qwen35_4BBytes)*weightsSlack + float64(vllmOverhead) + assert.GreaterOrEqual(t, utilOf(t, extra)*16*float64(gib), weightsAndOverhead, "the share must hold the weights and overhead: %v", extra) + assert.LessOrEqual(t, utilOf(t, extra), 14.0/16.0+0.001, "must leave a margin of the 15 GiB free") + assert.NotEmpty(t, flagOf(t, extra, "--max-model-len"), "without it vLLM would try the model's full 262144-token window") + assert.Contains(t, note, "sized from Hugging Face") +} + +func TestAutoContext_DoesNotFetchACachedModel(t *testing.T) { + withGPU(t, 48, 48) + cache := cachedModel(t, "Qwen/Qwen3-8B", qwen8BConfig, 16_400_000_000) + fetch := func() (hfPendingModel, bool) { + t.Fatal("fetched a model although it is cached") + return hfPendingModel{}, false + } + + extra, _ := planHFContext([]string{"--model", "Qwen/Qwen3.5-8B", "--model", "Qwen/Qwen3-8B"}, cache, 65536, fetch) + + assert.Equal(t, "65536", flagOf(t, extra, "--max-model-len")) +} + +// hfFetch is a metadata fetch against a fake Hugging Face API and a fake hf +// tool beside the engine binary. The tool records its arguments and HF_TOKEN +// and writes each file it is asked for, with config.json from FAKE_HF_CONFIG. +type hfFetch struct { + e *Executor + st *engineState + hub string + requests []*http.Request + mu sync.Mutex +} + +const fakeHFMetadataDownload = ` +id=$2; shift 2 +files="" +while [ $# -gt 0 ]; do + case "$1" in + --cache-dir) dir=$2; shift 2 ;; + --revision) rev=$2; shift 2 ;; + *) files="$files $1"; shift ;; + esac +done +echo "$id$files --revision $rev" >> "$(dirname "$dir")/hf-calls" +printf '%s' "${HF_TOKEN:-}" > "$(dirname "$dir")/hf-token" +repo="$dir/models--$(echo "$id" | sed 's#/#--#g')" +mkdir -p "$repo/refs" "$repo/snapshots/fetched1" +printf fetched1 > "$repo/refs/$rev" +for f in $files; do + case "$f" in + config.json) printf '%s' "$FAKE_HF_CONFIG" > "$repo/snapshots/fetched1/$f" ;; + model.safetensors.index.json) printf '%s' "$FAKE_HF_INDEX" > "$repo/snapshots/fetched1/$f" ;; + *) printf '{}' > "$repo/snapshots/fetched1/$f" ;; + esac +done +` + +func newHFFetch(t *testing.T, status int, siblings []hfSibling) *hfFetch { + t.Helper() + e, hub, _ := hfPullFixture(t, siblings, status, fakeHFMetadataDownload) + st, err := e.state("vllm") + require.NoError(t, err) + f := &hfFetch{e: e, st: st, hub: hub} + api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + f.mu.Lock() + f.requests = append(f.requests, r) + f.mu.Unlock() + if status != http.StatusOK { + http.Error(w, "x", status) + return + } + assert.NoError(t, json.NewEncoder(w).Encode(map[string]any{"siblings": siblings})) + })) + t.Cleanup(api.Close) + old := hfDefaultEndpoint + hfDefaultEndpoint = api.URL + t.Cleanup(func() { hfDefaultEndpoint = old }) + t.Setenv("FAKE_HF_CONFIG", qwen35_4BConfig) + t.Setenv("FAKE_HF_INDEX", "") + t.Setenv("HF_HUB_OFFLINE", "") + t.Setenv("TRANSFORMERS_OFFLINE", "") + t.Setenv("HF_HUB_DISABLE_IMPLICIT_TOKEN", "") + t.Setenv("HF_ENDPOINT", "") + t.Setenv("HF_TOKEN", "") + t.Setenv("HUGGING_FACE_HUB_TOKEN", "") + t.Setenv("HF_TOKEN_PATH", "") + // No token a real `hf auth login` stored may reach the fake server. + t.Setenv("HF_HOME", t.TempDir()) + return f +} + +func (f *hfFetch) fetch(t *testing.T, id, revision string) (hfPendingModel, bool) { + t.Helper() + return f.e.fetchHFMetadata(t.Context(), f.st, f.hub, id, revision, engineHFEnv(f.st)) +} + +func (f *hfFetch) setLaunchEnv(assignments ...string) { + f.st.mu.Lock() + f.st.plat.Runtime.LaunchEnv = assignments + f.st.mu.Unlock() +} + +func (f *hfFetch) toolCalls(t *testing.T) []string { + t.Helper() + raw, err := os.ReadFile(filepath.Join(filepath.Dir(f.hub), "hf-calls")) + if errors.Is(err, os.ErrNotExist) { + return nil + } + require.NoError(t, err) + return strings.Split(strings.TrimSpace(string(raw)), "\n") +} + +var qwen35Siblings = []hfSibling{ + {RFilename: "config.json", Size: 2_000}, + {RFilename: "model.safetensors-00001-of-00002.safetensors", Size: 5_000_000_000}, + {RFilename: "model.safetensors-00002-of-00002.safetensors", Size: 4_440_000_000}, + {RFilename: "tokenizer.json", Size: 8_000_000}, + {RFilename: "chat_template.jinja", Size: 7_000}, + {RFilename: "model.gguf", Size: 3_000_000_000}, // not downloaded, so not counted +} + +func TestFetchHFMetadata_PutsTheSmallFilesInTheCacheAtTheRevision(t *testing.T) { + f := newHFFetch(t, http.StatusOK, qwen35Siblings) + + pending, ok := f.fetch(t, "Qwen/Qwen3.5-4B", "v2") + + require.True(t, ok) + assert.Equal(t, int64(5_000_000_000+4_440_000_000), pending.weights, "the weights vLLM loads, not the other files") + assert.Equal(t, filepath.Join(f.hub, "models--Qwen--Qwen3.5-4B", "snapshots", "fetched1"), pending.snapshot) + assert.Equal(t, []string{"Qwen/Qwen3.5-4B config.json tokenizer.json chat_template.jinja --revision v2"}, f.toolCalls(t), "only the small files, at the revision") + require.Len(t, f.requests, 1) + assert.Equal(t, "/api/models/Qwen/Qwen3.5-4B/revision/v2", f.requests[0].URL.Path) + cfg, ok := readHFConfig(pending.snapshot) + require.True(t, ok) + native, layers, _, _, _ := cfg.text().kvShape() + assert.Equal(t, int64(262144), native) + assert.Equal(t, int64(32), layers) +} + +// A repository that ships its weights twice, as indexed shards and as one +// file, is sized by the shards its index lists, the copy vLLM loads. +func TestFetchHFMetadata_SizesOnlyTheWeightsTheIndexLists(t *testing.T) { + f := newHFFetch(t, http.StatusOK, []hfSibling{ + {RFilename: "config.json", Size: 2_000}, + {RFilename: "tokenizer.json", Size: 8_000_000}, + {RFilename: "model-00001-of-00002.safetensors", Size: 5_000_000_000}, + {RFilename: "model-00002-of-00002.safetensors", Size: 4_440_000_000}, + {RFilename: "model.safetensors.index.json", Size: 30_000}, + {RFilename: "model.safetensors", Size: 9_440_000_000}, // the same weights again + }) + t.Setenv("FAKE_HF_INDEX", `{"weight_map": {"a": "model-00001-of-00002.safetensors", "b": "model-00002-of-00002.safetensors"}}`) + + pending, ok := f.fetch(t, "Qwen/Qwen3.5-4B", "main") + + require.True(t, ok) + assert.Equal(t, int64(5_000_000_000+4_440_000_000), pending.weights, "the indexed shards, not both copies") + assert.Equal(t, []string{"Qwen/Qwen3.5-4B config.json tokenizer.json model.safetensors.index.json --revision main"}, f.toolCalls(t), "the index comes down with the metadata") +} + +// HF_HUB_OFFLINE or HF_TOKEN set in the engine's launch settings applies to +// PAIR's own Hugging Face requests and hf runs, as it does to vLLM's. +func TestFetchHFMetadata_UsesTheEnginesHuggingFaceSettings(t *testing.T) { + t.Run("offline in the engine settings", func(t *testing.T) { + f := newHFFetch(t, http.StatusOK, qwen35Siblings) + f.setLaunchEnv("HF_HUB_OFFLINE=1") + + _, ok := f.fetch(t, "Qwen/Qwen3.5-4B", "main") + + assert.False(t, ok) + assert.Empty(t, f.requests, "no request may reach Hugging Face") + assert.Empty(t, f.toolCalls(t), "the hf tool must not run") + }) + t.Run("TRANSFORMERS_OFFLINE in the engine settings", func(t *testing.T) { + f := newHFFetch(t, http.StatusOK, qwen35Siblings) + f.setLaunchEnv("TRANSFORMERS_OFFLINE=1") + + _, ok := f.fetch(t, "Qwen/Qwen3.5-4B", "main") + + assert.False(t, ok) + assert.Empty(t, f.requests, "no request may reach Hugging Face") + assert.Empty(t, f.toolCalls(t), "the hf tool must not run") + }) + t.Run("implicit token disabled in the engine settings", func(t *testing.T) { + f := newHFFetch(t, http.StatusOK, qwen35Siblings) + home := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(home, "token"), []byte("stored-token\n"), 0o600)) + f.setLaunchEnv("HF_HUB_DISABLE_IMPLICIT_TOKEN=1", "HF_TOKEN=engine-token", "HF_HOME="+home) + + _, ok := f.fetch(t, "Qwen/Qwen3.5-4B", "main") + + require.True(t, ok) + require.Len(t, f.requests, 1) + assert.Empty(t, f.requests[0].Header.Get("Authorization"), "neither HF_TOKEN nor a stored token is sent") + }) + t.Run("a token in the engine settings", func(t *testing.T) { + f := newHFFetch(t, http.StatusOK, qwen35Siblings) + f.setLaunchEnv("HF_TOKEN=engine-token") + + _, ok := f.fetch(t, "Qwen/Qwen3.5-4B", "main") + + require.True(t, ok) + require.Len(t, f.requests, 1) + assert.Equal(t, "Bearer engine-token", f.requests[0].Header.Get("Authorization")) + token, err := os.ReadFile(filepath.Join(filepath.Dir(f.hub), "hf-token")) + require.NoError(t, err) + assert.Equal(t, "engine-token", string(token), "the hf tool gets the engine's token") + }) + t.Run("a token stored by hf auth login", func(t *testing.T) { + f := newHFFetch(t, http.StatusOK, qwen35Siblings) + home := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(home, "token"), []byte("stored-token\n"), 0o600)) + f.setLaunchEnv("HF_HOME=" + home) + + _, ok := f.fetch(t, "Qwen/Qwen3.5-4B", "main") + + require.True(t, ok) + require.Len(t, f.requests, 1) + assert.Equal(t, "Bearer stored-token", f.requests[0].Header.Get("Authorization")) + }) + t.Run("a mirror in the engine settings", func(t *testing.T) { + f := newHFFetch(t, http.StatusOK, qwen35Siblings) + mirror := hfDefaultEndpoint + hfDefaultEndpoint = "http://127.0.0.1:1" // nothing listens: the Hub must not be asked + f.setLaunchEnv("HF_ENDPOINT=" + mirror + "/") + + _, ok := f.fetch(t, "Qwen/Qwen3.5-4B", "main") + + require.True(t, ok) + require.Len(t, f.requests, 1, "the listing came from the mirror") + assert.Equal(t, "/api/models/Qwen/Qwen3.5-4B", f.requests[0].URL.Path) + }) +} + +func TestFetchHFMetadata_FallsBackWhenHuggingFaceCannotAnswer(t *testing.T) { + test := func(name string, status int, offline string) { + t.Run(name, func(t *testing.T) { + f := newHFFetch(t, status, qwen35Siblings) + t.Setenv("HF_HUB_OFFLINE", offline) + + _, ok := f.fetch(t, "org/model", "main") + assert.False(t, ok) + assert.Empty(t, f.toolCalls(t)) + + // The start then falls back to the capped share, as before: 15 GiB + // free less the 1 GiB margin is 0.87 of 16 GiB. + withGPU(t, 16, 15) + extra, _ := planHFContext([]string{"--model", "org/model"}, f.hub, 65536, func() (hfPendingModel, bool) { + return f.fetch(t, "org/model", "main") + }) + assert.InDelta(t, 0.87, utilOf(t, extra), 0.0001) + assert.NotContains(t, extra, "--max-model-len") + }) + } + test("model not found", http.StatusNotFound, "") + test("gated model without a token", http.StatusUnauthorized, "") + test("server error", http.StatusInternalServerError, "") + test("HF_HUB_OFFLINE set for the service", http.StatusOK, "1") +} + +func TestApplyAutoLaunch_LimitsConcurrentRequestsUnlessTheUserSetsALimit(t *testing.T) { + withoutGPU(t) + test := func(name string, args []string, limit int, want string) { + t.Run(name, func(t *testing.T) { + rt := Runtime{AutoLaunch: &AutoLaunch{Kind: autoLaunchHFContext, ContextTokens: 65536, MaxNumSeqs: limit}} + + out := applyAutoLaunch(t.Context(), "vllm", rt, launchCommand{Bin: "vllm", Args: args}, t.TempDir(), nil) + + assert.Equal(t, want, flagOf(t, out.Args, "--max-num-seqs")) + assert.LessOrEqual(t, strings.Count(strings.Join(out.Args, " "), "--max-num-seqs"), 1) + }) + } + test("limit added", []string{"serve", "--model", "org/m"}, 32, "32") + test("the user's own limit wins", []string{"serve", "--model", "org/m", "--max-num-seqs", "4"}, 32, "4") + test("the user's own limit in = form wins", []string{"serve", "--model", "org/m", "--max-num-seqs=64"}, 32, "64") + test("a sharded launch gets it too", []string{"serve", "--model", "org/m", "--tensor-parallel-size", "2"}, 32, "32") + test("no limit in the manifest leaves vLLM's default", []string{"serve", "--model", "org/m"}, 0, "") +} + +func TestBundledVLLMLaunch_FirstStartOfTheDefaultModelFitsA16GiBGPU(t *testing.T) { + // The default model on a 16 GiB GPU with an empty cache, through the bundled + // manifest. Each added argument is one vLLM needs to start there: the share + // must hold the weights, the context must be bounded, and the concurrent + // requests must leave room for Qwen3.5's per-request state. + withGPU(t, 16, 15) + reg := loadWithOverrides(t, t.TempDir()) + manifest, ok := reg.Get("vllm") + require.True(t, ok, "vllm manifest not loaded") + rt := manifest.Platforms["linux/amd64"].Runtime + parent := t.TempDir() + hub := filepath.Join(parent, "hub") + in := launchCommand{Bin: "vllm", Args: []string{"serve", "--model", "Qwen/Qwen3.5-4B"}} + + out := applyAutoLaunch(t.Context(), "vllm", rt, in, hub, pendingModel(t, hub, "Qwen/Qwen3.5-4B", qwen35_4BConfig, qwen35_4BBytes)) + + weightsAndOverhead := float64(qwen35_4BBytes)*weightsSlack + float64(vllmOverhead) + assert.GreaterOrEqual(t, utilOf(t, out.Args)*16*float64(gib), weightsAndOverhead, "the share cannot hold the weights and overhead: %v", out.Args) + assert.NotEmpty(t, flagOf(t, out.Args, "--max-model-len")) + assert.Equal(t, "32", flagOf(t, out.Args, "--max-num-seqs"), "vLLM's default of 256 does not start on this GPU") +} + +// The probe reads the chat format vLLM will use: the tokenizer --tokenizer and +// --tokenizer-revision name, and the --chat-template the launch gives. When +// that cannot be read here, no parser is chosen rather than one read from the +// model's own template, which vLLM would not be using. +func TestPlanToolCalling_ProbesTheTokenizerAndTemplateVLLMWillUse(t *testing.T) { + cache := cachedModel(t, "org/m", qwen8BConfig, 1_000_000) + writeHFRepo(t, cache, "org/tok", map[string]string{"tokenizer.json": "{}", "tokenizer_config.json": "{}"}) + template := filepath.Join(t.TempDir(), "tools.jinja") + require.NoError(t, os.WriteFile(template, []byte("{{ messages }}"), 0o644)) + snapshot := func(id string) string { + return filepath.Join(cache, "models--"+strings.ReplaceAll(id, "/", "--"), "snapshots", "abc123") + } + probes := func(t *testing.T) *[][2]string { + t.Helper() + var seen [][2]string + old := probeToolParsers + probeToolParsers = func(_ context.Context, _, tokenizer, template string) (toolParserProbe, error) { + seen = append(seen, [2]string{tokenizer, template}) + return toolParserProbe{Matches: []string{"hermes"}}, nil + } + toolParserCache.Clear() + t.Cleanup(func() { + probeToolParsers = old + toolParserCache.Clear() + }) + return &seen + } + probed := func(name string, extra []string, wantTokenizer, wantTemplate string) { + t.Run(name, func(t *testing.T) { + seen := probes(t) + + tools, parser := planToolCalling(t.Context(), append([]string{"--model", "org/m"}, extra...), cache, "vllm", nil) + + assert.Equal(t, "hermes", parser) + assert.Contains(t, tools, "--tool-call-parser") + assert.Equal(t, [][2]string{{wantTokenizer, wantTemplate}}, *seen) + }) + } + noParser := func(name string, extra []string) { + t.Run(name, func(t *testing.T) { + seen := probes(t) + + tools, parser := planToolCalling(t.Context(), append([]string{"--model", "org/m"}, extra...), cache, "vllm", nil) + + assert.Empty(t, tools) + assert.Empty(t, parser) + assert.Empty(t, *seen, "nothing to probe") + }) + } + probed("the model's own tokenizer and template", nil, snapshot("org/m"), "") + probed("another cached tokenizer repository", []string{"--tokenizer", "org/tok"}, snapshot("org/tok"), "") + probed("the model's revision named as the tokenizer revision", []string{"--tokenizer-revision", "main"}, snapshot("org/m"), "") + probed("a tokenizer directory", []string{"--tokenizer", cache}, cache, "") + probed("a chat template file", []string{"--chat-template", template}, snapshot("org/m"), template) + probed("a chat template given inline", []string{"--chat-template", "{% for m in messages %}{{ m.content }}{% endfor %}"}, snapshot("org/m"), "{% for m in messages %}{{ m.content }}{% endfor %}") + noParser("a tokenizer repository that is not cached", []string{"--tokenizer", "org/not-cached"}) + noParser("a tokenizer revision that is not cached", []string{"--tokenizer-revision", "v9"}) + noParser("a chat template given as a relative path", []string{"--chat-template", "tools.jinja"}) +} diff --git a/services/nvpair-engine-manager/launchmodel.go b/services/nvpair-engine-manager/launchmodel.go new file mode 100644 index 000000000..83d368b64 --- /dev/null +++ b/services/nvpair-engine-manager/launchmodel.go @@ -0,0 +1,317 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "regexp" + "slices" + "strings" + + settings "nvpair-shared/enginesettings" +) + +// hfModelIDRe accepts a Hugging Face repository id, "name" or "org/name". It +// refuses anything that could be read as a flag, a path or a shell fragment, +// since the value becomes an engine argument. +var hfModelIDRe = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,95}(/[A-Za-z0-9][A-Za-z0-9._-]{0,95})?$`) + +// launchEdit is one change to an engine's saved launch arguments, made while +// holding the engine's operation lock from the first read to the last write. A +// settings edit, start or stop cannot land between reading the launch, changing +// it and putting it back, so a rollback restores exactly what this change +// replaced. Each caller keeps its own failure policy; launchEdit only supplies +// the steps. +type launchEdit struct { + e *Executor + st *engineState + engine string + // before is the launch state when the lock was taken. + before settings.LaunchState +} + +// editLaunch runs change with the engine's operation lock held. +func (e *Executor) editLaunch(engine string, change func(edit *launchEdit) error) error { + st, err := e.state(engine) + if err != nil { + return err + } + st.opMu.Lock() + defer st.opMu.Unlock() + return change(&launchEdit{e: e, st: st, engine: engine, before: e.launchStateLocked(engine, st)}) +} + +// save writes launch text. A running engine is restarted on it; a stopped one +// is started only when resume is set. +func (l *launchEdit) save(ctx context.Context, launchText string, resume bool) error { + _, err := l.e.configureLaunchLocked(ctx, l.st, settings.Configure{ + Engine: l.engine, + Settings: settings.Config{ServerPort: l.before.ServerPort, LaunchText: launchText}, + Resume: resume, + }, func() error { return nil }) + return err +} + +// restore puts back the launch text this change started from, and starts the +// engine again if it was running then, unless a Stop or Eject is waiting: the +// user asked for the engine to be off, so it is left stopped. It does not +// depend on the caller still waiting: a remote caller that gives up cancels ctx +// mid-start, and a restore under that ctx would fail at once and leave the +// engine down. +func (l *launchEdit) restore(ctx context.Context) error { + return l.save(context.WithoutCancel(ctx), l.before.LaunchText, l.before.Running && !l.e.stopWaiting(l.st)) +} + +// start starts the stopped engine on its saved launch and records that the +// user wants it on, as Start does. +func (l *launchEdit) start(ctx context.Context) error { + if err := l.e.doStart(ctx, l.st, l.engine, startOpts{}); err != nil { + return err + } + return l.e.setDesiredEnabled(l.engine, true) +} + +// switchLaunchModel is the launch_model action. An engine that serves one model +// per process has no load call: the model it serves is a launch argument, so +// loading a model means writing that argument into the saved launch settings +// and starting the engine on it. This is the same path the settings editor +// takes, so the saved launch text always shows what the engine is running. +// +// A running engine is restarted on the new model. A stopped engine, such as +// after Eject, is started on it. Either way, if the new model does not come up +// (it does not fit in GPU memory, or the id does not exist) the previous launch +// is put back, and a running engine is started on it again, so a failed switch +// does not leave the engine down or the settings pointing at a model that +// failed. +// +// A Stop or Eject during the Load ends it without starting anything: the engine +// stays stopped with the previous model selected. That is what the user asked +// for, not a failure, so the reply succeeds with stopped set, and only a restore +// that fails is an error. +func (e *Executor) switchLaunchModel(ctx context.Context, engine string, spec *ActionLaunchModel, params json.RawMessage) (json.RawMessage, error) { + var request struct { + Model string `json:"model"` + } + if err := json.Unmarshal(params, &request); err != nil { + return nil, errors.New(`launch_model params must be {"model": ""}`) + } + model := strings.TrimSpace(request.Model) + if !hfModelIDRe.MatchString(model) { + return nil, fmt.Errorf("%q is not a valid model id (want a Hugging Face id such as org/name)", request.Model) + } + + var restarted, stopped bool + err := e.editLaunch(engine, func(edit *launchEdit) error { + current := edit.before + if !current.Editable { + return fmt.Errorf("the launch settings of engine %q cannot be changed: %s", engine, current.Reason) + } + text, changed, err := withLaunchFlag(current.LaunchText, spec.Flag, model) + if err != nil { + return err + } + if !changed && current.Running { + return nil + } + restarted = true + if current.Running { + err = edit.save(ctx, text, false) + } else { + if changed { + if err := edit.save(ctx, text, false); err != nil { + return err + } + } + err = edit.start(ctx) + } + if err == nil { + return nil + } + if e.stopWaiting(edit.st) { + stopped, restarted = true, false + if !changed { + return nil + } + if restoreErr := edit.restore(ctx); restoreErr != nil { + return fmt.Errorf("loading %s was stopped, and restoring the previous model failed: %w", model, restoreErr) + } + return nil + } + if !changed { + return fmt.Errorf("starting %s failed: %w", model, err) + } + // The engine's own error carries a long log tail, so say what happened first. + if restoreErr := edit.restore(ctx); restoreErr != nil { + return fmt.Errorf("switching to %s failed and restoring the previous model also failed (%v): %w", model, restoreErr, err) + } + return fmt.Errorf("switching to %s failed; the previous model was restored: %w", model, err) + }) + if err != nil { + return nil, err + } + return json.Marshal(struct { + Model string `json:"model"` + Restarted bool `json:"restarted"` + Stopped bool `json:"stopped,omitempty"` + }{model, restarted, stopped}) +} + +// withLaunchFlag returns launch text with flag set to value. Every existing +// "flag value" or "flag=value" collapses into one "flag value" at the position +// of the first, and every other argument is kept in place; when the flag is +// absent it is added before any "--" terminator. changed is false when the +// arguments already say the same thing, whatever their quoting: the launch text +// shown to a user is quoted canonically, so comparing text would call an +// unchanged launch changed. +func withLaunchFlag(text, flag, value string) (result string, changed bool, err error) { + tokens, err := parseLaunchText(text) + if err != nil { + return "", false, err + } + rest, at := dropLaunchFlag(tokens, flag) + out := slices.Insert(rest, at, flag, value) + result, err = formatLaunchText(out) + return result, !slices.Equal(tokens, out), err +} + +// withoutLaunchFlag returns launch text with every "flag value" and +// "flag=value" before any "--" terminator removed, and whether any was. +func withoutLaunchFlag(text, flag string) (result string, changed bool, err error) { + tokens, err := parseLaunchText(text) + if err != nil { + return "", false, err + } + rest, _ := dropLaunchFlag(tokens, flag) + result, err = formatLaunchText(rest) + return result, len(rest) != len(tokens), err +} + +// dropLaunchFlag returns the arguments without any "flag value" or "flag=value" +// before a "--" terminator, and the index in that result where the flag first +// stood. Without the flag, the index is the terminator's, or the end. +// Arguments are a list rather than key/value pairs: flags repeat, may take no +// value, and keep their order, so the list is walked rather than looked up. +func dropLaunchFlag(tokens []string, flag string) (rest []string, at int) { + rest = make([]string, 0, len(tokens)+2) + at = -1 + for i := 0; i < len(tokens); i++ { + token := tokens[i] + switch { + case token == "--": + if at < 0 { + at = len(rest) + } + return append(rest, tokens[i:]...), at + case token == flag || strings.HasPrefix(token, flag+"="): + if at < 0 { + at = len(rest) + } + if token == flag { + i++ // its value + } + default: + rest = append(rest, token) + } + } + if at < 0 { + at = len(rest) + } + return rest, at +} + +// modelLaunchFlag is the flag an engine's launch_model action writes, or "" for +// an engine that does not choose its model by launch argument. +func modelLaunchFlag(m *Manifest) string { + for _, act := range m.Actions { + if act.LaunchModel != nil { + return act.LaunchModel.Flag + } + } + return "" +} + +// ejectByStopping ejects the served model by stopping the engine, which is the +// only way an engine with one model per process gives its GPU memory back. A +// stopped engine has nothing loaded, so that is a successful no-op. A model the +// engine reports it is not serving is refused, since stopping would eject the +// wrong one. +// The served model is asked of the engine itself, not read from the saved +// launch arguments, because an adopted engine may have been started on another. +func (e *Executor) ejectByStopping(ctx context.Context, engine string, params json.RawMessage) (json.RawMessage, error) { + var request struct { + Model string `json:"model"` + } + if len(params) > 0 && string(params) != "null" { + if err := json.Unmarshal(params, &request); err != nil { + return nil, errors.New(`unload params must be {"model": ""}`) + } + } + model := strings.TrimSpace(request.Model) + stopped := false + st, err := e.state(engine) + if err != nil { + return nil, err + } + // Eject is a stop: a start in flight, such as a Load restarting the engine, + // is cancelled rather than waited out. + done, cancelledStart := e.requestStop(st) + defer done() + err = e.editLaunch(engine, func(edit *launchEdit) error { + if !edit.before.Running { + // A start this Eject cancelled leaves the engine stopped; record that + // the user wants it off, as a Stop would, so it is not restarted. + if cancelledStart { + return e.setDesiredEnabled(engine, false) + } + return nil + } + // When the engine cannot say what it serves, it is stopped anyway: Eject + // exists to free the GPU, and stopping the engine is all it can do. + if model != "" { + if serving, err := e.servingModels(ctx, edit.st, engine); err == nil && !slices.Contains(serving, model) { + return fmt.Errorf("%s is not the model %s is serving (%s)", model, engine, strings.Join(serving, ", ")) + } + } + stopped = true + // Like Stop, this is explicit OFF intent; see Stop for why it is + // persisted even when the stop itself fails. + return errors.Join(e.doStop(edit.st, engine), e.setDesiredEnabled(engine, false)) + }) + if err != nil { + return nil, err + } + return json.Marshal(map[string]any{"model": request.Model, "stopped": stopped}) +} + +// servingModels asks a running engine which models it is serving, through its +// loaded_models action, and returns every name it reports for them: the name +// requests use and, for an engine that declares an inventory_field, the cached +// model id too. Callers hold st.opMu; this does not take it. +func (e *Executor) servingModels(ctx context.Context, st *engineState, engine string) ([]string, error) { + act, ok := st.manifest.Actions["loaded_models"] + if !ok || act.Result == nil { + return nil, fmt.Errorf("engine %q does not report the models it serves", engine) + } + raw, err := e.dispatchAction(ctx, st, engine, "loaded_models", act, nil) + if err != nil { + return nil, err + } + names, ok := extractStringsResult(raw, act.Result) + if !ok { + return nil, fmt.Errorf("engine %q returned an unreadable model list", engine) + } + if act.Result.InventoryField != "" { + inventory := *act.Result + inventory.Field = act.Result.InventoryField + ids, ok := extractStringsResult(raw, &inventory) + if !ok { + return nil, fmt.Errorf("engine %q returned an unreadable model list", engine) + } + names = append(names, ids...) + } + return names, nil +} diff --git a/services/nvpair-engine-manager/launchmodel_test.go b/services/nvpair-engine-manager/launchmodel_test.go new file mode 100644 index 000000000..dde3a11e6 --- /dev/null +++ b/services/nvpair-engine-manager/launchmodel_test.go @@ -0,0 +1,568 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "context" + "encoding/json" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + settings "nvpair-shared/enginesettings" +) + +func TestWithLaunchFlag_SetsTheFlagOnceAndKeepsTheRestOfTheText(t *testing.T) { + test := func(name, text, want string, wantChanged bool) { + t.Run(name, func(t *testing.T) { + got, changed, err := withLaunchFlag(text, "--model", "org/new") + + require.NoError(t, err) + assert.Equal(t, want, got) + assert.Equal(t, wantChanged, changed, "changed") + }) + } + test("added when absent", "--max-model-len 4096", "--max-model-len 4096 --model org/new", true) + test("added to empty text", "", "--model org/new", true) + test("replaces a separate value", "--model org/old --max-model-len 4096", "--model org/new --max-model-len 4096", true) + test("replaces an equals form", "--max-model-len 4096 --model=org/old --dtype auto", "--max-model-len 4096 --model org/new --dtype auto", true) + test("keeps leading environment assignments", "VLLM_LOGGING_LEVEL=DEBUG --model org/old", "VLLM_LOGGING_LEVEL=DEBUG --model org/new", true) + test("collapses a repeated flag", "--model a --x 1 --model=b", "--model org/new --x 1", true) + test("goes before a terminator", "--x 1 -- positional", "--x 1 --model org/new -- positional", true) + test("keeps quoted values", `--served-model-name "two words" --model org/old`, `--served-model-name "two words" --model org/new`, true) + test("leaves a prefix-sharing flag alone", "--model-impl vllm --model org/old", "--model-impl vllm --model org/new", true) + test("unchanged when already set", "--x 1 --model org/new", "--x 1 --model org/new", false) + test("unchanged despite different quoting", `OLLAMA_HOST="127.0.0.1:1" --model org/new`, "OLLAMA_HOST=127.0.0.1:1 --model org/new", false) +} + +func TestWithLaunchFlag_RejectsLaunchTextThatDoesNotParse(t *testing.T) { + _, _, err := withLaunchFlag(`--model "unterminated`, "--model", "x") + + assert.Error(t, err) +} + +// switchModelExecutor is a running fake engine with editable launch settings +// and a launch_model action. The fake ignores arguments it does not know, so +// the "model" flag is carried in the saved launch text without changing how it +// serves. +func switchModelExecutor(t *testing.T, flag string) *Executor { + t.Helper() + e := settingsExecutor(t, false) + m, ok := e.reg.Get("fake") + require.True(t, ok, "fake manifest missing") + m.Actions["load_model"] = Action{LaunchModel: &ActionLaunchModel{Flag: flag}} + // Like vLLM, the engine answers which model it serves, under the name + // requests use (id) and the model it loaded (root). + m.AdvertisedModels = advertisedModelsLoaded + m.Actions["loaded_models"] = Action{ + HTTP: &ActionHTTP{Method: "GET", Path: "/v1/models"}, + Result: &ActionResult{Array: "data", Field: "id", InventoryField: "root"}, + } + require.NoError(t, m.Validate()) + // An engine that chooses its model by launch argument needs one to start. + state, err := e.LaunchSettings("fake") + require.NoError(t, err) + text, _, err := withLaunchFlag(state.LaunchText, flag, "org/initial") + require.NoError(t, err) + _, err = e.ConfigureLaunch(t.Context(), settings.Configure{ + Engine: "fake", + Settings: settings.Config{ServerPort: state.ServerPort, LaunchText: text}, + }, func() error { return nil }) + require.NoError(t, err) + require.NoError(t, e.Start(t.Context(), "fake"), "start") + return e +} + +func loadModel(t *testing.T, e *Executor, model string) (json.RawMessage, error) { + t.Helper() + return e.Action(t.Context(), "fake", "load_model", json.RawMessage(`{"model":`+quoteJSON(model)+`}`)) +} + +func quoteJSON(s string) string { b, _ := json.Marshal(s); return string(b) } + +func launchTextOf(t *testing.T, e *Executor) string { + t.Helper() + return launchTextOfEngine(t, e, "fake") +} + +func launchTextOfEngine(t *testing.T, e *Executor, engine string) string { + t.Helper() + state, err := e.LaunchSettings(engine) + require.NoError(t, err) + return state.LaunchText +} + +func TestWithoutLaunchFlag_RemovesEveryFormOfTheFlag(t *testing.T) { + test := func(name, text, want string, wantChanged bool) { + t.Run(name, func(t *testing.T) { + got, changed, err := withoutLaunchFlag(text, "--model") + + require.NoError(t, err) + assert.Equal(t, want, got) + assert.Equal(t, wantChanged, changed, "changed") + }) + } + test("separate value", "--model org/a --port 8001", "--port 8001", true) + test("equals form", "--model=org/a --port 8001", "--port 8001", true) + test("repeated", "--model org/a --port 8001 --model=org/b", "--port 8001", true) + test("after the terminator is kept", "--port 8001 -- --model org/a", "--port 8001 -- --model org/a", false) + test("absent", "--port 8001", "--port 8001", false) +} + +// loadReply is the load_model result. +type loadReply struct { + Model string `json:"model"` + Restarted bool `json:"restarted"` + Stopped bool `json:"stopped,omitempty"` +} + +func decodeLoadReply(t *testing.T, raw json.RawMessage) loadReply { + t.Helper() + var reply loadReply + require.NoError(t, json.Unmarshal(raw, &reply)) + return reply +} + +func TestLaunchModel_SwitchesAndRestartsARunningEngine(t *testing.T) { + e := switchModelExecutor(t, "--model") + gen := startGeneration(t, e) + + raw, err := loadModel(t, e, "org/first") + + require.NoError(t, err, "load first") + assert.Equal(t, loadReply{Model: "org/first", Restarted: true}, decodeLoadReply(t, raw)) + assert.Contains(t, launchTextOf(t, e), "--model org/first", "the launch text carries the model") + after := mustStatus(t, e) + assert.True(t, after.Running, "the engine runs the new model") + assert.NotEqual(t, gen, startGeneration(t, e), "the engine was restarted") + + // A second switch replaces the first and keeps the rest of the launch text. + _, err = e.ConfigureLaunch(t.Context(), settings.Configure{Engine: "fake", Settings: settings.Config{ + ServerPort: after.Port, LaunchText: launchTextOf(t, e) + " --max-model-len 4096", + }}, func() error { return nil }) + require.NoError(t, err) + _, err = loadModel(t, e, "org/second") + require.NoError(t, err, "load second") + text := launchTextOf(t, e) + assert.Contains(t, text, "--model org/second") + assert.NotContains(t, text, "org/first") + assert.Contains(t, text, "--max-model-len 4096") + + // Asking for the model already being served does not restart anything. + gen = startGeneration(t, e) + raw, err = loadModel(t, e, "org/second") + require.NoError(t, err) + assert.False(t, decodeLoadReply(t, raw).Restarted, "a no-op switch reports no restart") + assert.Equal(t, gen, startGeneration(t, e), "a no-op switch restarted the engine") +} + +// startGeneration counts how many times the engine has been started. +func startGeneration(t *testing.T, e *Executor) int64 { + t.Helper() + st := settingsState(t, e) + st.mu.Lock() + defer st.mu.Unlock() + return st.gen +} + +func mustStatus(t *testing.T, e *Executor) EngineStatus { + t.Helper() + st, err := e.Status("fake") + require.NoError(t, err) + return st +} + +// Load on a stopped engine (after Eject, say) serves the model: it saves the +// choice and starts the engine, whether or not the model changed. +func TestLaunchModel_StartsAStoppedEngine(t *testing.T) { + e := switchModelExecutor(t, "--model") + require.NoError(t, e.Stop("fake")) + + raw, err := loadModel(t, e, "org/later") + + require.NoError(t, err, "load on a stopped engine") + assert.True(t, decodeLoadReply(t, raw).Restarted) + assert.True(t, mustStatus(t, e).Running, "loading a model into a stopped engine must start it") + assert.Contains(t, launchTextOf(t, e), "--model org/later") + + // The same model again, on an engine stopped once more, still starts it. + require.NoError(t, e.Stop("fake")) + _, err = loadModel(t, e, "org/later") + require.NoError(t, err) + assert.True(t, mustStatus(t, e).Running, "an unchanged model on a stopped engine must start it") +} + +func TestLaunchModel_RestoresTheChoiceWhenAStoppedEngineWillNotStart(t *testing.T) { + e := switchModelExecutor(t, "--model") + before := launchTextOf(t, e) + require.NoError(t, e.Stop("fake")) + + _, err := loadModel(t, e, "org/fails-to-load") + + assert.ErrorContains(t, err, "previous model was restored", "the failed start is reported with the restore") + assert.False(t, mustStatus(t, e).Running, "the engine must stay stopped when the model does not start") + assert.Equal(t, before, launchTextOf(t, e)) +} + +func ejectEngine(t *testing.T, e *Executor, model string) (json.RawMessage, error) { + t.Helper() + m, _ := e.reg.Get("fake") + m.Actions["unload_model"] = Action{StopEngine: true} + require.NoError(t, m.Validate()) + return e.Action(t.Context(), "fake", "unload_model", json.RawMessage(`{"model":`+quoteJSON(model)+`}`)) +} + +func TestEject_StopsTheEngineAndKeepsItsModelChoice(t *testing.T) { + e := switchModelExecutor(t, "--model") + _, err := loadModel(t, e, "org/served") + require.NoError(t, err) + type ejectReply struct { + Stopped bool `json:"stopped"` + } + eject := func(model string) ejectReply { + t.Helper() + raw, err := ejectEngine(t, e, model) + require.NoError(t, err) + var reply ejectReply + require.NoError(t, json.Unmarshal(raw, &reply)) + return reply + } + + _, err = ejectEngine(t, e, "org/other") + assert.ErrorContains(t, err, "not the model", "ejecting a model that is not served must be refused") + require.True(t, mustStatus(t, e).Running, "a refused eject stopped the engine") + + assert.Equal(t, ejectReply{Stopped: true}, eject("org/served")) + assert.False(t, mustStatus(t, e).Running, "eject must stop the engine") + assert.Contains(t, launchTextOf(t, e), "--model org/served", "eject must keep the model choice so Load can serve it again") + + assert.Equal(t, ejectReply{Stopped: false}, eject("org/served"), "ejecting from a stopped engine is a no-op") +} + +func TestLaunchModel_RefusesAnythingThatIsNotAModelID(t *testing.T) { + e := switchModelExecutor(t, "--model") + before := launchTextOf(t, e) + reject := func(name, params string) { + t.Run(name, func(t *testing.T) { + _, err := e.Action(t.Context(), "fake", "load_model", json.RawMessage(params)) + + assert.Error(t, err) + }) + } + rejectModel := func(name, model string) { + reject(name, `{"model":`+quoteJSON(model)+`}`) + } + accept := func(name, model string) { + t.Run(name, func(t *testing.T) { + _, err := loadModel(t, e, model) + + assert.NoError(t, err) + }) + } + rejectModel("empty", "") + rejectModel("blank", " ") + rejectModel("an option", "--help") + rejectModel("a short option", "-x") + rejectModel("an id followed by an option", "org/name --extra") + rejectModel("a space", "a b") + rejectModel("a relative path", "../etc/passwd") + rejectModel("an absolute path", "/abs/path") + rejectModel("an empty path segment", "org//name") + rejectModel("three path segments", "a/b/c") + rejectModel("a shell separator", "x;y") + rejectModel("a command substitution", "$(id)") + rejectModel("a newline", "org/na\nme") + reject("params without a model", `{"nope":1}`) + t.Run("missing params", func(t *testing.T) { + _, err := e.Action(t.Context(), "fake", "load_model", nil) + + assert.Error(t, err) + }) + require.Equal(t, before, launchTextOf(t, e), "a refused switch changed the launch text") + require.True(t, mustStatus(t, e).Running, "a refused switch stopped the engine") + + accept("a name without an organization", "gpt2") + accept("an organization and a name", "Qwen/Qwen2.5-0.5B-Instruct") + accept("dots, underscores and dashes", "org/name.v1_2-x") +} + +// If the new model does not come up, the engine goes back to the one it was +// serving instead of staying down. The fake engine exits at once when it is +// given the model org/fails-to-load, which stands in for a model that does not fit. +func TestLaunchModel_RestoresThePreviousModelWhenTheNewOneFails(t *testing.T) { + e := switchModelExecutor(t, "--model") + before := launchTextOf(t, e) + + _, err := loadModel(t, e, "org/fails-to-load") + + require.Error(t, err, "a model that failed to start was reported as loaded") + // The note comes first: the engine's own error that follows it is a long log tail. + assert.True(t, strings.HasPrefix(err.Error(), "switching to org/fails-to-load failed; the previous model was restored: "), + "error does not lead with the restore note: %v", err) + assert.Equal(t, before, launchTextOf(t, e)) + assert.True(t, mustStatus(t, e).Running, "the engine was left down after a failed switch") +} + +// A remote caller that stops waiting cancels the switch while the new model is +// still starting. The previous model must still come back: restoring it under +// the cancelled context would fail at once and leave the engine down. +func TestLaunchModel_RestoresThePreviousModelWhenTheCallerGivesUp(t *testing.T) { + e := switchModelExecutor(t, "--model") + // Every launch from here on takes a second to become ready, so the switch + // below is still starting when its caller gives up. + _, err := e.ConfigureLaunch(t.Context(), settings.Configure{Engine: "fake", Settings: settings.Config{ + ServerPort: mustStatus(t, e).Port, LaunchText: "FAKE_START_DELAY=1s " + launchTextOf(t, e), + }}, func() error { return nil }) + require.NoError(t, err) + before := launchTextOf(t, e) + ctx, cancel := context.WithTimeout(t.Context(), 300*time.Millisecond) + defer cancel() + + _, err = e.Action(ctx, "fake", "load_model", json.RawMessage(`{"model":"org/slow"}`)) + + require.Error(t, err, "a switch whose caller gave up was reported as done") + assert.ErrorContains(t, err, "the previous model was restored") + assert.Equal(t, before, launchTextOf(t, e)) + assert.True(t, mustStatus(t, e).Running, "the engine was left down after the caller gave up") +} + +func TestLaunchModelAction_AcceptsOnlyALongOptionFlag(t *testing.T) { + accept := func(name, flag string) { + t.Run(name, func(t *testing.T) { + assert.NoError(t, (&Action{LaunchModel: &ActionLaunchModel{Flag: flag}}).validate("load_model")) + }) + } + reject := func(name, flag string) { + t.Run(name, func(t *testing.T) { + assert.Error(t, (&Action{LaunchModel: &ActionLaunchModel{Flag: flag}}).validate("load_model")) + }) + } + accept("a long option", "--model") + reject("no dashes", "model") + reject("a space in the flag", "--mod el") + reject("a value in the flag", "--model=x") + reject("empty", "") +} + +// The restart needs editable launch settings and a readiness probe on every platform. +func TestLaunchModelAction_RequiresEditableLaunch(t *testing.T) { + m := testEngineManifest(fakeEngineBin) + m.Actions["load_model"] = Action{LaunchModel: &ActionLaunchModel{Flag: "--model"}} + + assert.ErrorContains(t, m.Validate(), "editable_launch") +} + +// slowStarts stops the engine and makes every later launch take 8 seconds to +// become ready (under the fake manifest's 10-second readiness limit), so a +// start is still in flight when the test acts on it. +func slowStarts(t *testing.T, e *Executor) { + t.Helper() + require.NoError(t, e.Stop("fake")) + _, err := e.ConfigureLaunch(t.Context(), settings.Configure{Engine: "fake", Settings: settings.Config{ + ServerPort: mustStatus(t, e).Port, LaunchText: "FAKE_START_DELAY=8s " + launchTextOf(t, e), + }}, func() error { return nil }) + require.NoError(t, err) +} + +// waitForStartInFlight returns once the engine's start has begun. +func waitForStartInFlight(t *testing.T, e *Executor) { + t.Helper() + st, err := e.state("fake") + require.NoError(t, err) + require.Eventually(t, func() bool { + st.mu.Lock() + defer st.mu.Unlock() + return st.startCancel != nil + }, 10*time.Second, 10*time.Millisecond, "the start never began") +} + +// Stopping an engine that is still loading cancels the start instead of +// waiting out its readiness window, and the user's stop is not reported as a +// failed start. +func TestStop_CancelsAStartThatIsStillLoading(t *testing.T) { + e := switchModelExecutor(t, "--model") + slowStarts(t, e) + started := make(chan error, 1) + go func() { started <- e.Start(t.Context(), "fake") }() + waitForStartInFlight(t, e) + + began := time.Now() + require.NoError(t, e.Stop("fake")) + + assert.Less(t, time.Since(began), 4*time.Second, "Stop waited out the start") + require.ErrorIs(t, <-started, errStartStopped, "the start says the user stopped it") + assert.False(t, mustStatus(t, e).Running) + for _, reported := range e.reporter.snapshot() { + assert.NotEqual(t, startFailedID("fake"), reported.ID, "a start the user stopped was reported as failed") + } +} + +// engine:start that a Stop cut short replies with the engine's status, stopped, +// not an error: the user asked for it, and the error would carry the cancelled +// engine's partial output. A caller that waits for the reply, such as the +// desktop's first-open start, would otherwise report a failed start. +func TestRunOp_AnswersAStartTheUserStoppedWithTheStoppedStatus(t *testing.T) { + e := switchModelExecutor(t, "--model") + slowStarts(t, e) + + status := stopDuringOp(t, e, "engine:start", `{"engine":"fake"}`) + + assert.False(t, status.Running) +} + +// The broker no longer holds a Stop back while a Restart is starting the engine +// again, so a Restart is cut short the same way and answers the same way: the +// desktop and the TUI wait for its reply. +func TestRunOp_AnswersARestartTheUserStoppedWithTheStoppedStatus(t *testing.T) { + e := switchModelExecutor(t, "--model") + slowStarts(t, e) + require.NoError(t, e.Start(t.Context(), "fake")) + + status := stopDuringOp(t, e, "engine:restart", `{"engine":"fake"}`) + + assert.False(t, status.Running) + enabled, known, err := e.desired.get("fake") + require.NoError(t, err) + assert.True(t, known) + assert.False(t, enabled, "the user turned the engine off") +} + +// stopDuringOp sends a lifecycle method through the manager, stops the engine once the +// start it begins is in flight, and returns the reply, which must be the +// engine's status, not an error: the user asked for the stop, and the error +// would carry the cancelled engine's partial output. +func stopDuringOp(t *testing.T, e *Executor, method, params string) EngineStatus { + t.Helper() + var out bytes.Buffer // read only after the reply, so no lock is needed + m := NewManager(NewCodec(&out), e, nil) + id := json.RawMessage("7") + replied := make(chan struct{}) + go func() { + m.runOp(t.Context(), &Message{JSONRPC: "2.0", ID: &id, Method: method, Params: json.RawMessage(params)}) + close(replied) + }() + waitForStartInFlight(t, e) + + began := time.Now() + require.NoError(t, e.Stop("fake")) + <-replied + + assert.Less(t, time.Since(began), 4*time.Second, "the reply waited out the start") + var reply struct { + Result *EngineStatus `json:"result"` + Error *json.RawMessage `json:"error"` + } + require.NoError(t, json.Unmarshal(out.Bytes(), &reply), out.String()) + require.Nil(t, reply.Error, "a start the user stopped is not an error: %s", out.String()) + require.NotNil(t, reply.Result) + return *reply.Result +} + +// loadOutcome is a load_model call's result and error. +type loadOutcome struct { + result json.RawMessage + err error +} + +// assertLoadStopped checks the reply to a Load that a Stop or Eject ended: the +// user asked for it, so it is not an error, and it says the Load was stopped. +func assertLoadStopped(t *testing.T, outcome loadOutcome) { + t.Helper() + require.NoError(t, outcome.err, "a Load the user stopped is not a failure") + assert.Equal(t, loadReply{Model: "org/slow", Stopped: true}, decodeLoadReply(t, outcome.result)) +} + +// Eject during a Load that is starting the engine stops it at once, puts the +// previous model back without starting it, and leaves the engine off. +func TestEject_DuringALoadStopsTheEngineAndKeepsThePreviousModel(t *testing.T) { + e := switchModelExecutor(t, "--model") + slowStarts(t, e) + before := launchTextOf(t, e) + loaded := make(chan loadOutcome, 1) + go func() { + result, err := loadModel(t, e, "org/slow") + loaded <- loadOutcome{result, err} + }() + waitForStartInFlight(t, e) + + began := time.Now() + _, err := ejectEngine(t, e, "") + require.NoError(t, err) + + assert.Less(t, time.Since(began), 4*time.Second, "Eject waited out the Load") + assertLoadStopped(t, <-loaded) + assert.Equal(t, before, launchTextOf(t, e)) + assert.False(t, mustStatus(t, e).Running, "a stopped Load must not start the previous model again") + enabled, known, err := e.desired.get("fake") + require.NoError(t, err) + assert.True(t, known) + assert.False(t, enabled, "the user turned the engine off") +} + +// The same while a Load is restarting a running engine: the restore that would +// normally bring the previous model back up must leave the engine stopped. +func TestEject_DuringALoadOfARunningEngineLeavesItStopped(t *testing.T) { + e := switchModelExecutor(t, "--model") + slowStarts(t, e) + require.NoError(t, e.Start(t.Context(), "fake")) + before := launchTextOf(t, e) + loaded := make(chan loadOutcome, 1) + go func() { + result, err := loadModel(t, e, "org/slow") + loaded <- loadOutcome{result, err} + }() + waitForStartInFlight(t, e) + + began := time.Now() + _, err := ejectEngine(t, e, "") + require.NoError(t, err) + + // Restarting the previous model first, only to stop it again, would hold + // the Eject for another 8-second start. + assert.Less(t, time.Since(began), 4*time.Second, "Eject waited for the previous model to start again") + assertLoadStopped(t, <-loaded) + assert.Equal(t, before, launchTextOf(t, e)) + assert.False(t, mustStatus(t, e).Running, "a stopped Load must not start the previous model again") +} + +// A start reports that it is starting for as long as it runs. A vLLM start +// can download and load for minutes, and a client that only knew "not +// running" showed the engine as off for all of that time. +func TestStart_ReportsStartingUntilTheEngineIsReady(t *testing.T) { + e := switchModelExecutor(t, "--model") + slowStarts(t, e) + assert.False(t, mustStatus(t, e).Starting, "a stopped engine is not starting") + + started := make(chan error, 1) + go func() { started <- e.Start(t.Context(), "fake") }() + waitForStartInFlight(t, e) + require.Eventually(t, func() bool { return mustStatus(t, e).Starting }, 5*time.Second, 20*time.Millisecond, + "the engine is reported starting while it loads") + assert.False(t, mustStatus(t, e).Running) + + require.NoError(t, <-started) + status := mustStatus(t, e) + assert.True(t, status.Running) + assert.False(t, status.Starting, "a running engine is no longer starting") +} + +// A start that is stopped, or fails, is no longer reported as starting. +func TestStart_StopsReportingStartingWhenTheStartEnds(t *testing.T) { + e := switchModelExecutor(t, "--model") + slowStarts(t, e) + started := make(chan error, 1) + go func() { started <- e.Start(t.Context(), "fake") }() + waitForStartInFlight(t, e) + require.Eventually(t, func() bool { return mustStatus(t, e).Starting }, 5*time.Second, 20*time.Millisecond) + + require.NoError(t, e.Stop("fake")) + require.Error(t, <-started) + status := mustStatus(t, e) + assert.False(t, status.Starting) + assert.False(t, status.Running) +} diff --git a/services/nvpair-engine-manager/lifecycle.go b/services/nvpair-engine-manager/lifecycle.go index 96ef0a959..5f5429dba 100644 --- a/services/nvpair-engine-manager/lifecycle.go +++ b/services/nvpair-engine-manager/lifecycle.go @@ -98,11 +98,26 @@ func (e *Executor) StartWith(ctx context.Context, engine string, opts startOpts) st.opMu.Lock() defer st.opMu.Unlock() if err := e.doStart(ctx, st, engine, opts); err != nil { - return err + return e.startStopped(st, err) } return e.setDesiredEnabled(engine, true) } +// errStartStopped marks a start that a Stop or Eject cut short. The user asked +// for that, so callers answer with the engine's status, stopped, rather than an +// error carrying the cancelled engine's partial output. +var errStartStopped = errors.New("the start was stopped") + +// startStopped wraps a failed start's error in errStartStopped when a Stop or +// Eject cut it short. Callers must still hold st.opMu: the Stop that cancelled +// the start waits for it, so it is still counted as waiting. +func (e *Executor) startStopped(st *engineState, err error) error { + if e.stopWaiting(st) && !e.shuttingDown.Load() { + return fmt.Errorf("%w: %w", errStartStopped, err) + } + return err +} + // doStart performs the start. Callers must hold st.opMu. func (e *Executor) doStart(ctx context.Context, st *engineState, engine string, opts startOpts) error { ctx, cancel := context.WithCancel(ctx) @@ -115,7 +130,7 @@ func (e *Executor) doStart(ctx context.Context, st *engineState, engine string, st.startCancel = nil st.mu.Unlock() }() - if e.shuttingDown.Load() { + if e.shuttingDown.Load() || e.stopWaiting(st) { return context.Canceled } @@ -128,6 +143,15 @@ func (e *Executor) doStart(ctx context.Context, st *engineState, engine string, } st.mu.Unlock() + // A start refused before bring-up is reported like one that fails during + // it. A desktop Start is sent without waiting for the reply and learns the + // outcome from engine:state-changed or a start-failed error, and a refusal + // changes no state, so without the report the engine stays shown as + // starting. + refuse := func(err error) error { + e.reportStartFailedUnlessShuttingDown(ctx, engine, err) + return err + } pathInstalled, _ := e.Detect(engine) rt := st.plat.Runtime port := st.port @@ -137,17 +161,17 @@ func (e *Executor) doStart(ctx context.Context, st *engineState, engine string, if port == 0 { p, err := freePort() if err != nil { - return fmt.Errorf("allocate port: %w", err) + return refuse(fmt.Errorf("allocate port: %w", err)) } port = p } if err := e.reservedPortError(port); err != nil { - return err + return refuse(err) } presence := e.reconcilePresence(ctx, engine, st, pathInstalled, port, true) if presence.Identified { if opts.RequireOwned { - return fmt.Errorf("cannot apply launch settings: another process already owns the server port") + return refuse(fmt.Errorf("cannot apply launch settings: another process already owns the server port")) } slog.Info("adopting already-running external engine", "engine", engine, "port", port) st.mu.Lock() @@ -170,11 +194,22 @@ func (e *Executor) doStart(ctx context.Context, st *engineState, engine string, return nil } if presence.Occupied && rt.modeOrDefault() == "process" { - return fmt.Errorf("cannot start engine %q on port %d: the port is occupied by a service that did not identify as %s", engine, port, st.manifest.DisplayName) + return refuse(fmt.Errorf("cannot start engine %q on port %d: the port is occupied by a service that did not identify as %s", engine, port, st.manifest.DisplayName)) } + // Not reported: a client offers Install, not Start, for an engine that is + // not installed, and an engine removed outside PAIR would otherwise raise an + // error on every launch that restores it. if !pathInstalled { return fmt.Errorf("engine %q is not installed", engine) } + // An engine whose model is a launch argument cannot start without one. vLLM + // would serve, and first download, a default model of its own instead. + if flag := modelLaunchFlag(st.manifest); flag != "" { + tokens, err := parseLaunchText(e.launchStateLocked(engine, st).LaunchText) + if model, _ := launchFlagValue(tokens, flag); err == nil && model == "" { + return refuse(fmt.Errorf("%s has no model selected; load a model to start it", st.manifest.DisplayName)) + } + } vars := st.pathVars() vars["host"] = effectiveBind(rt.Bind, opts.Bind) vars["port"] = strconv.Itoa(port) @@ -182,12 +217,22 @@ func (e *Executor) doStart(ctx context.Context, st *engineState, engine string, vars["cli"] = expandPath(rt.CLI) } + // Report the start while it runs, which for vLLM can be minutes of + // downloading and loading, so every client shows the engine as starting + // rather than off, whoever started it. st.mu.Lock() st.port = port st.stopping = false + st.starting = true st.mu.Unlock() + e.emitState(engine) - if err := e.bringUp(ctx, st, engine, rt, port, vars); err != nil { + err := e.bringUp(ctx, st, engine, rt, port, vars) + st.mu.Lock() + st.starting = false + st.mu.Unlock() + if err != nil { + e.emitState(engine) return err } @@ -230,15 +275,36 @@ func (e *Executor) bringUp(ctx context.Context, st *engineState, engine string, // bringUpProcess spawns the engine as a foreground process this service // owns, then waits for readiness. func (e *Executor) bringUpProcess(ctx context.Context, st *engineState, engine string, rt Runtime, port int, vars map[string]string) error { + // The readiness allowance counts from here, so the work before launch + // (vLLM's metadata fetch and tool-parser probe) is inside it, and a start + // stays within what a remote peer waits for. + began := time.Now() st.mu.Lock() binPath := st.binPath st.mu.Unlock() launch, err := resolveProcessLaunch(rt, binPath, vars) + if err == nil { + err = validateEffectiveLaunch(rt, launch, vars["host"], vars["port"]) + } if err != nil { + e.reportStartFailedUnlessShuttingDown(ctx, engine, err) return err } - if err := validateEffectiveLaunch(rt, launch, vars["host"], vars["port"]); err != nil { - return err + // The hub cache vLLM itself will use, from the environment it is started + // with, and the model it serves. Auto-launch adds sizing and parser flags + // but never changes --model or --revision, so they are read once here. + env := hfEnv(launch.Env) + hub := hfLaunchHub(launch.Args, env) + model, _ := launchFlagValue(launch.Args, "--model") + revision := hfRevision(launch.Args) + fetch := sync.OnceValues(func() (hfPendingModel, bool) { + return e.fetchHFMetadata(ctx, st, hub, model, revision, env) + }) + launch = applyAutoLaunch(ctx, engine, rt, launch, hub, fetch) + if err := checkLaunchWeights(rt.UnsupportedWeights, launch.Args, hub, fetch); err != nil { + werr := fmt.Errorf("start %s: %w", engine, err) + e.reportStartFailedUnlessShuttingDown(ctx, engine, werr) + return werr } diagnostics := newStartupOutput(launchDiagnosticArgs(rt), launch.Env) @@ -269,7 +335,7 @@ func (e *Executor) bringUpProcess(ctx context.Context, st *engineState, engine s case <-readyCtx.Done(): } }() - readyErr := e.waitReady(readyCtx, rt.Ready, port) + readyErr := e.waitReady(readyCtx, rt.Ready, port, began) select { case <-proc.done: readyErr = fmt.Errorf("process exited before readiness; check the launch options") @@ -280,7 +346,7 @@ func (e *Executor) bringUpProcess(ctx context.Context, st *engineState, engine s st.mu.Lock() st.stopping = true st.mu.Unlock() - proc.stop(stopGrace(rt)) + proc.stop(e.stopGraceFor(rt)) st.mu.Lock() st.proc = nil st.mu.Unlock() @@ -325,6 +391,7 @@ func (e *Executor) bringUpCommand(ctx context.Context, st *engineState, engine s } if rt.EditableLaunch != nil && index == rt.EditableLaunch.StartIndex { if err := validateEffectiveLaunch(rt, launch, vars["host"], vars["port"]); err != nil { + e.reportStartFailedUnlessShuttingDown(ctx, engine, err) return cleanup(err) } } @@ -348,7 +415,7 @@ func (e *Executor) bringUpCommand(ctx context.Context, st *engineState, engine s // official stop command before Start releases the lifecycle lock. started = true } - if err := e.waitReady(ctx, rt.Ready, port); err != nil { + if err := e.waitReady(ctx, rt.Ready, port, time.Now()); err != nil { werr := fmt.Errorf("engine %q did not become ready: %w", engine, err) e.reportStartFailedUnlessShuttingDown(ctx, engine, werr) e.emitState(engine) @@ -414,6 +481,8 @@ func (e *Executor) Stop(engine string) error { if err != nil { return err } + done, _ := e.requestStop(st) + defer done() st.opMu.Lock() defer st.opMu.Unlock() stopErr := e.doStop(st, engine) @@ -455,7 +524,7 @@ func (e *Executor) doStop(st *engineState, engine string) error { // we manage :11435) is never touched; we decline only a process // whose image we can't confirm is ours. pid, image, ok := pidOnPort(port) - if ok && isManagedInstallPath(binPath, st.installDir) && isOurEngineImage(image, binPath) { + if ok && isManagedInstallPath(binPath, st.installDir) && (isOurEngineImage(image, binPath) || procRunsScript(pid, binPath)) { st.mu.Lock() st.stopping = true if st.healthStop != nil { @@ -463,7 +532,7 @@ func (e *Executor) doStop(st *engineState, engine string) error { st.healthStop = nil } st.mu.Unlock() - grace := stopGrace(rt) + grace := e.stopGraceFor(rt) terminatePID(pid, grace) if !pidAlive(pid) { e.markStopped(st, engine) @@ -499,7 +568,7 @@ func (e *Executor) doStop(st *engineState, engine string) error { return e.reconcileFailedCommandStop(st, engine, rt.Ready == nil || !e.waitUnavailable(rt.Ready, port, time.Second), err) } } else if proc != nil { - proc.stop(stopGrace(rt)) + proc.stop(e.stopGraceFor(rt)) } e.markStopped(st, engine) @@ -550,6 +619,26 @@ func (e *Executor) runCommandStop(st *engineState, engine string, rt Runtime, po // stopGrace is how long to wait after a graceful stop before escalating to a // forced kill, taken from the manifest's stop spec: default 5s, the spec's // grace_s when set, or 0 when the spec asks for an immediate kill. +// shutdownStopGraceCap is the longest an engine is given to exit on a polite +// signal once the whole service is shutting down. The parent bounds the engine +// stop to a few seconds (the broker allows engine-manager 5s for it) and then +// kills engine-manager, which cannot escalate to a forced kill after that: the +// engine, running in its own process group, would be left behind. So at +// shutdown the force kill has to come inside that window, however generous the +// manifest's grace is for an ordinary stop. An engine that is still loading +// (vLLM with its weights half read) does not answer a polite signal at all. +const shutdownStopGraceCap = 3 * time.Second + +// stopGraceFor is the grace for stopping an engine now: the manifest's, capped +// during shutdown. +func (e *Executor) stopGraceFor(rt Runtime) time.Duration { + grace := stopGrace(rt) + if e.shuttingDown.Load() && grace > shutdownStopGraceCap { + return shutdownStopGraceCap + } + return grace +} + func stopGrace(rt Runtime) time.Duration { grace := 5 * time.Second if sp := rt.Stop; sp != nil { @@ -744,7 +833,7 @@ func (e *Executor) Restart(ctx context.Context, engine string) error { return err } if err := e.doStart(ctx, st, engine, startOpts{}); err != nil { - return err + return e.startStopped(st, err) } return e.setDesiredEnabled(engine, true) } @@ -890,7 +979,9 @@ func (e *Executor) runHealth(ctx context.Context, st *engineState, engine string } } -func (e *Executor) waitReady(ctx context.Context, p *Probe, port int) error { +// waitReady polls the readiness probe until it passes or the manifest's +// allowance, counted from began, runs out. +func (e *Executor) waitReady(ctx context.Context, p *Probe, port int, began time.Time) error { if p == nil { return nil } @@ -898,7 +989,7 @@ func (e *Executor) waitReady(ctx context.Context, p *Probe, port int) error { if p.TimeoutS > 0 { timeout = time.Duration(p.TimeoutS) * time.Second } - deadline := time.Now().Add(timeout) + deadline := began.Add(timeout) for { if e.probe(ctx, p, port) { return nil @@ -991,12 +1082,45 @@ func (e *Executor) reportStartFailedUnlessShuttingDown(ctx context.Context, engi } // CommandContext can surface an exit status instead of context.Canceled on // Windows, so the canceled operation context is the authoritative signal. - if e.shuttingDown.Load() && ctx.Err() != nil { + // A start the user stopped did not fail either. + if ctx.Err() != nil && (e.shuttingDown.Load() || e.stopWaitingFor(engine)) { return } e.reportStartFailed(engine, err) } +// requestStop cancels a start in flight, as StopAll does, so a Stop or Eject +// does not wait out a readiness window of minutes for an engine it is about to +// stop, and keeps any other start from beginning until the caller is done. The +// returned function ends the request; call it after the stop, once opMu is +// released or about to be. cancelled reports whether a start was in flight. +func (e *Executor) requestStop(st *engineState) (done func(), cancelled bool) { + st.mu.Lock() + st.stopsWaiting++ + cancel := st.startCancel + st.mu.Unlock() + if cancel != nil { + cancel() + } + return func() { + st.mu.Lock() + st.stopsWaiting-- + st.mu.Unlock() + }, cancel != nil +} + +// stopWaiting reports whether a Stop or Eject is waiting for the engine. +func (e *Executor) stopWaiting(st *engineState) bool { + st.mu.Lock() + defer st.mu.Unlock() + return st.stopsWaiting > 0 +} + +func (e *Executor) stopWaitingFor(engine string) bool { + st, err := e.state(engine) + return err == nil && e.stopWaiting(st) +} + func freePort() (int, error) { l, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { diff --git a/services/nvpair-engine-manager/loadedwatch.go b/services/nvpair-engine-manager/loadedwatch.go index 16e8bfaec..982279eb6 100644 --- a/services/nvpair-engine-manager/loadedwatch.go +++ b/services/nvpair-engine-manager/loadedwatch.go @@ -5,12 +5,14 @@ package main import ( "context" + "maps" "sort" "time" ) -// modelsChangedMethod is the push emitted whenever an engine's set of loaded -// (in-memory) models changes. +// modelsChangedMethod is the push emitted whenever one of an engine's model +// lists changes: what it serves, what is loaded in memory, or what it has +// downloaded. const modelsChangedMethod = "engine:models-changed" // defaultLoadedPollSeconds is the default cadence of the loaded-model watcher, @@ -19,7 +21,7 @@ const modelsChangedMethod = "engine:models-changed" const defaultLoadedPollSeconds = 5 // modelsChangedParams is the engine:models-changed payload: which engine's -// loaded set changed (the trigger) plus the full current model result — the same +// lists changed (the trigger) plus the full current model result — the same // shape engine:models returns — so a push-driven consumer swaps its whole // snapshot instead of applying a delta. type modelsChangedParams struct { @@ -28,23 +30,23 @@ type modelsChangedParams struct { } // residencyNeutralActions are engine actions that only read state and so can -// never change what's resident in memory. pokeLoaded is skipped after them (see +// never change a model list. pokeLoaded is skipped after them (see // Manager.runAction) to avoid a wasted out-of-cycle sweep. Anything not listed // here — load_model/unload_model/run_model/pull_model/delete_model, and chat, -// which can JIT-load a model — is treated as potentially residency-affecting. +// which can JIT-load a model — is treated as potentially changing one. var residencyNeutralActions = map[string]bool{ "list_models": true, "list_downloaded": true, "loaded_models": true, } -// watchLoaded polls the loaded-model set of every running engine and emits -// engine:models-changed whenever it changes — covering explicit load/unload, LM -// Studio JIT auto-load, and TTL/idle auto-eviction. The engines themselves push -// nothing, so this internal poll is the only signal; keeping it here means a -// push-driven consumer never has to poll. It runs until ctx is cancelled and is -// pokeable (see Executor.pokeLoaded) for sub-tick latency after an explicit -// action. The first sweep only seeds the baseline (no emit) so a consumer's +// watchLoaded polls the model lists of every engine and emits +// engine:models-changed whenever one changes — covering explicit load/unload, LM +// Studio JIT auto-load, TTL/idle auto-eviction, and a model pulled or deleted. +// The engines themselves push nothing, so this internal poll is the only signal; +// keeping it here means a push-driven consumer never has to poll. It runs until +// ctx is cancelled and is pokeable (see Executor.pokeLoaded) for sub-tick +// latency after an explicit action. The first sweep only seeds the baseline (no emit) so a consumer's // initial engine:models fetch isn't immediately shadowed by a redundant push. func (e *Executor) watchLoaded(ctx context.Context) { interval := e.loadedPollInterval @@ -54,7 +56,7 @@ func (e *Executor) watchLoaded(ctx context.Context) { t := time.NewTicker(interval) defer t.Stop() - var prev map[string][]string + var prev modelLists seeded := false for { select { @@ -63,7 +65,7 @@ func (e *Executor) watchLoaded(ctx context.Context) { case <-t.C: case <-e.loadedPoke: } - changed, next, res := e.sweepLoaded(ctx, prev) + changed, next, res := e.sweepModels(ctx, prev) prev = next if !seeded { seeded = true @@ -75,30 +77,69 @@ func (e *Executor) watchLoaded(ctx context.Context) { } } -// sweepLoaded performs one loaded-set poll: it fetches the current model result -// and computes which engines' loaded sets changed relative to prevLoaded. It -// returns those engine names, the refreshed baseline to carry into the next -// sweep, and the full result (the push payload). The baseline is prevLoaded with -// every engine queried this sweep overwritten by its fresh value; an engine -// absent from this sweep keeps its last-good entry (see changedEngines) so a -// transient loaded_models miss isn't mistaken for an unload. It never mutates -// prevLoaded. Split out of watchLoaded so the diff logic is unit-testable -// without goroutines or timers. -func (e *Executor) sweepLoaded(ctx context.Context, prevLoaded map[string][]string) (changed []string, next map[string][]string, res ModelsResult) { +// modelLists is the watcher's baseline: the last-good value of each per-engine +// list in a ModelsResult, and the engines whose downloaded list the last sweep +// could not read. +type modelLists struct { + advertised, loaded, downloaded map[string][]string + downloadedMissed map[string]bool +} + +// sweepModels performs one poll: it fetches the current model result and +// computes which engines changed relative to prev in any of the lists a consumer +// shows — what is served (ByEngine), loaded (LoadedByEngine) or downloaded +// (DownloadedByEngine). The loaded set alone misses a pull or delete, and for an +// engine that advertises only what it serves (vLLM) the downloaded list changes +// on its own. It returns those engine names, the refreshed baseline to carry +// into the next sweep, and the full result (the push payload). In each list an +// engine absent from this sweep keeps its last-good entry (see changedEngines), +// so a transient miss isn't mistaken for a change. It never mutates prev. Split +// out of watchLoaded so the diff logic is unit-testable without goroutines or +// timers. +func (e *Executor) sweepModels(ctx context.Context, prev modelLists) (changed []string, next modelLists, res ModelsResult) { res = e.ModelsResult(ctx) - changed = changedEngines(prevLoaded, res.LoadedByEngine) - next = make(map[string][]string, len(prevLoaded)+len(res.LoadedByEngine)) - for name, ld := range prevLoaded { - next[name] = ld + seen := map[string]bool{} + sweep := func(prev, cur map[string][]string) map[string][]string { + for _, name := range changedEngines(prev, cur) { + if !seen[name] { + seen[name] = true + changed = append(changed, name) + } + } + next := make(map[string][]string, len(prev)+len(cur)) + maps.Copy(next, prev) + maps.Copy(next, cur) + return next } - for name, ld := range res.LoadedByEngine { - next[name] = ld + next = modelLists{ + advertised: sweep(prev.advertised, res.ByEngine), + loaded: sweep(prev.loaded, res.LoadedByEngine), + downloaded: sweep(prev.downloaded, res.DownloadedByEngine), + } + // A push made while an engine's downloaded list could not be read carries + // no list for it, so a consumer drops its rows. When the list reads again, + // push once even if it matches the last-good one, or those rows stay + // missing until something else changes. + for name := range prev.downloaded { + if _, ok := res.DownloadedByEngine[name]; !ok { + if next.downloadedMissed == nil { + next.downloadedMissed = map[string]bool{} + } + next.downloadedMissed[name] = true + } } + for name := range prev.downloadedMissed { + if _, ok := res.DownloadedByEngine[name]; ok && !seen[name] { + seen[name] = true + changed = append(changed, name) + } + } + sort.Strings(changed) return changed, next, res } -// pokeLoaded requests an immediate loaded-set check (non-blocking, coalesced). -// Called after an action that may change residency (load/unload/pull/…) so the +// pokeLoaded requests an immediate model-list check (non-blocking, coalesced). +// Called after an action that may change a list (load/unload/pull/delete/…) so the // engine:models-changed push doesn't wait up to a full poll interval. func (e *Executor) pokeLoaded() { select { @@ -107,7 +148,7 @@ func (e *Executor) pokeLoaded() { } } -// changedEngines returns, sorted, the engine names whose loaded set was reported +// changedEngines returns, sorted, the engine names whose list was reported // this sweep (present in cur) and differs from the baseline prev — either a new // key or a key whose set changed order-insensitively. // @@ -117,7 +158,7 @@ func (e *Executor) pokeLoaded() { // consumers via engine:state-changed. Treating every disappearance as a change // would emit spurious engine:models-changed churn — a drop then a re-add — on a // single transient failure, so the watcher keeps such keys as last-good (see -// sweepLoaded) and waits for a definitive reading. +// sweepModels) and waits for a definitive reading. func changedEngines(prev, cur map[string][]string) []string { var changed []string for name, cv := range cur { diff --git a/services/nvpair-engine-manager/main_test.go b/services/nvpair-engine-manager/main_test.go index 0368a5895..f6a00fc0b 100644 --- a/services/nvpair-engine-manager/main_test.go +++ b/services/nvpair-engine-manager/main_test.go @@ -34,6 +34,15 @@ func TestMain(m *testing.M) { panic("build fake-engine: " + err.Error() + "\n" + string(out)) } + // The Hugging Face tool an hf_cache delete runs, beside the fake engine as + // `hf` sits beside `vllm` in a vLLM virtualenv. It is a shell script, so the + // tests that need it skip on Windows, where vLLM has no install either. + if runtime.GOOS != "windows" { + if err := os.WriteFile(filepath.Join(tmp, "hf"), []byte(fakeHFCacheRm), 0o755); err != nil { + panic("write fake hf: " + err.Error()) + } + } + managerBin = filepath.Join(tmp, "nvpair-engine-manager"+suffix) if out, err := exec.Command("go", "build", "-o", managerBin, ".").CombinedOutput(); err != nil { panic("build engine-manager: " + err.Error() + "\n" + string(out)) diff --git a/services/nvpair-engine-manager/manager.go b/services/nvpair-engine-manager/manager.go index af41f9ddb..b6919906b 100644 --- a/services/nvpair-engine-manager/manager.go +++ b/services/nvpair-engine-manager/manager.go @@ -119,13 +119,22 @@ func (m *Manager) Run(ctx context.Context) error { slog.Warn("failed to subscribe to discovery relay for ec peers", "err", err) } - // Watch each running engine's loaded-model set and emit engine:models-changed - // on change (load/unload, JIT auto-load, TTL eviction). Bound to ctx: stops + // Watch each engine's model lists and emit engine:models-changed on change + // (load/unload, JIT auto-load, TTL eviction, pull, delete). Bound to ctx: stops // when the read loop exits below. go m.exec.watchLoaded(ctx) go m.watchPeerSettings(ctx) - err := m.readLoop(ctx) + // The read blocks on stdin until the parent sends something or closes it, so + // a SIGTERM that cancels ctx would otherwise never reach the shutdown below: + // the process would sit there, with its engines running, until stdin closed. + readDone := make(chan error, 1) + go func() { readDone <- m.readLoop(ctx) }() + var err error + select { + case err = <-readDone: + case <-ctx.Done(): + } // Cancel first so any in-flight start aborts at its readiness wait, // then stop engines so the parent's exit doesn't orphan them. cancel() @@ -334,6 +343,11 @@ func (m *Manager) runOp(ctx context.Context, msg *Message) { case "engine:restart": err = m.exec.Restart(ctx, p.Engine) } + // A start, restart or install-and-start that the user's Stop cut short + // answers with the status, stopped. + if errors.Is(err, errStartStopped) { + err = nil + } if err != nil { m.codec.RespondError(msg.ID, -32000, err.Error()) return @@ -385,6 +399,9 @@ func (m *Manager) runCatalog(ctx context.Context, msg *Message) { m.codec.RespondError(msg.ID, -32603, err.Error()) return } + // A target whose engine cannot load some weights is not offered them, as + // filterForTarget does for MLX. The rule is the target's, not this host's. + res.Models = withoutUnsupportedWeights(res.Models, m.exec.unsupportedWeightsFor(normalizeCatalogEngine(p.Engine), res.Platform, res.Arch)) m.codec.Respond(msg.ID, res) } @@ -424,11 +441,11 @@ func (m *Manager) runAction(ctx context.Context, msg *Message) { m.codec.RespondError(msg.ID, -32000, err.Error()) return } - // An action may have changed what's resident (load/unload/pull/delete/run, or - // a chat that JIT-loads); nudge the loaded watcher so engine:models-changed - // fires promptly instead of waiting for the next poll tick. Pure reads - // (list_models/loaded_models/list_downloaded) can't move residency, so skip - // the nudge and its wasted sweep for them. + // An action may have changed a model list (load/unload/pull/delete/run, or + // a chat that JIT-loads); nudge the watcher so engine:models-changed fires + // promptly instead of waiting for the next poll tick. Pure reads + // (list_models/loaded_models/list_downloaded) can't change one, so skip the + // nudge and its wasted sweep for them. if !residencyNeutralActions[p.Action] { m.exec.pokeLoaded() } diff --git a/services/nvpair-engine-manager/manager_test.go b/services/nvpair-engine-manager/manager_test.go index 75685f8ac..eb471b321 100644 --- a/services/nvpair-engine-manager/manager_test.go +++ b/services/nvpair-engine-manager/manager_test.go @@ -8,8 +8,8 @@ import ( "bytes" "context" "encoding/json" + "io" "net" - "testing" "time" @@ -130,3 +130,36 @@ func mustContain(t *testing.T, s, sub string) { t.Helper() require.Contains(t, s, sub) } + +// A SIGTERM cancels the manager's context while the read loop is blocked on +// stdin, which stays open and silent. Shutdown must not wait for the next +// message: it has to stop the engines and return, or the process hangs with its +// engines running until the parent closes the pipe. +func TestRunShutsDownWhenCancelledWhileStdinIsIdle(t *testing.T) { + ex := newTestExecutor(t, testEngineManifest(fakeEngineBin)) + if err := ex.Start(t.Context(), "fake"); err != nil { + t.Fatalf("start: %v", err) + } + + stdinRead, stdinWrite := io.Pipe() // never written to, never closed + t.Cleanup(func() { _ = stdinWrite.Close() }) + m := NewManager(NewCodec(struct { + io.Reader + io.Writer + }{stdinRead, io.Discard}), ex, nil) + + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { done <- m.Run(ctx) }() + time.Sleep(200 * time.Millisecond) // let it reach the blocking read + cancel() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("Run did not return after its context was cancelled") + } + if status, _ := ex.Status("fake"); status.Running { + t.Fatal("the engine was left running by the shutdown") + } +} diff --git a/services/nvpair-engine-manager/manifests/vllm.json b/services/nvpair-engine-manager/manifests/vllm.json new file mode 100644 index 000000000..b5d28afcf --- /dev/null +++ b/services/nvpair-engine-manager/manifests/vllm.json @@ -0,0 +1,107 @@ +{ + "$schema": "../manifest.schema.json", + "engine": "vllm", + "display_name": "vLLM", + "manifest_version": 1, + "advertised_models": "loaded", + "install": { "mode": "user" }, + "models_dir": "~/.cache/huggingface", + "uninstall": { "remove": ["{install_dir}"] }, + "runtime": { + "bin": "{install_dir}/venv/bin/vllm", + "editable_launch": { + "fixed_args": ["serve"], + "controls": [ + { "flags": ["--host"], "value": "{server.host}" }, + { "flags": ["--port"], "value": "{server.port}" } + ] + }, + "args": ["serve", "--model", "Qwen/Qwen3.5-4B", "--host", "{host}", "--port", "{port}"], + "auto_launch": { "kind": "hf_context", "context_tokens": 65536, "max_num_seqs": 32 }, + "env": { "HF_HOME": "{models_dir}", "VLLM_USE_FLASHINFER_SAMPLER": "0" }, + "bind": "127.0.0.1", + "port": 8001, + "ready": { "http": "http://127.0.0.1:{port}/health", "status": 200, "timeout_s": 480 }, + "stop": { "signal": "term", "grace_s": 15 }, + "health": { "http": "http://127.0.0.1:{port}/health", "status": 200, "interval_s": 5 } + }, + "platforms": { + "linux/amd64": { + "detect": ["{install_dir}/venv/bin/vllm", "~/.local/bin/vllm"], + "install": { + "fetch": { + "url": "https://github.com/astral-sh/uv/releases/download/0.13.1/uv-x86_64-unknown-linux-gnu.tar.gz", + "sha256": "eadd9d27c134ab5ac1ffd4473b97a741407f3ea238288f03121fc4deddf90220" + }, + "run": [ + "sh", "-c", + "set -eu; command -v nvidia-smi >/dev/null 2>&1 || { echo 'vLLM needs an NVIDIA GPU and driver (nvidia-smi not found)' >&2; exit 1; }; cuda_major=$(nvidia-smi -q 2>/dev/null | awk -F' *: *' '/^CUDA (UMD )?Version/ {split($2, v, \".\"); print v[1]; exit}'); [ \"${cuda_major:-0}\" -ge 13 ] 2>/dev/null || { echo \"vLLM needs an NVIDIA driver that supports CUDA 13 or newer (found: ${cuda_major:-none}). Update your NVIDIA driver and try again.\" >&2; exit 1; }; free_kb=$(df -Pk \"$2\" | awk 'NR==2 {print $4}'); [ \"${free_kb:-0}\" -ge 26214400 ] || { echo 'vLLM needs at least 25 GB of free disk space (the engine and its default model)' >&2; exit 1; }; mkdir -p \"$2/uv\"; tar -xzf \"$1\" -C \"$2/uv\" --strip-components=1; export UV_PYTHON_INSTALL_DIR=\"$2/python\" UV_PYTHON_PREFERENCE=only-managed; \"$2/uv/uv\" venv --no-cache --python 3.12 \"$2/venv\"; \"$2/uv/uv\" pip install --no-cache --python \"$2/venv/bin/python\" \"vllm==0.32.0\" \"mistral_common==1.12.0\"", + "vllm-install", "{download}", "{install_dir}" + ] + } + }, + "linux/arm64": { + "detect": ["{install_dir}/venv/bin/vllm", "~/.local/bin/vllm"], + "install": { + "fetch": { + "url": "https://github.com/astral-sh/uv/releases/download/0.13.1/uv-aarch64-unknown-linux-gnu.tar.gz", + "sha256": "9041c73c99fdc5892cf128140beb354eca9231d502068cc36e6cb53a8e59ad14" + }, + "run": [ + "sh", "-c", + "set -eu; command -v nvidia-smi >/dev/null 2>&1 || { echo 'vLLM needs an NVIDIA GPU and driver (nvidia-smi not found)' >&2; exit 1; }; cuda_major=$(nvidia-smi -q 2>/dev/null | awk -F' *: *' '/^CUDA (UMD )?Version/ {split($2, v, \".\"); print v[1]; exit}'); [ \"${cuda_major:-0}\" -ge 13 ] 2>/dev/null || { echo \"vLLM needs an NVIDIA driver that supports CUDA 13 or newer (found: ${cuda_major:-none}). Update your NVIDIA driver and try again.\" >&2; exit 1; }; free_kb=$(df -Pk \"$2\" | awk 'NR==2 {print $4}'); [ \"${free_kb:-0}\" -ge 26214400 ] || { echo 'vLLM needs at least 25 GB of free disk space (the engine and its default model)' >&2; exit 1; }; mkdir -p \"$2/uv\"; tar -xzf \"$1\" -C \"$2/uv\" --strip-components=1; export UV_PYTHON_INSTALL_DIR=\"$2/python\" UV_PYTHON_PREFERENCE=only-managed; \"$2/uv/uv\" venv --no-cache --python 3.12 \"$2/venv\"; \"$2/uv/uv\" pip install --no-cache --python \"$2/venv/bin/python\" \"vllm==0.32.0\" \"mistral_common==1.12.0\"", + "vllm-install", "{download}", "{install_dir}" + ] + } + }, + "darwin/arm64": { + "detect": ["{install_dir}/venv/bin/vllm", "~/.venv-vllm-metal/bin/vllm"], + "runtime": { + "unsupported_weights": { + "formats": ["fp8"], + "except_model_types": ["qwen3_5", "qwen3_5_moe"], + "reason": "vLLM on Apple Silicon (vLLM-Metal) cannot load FP8 weights or an FP8 KV cache. Choose the model's BF16 version instead." + } + }, + "install": { + "fetch": { + "url": "https://github.com/astral-sh/uv/releases/download/0.13.1/uv-aarch64-apple-darwin.tar.gz", + "sha256": "cfe454573412ff7380d15fc7dae5997ab53ffce72fce0a78e91ddbc84001acd2" + }, + "run": [ + "sh", "-c", + "set -eu; [ \"$(uname -m)\" = arm64 ] || { echo 'vLLM on macOS needs an Apple Silicon Mac (arm64)' >&2; exit 1; }; macos_major=$(sw_vers -productVersion | cut -d. -f1); [ \"${macos_major:-0}\" -ge 15 ] 2>/dev/null || { echo \"vLLM needs macOS 15 or newer (found: ${macos_major:-unknown}). Update macOS and try again.\" >&2; exit 1; }; free_kb=$(df -Pk \"$2\" | awk 'NR==2 {print $4}'); [ \"${free_kb:-0}\" -ge 20971520 ] || { echo 'vLLM needs at least 20 GB of free disk space (the engine and its default model)' >&2; exit 1; }; mkdir -p \"$2/uv\" \"$2/wheels\"; tar -xzf \"$1\" -C \"$2/uv\" --strip-components=1; export UV_PYTHON_INSTALL_DIR=\"$2/python\" UV_PYTHON_PREFERENCE=only-managed; \"$2/uv/uv\" venv --no-cache --python 3.12 \"$2/venv\"; core_whl=\"$2/wheels/vllm-0.32.0+cpu-cp38-abi3-macosx_11_0_arm64.whl\"; metal_whl=\"$2/wheels/vllm_metal-0.32.0-cp312-cp312-macosx_15_0_arm64.whl\"; curl -fsSL --retry 3 -o \"${core_whl:-}\" 'https://github.com/vllm-project/vllm/releases/download/v0.32.0/vllm-0.32.0%2Bcpu-cp38-abi3-macosx_11_0_arm64.whl'; echo \"c28cf0fede58f02c4886d2794217ec3c7e4bbd22dfeb16b4b32586a8f230bb75 ${core_whl:-}\" | shasum -a 256 -c -; curl -fsSL --retry 3 -o \"${metal_whl:-}\" 'https://github.com/vllm-project/vllm-metal/releases/download/v0.32.0/vllm_metal-0.32.0-cp312-cp312-macosx_15_0_arm64.whl'; echo \"458d34157eb02a6bfb0e293c70485df18f72b7eb1d1ce8f1eefea95ba47ab63c ${metal_whl:-}\" | shasum -a 256 -c -; \"$2/uv/uv\" pip install --no-cache --python \"$2/venv/bin/python\" \"${core_whl:-}\" \"mistral_common==1.12.0\"; \"$2/uv/uv\" pip install --no-cache --python \"$2/venv/bin/python\" \"${metal_whl:-}\"; rm -rf \"$2/wheels\"", + "vllm-install", "{download}", "{install_dir}" + ] + } + } + }, + "actions": { + "list_models": { + "description": "List the models in the Hugging Face cache that vLLM could serve (complete downloads with a config, tokenizer and weights). Needs no running engine.", + "hf_cache": { "dir": "{hf_hub_cache}" }, + "result": { "array": "data", "field": "id" } + }, + "load_model": { + "description": "Serve a different model (params: {\"model\": \"\"}). vLLM serves one model per process, so this saves the model into the launch arguments and restarts the engine on it; if the new model does not start, the previous one is restored.", + "launch_model": { "flag": "--model" } + }, + "unload_model": { + "description": "Eject the model vLLM is serving (params: {\"model\": \"\"}). vLLM cannot drop its model and keep running, so this stops the engine, which frees all of its GPU memory. Load starts it again.", + "stop_engine": true + }, + "delete_model": { + "description": "Delete a downloaded model from the Hugging Face cache (params: {\"model\": \"\"}). The model a running vLLM is serving is refused (eject it first). Deleting the model a stopped vLLM is set to serve leaves it with no model selected, and it will not start until a model is loaded.", + "hf_cache": { "dir": "{hf_hub_cache}", "delete": true, "tool": "hf" } + }, + "pull_model": { + "description": "Download the files vLLM loads from a Hugging Face model into the shared cache (params: {\"model\": \"\"}): weights in one format, config and tokenizer. Other formats in the repository, such as GGUF or original/ checkpoints, are not fetched. Needs no running engine. Once downloaded it appears in list_models and can be served with load_model.", + "hf_pull": { "tool": "hf", "dir": "{hf_hub_cache}" } + }, + "loaded_models": { + "description": "List the model this vLLM server is serving (it serves exactly one). Its served name (id) is what peers and the proxy are told it can answer for; its --model (root) marks the downloaded model as loaded.", + "http": { "method": "GET", "path": "/v1/models" }, + "result": { "array": "data", "field": "id", "inventory_field": "root" } + } + } +} diff --git a/services/nvpair-engine-manager/models.go b/services/nvpair-engine-manager/models.go index 31db17aa8..b9922989e 100644 --- a/services/nvpair-engine-manager/models.go +++ b/services/nvpair-engine-manager/models.go @@ -39,10 +39,24 @@ const modelsTimeout = 5 * time.Second // loaded_models succeeds, a loaded name can appear here without a matching // ByEngine/Models entry until the next successful sweep. Its own omitempty drops // the whole map when no engine reports loaded state. +// +// DownloadedByEngine is the model inventory of an engine that advertises only +// what it is serving (advertised_models "loaded", such as vLLM): ByEngine holds +// the names a request may be routed to, and this holds every model the engine +// can load, so a model manager keeps its rows for models that are on disk but +// not being served. It is read whenever the engine is installed, running or +// not, because such an engine's list_models needs no running server. Its +// LoadedByEngine entry is present and empty while it is stopped, and names +// models in this list's terms (loaded_models' +// result.inventory_field), so a model served under an alias still marks its +// cached row as loaded. An empty list means the inventory was read and is empty; +// a missing key means it was not read (not installed, or the read failed). An +// engine whose advertised list is its inventory has no entry here. type ModelsResult struct { - Models []string `json:"models"` - ByEngine map[string][]string `json:"modelsByEngine,omitempty"` - LoadedByEngine map[string][]string `json:"loadedByEngine,omitempty"` + Models []string `json:"models"` + ByEngine map[string][]string `json:"modelsByEngine,omitempty"` + LoadedByEngine map[string][]string `json:"loadedByEngine,omitempty"` + DownloadedByEngine map[string][]string `json:"downloadedByEngine,omitempty"` } // Models returns the union of model names served by every installed, running @@ -53,14 +67,24 @@ func (e *Executor) Models(ctx context.Context) []string { return e.ModelsResult(ctx).Models } +// engineModels is one engine's answer within a ModelsResult sweep. Each list +// has an ok flag because an empty list read successfully and a list that could +// not be read mean different things on the wire. +type engineModels struct { + advertised, loaded, downloaded []string + advertisedOK, loadedOK, downloadedOK bool +} + // ModelsResult queries each installed, running engine's loopback list_models // endpoint (reusing the manifest action machinery), normalizes the // engine-specific shapes, and returns both the flat de-duplicated union and the // per-engine attribution — from a single sweep, so a caller needing both // (the /v1/models HTTP surface, the engine:models RPC) never sweeps twice. // Stopped engines contribute nothing — their model list isn't queryable while -// the server is down. Best-effort: a per-engine failure is logged and skipped -// so the engines that did answer still surface their models. +// the server is down — except the DownloadedByEngine inventory of an engine +// that advertises its loaded set, which is read from disk. Best-effort: a +// per-engine failure is logged and skipped so the engines that did answer +// still surface their models. // // The per-engine queries run CONCURRENTLY. A host running several engines (e.g. // Ollama + LM Studio) must not have the later engines squeezed out of a single @@ -76,10 +100,7 @@ func (e *Executor) ModelsResult(ctx context.Context) ModelsResult { engineNames := e.reg.Names() // Results indexed by engine position so the flattened output stays // deterministic (registration order) despite concurrent completion. - perEngine := make([][]string, len(engineNames)) - listOK := make([]bool, len(engineNames)) - loaded := make([][]string, len(engineNames)) - loadedOK := make([]bool, len(engineNames)) + perEngine := make([]engineModels, len(engineNames)) var wg sync.WaitGroup for i, name := range engineNames { mf, ok := e.reg.Get(name) @@ -96,76 +117,124 @@ func (e *Executor) ModelsResult(ctx context.Context) ModelsResult { if listSpec == nil && loadedSpec == nil { continue } - wg.Add(1) - go func(i int, name string, listSpec, loadedSpec *ActionResult) { - defer wg.Done() + wg.Go(func() { + // An engine that can only answer for what it is running advertises its + // loaded set. Its list_models is a view of what could be loaded, which + // is not what a peer or the proxy may route to, so it is reported apart. + if mf.AdvertisedModels == advertisedModelsLoaded && loadedSpec != nil { + perEngine[i] = e.loadedOnlyEngineModels(ctx, name, listSpec, loadedSpec) + return + } st, err := e.Status(name) if err != nil || !st.Running { return } + var r engineModels if listSpec != nil { - if raw, err := e.Action(ctx, name, "list_models", nil); err != nil { - slog.Debug("engine:models list_models failed", "engine", name, "err", err) - } else if models, ok := extractStringsResult(raw, listSpec); ok { - perEngine[i] = models - listOK[i] = true - } else { - slog.Debug("engine:models list_models returned an invalid inventory", "engine", name) - } + r.advertised, r.advertisedOK = e.queryModelList(ctx, name, "list_models", listSpec) } if loadedSpec != nil { - if raw, err := e.Action(ctx, name, "loaded_models", nil); err != nil { - slog.Debug("engine:models loaded_models failed", "engine", name, "err", err) - } else if models, ok := extractStringsResult(raw, loadedSpec); ok { - // A successful query — even an empty result — records the key - // so a consumer can tell "running, nothing loaded" apart from - // "unknown". An action or shape error leaves loadedOK[i] false. - loaded[i] = models - loadedOK[i] = true - } else { - slog.Debug("engine:models loaded_models returned an invalid inventory", "engine", name) - } + // A successful query — even an empty result — records the key so a + // consumer can tell "running, nothing loaded" apart from "unknown". + r.loaded, r.loadedOK = e.queryModelList(ctx, name, "loaded_models", loadedSpec) } - }(i, name, listSpec, loadedSpec) + perEngine[i] = r + }) } wg.Wait() res := ModelsResult{Models: []string{}} seen := map[string]bool{} - for i, models := range perEngine { - if listOK[i] { - if res.ByEngine == nil { - res.ByEngine = make(map[string][]string) - } - // Normalize nil -> [] so the wire distinguishes a successful empty - // inventory from a missing/not-queryable engine. - if models == nil { - models = []string{} - } - res.ByEngine[engineNames[i]] = models - for _, m := range models { + for i, r := range perEngine { + name := engineNames[i] + if r.advertisedOK { + res.ByEngine = withModels(res.ByEngine, name, r.advertised) + for _, m := range r.advertised { if m != "" && !seen[m] { seen[m] = true res.Models = append(res.Models, m) } } } - if loadedOK[i] { - if res.LoadedByEngine == nil { - res.LoadedByEngine = make(map[string][]string) - } - // Normalize nil -> [] so the wire carries an empty array ("queried, - // nothing resident"), never null. - ld := loaded[i] - if ld == nil { - ld = []string{} - } - res.LoadedByEngine[engineNames[i]] = ld + if r.loadedOK { + res.LoadedByEngine = withModels(res.LoadedByEngine, name, r.loaded) + } + if r.downloadedOK { + res.DownloadedByEngine = withModels(res.DownloadedByEngine, name, r.downloaded) } } return res } +// loadedOnlyEngineModels answers for an engine that advertises its loaded set. +// Its inventory comes from list_models whenever it is installed, since that +// list is read from disk and needs no running server, and while it is stopped +// it reports nothing loaded. While it runs, one +// loaded_models response gives both the names it answers to (result.field, +// such as vLLM's served name) and, when result.inventory_field is set, the same +// models in the inventory's terms (vLLM's --model). +func (e *Executor) loadedOnlyEngineModels(ctx context.Context, name string, listSpec, loadedSpec *ActionResult) engineModels { + var r engineModels + st, err := e.Status(name) + if err != nil || !st.Installed { + return r + } + if listSpec != nil { + r.downloaded, r.downloadedOK = e.queryModelList(ctx, name, "list_models", listSpec) + } + if !st.Running { + // A stopped engine of this kind has nothing loaded, which is known, not + // unknown: its rows show as not loaded rather than unreported. + r.loadedOK = true + return r + } + raw, err := e.Action(ctx, name, "loaded_models", nil) + if err != nil { + slog.Debug("engine:models loaded_models failed", "engine", name, "err", err) + return r + } + r.advertised, r.advertisedOK = extractStringsResult(raw, loadedSpec) + r.loaded, r.loadedOK = r.advertised, r.advertisedOK + if loadedSpec.InventoryField != "" { + inventorySpec := *loadedSpec + inventorySpec.Field = loadedSpec.InventoryField + r.loaded, r.loadedOK = extractStringsResult(raw, &inventorySpec) + } + if !r.advertisedOK || !r.loadedOK { + slog.Debug("engine:models loaded_models returned an invalid inventory", "engine", name) + } + return r +} + +// queryModelList runs one model-list action and extracts its names. ok is +// false when the action failed or its response was not a valid list. +func (e *Executor) queryModelList(ctx context.Context, name, action string, spec *ActionResult) ([]string, bool) { + raw, err := e.Action(ctx, name, action, nil) + if err != nil { + slog.Debug("engine:models action failed", "engine", name, "action", action, "err", err) + return nil, false + } + models, ok := extractStringsResult(raw, spec) + if !ok { + slog.Debug("engine:models action returned an invalid inventory", "engine", name, "action", action) + } + return models, ok +} + +// withModels records one engine's list in a per-engine map, creating the map on +// first use. A nil list is stored as [] so the wire carries an empty array +// ("read, nothing there"), never null. +func withModels(m map[string][]string, engine string, models []string) map[string][]string { + if m == nil { + m = make(map[string][]string) + } + if models == nil { + models = []string{} + } + m[engine] = models + return m +} + // extractStrings is the value-only wrapper used by focused extractor tests. // Callers that must distinguish a valid empty inventory from an invalid // response use extractStringsResult. diff --git a/services/nvpair-engine-manager/models_test.go b/services/nvpair-engine-manager/models_test.go index 6a8c71474..7e8e789be 100644 --- a/services/nvpair-engine-manager/models_test.go +++ b/services/nvpair-engine-manager/models_test.go @@ -9,7 +9,9 @@ import ( "encoding/json" "fmt" "net/http" - + "os" + "path/filepath" + "runtime" "testing" "github.com/stretchr/testify/assert" @@ -251,44 +253,117 @@ func TestModelsResultNoLoadedAction_OmitsKey(t *testing.T) { require.Nil(t, ex.ModelsResult(ctx).LoadedByEngine) } -// TestSweepLoadedSeedsThenEmitsOnChange covers the watcher's core diff without +// TestSweepModels_SeedsThenEmitsOnChange covers the watcher's core diff without // goroutines or timers: the first sweep only seeds a baseline (the caller, // watchLoaded, suppresses its emit), a subsequent identical sweep reports no // change, and a sweep after an eviction reports the engine as changed and // carries the fresh full model result as the push payload. -func TestSweepLoaded_SeedsThenEmitsOnChange(t *testing.T) { +func TestSweepModels_SeedsThenEmitsOnChange(t *testing.T) { ex := newTestExecutor(t, loadedActionManifest()) - ctx := context.Background() + ctx := t.Context() t.Cleanup(func() { _ = ex.Stop("fake") }) require.NoError(t, ex.Start(ctx, "fake"), "start") - // Seed sweep: baseline is {fake:[llama3.2:1b]}. - _, prev, _ := ex.sweepLoaded(ctx, nil) - require.Equal(t, map[string][]string{"fake": {"llama3.2:1b"}}, prev, "seed baseline") + // Seed sweep: the loaded baseline is {fake:[llama3.2:1b]}. + _, prev, _ := ex.sweepModels(ctx, modelLists{}) + require.Equal(t, map[string][]string{"fake": {"llama3.2:1b"}}, prev.loaded, "seed baseline") - // No residency change -> no engine reported changed. - changed, prev, _ := ex.sweepLoaded(ctx, prev) + // No change -> no engine reported changed. + changed, prev, _ := ex.sweepModels(ctx, prev) require.Empty(t, changed, "unchanged sweep reported") // Evict everything -> fake changes; payload carries the empty loaded set. setLoaded(t, ex, nil) - changed, _, res := ex.sweepLoaded(ctx, prev) + changed, _, res := ex.sweepModels(ctx, prev) require.Equal(t, []string{"fake"}, changed, "changed") require.Equal(t, map[string][]string{"fake": {}}, res.LoadedByEngine, "pushed LoadedByEngine") } -// TestSweepLoadedRetainsLastGoodOnTransientMiss covers the anti-churn guard: an +// A pull or delete changes what an engine has, not what it has loaded. The push +// has to follow it too, or a client that waits for pushes (the terminal +// interface) keeps showing the old list until it asks again. +func TestSweepModels_EmitsWhenAModelIsPulledOrDeleted(t *testing.T) { + test := func(name string, manifest func(t *testing.T) *Manifest, start bool, list func(ModelsResult) []string) { + t.Run(name, func(t *testing.T) { + m := manifest(t) + cache := m.Actions["list_models"].HFCache.Dir + ex := newTestExecutor(t, m) + ctx := t.Context() + if start { + t.Cleanup(func() { _ = ex.Stop("fake") }) + require.NoError(t, ex.Start(ctx, "fake")) + } + _, prev, _ := ex.sweepModels(ctx, modelLists{}) + + writeHFRepo(t, cache, "org/pulled-model", llmFiles(map[string]string{"model.safetensors": "w"})) + changed, prev, res := ex.sweepModels(ctx, prev) + assert.Equal(t, []string{"fake"}, changed, "after the pull") + assert.Contains(t, list(res), "org/pulled-model", "the pushed list") + + require.NoError(t, os.RemoveAll(filepath.Join(cache, "models--org--pulled-model"))) + changed, _, res = ex.sweepModels(ctx, prev) + assert.Equal(t, []string{"fake"}, changed, "after the delete") + assert.NotContains(t, list(res), "org/pulled-model", "the pushed list") + }) + } + test("a stopped engine that advertises only what it serves, such as vLLM", + func(t *testing.T) *Manifest { return loadedOnlyManifest(t, fakeEngineBin) }, false, + func(r ModelsResult) []string { return r.DownloadedByEngine["fake"] }) + test("a running engine that advertises what it has", + func(t *testing.T) *Manifest { + m := loadedOnlyManifest(t, fakeEngineBin) + m.AdvertisedModels = "" + delete(m.Actions, "loaded_models") + require.NoError(t, m.Validate()) + return m + }, true, + func(r ModelsResult) []string { return r.ByEngine["fake"] }) +} + +// While vLLM's cache cannot be read, a push caused by another engine carries no +// vLLM list, and a consumer drops those rows. When the cache reads again the +// list matches the last-good one, but it still has to be pushed once. +func TestSweepModels_PushesADownloadedListThatReadsAgainUnchanged(t *testing.T) { + m := loadedOnlyManifest(t, fakeEngineBin) + cache := m.Actions["list_models"].HFCache.Dir + ex := newTestExecutor(t, m) + ctx := t.Context() + _, prev, _ := ex.sweepModels(ctx, modelLists{}) + require.Equal(t, []string{"other/cached-model"}, prev.downloaded["fake"], "seed baseline") + + require.NoError(t, os.Chmod(cache, 0)) + t.Cleanup(func() { _ = os.Chmod(cache, 0o755) }) + changed, prev, res := ex.sweepModels(ctx, prev) + if _, read := res.DownloadedByEngine["fake"]; read { + t.Skip("this system can still read a directory with no permissions") + } + assert.Empty(t, changed, "a cache that cannot be read is not a change") + + require.NoError(t, os.Chmod(cache, 0o755)) + changed, prev, res = ex.sweepModels(ctx, prev) + assert.Equal(t, []string{"fake"}, changed, "the list read again") + assert.Equal(t, []string{"other/cached-model"}, res.DownloadedByEngine["fake"], "the pushed list") + + changed, _, _ = ex.sweepModels(ctx, prev) + assert.Empty(t, changed, "pushed once, not on every sweep after") +} + +// TestSweepModels_RetainsLastGoodOnTransientMiss covers the anti-churn guard: an // engine that drops out of a sweep (a transient loaded_models miss, indistinct // from a stop) is NOT reported as changed and keeps its last-good baseline, so a // single blip can't emit a spurious drop-then-re-add pair. -func TestSweepLoaded_RetainsLastGoodOnTransientMiss(t *testing.T) { +func TestSweepModels_RetainsLastGoodOnTransientMiss(t *testing.T) { ex := newTestExecutor(t, loadedActionManifest()) - prev := map[string][]string{"fake": {"llama3.2:1b"}} + prev := modelLists{ + advertised: map[string][]string{"fake": {"llama3.2:1b"}}, + loaded: map[string][]string{"fake": {"llama3.2:1b"}}, + downloaded: map[string][]string{}, + } // The engine isn't started, so ModelsResult reports it neither running nor - // queryable: LoadedByEngine has no "fake" key this sweep. - changed, next, _ := ex.sweepLoaded(context.Background(), prev) + // queryable: no list has a "fake" key this sweep. + changed, next, _ := ex.sweepModels(t.Context(), prev) require.Empty(t, changed, "a disappeared engine reported") - require.Equal(t, map[string][]string{"fake": {"llama3.2:1b"}}, next, "baseline after miss") + require.Equal(t, prev, next, "baseline after miss") } func TestChangedEngines_IdentifiesInventoryDifferences(t *testing.T) { @@ -320,3 +395,147 @@ func TestSameStringSet_IgnoresOrdering(t *testing.T) { test("duplicates matter", []string{"x", "x"}, []string{"x", "y"}, false) test("same multiset with dups", []string{"x", "x", "y"}, []string{"y", "x", "x"}, true) } + +// loadedOnlyManifest is a vLLM-like engine: it advertises only what it is +// serving (loaded_models on the OpenAI-shape /v1/models), and lists what it has +// downloaded from a Hugging Face cache, which needs no running server. +func loadedOnlyManifest(t *testing.T, bin string) *Manifest { + t.Helper() + cache := t.TempDir() + writeHFRepo(t, cache, "other/cached-model", llmFiles(map[string]string{"model.safetensors": "w"})) + m := testEngineManifest(bin) + m.AdvertisedModels = advertisedModelsLoaded + m.Actions["list_models"] = Action{ + HFCache: &ActionHFCache{Dir: cache}, + Result: &ActionResult{Array: "data", Field: "id"}, + } + m.Actions["loaded_models"] = Action{ + HTTP: &ActionHTTP{Method: "GET", Path: "/v1/models"}, + Result: &ActionResult{Array: "data", Field: "id", InventoryField: "root"}, + } + require.NoError(t, m.Validate()) + return m +} + +// An engine that serves one model per process can list many models on disk but +// answers for only the one it is running. The proxy routes to any node that +// advertises a model, so advertising the on-disk list would send requests to an +// engine that cannot answer them. advertised_models: "loaded" advertises the +// served model and reports the on-disk list apart, for model managers. +func TestModelsResult_AdvertisesOnlyTheServedModelOfALoadedOnlyEngine(t *testing.T) { + ex := newTestExecutor(t, loadedOnlyManifest(t, fakeEngineBin)) + ctx := t.Context() + t.Cleanup(func() { _ = ex.Stop("fake") }) + require.NoError(t, ex.Start(ctx, "fake")) + + res := ex.ModelsResult(ctx) + served := []string{"llama3.2:1b"} // what the fake engine is running + assert.Equal(t, map[string][]string{"fake": served}, res.ByEngine, "advertised per engine") + assert.Equal(t, served, res.Models, "advertised union") + assert.Equal(t, map[string][]string{"fake": served}, res.LoadedByEngine, "loaded") + assert.Equal(t, map[string][]string{"fake": {"other/cached-model"}}, res.DownloadedByEngine, "downloaded") +} + +// Ejecting a vLLM model stops the engine. Its downloaded models must still be +// reported, so a model manager on this node or a peer keeps the rows it needs to +// Load or Delete them, while nothing is advertised for routing. +func TestModelsResult_ReportsTheDownloadedModelsOfAStoppedLoadedOnlyEngine(t *testing.T) { + ex := newTestExecutor(t, loadedOnlyManifest(t, fakeEngineBin)) + ctx := t.Context() + + res := ex.ModelsResult(ctx) + assert.Equal(t, map[string][]string{"fake": {"other/cached-model"}}, res.DownloadedByEngine, "downloaded") + assert.Empty(t, res.ByEngine, "advertised per engine") + assert.Empty(t, res.Models, "advertised union") + assert.Equal(t, map[string][]string{"fake": {}}, res.LoadedByEngine, "nothing loaded, and known to be") +} + +// A model the engine cannot read is unknown, not an empty cache: an engine that +// is not installed reports no downloaded key at all. +func TestModelsResult_OmitsTheDownloadedModelsOfAnEngineThatIsNotInstalled(t *testing.T) { + missing := filepath.Join(t.TempDir(), "not-installed") + ex := newTestExecutor(t, loadedOnlyManifest(t, missing)) + + res := ex.ModelsResult(t.Context()) + assert.Nil(t, res.DownloadedByEngine) +} + +// With vLLM's --served-model-name, requests use the alias, but the model list +// names the cached repository. The alias is advertised for routing and the +// loaded marker uses the repository, so the cached row shows as loaded. +func TestModelsResult_MarksTheCachedModelLoadedWhenServedUnderAnAlias(t *testing.T) { + m := loadedOnlyManifest(t, fakeEngineBin) + key := runtime.GOOS + "/" + runtime.GOARCH + plat := m.Platforms[key] + plat.Runtime.Env = map[string]string{"OLLAMA_HOST": "127.0.0.1:{port}", "FAKE_SERVED_MODEL_NAME": "qwen-api"} + m.Platforms[key] = plat + ex := newTestExecutor(t, m) + ctx := t.Context() + t.Cleanup(func() { _ = ex.Stop("fake") }) + require.NoError(t, ex.Start(ctx, "fake")) + + res := ex.ModelsResult(ctx) + assert.Equal(t, map[string][]string{"fake": {"qwen-api"}}, res.ByEngine, "advertised per engine") + assert.Equal(t, []string{"qwen-api"}, res.Models, "advertised union") + assert.Equal(t, map[string][]string{"fake": {"llama3.2:1b"}}, res.LoadedByEngine, "loaded") +} + +// An engine whose advertised list is its inventory reports nothing apart. +func TestModelsResult_OmitsDownloadedForAnEngineThatAdvertisesItsList(t *testing.T) { + ex := newTestExecutor(t, loadedActionManifest()) + ctx := t.Context() + t.Cleanup(func() { _ = ex.Stop("fake") }) + require.NoError(t, ex.Start(ctx, "fake")) + + assert.Nil(t, ex.ModelsResult(ctx).DownloadedByEngine) +} + +func TestInventoryField_IsOnlyForTheLoadedModelsOfALoadedOnlyEngine(t *testing.T) { + manifestWith := func(advertised, action, field string) *Manifest { + m := loadedActionManifest() + m.AdvertisedModels = advertised + act := m.Actions[action] + result := *act.Result + result.InventoryField = field + act.Result = &result + m.Actions[action] = act + return m + } + accept := func(name, advertised, action, field string) { + t.Run(name, func(t *testing.T) { + assert.NoError(t, manifestWith(advertised, action, field).Validate()) + }) + } + reject := func(name, advertised, action, field string) { + t.Run(name, func(t *testing.T) { + assert.Error(t, manifestWith(advertised, action, field).Validate()) + }) + } + accept("loaded_models of a loaded-only engine", advertisedModelsLoaded, "loaded_models", "root") + reject("engine that advertises its list", advertisedModelsList, "loaded_models", "root") + reject("action other than loaded_models", advertisedModelsLoaded, "list_models", "root") + reject("object path instead of a field name", advertisedModelsLoaded, "loaded_models", "a.b") +} + +func TestAdvertisedModels_AcceptsOnlyKnownValuesItCanServe(t *testing.T) { + accept := func(name, value string) { + t.Run(name, func(t *testing.T) { + m := loadedActionManifest() + m.AdvertisedModels = value + + assert.NoError(t, m.Validate()) + }) + } + reject := func(name string, m *Manifest, value string) { + t.Run(name, func(t *testing.T) { + m.AdvertisedModels = value + + assert.Error(t, m.Validate()) + }) + } + accept("unset", "") + accept("the list", advertisedModelsList) + accept("loaded, with a loaded_models action", advertisedModelsLoaded) + reject("loaded without a loaded_models action", testEngineManifest(fakeEngineBin), advertisedModelsLoaded) + reject("an unknown value", testEngineManifest(fakeEngineBin), "everything") +} diff --git a/services/nvpair-engine-manager/proc.go b/services/nvpair-engine-manager/proc.go index a1983ec0e..17206fdec 100644 --- a/services/nvpair-engine-manager/proc.go +++ b/services/nvpair-engine-manager/proc.go @@ -90,8 +90,13 @@ func (mp *managedProc) stop(grace time.Duration) { return // already exited default: } + pid := mp.cmd.Process.Pid + // However the process we started exits, an engine that forks helpers (vLLM's + // EngineCore) can leave them behind, holding GPU memory, when its main + // process goes down mid-start. Nothing else will ever clean them up. + defer killGroupRemnants(pid) if grace <= 0 { - _ = signalPID(mp.cmd.Process.Pid, true) + _ = signalPID(pid, true) <-mp.done return } @@ -103,7 +108,7 @@ func (mp *managedProc) stop(grace time.Duration) { return case <-timer.C: } - _ = signalPID(mp.cmd.Process.Pid, true) + _ = signalPID(pid, true) <-mp.done } diff --git a/services/nvpair-engine-manager/proc_group_unix_test.go b/services/nvpair-engine-manager/proc_group_unix_test.go new file mode 100644 index 000000000..7b2520391 --- /dev/null +++ b/services/nvpair-engine-manager/proc_group_unix_test.go @@ -0,0 +1,223 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build linux + +package main + +import ( + "os" + "runtime" + "strconv" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// processGone reports whether pid is gone. A killed process whose parent has +// not reaped it yet is a zombie, which is as dead as a process gets. +func processGone(pid int) bool { + data, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat") + if err != nil { + return true + } + // "pid (comm) S ..." — comm may contain spaces and parentheses, so cut after the last ")". + fields := strings.Fields(string(data[strings.LastIndex(string(data), ")")+1:])) + return len(fields) > 0 && fields[0] == "Z" +} + +func waitGone(pid int, within time.Duration) bool { + deadline := time.Now().Add(within) + for time.Now().Before(deadline) { + if processGone(pid) { + return true + } + time.Sleep(20 * time.Millisecond) + } + return processGone(pid) +} + +// vLLM's main process can exit on a polite stop while the helper it forked +// (EngineCore, holding the GPU) is still alive, or never gets to answer one. +// Stopping an engine has to take the helpers with it: they are in the +// engine's process group and no one else will ever clean them up. +func TestManagedProcStop_KillsHelpersTheMainProcessLeftBehind(t *testing.T) { + // The main process exits on SIGTERM; its helper ignores it and has closed the + // engine's output pipes, so the stop sees the main process end and nothing + // else would ever notice the helper. + script := `trap 'exit 0' TERM; (trap '' TERM; exec sleep 300 >/dev/null 2>&1) & echo $!; while true; do sleep 0.1; done` + childPID := make(chan int, 1) + proc, err := startManagedProc("sh", []string{"-c", script}, nil, func(stream, line string) { + if pid, err := strconv.Atoi(strings.TrimSpace(line)); err == nil && stream == "stdout" { + select { + case childPID <- pid: + default: + } + } + }) + require.NoError(t, err) + var helper int + select { + case helper = <-childPID: + case <-time.After(5 * time.Second): + require.FailNow(t, "the helper process never started") + } + t.Cleanup(func() { killGroupRemnants(proc.cmd.Process.Pid) }) + require.False(t, processGone(helper), "the helper died before the stop") + + started := time.Now() + proc.stop(5 * time.Second) + + // The main process answered the polite signal, so this must not have waited + // out the grace: what is under test is the cleanup after it exits. + assert.LessOrEqual(t, time.Since(started), 2*time.Second, "the main process should have exited on the graceful signal") + assert.True(t, waitGone(helper, 3*time.Second), "helper process %d survived the engine being stopped", helper) +} + +// A process that does not answer the polite signal is force-killed after the +// grace, together with its helpers. +func TestManagedProcStop_ForceKillsAProcessThatIgnoresTheGracefulSignal(t *testing.T) { + script := `trap '' TERM; (trap '' TERM; exec sleep 300) & echo $!; while true; do sleep 1; done` + childPID := make(chan int, 1) + proc, err := startManagedProc("sh", []string{"-c", script}, nil, func(stream, line string) { + if pid, err := strconv.Atoi(strings.TrimSpace(line)); err == nil && stream == "stdout" { + select { + case childPID <- pid: + default: + } + } + }) + require.NoError(t, err) + helper := <-childPID + t.Cleanup(func() { killGroupRemnants(proc.cmd.Process.Pid) }) + + started := time.Now() + proc.stop(500 * time.Millisecond) + + assert.LessOrEqual(t, time.Since(started), 3*time.Second, "the stop must not wait much past the grace") + assert.True(t, waitGone(proc.cmd.Process.Pid, time.Second), "the process survived the forced stop") + assert.True(t, waitGone(helper, 3*time.Second), "the helper survived the forced stop") +} + +// Shutdown is bounded by the parent, so the force kill has to come well inside +// that bound however long the manifest's grace is. +func TestStopGraceFor_CapsTheManifestGraceDuringShutdown(t *testing.T) { + test := func(name string, shuttingDown bool, stop StopSpec, want time.Duration) { + t.Run(name, func(t *testing.T) { + e := &Executor{} + e.shuttingDown.Store(shuttingDown) + + assert.Equal(t, want, e.stopGraceFor(Runtime{Stop: &stop})) + }) + } + test("an ordinary stop gets the manifest's grace", false, StopSpec{Signal: "term", GraceS: 30}, 30*time.Second) + test("a shutdown stop is capped", true, StopSpec{Signal: "term", GraceS: 30}, shutdownStopGraceCap) + test("a shutdown keeps a grace below the cap", true, StopSpec{Signal: "term", GraceS: 1}, time.Second) + test("an immediate-kill spec has no grace", true, StopSpec{Signal: "kill"}, 0) +} + +// vLLM that is still reading its weights does not answer a polite stop, and its +// manifest allows a generous grace for an ordinary stop. When the whole service +// shuts down mid-start the parent allows only a few seconds before it kills +// engine-manager, so the engine has to be force-killed inside that time or it is +// orphaned for good (it runs in its own process group, so nothing else will). +func TestStopAll_ForceKillsAStartingEngineThatIgnoresTheGracefulSignal(t *testing.T) { + // A shell that ignores SIGTERM and forks a helper that does too, like vLLM's + // API server and its EngineCore while loading. + script := `trap '' TERM; (trap '' TERM; exec sleep 300) & echo $! > ` + "%s" + `; while true; do sleep 1; done` + pidFile := t.TempDir() + "/helper.pid" + + m := testEngineManifest("/bin/sh") + key := runtime.GOOS + "/" + runtime.GOARCH + p := m.Platforms[key] + p.Runtime.Args = []string{"-c", strings.Replace(script, "%s", pidFile, 1)} + p.Runtime.Env = nil + p.Runtime.Stop = &StopSpec{Signal: "term", GraceS: 30} // far longer than a shutdown can wait + p.Runtime.Ready = &Probe{HTTP: "http://127.0.0.1:{port}/", Status: 200, TimeoutS: 60} + p.Runtime.Health = nil + m.Platforms[key] = p + m.Actions = nil + ex := newTestExecutor(t, m) + + started := make(chan error, 1) + go func() { started <- ex.Start(t.Context(), "fake") }() + + var helper int + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if data, err := os.ReadFile(pidFile); err == nil { + if pid, err := strconv.Atoi(strings.TrimSpace(string(data))); err == nil { + helper = pid + break + } + } + time.Sleep(20 * time.Millisecond) + } + require.NotZero(t, helper, "the engine never started") + st, _ := ex.state("fake") + for i := 0; i < 100; i++ { + st.mu.Lock() + proc := st.proc + st.mu.Unlock() + if proc != nil { + t.Cleanup(func() { killGroupRemnants(proc.cmd.Process.Pid) }) + break + } + time.Sleep(20 * time.Millisecond) + } + + began := time.Now() + ex.StopAll() + + assert.LessOrEqual(t, time.Since(began), shutdownStopGraceCap+3*time.Second, "the parent only waits a few seconds for shutdown") + assert.True(t, waitGone(helper, 2*time.Second), "the engine's helper process %d was left running", helper) + select { + case <-started: + case <-time.After(3 * time.Second): + assert.Fail(t, "the cancelled start never returned") + } +} + +// (The script ends in a command after the sleep so the shell is not replaced by it.) +// A Python-packaged engine (vLLM) runs as "interpreter script args", so its +// executable never matches the engine path. The script in the command line is +// what proves a leftover process is the engine this service manages. +func TestProcRunsScript_RecognisesAnInterpreterRunningTheEngineScript(t *testing.T) { + dir := t.TempDir() + script := dir + "/engine" + require.NoError(t, os.WriteFile(script, []byte("#!/bin/sh\nsleep 30\ntrue\n"), 0o755)) + proc, err := startManagedProc("/bin/sh", []string{script, "serve", "--port", "1"}, nil, nil) + require.NoError(t, err) + t.Cleanup(func() { proc.stop(0) }) + pid := proc.cmd.Process.Pid + // /proc//cmdline reads empty for a moment while a new process finishes + // exec'ing. A leftover engine has been running far longer than that. + for i := 0; i < 200; i++ { + if raw, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/cmdline"); err == nil && len(raw) > 0 { + break + } + time.Sleep(10 * time.Millisecond) + } + + recognised := func(name string, pid int, path string) { + t.Run(name, func(t *testing.T) { + assert.True(t, procRunsScript(pid, path)) + }) + } + notRecognised := func(name string, pid int, path string) { + t.Run(name, func(t *testing.T) { + assert.False(t, procRunsScript(pid, path)) + }) + } + recognised("the engine script", pid, script) + recognised("the engine script through a dot segment", pid, dir+"/./engine") + notRecognised("an unrelated script", pid, dir+"/other-engine") + // A bare word is not an engine path, even though it is one of the arguments. + notRecognised("a relative path among the arguments", pid, "serve") + notRecognised("pid 0", 0, script) + notRecognised("an empty path", pid, "") + notRecognised("a pid that does not exist", 999999999, script) +} diff --git a/services/nvpair-engine-manager/proc_unix.go b/services/nvpair-engine-manager/proc_unix.go index 12c9b190d..47186a6ae 100644 --- a/services/nvpair-engine-manager/proc_unix.go +++ b/services/nvpair-engine-manager/proc_unix.go @@ -10,6 +10,7 @@ import ( "errors" "os" "os/exec" + "path/filepath" "strconv" "strings" "syscall" @@ -108,6 +109,34 @@ func procImage(pid int) string { return "" } +// procRunsScript reports whether pid was started as an interpreter running the +// engine script at binPath. An engine installed as a Python package (vLLM's +// venv/bin/vllm) is a script, so the process's executable is the interpreter and +// procImage never matches binPath; what does match is the script named in the +// command line, as one of its first arguments ("python /path/venv/bin/vllm +// serve ..."). Linux only: elsewhere there is no /proc to read, and the answer +// is false, which keeps the caller failing closed. +func procRunsScript(pid int, binPath string) bool { + if pid <= 0 || binPath == "" { + return false + } + raw, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/cmdline") + if err != nil { + return false + } + want := filepath.Clean(binPath) + if !filepath.IsAbs(want) { + return false // an engine path is absolute; a bare word could match any argument + } + args := strings.Split(strings.TrimRight(string(raw), "\x00"), "\x00") + for i := 0; i < len(args) && i < 3; i++ { + if args[i] != "" && filepath.Clean(args[i]) == want { + return true + } + } + return false +} + // signalPID sends SIGTERM (or SIGKILL when force) to the process group when // possible. The owned-process escalation and orphan reclaim share it so forced // stops reach forked helpers (model runners, etc.). @@ -125,6 +154,16 @@ func signalPID(pid int, force bool) error { return syscall.Kill(pid, sig) } +// killGroupRemnants force-kills whatever is left in the process group that +// configureSysProcAttr created for pid, which is the group's leader and so its +// id. It is called right after the leader has exited and been reaped; an empty +// group is the usual case and is not an error. +func killGroupRemnants(pid int) { + if pid > 1 { + _ = syscall.Kill(-pid, syscall.SIGKILL) + } +} + // pidAlive reports whether the PID still exists (signal 0 probes existence // without delivering a signal; EPERM means it exists but isn't ours). func pidAlive(pid int) bool { diff --git a/services/nvpair-engine-manager/proc_windows.go b/services/nvpair-engine-manager/proc_windows.go index 77fd21811..e9d6ac09a 100644 --- a/services/nvpair-engine-manager/proc_windows.go +++ b/services/nvpair-engine-manager/proc_windows.go @@ -54,6 +54,14 @@ func taskkill(cmd *exec.Cmd, force bool) error { return kc.Run() } +// procRunsScript is unused on Windows, where the engines are real executables +// and procImage identifies them. +func procRunsScript(int, string) bool { return false } + +// killGroupRemnants has nothing to do on Windows: every stop is taskkill /T, +// which already takes the whole process tree. +func killGroupRemnants(int) {} + var ( iphlpapi = windows.NewLazySystemDLL("iphlpapi.dll") procGetExtendedTcpTable = iphlpapi.NewProc("GetExtendedTcpTable") diff --git a/services/nvpair-engine-manager/pull.go b/services/nvpair-engine-manager/pull.go index 97d8fb565..5d464536b 100644 --- a/services/nvpair-engine-manager/pull.go +++ b/services/nvpair-engine-manager/pull.go @@ -122,6 +122,13 @@ func (e *Executor) PullModelStream(ctx context.Context, engine, model string, pa params, _ = json.Marshal(map[string]string{"name": model, "model": model}) } + // Hugging Face repository (vLLM): download with the engine's own `hf` tool and + // report progress from the bytes landing in the cache. It needs no running + // engine, so a model can be fetched before it is first served. + if act.HFPull != nil { + return e.pullHFModel(ctx, st, engine, model, act) + } + // CLI action (e.g. lms get): no structured line progress; emit a start // marker and return the final result via the existing runner. if len(act.Cmd) > 0 { diff --git a/services/nvpair-engine-manager/registry.go b/services/nvpair-engine-manager/registry.go index 3b3278871..ecf04eccf 100644 --- a/services/nvpair-engine-manager/registry.go +++ b/services/nvpair-engine-manager/registry.go @@ -44,7 +44,17 @@ var allowedPlaceholders = map[string]bool{ "models_dir": true, } +// hfHubCachePlaceholder is the Hugging Face hub cache the engine itself uses, +// resolved from its launch environment the way huggingface_hub resolves it +// (see hfEnv.hubCache). It is available only in the dir of an hf_cache or +// hf_pull action, whose files must be the ones the engine reads. +const hfHubCachePlaceholder = "hf_hub_cache" + var placeholderRe = regexp.MustCompile(`\{([a-zA-Z_][a-zA-Z0-9_]*)\}`) + +// resultFieldNameRe matches one element field name, such as "root". +var resultFieldNameRe = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_-]*$`) + var resultMatchFieldPathRe = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_-]*(?:\.[a-zA-Z_][a-zA-Z0-9_-]*)*$`) var artifactNameRe = regexp.MustCompile(`^[a-z][a-z0-9_]{0,31}$`) var sha256Re = regexp.MustCompile(`^[a-fA-F0-9]{64}$`) @@ -63,6 +73,12 @@ type Manifest struct { ManifestVersion int `json:"manifest_version"` Platforms map[string]Platform `json:"platforms"` Actions map[string]Action `json:"actions,omitempty"` + // AdvertisedModels says which models this node tells its peers and its own + // proxy it can serve: "list" (the default) advertises list_models, "loaded" + // advertises loaded_models. An engine that serves one model per process + // can list many models it has on disk but only answer for the one it is + // running, and the proxy sends a request to any node advertising its model. + AdvertisedModels string `json:"advertised_models,omitempty"` } // Platform is the per-`/` block. Variance lives here @@ -150,6 +166,14 @@ type Runtime struct { // a non-nil empty slice keeps only fixed and managed arguments. Use // omitzero to omit nil while preserving explicit empty arrays in JSON. LaunchArgs []string `json:"launch_args,omitzero"` + // AutoLaunch lets the manager add startup arguments that depend on the model + // being served and the machine, at every start. Arguments the user saved + // always win, and nothing added here is written back to saved settings. + AutoLaunch *AutoLaunch `json:"auto_launch,omitempty"` + // UnsupportedWeights names weight formats this platform's engine cannot + // load. A start or download of such a model is refused with its reason, + // before the engine spends minutes downloading or loading it only to fail. + UnsupportedWeights *UnsupportedWeights `json:"unsupported_weights,omitempty"` // Mode selects the lifecycle model: "process" means PAIR owns the launched // foreground process; "command" means PAIR invokes commands that manage a // separate daemon. Empty defaults to "process". @@ -188,6 +212,39 @@ type Runtime struct { Health *Probe `json:"health,omitempty"` } +// AutoLaunch configures start-time launch arguments. The only kind is +// "hf_context": read the model's config.json from the Hugging Face cache and +// set the context length, any rope scaling it needs, and the GPU memory share. +type AutoLaunch struct { + Kind string `json:"kind"` + // ContextTokens is the context window to aim for. Models that natively + // support less are extended where that is safe; models that support more are + // capped here so the KV cache stays affordable. + ContextTokens int `json:"context_tokens"` + // MaxNumSeqs caps how many requests the engine runs at once + // (--max-num-seqs). vLLM's default of 256 suits a server; on a desktop GPU a + // hybrid model such as Qwen3.5 cannot fit one state block per sequence for + // 256 of them and refuses to start. 0 leaves the engine's default. + MaxNumSeqs int `json:"max_num_seqs,omitempty"` +} + +const autoLaunchHFContext = "hf_context" + +// UnsupportedWeights is a weight format a platform's engine cannot load, read +// from a Hugging Face checkpoint the way vLLM reads it (see weightFormat). +type UnsupportedWeights struct { + // Formats are the refused formats. The only one is "fp8": FP8 weights or an + // FP8 KV cache. + Formats []string `json:"formats"` + // ExceptModelTypes are config.json model_type values the engine loads in + // these formats anyway. + ExceptModelTypes []string `json:"except_model_types,omitempty"` + // Reason is what the refusal tells the user, including what to choose instead. + Reason string `json:"reason"` +} + +const weightFormatFP8 = "fp8" + type EditableLaunch struct { FixedArgs []string `json:"fixed_args"` Controls []LaunchControl `json:"controls"` @@ -257,6 +314,21 @@ type Action struct { HTTP *ActionHTTP `json:"http,omitempty"` Cmd []string `json:"cmd,omitempty"` RemovePath *ActionRemovePath `json:"remove_path,omitempty"` + // HFCache lists the models in a Hugging Face hub cache directory. It reads + // the disk and needs no running engine, so an engine that serves one model + // per process can still show everything it could be switched to. + HFCache *ActionHFCache `json:"hf_cache,omitempty"` + // LaunchModel makes the action switch the model an engine is serving, for an + // engine that has no load call because the model is a launch argument. + LaunchModel *ActionLaunchModel `json:"launch_model,omitempty"` + // HFPull downloads a Hugging Face repository into the hub cache with the + // engine's own `hf` tool, reporting byte progress. Only pull_model uses it. + HFPull *ActionHFPull `json:"hf_pull,omitempty"` + // StopEngine makes the action stop the engine. It is how an engine that serves + // one model per process ejects that model: the process cannot drop its model + // and keep running, but stopping it frees all of its GPU memory. Only + // unload_model may set it. + StopEngine bool `json:"stop_engine,omitempty"` // ProgressProtocol selects a narrowly defined streaming adapter for an // asynchronous pull HTTP action. Empty keeps the ordinary response-stream // behavior; named protocols are validated so a typo cannot silently fall @@ -293,6 +365,13 @@ type ActionResult struct { // field) rather than returning a presence-only list (Ollama's /api/ps, // every row of which is resident and so needs no Match). Match *ResultMatch `json:"match,omitempty"` + // InventoryField, valid only on the loaded_models action of an engine with + // advertised_models "loaded", names the element field that identifies each + // loaded model the way list_models does. vLLM's /v1/models gives the name a + // request uses as "id" (a --served-model-name alias when one is set) and the + // model it loaded as "root" (its --model): peers and the proxy are told the + // id, and the loaded marker on the model list uses the root. + InventoryField string `json:"inventory_field,omitempty"` } // ResultMatch is the optional row filter on an ActionResult. Exactly one of @@ -314,6 +393,36 @@ type ActionRemovePath struct { Root string `json:"root"` } +// ActionHFCache names the Hugging Face hub cache directory to list. Dir is +// templated, normally "{hf_hub_cache}", the cache the engine itself uses. The result is {"data":[{"id","size"}]}, +// which the ordinary result extractor reads with array "data" and field "id". +type ActionHFCache struct { + Dir string `json:"dir"` + // Delete turns the action into the removal of one model, named by the + // caller's {"model": ""} param, from that cache. Only a + // delete_model action may set it, and it needs Tool. + Delete bool `json:"delete,omitempty"` + // Tool is the file name of the Hugging Face CLI that performs a delete + // (`hf cache rm`), looked up beside the engine binary as for hf_pull. + Tool string `json:"tool,omitempty"` +} + +// ActionHFPull downloads a model repository. Tool is the file name of the +// Hugging Face CLI, looked up in the directory that holds the engine binary +// (a vLLM virtualenv ships `hf` beside `vllm`). Dir is the hub cache directory, +// normally "{hf_hub_cache}", the cache the engine itself uses. +type ActionHFPull struct { + Tool string `json:"tool"` + Dir string `json:"dir"` +} + +// ActionLaunchModel names the launch flag that selects the engine's model, such +// as "--model". The action takes {"model": ""}, writes it into the saved +// launch settings and restarts a running engine on it. +type ActionLaunchModel struct { + Flag string `json:"flag"` +} + // ActionHTTP is a templated call against the engine's loopback base URL. // ParamsIn selects whether the caller's params are sent as the JSON request // body (the default) or as URL-encoded query parameters. BodySchema is @@ -716,6 +825,19 @@ func (m *Manifest) Validate() error { // process is spawned, the reply races the engine's own startup, and the // caller can still be handed the pre-restart state. Fail at load rather // than ship a guarantee the manifest cannot keep. + // launch_model rewrites the saved launch settings and restarts the engine, + // so every platform needs editable launch settings and a readiness probe + // for the same reason as restart_after. + if a.LaunchModel != nil { + for key, p := range m.Platforms { + if p.Runtime.EditableLaunch == nil { + return fmt.Errorf("action %q: launch_model requires platform %q to declare runtime.editable_launch", name, key) + } + if p.Runtime.Ready == nil { + return fmt.Errorf("action %q: launch_model requires platform %q to declare runtime.ready", name, key) + } + } + } if a.RestartAfter { for key, p := range m.Platforms { if p.Runtime.Ready == nil { @@ -724,10 +846,49 @@ func (m *Manifest) Validate() error { } } } + switch m.AdvertisedModels { + case "", advertisedModelsList: + case advertisedModelsLoaded: + if _, ok := m.Actions["loaded_models"]; !ok { + return fmt.Errorf("advertised_models %q requires a loaded_models action", m.AdvertisedModels) + } + default: + return fmt.Errorf("advertised_models %q invalid (want %q or %q)", m.AdvertisedModels, advertisedModelsList, advertisedModelsLoaded) + } + if act, ok := m.Actions["loaded_models"]; ok && act.Result != nil && act.Result.InventoryField != "" && m.AdvertisedModels != advertisedModelsLoaded { + return fmt.Errorf("action %q: result.inventory_field requires advertised_models %q", "loaded_models", advertisedModelsLoaded) + } return m.validatePlaceholders() } +// advertisedModels values for Manifest.AdvertisedModels. +const ( + advertisedModelsList = "list" + advertisedModelsLoaded = "loaded" +) + func (p *Platform) validate(key string) error { + if auto := p.Runtime.AutoLaunch; auto != nil { + if auto.Kind != autoLaunchHFContext || auto.ContextTokens < 1024 { + return fmt.Errorf("platform %q: auto_launch needs kind %q and context_tokens of at least 1024", key, autoLaunchHFContext) + } + if auto.MaxNumSeqs < 0 { + return fmt.Errorf("platform %q: auto_launch max_num_seqs must not be negative", key) + } + if p.Runtime.EditableLaunch == nil || p.Runtime.modeOrDefault() != "process" { + return fmt.Errorf("platform %q: auto_launch requires a process-mode runtime with editable_launch", key) + } + } + if rule := p.Runtime.UnsupportedWeights; rule != nil { + if len(rule.Formats) == 0 || strings.TrimSpace(rule.Reason) == "" { + return fmt.Errorf("platform %q: unsupported_weights needs formats and a reason", key) + } + for _, format := range rule.Formats { + if format != weightFormatFP8 { + return fmt.Errorf("platform %q: unsupported_weights format %q is not one of [%s]", key, format, weightFormatFP8) + } + } + } if policy := p.Runtime.EditableLaunch; policy != nil { if policy.StartIndex < 0 || (p.Runtime.modeOrDefault() == "command" && policy.StartIndex >= len(p.Runtime.Start)) { return fmt.Errorf("platform %q: invalid editable launch definition", key) @@ -883,6 +1044,9 @@ func (a *Action) validate(name string) error { hasHTTP := a.HTTP != nil hasCmd := len(a.Cmd) > 0 hasRemovePath := a.RemovePath != nil + hasHFCache := a.HFCache != nil + hasLaunchModel := a.LaunchModel != nil + hasHFPull := a.HFPull != nil kinds := 0 if hasHTTP { kinds++ @@ -893,8 +1057,44 @@ func (a *Action) validate(name string) error { if hasRemovePath { kinds++ } + if hasHFCache { + kinds++ + } + if hasLaunchModel { + kinds++ + } + if hasHFPull { + kinds++ + } + if a.StopEngine { + kinds++ + } if kinds != 1 { - return fmt.Errorf("action %q: exactly one of http, cmd, or remove_path is required", name) + return fmt.Errorf("action %q: exactly one of http, cmd, remove_path, hf_cache, launch_model, hf_pull, or stop_engine is required", name) + } + if a.StopEngine && name != "unload_model" { + return fmt.Errorf("action %q: stop_engine belongs to unload_model", name) + } + if hasHFPull { + tool := a.HFPull.Tool + if name != pullModelAction || tool == "" || tool != filepath.Base(tool) || strings.TrimSpace(a.HFPull.Dir) == "" { + return fmt.Errorf("action %q: hf_pull belongs to pull_model and needs a tool file name and a dir", name) + } + } + if hasLaunchModel { + flag := a.LaunchModel.Flag + if !strings.HasPrefix(flag, "-") || strings.ContainsAny(flag, " \t=") { + return fmt.Errorf("action %q: launch_model.flag %q must be a flag such as \"--model\"", name, flag) + } + } + if hasHFCache && strings.TrimSpace(a.HFCache.Dir) == "" { + return fmt.Errorf("action %q: hf_cache.dir is required", name) + } + if hasHFCache && a.HFCache.Delete && name != "delete_model" { + return fmt.Errorf("action %q: hf_cache.delete belongs to delete_model", name) + } + if hasHFCache && a.HFCache.Delete != (strings.TrimSpace(a.HFCache.Tool) != "") { + return fmt.Errorf("action %q: hf_cache.delete and hf_cache.tool go together", name) } if hasRemovePath { if strings.TrimSpace(a.RemovePath.Path) == "" || strings.TrimSpace(a.RemovePath.Root) == "" { @@ -922,6 +1122,14 @@ func (a *Action) validate(name string) error { if a.Result != nil && (strings.TrimSpace(a.Result.Array) == "" || strings.TrimSpace(a.Result.Field) == "") { return fmt.Errorf("action %q: result.array and result.field are required when result is set", name) } + if a.Result != nil && a.Result.InventoryField != "" { + if name != "loaded_models" { + return fmt.Errorf("action %q: result.inventory_field is only valid on loaded_models", name) + } + if !resultFieldNameRe.MatchString(a.Result.InventoryField) { + return fmt.Errorf("action %q: result.inventory_field %q is not a valid field name", name, a.Result.InventoryField) + } + } if a.Result != nil && a.Result.Match != nil { m := a.Result.Match if strings.TrimSpace(m.Field) == "" { @@ -972,7 +1180,29 @@ func (m *Manifest) validatePlaceholders() error { return fmt.Errorf("platform %q: %w", key, err) } } - return validatePlaceholderStrings(m.actionTemplatedStrings(), allowedPlaceholders) + if err := validatePlaceholderStrings(m.actionTemplatedStrings(), allowedPlaceholders); err != nil { + return err + } + hfAllowed := map[string]bool{hfHubCachePlaceholder: true} + for name := range allowedPlaceholders { + hfAllowed[name] = true + } + return validatePlaceholderStrings(m.hfActionDirs(), hfAllowed) +} + +// hfActionDirs are the dirs of the hf_cache and hf_pull actions, which may also +// name {hf_hub_cache}. +func (m *Manifest) hfActionDirs() []string { + var out []string + for _, act := range m.Actions { + if act.HFPull != nil { + out = append(out, act.HFPull.Dir) + } + if act.HFCache != nil { + out = append(out, act.HFCache.Dir) + } + } + return out } func validatePlaceholderStrings(templates []string, allowed map[string]bool) error { @@ -1032,7 +1262,7 @@ func (m *Manifest) actionTemplatedStrings() []string { // Path may template caller params (e.g. {model}); validated at call time. } } - // Action http.path / cmd are resolved from runtime params (e.g. + // hf_cache and hf_pull dirs are checked by hfActionDirs. Action http.path / cmd are resolved from runtime params (e.g. // {model}), so they're validated at call time, not statically here. return out } diff --git a/services/nvpair-engine-manager/registry_test.go b/services/nvpair-engine-manager/registry_test.go index fedd89851..b8bef4df7 100644 --- a/services/nvpair-engine-manager/registry_test.go +++ b/services/nvpair-engine-manager/registry_test.go @@ -317,7 +317,7 @@ func TestValidate_RejectsInvalidManifests(t *testing.T) { }, "json_match.value is required"}, {"action without http or cmd", func(m *Manifest) { m.Actions = map[string]Action{"x": {Description: "neither"}} - }, "exactly one of http, cmd, or remove_path"}, + }, "exactly one of http, cmd, remove_path, hf_cache, launch_model, hf_pull, or stop_engine"}, {"action missing method", func(m *Manifest) { m.Actions = map[string]Action{"x": {HTTP: &ActionHTTP{Path: "/p"}}} }, "http.method and http.path"}, @@ -551,6 +551,23 @@ func TestBundledOllama_AllowsConfiguredStartupTime(t *testing.T) { } } +// A vLLM start fetches model metadata and probes for a tool parser before the +// engine launches, all inside its readiness allowance. That whole allowance +// must fit what a remote peer waits for, and the work before launch must leave +// the engine time to come up. +func TestBundledVLLM_StartFitsTheRemoteReadinessBudget(t *testing.T) { + reg := NewRegistry() + require.NoError(t, reg.LoadFS(bundledManifests, "manifests")) + m, ok := reg.Get("vllm") + require.True(t, ok, "vllm manifest not loaded") + for key, p := range m.Platforms { + require.NotNil(t, p.Runtime.Ready, "readiness probe missing (%v)", key) + startAllowance := time.Duration(p.Runtime.Ready.TimeoutS) * time.Second + assert.Greater(t, remoteReadyResponseHeaderTimeout, startAllowance, "the whole start, work before launch included (%v)", key) + assert.Greater(t, startAllowance, hfMetadataTimeout+toolParserProbeTimeout, "the work before launch leaves the engine time to come up (%v)", key) + } +} + func writeManifest(t *testing.T, dir, name string, m Manifest) { t.Helper() data, err := json.MarshalIndent(m, "", " ") @@ -677,3 +694,11 @@ func TestLMStudioInstall_UsesSafeBootstrapCommands(t *testing.T) { assert.Equal(t, []string{"bash", "{download}"}, p.Install.Run, " (%v)", key) } } + +func TestValidate_RejectsANegativeAutoLaunchRequestLimit(t *testing.T) { + m := validManifest() + platform := m.Platforms["linux/amd64"] + platform.Runtime.AutoLaunch = &AutoLaunch{Kind: autoLaunchHFContext, ContextTokens: 65536, MaxNumSeqs: -1} + m.Platforms["linux/amd64"] = platform + require.ErrorContains(t, m.Validate(), "max_num_seqs", "a negative request limit was accepted") +} diff --git a/services/nvpair-engine-manager/remediation_test.go b/services/nvpair-engine-manager/remediation_test.go index 194b667a5..fc2ea64c1 100644 --- a/services/nvpair-engine-manager/remediation_test.go +++ b/services/nvpair-engine-manager/remediation_test.go @@ -7,6 +7,7 @@ import ( "context" "encoding/json" "fmt" + "net" "net/http" "net/http/httptest" "net/url" @@ -475,3 +476,33 @@ func TestStop_ClearsUnhealthy(t *testing.T) { } require.False(t, hasErr(ex.Errors(), unhealthyID("d")), "stop must clear the lingering unhealthy error") } + +// A start refused before the engine is launched changes no state, so a client +// that sends Start without waiting for the reply, as the desktop does, learns of +// the refusal only from the start-failed error. +func TestStart_ReportsARefusalAsAFailedStart(t *testing.T) { + test := func(name string, executor func(t *testing.T) *Executor, wantErr string) { + t.Run(name, func(t *testing.T) { + e := executor(t) + require.ErrorContains(t, e.Start(t.Context(), "fake"), wantErr) + assert.True(t, hasErr(e.Errors(), startFailedID("fake")), "the refusal is reported as a failed start") + }) + } + test("an engine whose model is a launch argument, with no model selected", func(t *testing.T) *Executor { + e := settingsExecutor(t, false) + m, _ := e.reg.Get("fake") + m.Actions["load_model"] = Action{LaunchModel: &ActionLaunchModel{Flag: "--model"}} + require.NoError(t, m.Validate()) + return e + }, "no model selected") + test("a port taken by a service that is not the engine", func(t *testing.T) *Executor { + l, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { _ = l.Close() }) + e := settingsExecutor(t, false) + st, err := e.state("fake") + require.NoError(t, err) + st.port = l.Addr().(*net.TCPAddr).Port + return e + }, "occupied") +} diff --git a/services/nvpair-engine-manager/remoteclient.go b/services/nvpair-engine-manager/remoteclient.go index 3a19bdf07..e54cdd33a 100644 --- a/services/nvpair-engine-manager/remoteclient.go +++ b/services/nvpair-engine-manager/remoteclient.go @@ -38,8 +38,8 @@ type remoteClient struct { } // waitsForEngineReadiness reports whether a peer may only answer after an -// engine is healthy or an Ollama model is loaded, which the ordinary 30s -// response-header budget cannot cover. +// engine is healthy or an Ollama or vLLM model is loaded, which the ordinary +// 30s response-header budget cannot cover. // // Cutting such a call off is worse than slow. The initiator's cancellation // propagates into the peer's in-flight handler, so bringUpCommand tears the @@ -51,8 +51,10 @@ func waitsForEngineReadiness(path, engine string) bool { } // controlDeletePath: LM Studio's delete_model declares restart_after, so the // peer replies only after the post-delete restart is ready. + // vLLM's load_model restarts the engine on the new model and replies once + // it is ready. if path == controlLoadPath { - return engine == "ollama" + return engine == "ollama" || engine == "vllm" } return path == controlStartPath || path == controlDeletePath } diff --git a/services/nvpair-engine-manager/remoteclient_test.go b/services/nvpair-engine-manager/remoteclient_test.go index 36f8939b7..d685e0502 100644 --- a/services/nvpair-engine-manager/remoteclient_test.go +++ b/services/nvpair-engine-manager/remoteclient_test.go @@ -74,6 +74,7 @@ func TestRemoteReadinessBudget_CoversEngineStartupAllowance(t *testing.T) { test("start waits", controlStartPath, "ollama", true) test("delete waits", controlDeletePath, "lmstudio", true) test("Ollama load waits", controlLoadPath, "ollama", true) + test("vLLM load waits", controlLoadPath, "vllm", true) test("LM Studio load does not wait", controlLoadPath, "lmstudio", false) test("stop does not wait", controlStopPath, "ollama", false) test("unload does not wait", controlUnloadPath, "ollama", false) diff --git a/services/nvpair-engine-manager/settings.go b/services/nvpair-engine-manager/settings.go index 568915107..c1414bcfe 100644 --- a/services/nvpair-engine-manager/settings.go +++ b/services/nvpair-engine-manager/settings.go @@ -398,6 +398,13 @@ func (e *Executor) ConfigureLaunch(ctx context.Context, p settings.Configure, re } st.opMu.Lock() defer st.opMu.Unlock() + return e.configureLaunchLocked(ctx, st, p, rebind) +} + +// configureLaunchLocked is ConfigureLaunch for a caller that already holds +// st.opMu, so a model switch or a delete can change the launch, start or roll +// back without another operation slipping in between. +func (e *Executor) configureLaunchLocked(ctx context.Context, st *engineState, p settings.Configure, rebind func() error) (settings.LaunchState, error) { if err := ctx.Err(); err != nil { return settings.LaunchState{}, err } diff --git a/services/nvpair-engine-manager/shutdown_test.go b/services/nvpair-engine-manager/shutdown_test.go index d30c56f9c..fc483f3dd 100644 --- a/services/nvpair-engine-manager/shutdown_test.go +++ b/services/nvpair-engine-manager/shutdown_test.go @@ -213,7 +213,7 @@ func TestStopAll_RetriesAFailedStopCommand(t *testing.T) { }, }, }) - require.NoError(t, ex.Start(context.Background(), "daemon")) + require.NoError(t, ex.Start(t.Context(), "daemon")) ex.StopAll() st, err := ex.Status("daemon") require.NoError(t, err) diff --git a/services/nvpair-engine-manager/spec.md b/services/nvpair-engine-manager/spec.md index ac2ec2c5f..5edb97e0a 100644 --- a/services/nvpair-engine-manager/spec.md +++ b/services/nvpair-engine-manager/spec.md @@ -49,7 +49,7 @@ extensibility story for an open-source product. and must not be started or adopted because its `/models` contract does not expose router residency state. - Per-engine stdout/stderr log capture and structured operational error records, surfaced via the errors pipeline. -- A normalized node-level model list (`engine:models`): union of every running engine's `list_models`, name-extracted via each action's declarative `result` spec, plus the per-engine set of models loaded in memory (`loadedByEngine`, from each engine's `loaded_models` action). A successful explicit empty inventory remains an engine key with `[]`; a missing/malformed/failed inventory omits that engine key instead of being mislabeled as authoritative empty. A watcher polls the loaded set and pushes `engine:models-changed` when it changes (explicit load/unload, JIT auto-load, TTL/idle eviction). +- A normalized node-level model list (`engine:models`): union of every running engine's `list_models`, name-extracted via each action's declarative `result` spec, plus the per-engine set of models loaded in memory (`loadedByEngine`, from each engine's `loaded_models` action). An engine with `advertised_models: "loaded"` (vLLM) advertises only its served model in `modelsByEngine` and reports its `list_models` inventory as `downloadedByEngine`, read even while it is stopped; its `loadedByEngine` names models the way that inventory does. A successful explicit empty inventory remains an engine key with `[]`; a missing/malformed/failed inventory omits that engine key instead of being mislabeled as authoritative empty. A watcher polls these lists and pushes `engine:models-changed` when one changes (explicit load/unload, JIT auto-load, TTL/idle eviction, a pull or a delete). - Expose all of the above over the `engine:*` JSON-RPC surface to whatever orchestrates the service, plus an optional plain-HTTP LAN endpoint (`--http-port`, `GET /v1/models`) that serves the model list to a peer's discovery daemon (the list moved off the size-limited mDNS TXT onto HTTP). **Out of scope** @@ -82,7 +82,7 @@ extensibility story for an open-source product. **Functional** - Load + validate per-engine JSON manifests (bundled + user dir); select the host `/` block; resolve placeholders (`{bin}`, `{cli}`, `{port}`, `{download}`, `{install_dir}`). Install commands also receive resolved download and destination paths in child-scoped `NVPAIR_INSTALL_*` environment variables so shell reparsing cannot corrupt them. - Support both `process` (owned foreground) and `command` (daemon + control-CLI) runtimes; execute detect / install / uninstall / start / stop / restart / status / health and HTTP **or** CLI actions; emit `engine:*` results and notifications. -- Bound process-mode stops: on Unix send SIGTERM to the owned process group, then SIGKILL after `runtime.stop.grace_s` (five seconds by default); `signal:"kill"` skips the grace. On Windows, windowless managed engines require immediate `taskkill /T /F`. +- Bound process-mode stops: on Unix send SIGTERM to the owned process group, then SIGKILL after `runtime.stop.grace_s` (five seconds by default, capped at three seconds while the service shuts down); the SIGKILL also takes any helper processes left in the group; `signal:"kill"` skips the grace. On Windows, windowless managed engines require immediate `taskkill /T /F`. - Emit `errors:report` / `errors:clear` on its stdio for the Broker to forward to `nvpair-errors`. **Non-functional** @@ -136,9 +136,9 @@ Requests (caller → service): | Method | Params | Result | |--------|--------|--------| -| `engine:get-installed` | — | `{ engines: [EngineStatus] }` | +| `engine:get-installed` | — | `{ engines: [EngineStatus] }` for every engine this host can run; an engine whose manifest has no block for the host's OS and CPU is left out | | `engine:describe` | `{ engine }` | the engine's manifest | -| `engine:status` | `{ engine }` | `EngineStatus` | +| `engine:status` | `{ engine }` | `EngineStatus`; an error for an engine this host cannot run | | `engine:install` | `{ engine }` | `EngineStatus` (after install) | | `engine:uninstall` | `{ engine }` | `EngineStatus` (after removal) | | `engine:uninstall-managed` | — | `{ engines: [{ engine, removed, error? }] }` — removes the engines this service installed, selected by their install marker, each through the `engine:uninstall` path. Backs both clients' "reset all data"; the platform uninstallers reach the same code through `--uninstall-managed` because they have no broker. Neither `removed` nor `error` means the engine was left alone as not ours. | @@ -161,7 +161,7 @@ Requests (caller → service): | `shutdown` | — | `null` | | `log/set-level` | `{ level }` | `{ level }` | -Notifications (service → caller): `ready{version}`, `engine:state-changed{EngineStatus}`, `engine:models-changed{engine, models}` (pushed when an engine's loaded-in-memory model set changes; `models` is the full `engine:models` shape incl. `loadedByEngine`), `engine:install-progress{engine, stage, percent?, error?}` (`error` on the terminal `failed` frame), `engine:pull-progress{engine, op, stage, percent?, message}` (live progress for a local model pull driven via `engine:action{action:"pull_model"}` — the local counterpart of `engine:remote-progress`), `engine:remote-progress{opId, node, engine, op, stage, percent?, message}` (live progress relayed from a remote install/pull), and `errors:report` / `errors:clear` (consumed by `nvpair-errors` via the Broker). `install` / `start` / `stop` / `restart` / `action` / `remote-*` each run in their own goroutine so the read loop never blocks; their responses arrive when the op completes. +Notifications (service → caller): `ready{version}`, `engine:state-changed{EngineStatus}`, `engine:models-changed{engine, models}` (pushed when an engine's served, loaded-in-memory or downloaded model list changes; `models` is the full `engine:models` shape), `engine:install-progress{engine, stage, percent?, error?}` (`error` on the terminal `failed` frame), `engine:pull-progress{engine, op, stage, percent?, message}` (live progress for a local model pull driven via `engine:action{action:"pull_model"}` — the local counterpart of `engine:remote-progress`), `engine:remote-progress{opId, node, engine, op, stage, percent?, message}` (live progress relayed from a remote install/pull), and `errors:report` / `errors:clear` (consumed by `nvpair-errors` via the Broker). `install` / `start` / `stop` / `restart` / `action` / `remote-*` each run in their own goroutine so the read loop never blocks; their responses arrive when the op completes. `percent` on all three progress notifications follows one rule: it is present only when the step is measurable (download bytes, byte-progress pulls) or terminal (`100` on install `done` / `already-installed`, `-1` on install `failed` and pull `error`). An indeterminate step omits it — the install `verified` and `installing` stages always do, because the install command is opaque — and a consumer renders the stage alone rather than a percentage. @@ -183,7 +183,7 @@ Example `engine:pull-progress` (stdout): ### 7.2 Remote engine management (the `ec` surface) With `--control-port` and a clustered `--cluster-dir`, engine-manager serves a cluster-scoped remote-control surface over pin-based mTLS (`nvpair-shared/clustertrust`): it presents this node's cluster leaf, requires a client cert, and `403`s any caller that isn't a byte-for-byte pinned cluster peer. The listener is bound whenever `--control-port` is set and admits callers by live membership — its leaf is resolved per handshake, so an unclustered node presents none and every handshake is refused — so no restart is needed on `cluster:identity-changed`; the broker registers `ec` whenever a cluster dir is configured. Routes under `/v1`: `GET /engines`, streaming `POST /engines/install` and `/models/pull` (chunked NDJSON — zero+ `{"type":"progress"}` frames then one terminal `{"type":"result"}`/`{"type":"error"}` frame), non-streaming `POST /models/{load,unload,delete}`, and non-streaming `POST /engines/{start,stop}` returning `EngineStatus`. -The `engine:remote-*` methods are the client half: engine-manager resolves the target `node` in an `ec` peer directory (fed by its own `discovery:subscribe{services:[ec]}` to the broker relay), dials the peer's `ec` surface with the same pinned identity, relays each streamed progress frame up as `engine:remote-progress` (keyed by a minted `opId`), and settles the request on the terminal frame. Remote install/pull run for the operation's full duration (no broker-imposed timeout); remote stop and model unload/delete are fast request/response, remote Ollama model load uses the readiness-sized response budget, and remote start waits for the engine's bounded readiness result. They error if this node isn't clustered or the target isn't a pinned peer. +The `engine:remote-*` methods are the client half: engine-manager resolves the target `node` in an `ec` peer directory (fed by its own `discovery:subscribe{services:[ec]}` to the broker relay), dials the peer's `ec` surface with the same pinned identity, relays each streamed progress frame up as `engine:remote-progress` (keyed by a minted `opId`), and settles the request on the terminal frame. Remote install/pull run for the operation's full duration (no broker-imposed timeout); remote stop and model unload/delete are fast request/response, remote Ollama and vLLM model loads use the readiness-sized response budget (a vLLM load restarts the engine on the new model), and remote start waits for the engine's bounded readiness result. They error if this node isn't clustered or the target isn't a pinned peer. ## 8. Dependencies - **Upstream**: the orchestrator that spawns it and forwards `engine:*` (`nvpair-ui-broker`). diff --git a/services/nvpair-engine-manager/status.go b/services/nvpair-engine-manager/status.go index 322e02903..c4c93e54f 100644 --- a/services/nvpair-engine-manager/status.go +++ b/services/nvpair-engine-manager/status.go @@ -59,15 +59,16 @@ func (e *Executor) Status(engine string) (EngineStatus, error) { func (e *Executor) StatusAtPort(engine string, probePort int) (EngineStatus, error) { st, err := e.state(engine) if err != nil { - // Known engine with no host-platform block: return a shell status - // (matching get-installed) rather than erroring, so status is - // consistent for the same engine across methods. - if m, ok := e.reg.Get(engine); ok { - return EngineStatus{Engine: engine, DisplayName: m.DisplayName}, nil - } return EngineStatus{}, err } - st.opMu.Lock() + // A start, stop, install or model switch holds opMu for as long as it takes, + // and an engine like vLLM takes a minute or more to become ready. A status + // read must not wait that out: callers give it seconds, and the desktop + // treats it as a fast read. While an operation runs, what is already known + // about the engine is the answer, and the operation reports its own result. + if !st.opMu.TryLock() { + return e.snapshot(engine, st), nil + } defer st.opMu.Unlock() pathInstalled, _ := e.Detect(engine) st.mu.Lock() @@ -80,22 +81,22 @@ func (e *Executor) StatusAtPort(engine string, probePort int) (EngineStatus, err return e.snapshot(engine, st), nil } -// GetInstalled lists every known engine with its current status. +// GetInstalled lists every engine this host can run with its current status. func (e *Executor) GetInstalled() []EngineStatus { var out []EngineStatus for _, name := range e.reg.Names() { st, err := e.state(name) if err != nil { - // Known engine with no block for this host: surface a shell - // status so the UI can still list it as unavailable here. - dn := name - if m, ok := e.reg.Get(name); ok { - dn = m.DisplayName - } - out = append(out, EngineStatus{Engine: name, DisplayName: dn}) + // No block for this host, such as vLLM on Windows: the engine + // cannot be installed or run here, so it is not offered at all. + continue + } + // Same reasoning as StatusAtPort: one engine mid-start must not hold up + // the listing of every other engine. + if !st.opMu.TryLock() { + out = append(out, e.snapshot(name, st)) continue } - st.opMu.Lock() pathInstalled, _ := e.Detect(name) st.mu.Lock() port := st.port @@ -129,6 +130,7 @@ func (e *Executor) snapshot(engine string, st *engineState) EngineStatus { DisplayName: st.manifest.DisplayName, Installed: st.installed, Running: st.running, + Starting: st.starting && !st.running, Healthy: st.healthy, Port: st.port, } diff --git a/services/nvpair-engine-manager/testdata/fakeengine/main.go b/services/nvpair-engine-manager/testdata/fakeengine/main.go index e58973ce8..a86c364ab 100644 --- a/services/nvpair-engine-manager/testdata/fakeengine/main.go +++ b/services/nvpair-engine-manager/testdata/fakeengine/main.go @@ -81,12 +81,32 @@ func main() { os.Exit(2) } } + // "--fail-launch", or a model named org/fails-to-load, stands in for a launch + // the engine rejects, such as a model that does not fit in GPU memory. for _, arg := range os.Args[1:] { - if arg == "--fail-launch" { + if arg == "--fail-launch" || arg == "org/fails-to-load" { fmt.Fprintln(os.Stderr, "invalid launch", os.Args[1:]) os.Exit(2) } } + // Like vLLM, a server started with --model serves that one model and no + // other. + for i, arg := range os.Args[1:] { + served, ok := strings.CutPrefix(arg, "--model=") + if arg == "--model" && i+2 < len(os.Args) { + served, ok = os.Args[i+2], true + } + if ok { + models = map[string]bool{served: true} + loaded = map[string]bool{served: true} + } + } + // FAKE_SERVE_MODEL stands in for a server someone started on another model + // than the saved launch arguments name, which PAIR then adopted. + if served := os.Getenv("FAKE_SERVE_MODEL"); served != "" { + models = map[string]bool{served: true} + loaded = map[string]bool{served: true} + } if path := os.Getenv("FAKE_START_LOG"); path != "" { file, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600) if err != nil { @@ -287,11 +307,17 @@ func main() { w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(map[string]any{"models": ms}) }) - // LM Studio-shape (OpenAI-compatible). + // LM Studio-shape (OpenAI-compatible). Each row also carries vLLM's root + // (the model it loaded). FAKE_SERVED_MODEL_NAME stands in for vLLM's + // --served-model-name: the id becomes that alias while root keeps the model. mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) { data := make([]map[string]string, 0) for _, n := range modelNames() { - data = append(data, map[string]string{"id": n, "object": "model"}) + id := n + if alias := os.Getenv("FAKE_SERVED_MODEL_NAME"); alias != "" { + id = alias + } + data = append(data, map[string]string{"id": id, "root": n, "object": "model"}) } w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": data}) diff --git a/services/nvpair-engine-manager/toolparser.go b/services/nvpair-engine-manager/toolparser.go new file mode 100644 index 000000000..4f06e8efa --- /dev/null +++ b/services/nvpair-engine-manager/toolparser.go @@ -0,0 +1,143 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "context" + "crypto/sha256" + _ "embed" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "slices" + "sync" + "time" +) + +// toolParserProbeScript asks the installed vLLM which of its tool-call parsers +// read a model's tool-call format. See the script for how. +// +//go:embed toolparser_probe.py +var toolParserProbeScript string + +// toolParserProbeTimeout bounds one probe. It imports transformers and vLLM's +// parsers, which takes a few seconds once their files are in the page cache. +const toolParserProbeTimeout = 2 * time.Minute + +// toolParserProbe is what the probe found for one model. +type toolParserProbe struct { + // HF is true when the tokenizer ships response_template metadata, which + // vLLM's hf parser reads. + HF bool `json:"hf"` + // Matches are the parsers that extracted the model's example tool call. + Matches []string `json:"matches"` +} + +// probeToolParsers runs the probe with the Python next to the vLLM binary. It +// is a variable so tests can supply results without a vLLM install. +var probeToolParsers = runToolParserProbe + +// toolParserCache keeps each probe's result for the life of the service, keyed +// by vLLM binary, snapshot directory and chat template (with a template file's +// contents), so a model is probed once rather than on every start, wake and +// switch. A different vLLM, model revision or template is a different key. +var toolParserCache sync.Map + +// preferredToolParsers breaks ties when several parsers read a format equally +// well: Qwen's JSON is read alike by hermes, ernie45, granite4 and +// qwen3_json, and its XML form by qwen3_coder, qwen3_xml, seed_oss and step3p5. +// These are vLLM's general-purpose parsers for the common formats, tried in +// this order; any other match is taken in name order. +var preferredToolParsers = []string{"hermes", "qwen3_coder", "llama3_json", "mistral", "pythonic"} + +// chooseToolParser picks the parser to launch with, or "" when none fits. +func chooseToolParser(p toolParserProbe) string { + if p.HF { + return "hf" + } + for _, name := range preferredToolParsers { + if slices.Contains(p.Matches, name) { + return name + } + } + if len(p.Matches) == 0 { + return "" + } + return slices.Min(p.Matches) +} + +// vllmPython is the interpreter of the environment vLLM is installed in. +func vllmPython(bin string) (string, error) { + if resolved, err := filepath.EvalSymlinks(bin); err == nil { + bin = resolved + } + name := "python" + if runtime.GOOS == "windows" { + name = "python.exe" + } + python := filepath.Join(filepath.Dir(bin), name) + if _, err := os.Stat(python); err != nil { + return "", fmt.Errorf("no Python next to %s", bin) + } + return python, nil +} + +// runToolParserProbe runs the probe under ctx, so a start that is cancelled +// (Stop, a remote caller giving up, service shutdown) stops it too. tokenizer +// is the directory of the tokenizer vLLM loads, and template the launch's +// --chat-template, or "" for the tokenizer's own. +func runToolParserProbe(ctx context.Context, bin, tokenizer, template string) (toolParserProbe, error) { + python, err := vllmPython(bin) + if err != nil { + return toolParserProbe{}, err + } + ctx, cancel := context.WithTimeout(ctx, toolParserProbeTimeout) + defer cancel() + // -I keeps the working directory and user site off the import path. + args := []string{"-I", "-c", toolParserProbeScript, tokenizer} + if template != "" { + args = append(args, template) + } + cmd := exec.CommandContext(ctx, python, args...) + cmd.Dir = os.TempDir() + out, err := cmd.Output() + if err != nil { + return toolParserProbe{}, fmt.Errorf("tool-call parser probe: %w", err) + } + lines := bytes.Split(bytes.TrimSpace(out), []byte("\n")) + var probe toolParserProbe + if err := json.Unmarshal(lines[len(lines)-1], &probe); err != nil { + return toolParserProbe{}, errors.New("tool-call parser probe printed no result") + } + return probe, nil +} + +// detectToolParser returns the parser for the tokenizer in the tokenizer +// directory, rendered with template when the launch overrides it, probing it +// the first time. A failed probe is not cached, so the next start tries again. +func detectToolParser(ctx context.Context, bin, tokenizer, template string) (string, []string, error) { + key := bin + "\x00" + tokenizer + "\x00" + template + // A --chat-template file can be edited between starts under the same path, + // so its contents are part of the key. vLLM reads the file on every start. + if raw, err := os.ReadFile(template); template != "" && err == nil { + sum := sha256.Sum256(raw) + key += "\x00" + hex.EncodeToString(sum[:]) + } + if cached, ok := toolParserCache.Load(key); ok { + probe := cached.(toolParserProbe) + return chooseToolParser(probe), probe.Matches, nil + } + probe, err := probeToolParsers(ctx, bin, tokenizer, template) + if err != nil { + return "", nil, err + } + toolParserCache.Store(key, probe) + return chooseToolParser(probe), probe.Matches, nil +} diff --git a/services/nvpair-engine-manager/toolparser_probe.py b/services/nvpair-engine-manager/toolparser_probe.py new file mode 100644 index 000000000..938f34f0d --- /dev/null +++ b/services/nvpair-engine-manager/toolparser_probe.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Find the vLLM tool-call parsers that read a model's tool-call format. + +nvpair-engine-manager runs this with the Python of the vLLM it is about to +launch, passing the directory of the tokenizer vLLM will load and, when the +launch overrides it, the --chat-template value: + + python -I -c [] + +A model's chat template writes an earlier tool call back into the conversation +in the same format the model uses when it calls a tool. So the template is +rendered with one example call, and that text is offered to every parser the +installed vLLM provides; a parser matches when it extracts exactly that call. + +The last line printed is JSON: {"hf": bool, "matches": [parser names]}, where +hf means the tokenizer ships response_template metadata for vLLM's hf parser. +Nothing is downloaded, and no code from the model repository runs. +""" + +import json +import os +import sys + +os.environ["HF_HUB_OFFLINE"] = "1" +os.environ["TRANSFORMERS_OFFLINE"] = "1" + +ARGS = {"city": "Paris"} +TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather in a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} +USER = {"role": "user", "content": "What is the weather in Paris?"} +RESULT = {"role": "tool", "tool_call_id": "call_1", "content": "sunny"} + + +def rendered_call(tok): + """Return the example call as the template writes it, or "" if it does not.""" + try: + prompt = tok.apply_chat_template([USER], tools=[TOOL], add_generation_prompt=True, tokenize=False) + base = tok.apply_chat_template([USER], tools=[TOOL], tokenize=False) + except Exception: + return "" + # Templates differ on whether arguments are an object or a JSON string. + for arguments in (ARGS, json.dumps(ARGS)): + call = { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": arguments}}], + } + # Some templates render a final assistant turn as the generation prompt + # and drop its calls; a tool result after it makes them write the call. + for conversation in ([USER, call], [USER, call, RESULT]): + try: + full = tok.apply_chat_template(conversation, tools=[TOOL], tokenize=False) + except Exception: + continue + for prefix in (prompt, base, os.path.commonprefix([prompt, full])): + if full.startswith(prefix): + text = full[len(prefix):] + break + if "get_weather" in text and "Paris" in text: + return text + return "" + + +def main(): + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(sys.argv[1], trust_remote_code=False) + if len(sys.argv) > 2: + # The template the launch names replaces the tokenizer's, read the way + # vLLM reads --chat-template: a file path or the template itself. + try: + from vllm.renderers.chat_utils import load_chat_template + except ImportError: + from vllm.entrypoints.chat_utils import load_chat_template + tok.chat_template = load_chat_template(sys.argv[2]) + + hf = False + try: + from vllm.parser.response_template import validate_tokenizer_response_template + + validate_tokenizer_response_template(tok, reasoning=False, tools=True) + hf = True + except Exception: + pass + + matches = [] + text = rendered_call(tok) + if text: + try: + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + except ImportError: + from vllm.entrypoints.openai.protocol import ChatCompletionRequest + from vllm.tool_parsers import ToolParserManager + + request = ChatCompletionRequest(model="probe", messages=[USER], tools=[TOOL]) + for name in sorted(ToolParserManager.list_registered()): + if name == "hf": + continue + try: + parser = ToolParserManager.get_tool_parser(name)(tok, tools=request.tools) + result = parser.extract_tool_calls(text, request) + calls = result.tool_calls if result.tools_called else [] + if ( + len(calls) == 1 + and calls[0].function.name == "get_weather" + and json.loads(calls[0].function.arguments) == ARGS + ): + matches.append(name) + except Exception: + continue + + print(json.dumps({"hf": hf, "matches": matches})) + + +main() diff --git a/services/nvpair-engine-manager/toolparser_test.go b/services/nvpair-engine-manager/toolparser_test.go new file mode 100644 index 000000000..180a0f28e --- /dev/null +++ b/services/nvpair-engine-manager/toolparser_test.go @@ -0,0 +1,157 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// withToolParserProbe replaces the probe of the installed vLLM, and empties the +// cache of earlier results, for one test. +func withToolParserProbe(t *testing.T, probe func(bin, snapshot string) (toolParserProbe, error)) { + t.Helper() + old := probeToolParsers + probeToolParsers = func(_ context.Context, bin, snapshot, _ string) (toolParserProbe, error) { + return probe(bin, snapshot) + } + toolParserCache.Clear() + t.Cleanup(func() { + probeToolParsers = old + toolParserCache.Clear() + }) +} + +func TestChooseToolParser_PrefersGeneralPurposeParsersAmongEqualMatches(t *testing.T) { + test := func(name string, probe toolParserProbe, want string) { + t.Run(name, func(t *testing.T) { + assert.Equal(t, want, chooseToolParser(probe)) + }) + } + test("Qwen JSON in tool_call tags", toolParserProbe{Matches: []string{"ernie45", "granite4", "hermes", "qwen3_json"}}, "hermes") + test("Qwen XML function tags", toolParserProbe{Matches: []string{"qwen3_coder", "qwen3_xml", "seed_oss", "step3p5"}}, "qwen3_coder") + test("only uncommon parsers", toolParserProbe{Matches: []string{"step3p5", "seed_oss"}}, "seed_oss") + test("response template wins", toolParserProbe{HF: true, Matches: []string{"hermes"}}, "hf") + test("nothing matched", toolParserProbe{}, "") +} + +func TestDetectToolParser_ProbesEachModelOnce(t *testing.T) { + probes := 0 + withToolParserProbe(t, func(bin, snapshot string) (toolParserProbe, error) { + probes++ + return toolParserProbe{Matches: []string{"hermes"}}, nil + }) + + for range 3 { + parser, _, err := detectToolParser(t.Context(), "vllm", "/hub/a", "") + require.NoError(t, err) + require.Equal(t, "hermes", parser) + } + require.Equal(t, 1, probes, "the same model is probed once") + _, _, err := detectToolParser(t.Context(), "vllm", "/hub/c", "") + require.NoError(t, err) + assert.Equal(t, 2, probes, "another model must be probed on its own") +} + +// Editing a --chat-template file and restarting must probe the new template, +// not reuse what the old one matched. +func TestDetectToolParser_ProbesAgainWhenTheTemplateFileChanges(t *testing.T) { + probes := 0 + withToolParserProbe(t, func(bin, snapshot string) (toolParserProbe, error) { + probes++ + return toolParserProbe{Matches: []string{"hermes"}}, nil + }) + template := filepath.Join(t.TempDir(), "chat.jinja") + require.NoError(t, os.WriteFile(template, []byte("{{ messages }}"), 0o644)) + + for range 2 { + _, _, err := detectToolParser(t.Context(), "vllm", "/hub/a", template) + require.NoError(t, err) + } + require.Equal(t, 1, probes, "an unchanged template is probed once") + + require.NoError(t, os.WriteFile(template, []byte("{{ messages }}{{ tools }}"), 0o644)) + _, _, err := detectToolParser(t.Context(), "vllm", "/hub/a", template) + require.NoError(t, err) + assert.Equal(t, 2, probes, "the edited template is probed again") +} + +func TestDetectToolParser_RetriesAFailedProbe(t *testing.T) { + probes := 0 + withToolParserProbe(t, func(bin, snapshot string) (toolParserProbe, error) { + probes++ + if probes == 1 { + return toolParserProbe{}, errors.New("timed out") + } + return toolParserProbe{Matches: []string{"hermes"}}, nil + }) + + _, _, err := detectToolParser(t.Context(), "vllm", "/hub/a", "") + require.Error(t, err, "the first probe's failure was not reported") + + parser, _, err := detectToolParser(t.Context(), "vllm", "/hub/a", "") + require.NoError(t, err, "the next start must probe again") + assert.Equal(t, "hermes", parser) +} + +// The probe run against the real vLLM install and the models in the local +// Hugging Face cache. NVPAIR_LIVE_VLLM_BIN names the vllm binary, and +// NVPAIR_LIVE_VLLM_TOOL_MODELS lists model=parser pairs to check, with an empty +// parser for a model no parser reads, for example +// "Qwen/Qwen3-8B=hermes,nvidia/NVIDIA-Nemotron-Nano-9B-v2=". +func TestToolParserProbe_FindsTheParserWithTheInstalledVLLM(t *testing.T) { + bin := os.Getenv("NVPAIR_LIVE_VLLM_BIN") + if bin == "" { + t.Skip("set NVPAIR_LIVE_VLLM_BIN and NVPAIR_LIVE_VLLM_TOOL_MODELS to probe real models with the installed vLLM") + } + home, err := os.UserHomeDir() + require.NoError(t, err) + hub := filepath.Join(home, ".cache", "huggingface", "hub") + for _, pair := range strings.Split(os.Getenv("NVPAIR_LIVE_VLLM_TOOL_MODELS"), ",") { + model, want, _ := strings.Cut(pair, "=") + t.Run(model, func(t *testing.T) { + found, ok := findLaunchModel([]string{"--model", model}, hub, nil) + require.True(t, ok, "%s is not in %s", model, hub) + + probe, err := runToolParserProbe(t.Context(), bin, found.snapshot, "") + + require.NoError(t, err) + assert.Equal(t, want, chooseToolParser(probe), "chosen from %+v", probe) + }) + } +} + +// The probe runs under the start's context, so cancelling the start (Stop, a +// caller giving up, shutdown) stops the probe instead of leaving it to run out +// its own timeout. +func TestDetectToolParser_StopsTheProbeWhenTheStartIsCancelled(t *testing.T) { + old := probeToolParsers + t.Cleanup(func() { + probeToolParsers = old + toolParserCache.Clear() + }) + toolParserCache.Clear() + probeToolParsers = func(ctx context.Context, bin, snapshot, template string) (toolParserProbe, error) { + select { + case <-ctx.Done(): + return toolParserProbe{}, ctx.Err() + case <-time.After(5 * time.Second): + return toolParserProbe{}, errors.New("the probe ran on after its start was cancelled") + } + } + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, _, err := detectToolParser(ctx, "vllm", "/hub/cancelled", "") + + require.ErrorIs(t, err, context.Canceled) +} diff --git a/services/nvpair-engine-manager/vllm_manifest_test.go b/services/nvpair-engine-manager/vllm_manifest_test.go new file mode 100644 index 000000000..be7f6ac0d --- /dev/null +++ b/services/nvpair-engine-manager/vllm_manifest_test.go @@ -0,0 +1,224 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// vLLM is Linux (NVIDIA) and Apple Silicon only, installs through a pinned uv +// bootstrap, serves exactly one model per process and has no load, unload, pull or delete API. These +// are the properties the rest of the stack (UI capabilities, the model wire, +// the install gate) relies on, so a manifest edit that drops one fails here. +func TestBundledVLLMManifest_PinsTheInstallAndSizesTheLaunchOnEachPlatform(t *testing.T) { + reg := loadWithOverrides(t, t.TempDir()) + manifest, ok := reg.Get("vllm") + require.True(t, ok, "vllm manifest not loaded") + require.Len(t, manifest.Platforms, 3, "want linux/amd64, linux/arm64 and darwin/arm64 only") + + // wantPin is the version the platform's install script must pin: Linux + // installs vLLM itself from PyPI, macOS installs the vLLM core and + // vLLM-Metal wheels. + test := func(key, wantPin string) { + t.Run(key, func(t *testing.T) { + platform, ok := manifest.Platforms[key] + require.True(t, ok, "vllm has no %s platform block", key) + + // Our managed install must win over an external one. + require.GreaterOrEqual(t, len(platform.Detect), 2, "detect must list the managed install and an external one") + assert.True(t, strings.HasPrefix(platform.Detect[0], "{install_dir}/"), "detect order = %v, want the managed install first", platform.Detect) + require.NotNil(t, platform.Install) + require.NotNil(t, platform.Install.Fetch, "the uv bootstrap must be a fetch") + require.NotEmpty(t, platform.Install.Fetch.SHA256, "the uv bootstrap must be sha256-pinned") + assert.Empty(t, platform.Install.Script, "install must not use the unpinned script escape hatch") + script := strings.Join(platform.Install.Run, " ") + assert.Contains(t, script, wantPin, "the vLLM install must be pinned") + // vLLM 0.32 asks only for mistral_common>=1.11.6. Release 1.13 made an + // argument of get_lark_from_jinja required that vLLM 0.32 does not pass, + // so with the Mistral tool-call parser every request to a Mistral model + // fails. vLLM 0.33 makes the same call; drop the pin only with a vLLM + // release that itself requires mistral_common 1.13. + assert.Contains(t, script, "mistral_common==1.12.0", "the install must pin mistral_common==1.12.0 for vLLM 0.32") + if strings.HasPrefix(key, "darwin/") { + // The macOS wheels are downloaded by the script itself, so each one + // has to be checked against a digest before it is installed. + assert.Equal(t, 2, strings.Count(script, "shasum -a 256 -c"), "want one wheel digest check for the vLLM core and one for vLLM-Metal") + } + + rt := platform.Runtime + require.NotNil(t, rt.EditableLaunch) + assert.Equal(t, []string{"serve"}, rt.EditableLaunch.FixedArgs, "launch must start with the fixed serve subcommand") + assert.Equal(t, "~/.cache/huggingface", platform.ModelsDir, "models_dir must be the Hugging Face cache") + // vLLM compiles FlashInfer sampling kernels at first start and that build + // cannot cope with the spaces in PAIR's app-data path + // ("Nvidia Corporation/Personal AI Router"), so it must stay off. It also + // keeps a system CUDA toolkit and C++ compiler from being required. + assert.Equal(t, "0", rt.Env["VLLM_USE_FLASHINFER_SAMPLER"], "the FlashInfer sampler must be disabled") + // vLLM reserves 92% of GPU memory by default and refuses a context above + // the model's own limit. Both are decided per model at start, so the + // manifest must not pin them in the saved launch. + require.NotNil(t, rt.AutoLaunch, "vllm must size its context and GPU share automatically") + assert.Equal(t, autoLaunchHFContext, rt.AutoLaunch.Kind) + assert.Equal(t, 65536, rt.AutoLaunch.ContextTokens, "vllm must aim for a 64K window") + // vLLM's default of 256 concurrent requests does not fit a hybrid model + // such as the default Qwen3.5-4B on a 16 GiB GPU. + assert.Equal(t, 32, rt.AutoLaunch.MaxNumSeqs, "vllm must run at most 32 requests at a time") + assert.NotContains(t, rt.Args, "--gpu-memory-utilization", "must not be pinned in the default launch") + assert.NotContains(t, rt.Args, "--max-model-len", "must not be pinned in the default launch") + assert.Equal(t, 8001, rt.Port, "the facade owns 8000") + }) + } + test("linux/amd64", "vllm==") + test("linux/arm64", "vllm==") + test("darwin/arm64", "vllm-0.32.0") +} + +func TestBundledVLLMManifest_DeclaresTheModelActionsTheStackReliesOn(t *testing.T) { + reg := loadWithOverrides(t, t.TempDir()) + manifest, ok := reg.Get("vllm") + require.True(t, ok, "vllm manifest not loaded") + + assert.Contains(t, manifest.Actions, "list_models") + assert.Contains(t, manifest.Actions, "loaded_models") + // list_models is what is on disk; only the model being served is advertised. + // Listing, Delete and Pull use the hub cache vLLM itself uses, resolved from + // its launch environment ({hf_hub_cache}). + if list := manifest.Actions["list_models"]; assert.NotNil(t, list.HFCache, "list_models must list the Hugging Face cache") { + assert.Equal(t, "{hf_hub_cache}", list.HFCache.Dir) + } + assert.Equal(t, advertisedModelsLoaded, manifest.AdvertisedModels, "peers and the proxy must only be offered the served model") + if loaded := manifest.Actions["loaded_models"]; assert.NotNil(t, loaded.HTTP, "loaded_models must read vLLM's /v1/models") { + assert.Equal(t, "/v1/models", loaded.HTTP.Path) + } + // Switching the model is a launch-argument change plus a restart, since the + // engine has no load call. + if load := manifest.Actions["load_model"]; assert.NotNil(t, load.LaunchModel, "load_model must switch the model through a launch argument") { + assert.Equal(t, "--model", load.LaunchModel.Flag) + } + // vLLM cannot drop its model and keep running, so Eject stops the engine. + assert.True(t, manifest.Actions["unload_model"].StopEngine, "unload_model must stop the engine") + // Deleting is a removal from the shared cache that needs no running engine. + if del := manifest.Actions["delete_model"]; assert.NotNil(t, del.HFCache, "delete_model must remove from the hub cache") { + assert.True(t, del.HFCache.Delete) + assert.Equal(t, "{hf_hub_cache}", del.HFCache.Dir) + } + // Downloads go through the hf tool that ships in the engine's virtualenv. + if pull := manifest.Actions["pull_model"]; assert.NotNil(t, pull.HFPull, "pull_model must download through hf into the hub cache") { + assert.Equal(t, "hf", pull.HFPull.Tool) + assert.Equal(t, "{hf_hub_cache}", pull.HFPull.Dir) + } +} + +// The vLLM wheel is a CUDA 13 build, so an older driver installs cleanly and +// then fails at startup, after a multi-GB download. The install script has to +// refuse before downloading anything. This runs the manifest's own script, with +// its placeholders resolved the way the executor resolves them, against a fake +// nvidia-smi. Recent drivers print "CUDA UMD Version" and mark the older +// "CUDA Version" field deprecated, so both spellings have to be understood. +func TestVLLMInstall_RefusesDriversWithoutCUDA13(t *testing.T) { + if runtime.GOOS != "linux" { + t.Skip("the vLLM install script is Linux only") + } + tools := map[string]string{} + for _, name := range []string{"sh", "awk", "df", "mkdir", "tar"} { + path, err := exec.LookPath(name) + if err != nil { + t.Skipf("%s is not available", name) + } + tools[name] = path + } + reg := loadWithOverrides(t, t.TempDir()) + manifest, ok := reg.Get("vllm") + require.True(t, ok, "vllm manifest not loaded") + + // runInstall runs the Linux install script with only the tools above and, + // unless smiOutput is nil, a fake nvidia-smi that prints it. + runInstall := func(t *testing.T, smiOutput *string) (text string, runErr error, installDir string) { + t.Helper() + bin := t.TempDir() + for name, path := range tools { + if name != "sh" { + require.NoError(t, os.Symlink(path, filepath.Join(bin, name))) + } + } + if smiOutput != nil { + script := "#!/bin/sh\nprintf '%s' '" + *smiOutput + "'\n" + require.NoError(t, os.WriteFile(filepath.Join(bin, "nvidia-smi"), []byte(script), 0o755)) + } + installDir = t.TempDir() + argv, err := resolveArgs(manifest.Platforms["linux/amd64"].Install.Run, map[string]string{ + "download": filepath.Join(installDir, "missing-uv.tar.gz"), + "install_dir": installDir, + }) + require.NoError(t, err) + // The executor expands OS environment references in every install + // argument after resolving placeholders (install.go), so "$name" in + // the script would be replaced before the shell ever saw it. Apply + // the same expansion, and require it to be a no-op: only positional + // parameters and "${name:-default}" forms survive it. + for i := range argv { + require.Equal(t, argv[i], expandPath(argv[i]), "the executor rewrites install argument %d, which breaks the script", i) + } + cmd := exec.Command(tools["sh"], argv[1:]...) + cmd.Env = []string{"PATH=" + bin} + out, runErr := cmd.CombinedOutput() + return string(out), runErr, installDir + } + accepted := func(name, smiOutput string) { + t.Run(name, func(t *testing.T) { + text, _, _ := runInstall(t, &smiOutput) + + // The driver check passed. The script goes on to fail later, + // because the archive does not exist; it must not blame the driver. + assert.NotContains(t, text, "NVIDIA driver", "a CUDA 13 driver was refused") + }) + } + refused := func(name string, smiOutput *string, wantInText string) { + t.Run(name, func(t *testing.T) { + text, runErr, installDir := runInstall(t, smiOutput) + + require.Error(t, runErr, "install was not refused: %s", text) + assert.Contains(t, text, wantInText) + assert.NoFileExists(t, filepath.Join(installDir, "uv"), "the script started installing before it refused") + }) + } + accepted("current driver format", "CUDA Version : 13.4 [Deprecated; will be removed in CUDA 14.0. Use CUDA UMD Version instead]\nCUDA UMD Version : 13.4\n") + accepted("older CUDA 13 driver format", "CUDA Version : 13.0\n") + accepted("UMD only", "CUDA UMD Version : 14.0\n") + refused("no nvidia-smi", nil, "nvidia-smi not found") + refused("old driver", ptr("CUDA Version : 12.4\n"), "CUDA 13 or newer (found: 12)") + refused("unrecognised output", ptr("Driver Version : 1.2.3\n"), "CUDA 13 or newer (found: none)") +} + +// The executor expands OS environment references in every install argument +// after resolving placeholders, so a "$name" in a script would be replaced +// before the shell ran it. Each platform's script must survive that untouched. +func TestVLLMInstallScripts_SurviveArgumentExpansion(t *testing.T) { + reg := loadWithOverrides(t, t.TempDir()) + manifest, ok := reg.Get("vllm") + require.True(t, ok, "vllm manifest not loaded") + for key, platform := range manifest.Platforms { + t.Run(key, func(t *testing.T) { + argv, err := resolveArgs(platform.Install.Run, map[string]string{ + "download": "/tmp/uv.tar.gz", + "install_dir": "/tmp/vllm", + }) + + require.NoError(t, err) + for i := range argv { + assert.Equal(t, argv[i], expandPathForOS(argv[i], "linux"), "the executor rewrites install argument %d, which breaks the script", i) + } + }) + } +} + +func ptr(s string) *string { return &s } diff --git a/services/nvpair-engine-manager/weightformat.go b/services/nvpair-engine-manager/weightformat.go new file mode 100644 index 000000000..2a44476cc --- /dev/null +++ b/services/nvpair-engine-manager/weightformat.go @@ -0,0 +1,201 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "regexp" + "slices" + "strings" +) + +// hfQuantConfigFile is where NVIDIA ModelOpt checkpoints declare their +// quantization; vLLM reads it beside config.json. +const hfQuantConfigFile = "hf_quant_config.json" + +// checkpointQuantization is the part of config.json and hf_quant_config.json +// that says how a checkpoint's weights and KV cache are stored. +type checkpointQuantization struct { + ModelType string `json:"model_type"` + QuantizationConfig *struct { + QuantMethod string `json:"quant_method"` + QuantAlgo string `json:"quant_algo"` + ConfigGroups map[string]struct { + Weights *floatScheme `json:"weights"` + } `json:"config_groups"` + KVCacheScheme *floatScheme `json:"kv_cache_scheme"` + } `json:"quantization_config"` + // Quantization is ModelOpt's hf_quant_config.json body. + Quantization *struct { + QuantAlgo string `json:"quant_algo"` + KVCacheQuantAlgo string `json:"kv_cache_quant_algo"` + } `json:"quantization"` +} + +// floatScheme is a compressed-tensors scheme: FP8 is type float, 8 bits. +type floatScheme struct { + Type string `json:"type"` + NumBits int `json:"num_bits"` +} + +func (s *floatScheme) fp8() bool { return s != nil && s.Type == "float" && s.NumBits == 8 } + +// weightFormat returns "fp8" for a checkpoint whose weights or KV cache vLLM +// would load as FP8, and its model type. It reads the three ways checkpoints +// say so: quant_method fp8 (or fbgemm_fp8) in config.json, a compressed-tensors +// float 8-bit scheme, and ModelOpt's FP8 quant_algo or kv_cache_quant_algo, +// in config.json or hf_quant_config.json. hfQuant is nil when there is none. +func weightFormat(config, hfQuant []byte) (format, modelType string, err error) { + var cfg checkpointQuantization + if err := json.Unmarshal(config, &cfg); err != nil { + return "", "", fmt.Errorf("read config.json: %w", err) + } + if len(hfQuant) > 0 { + var modelopt checkpointQuantization + if err := json.Unmarshal(hfQuant, &modelopt); err != nil { + return "", "", fmt.Errorf("read %s: %w", hfQuantConfigFile, err) + } + cfg.Quantization = modelopt.Quantization + } + isFP8 := func(algo string) bool { return strings.Contains(strings.ToUpper(algo), "FP8") } + if q := cfg.Quantization; q != nil && (isFP8(q.QuantAlgo) || isFP8(q.KVCacheQuantAlgo)) { + return weightFormatFP8, cfg.ModelType, nil + } + if q := cfg.QuantizationConfig; q != nil { + switch strings.ToLower(q.QuantMethod) { + case "fp8", "fbgemm_fp8": + return weightFormatFP8, cfg.ModelType, nil + case "modelopt": + if isFP8(q.QuantAlgo) { + return weightFormatFP8, cfg.ModelType, nil + } + case "compressed-tensors": + if q.KVCacheScheme.fp8() { + return weightFormatFP8, cfg.ModelType, nil + } + for _, group := range q.ConfigGroups { + if group.Weights.fp8() { + return weightFormatFP8, cfg.ModelType, nil + } + } + } + } + return "", cfg.ModelType, nil +} + +// refuses reports whether the rule refuses a checkpoint of this format and +// model type. +func (u *UnsupportedWeights) refuses(format, modelType string) bool { + return u != nil && format != "" && slices.Contains(u.Formats, format) && !slices.Contains(u.ExceptModelTypes, modelType) +} + +// checkSnapshotWeights refuses the model in snapshot when rule names its +// format. A snapshot whose config cannot be read is let through: the engine +// then reports its own error. +func checkSnapshotWeights(rule *UnsupportedWeights, model, snapshot string) error { + if rule == nil { + return nil + } + config, err := os.ReadFile(filepath.Join(snapshot, "config.json")) + if err != nil { + return nil + } + hfQuant, err := os.ReadFile(filepath.Join(snapshot, hfQuantConfigFile)) + if err != nil { + hfQuant = nil + } + return checkWeights(rule, model, config, hfQuant) +} + +func checkWeights(rule *UnsupportedWeights, model string, config, hfQuant []byte) error { + format, modelType, err := weightFormat(config, hfQuant) + if err != nil || !rule.refuses(format, modelType) { + return nil + } + return fmt.Errorf("%s is an FP8 model: %s", model, rule.Reason) +} + +// checkLaunchWeights refuses a start whose model the platform's engine cannot +// load, reading the model at the revision vLLM loads, from the cache or, for a +// model this start would download, through fetch. A model whose files cannot +// be read is let through. +func checkLaunchWeights(rule *UnsupportedWeights, args []string, hubDir string, fetch hfModelFetcher) error { + if rule == nil { + return nil + } + model, ok := launchFlagValue(args, "--model") + if !ok || model == "" { + return nil + } + revision := hfRevision(args) + snapshot, ok := hfSnapshotDir(hubDir, model, revision) + if !ok { + pending, fetched := fetch() + if !fetched { + return nil + } + snapshot = pending.snapshot + } + return checkSnapshotWeights(rule, model, snapshot) +} + +// checkRepoWeights refuses a download the platform's engine could not load, +// from the repository's config.json and hf_quant_config.json at main. A +// repository whose config cannot be read is let through. +func (e *Executor) checkRepoWeights(ctx context.Context, rule *UnsupportedWeights, id string, files []string, env hfEnv) error { + if rule == nil { + return nil + } + config, err := e.hfGet(ctx, env.endpoint()+"/"+id+"/resolve/main/config.json", maxHFConfigBody, env) + if err != nil { + return nil + } + var hfQuant []byte + if slices.Contains(files, hfQuantConfigFile) { + if hfQuant, err = e.hfGet(ctx, env.endpoint()+"/"+id+"/resolve/main/"+hfQuantConfigFile, maxHFConfigBody, env); err != nil { + hfQuant = nil + } + } + return checkWeights(rule, id, config, hfQuant) +} + +// unsupportedWeightsFor is the weight rule an engine's manifest declares for a +// target platform, or nil. +func (e *Executor) unsupportedWeightsFor(engine, platform, arch string) *UnsupportedWeights { + m, ok := e.reg.Get(engine) + if !ok { + return nil + } + p, ok := m.Platforms[platform+"/"+arch] + if !ok { + return nil + } + return p.Runtime.UnsupportedWeights +} + +// fp8NameRe matches FP8 as a token of a repository id, as in +// nvidia/Llama-3.1-8B-Instruct-FP8. +var fp8NameRe = regexp.MustCompile(`(?i)(^|[-_./])fp8($|[-_.])`) + +// withoutUnsupportedWeights drops catalog entries a rule refuses. A listing +// carries no config, so FP8 is read from the repository's tags and name; a +// model type the rule excepts cannot be told from a listing, so those entries +// are dropped too and stay available to a pull by id. +func withoutUnsupportedWeights(models []CatalogModel, rule *UnsupportedWeights) []CatalogModel { + if rule == nil || !slices.Contains(rule.Formats, weightFormatFP8) { + return models + } + out := make([]CatalogModel, 0, len(models)) + for _, m := range models { + if fp8NameRe.MatchString(m.ID) || slices.ContainsFunc(m.Tags, func(tag string) bool { return strings.EqualFold(tag, weightFormatFP8) }) { + continue + } + out = append(out, m) + } + return out +} diff --git a/services/nvpair-engine-manager/weightformat_test.go b/services/nvpair-engine-manager/weightformat_test.go new file mode 100644 index 000000000..aecfcccbf --- /dev/null +++ b/services/nvpair-engine-manager/weightformat_test.go @@ -0,0 +1,259 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// The checkpoints below are trimmed to the fields that decide the format, as +// published on Hugging Face. +const ( + // nvidia/Llama-3.1-8B-Instruct-FP8: ModelOpt declares FP8 weights and an FP8 + // KV cache in hf_quant_config.json; config.json says nothing. + modelOptFP8Config = `{"model_type":"llama","torch_dtype":"bfloat16"}` + modelOptFP8HFQuant = `{"producer":{"name":"modelopt"},"quantization":{"quant_algo":"FP8","kv_cache_quant_algo":"FP8","exclude_modules":["lm_head"]}}` + // Qwen/Qwen3-8B-FP8. + qwenFP8Config = `{"model_type":"qwen3","quantization_config":{"quant_method":"fp8","weight_block_size":[128,128]}}` + // Qwen/Qwen3.5-4B-FP8 style: vLLM-Metal dequantizes this family itself. + qwen35FP8Config = `{"model_type":"qwen3_5","quantization_config":{"quant_method":"fp8","weight_block_size":[128,128]}}` + // RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8. + compressedFP8Config = `{"model_type":"llama","quantization_config":{"quant_method":"compressed-tensors","config_groups":{"group_0":{"weights":{"type":"float","num_bits":8}}}}}` + // A compressed-tensors INT4 checkpoint. + compressedInt4Config = `{"model_type":"llama","quantization_config":{"quant_method":"compressed-tensors","config_groups":{"group_0":{"weights":{"type":"int","num_bits":4}}}}}` + bf16Config = `{"model_type":"llama","torch_dtype":"bfloat16"}` +) + +var macFP8Rule = &UnsupportedWeights{ + Formats: []string{"fp8"}, + ExceptModelTypes: []string{"qwen3_5", "qwen3_5_moe"}, + Reason: "choose the model's BF16 version instead", +} + +func TestWeightFormat_ReadsFP8TheWayVLLMDoes(t *testing.T) { + test := func(name, config, hfQuant, wantFormat string) { + t.Run(name, func(t *testing.T) { + var quant []byte + if hfQuant != "" { + quant = []byte(hfQuant) + } + + format, _, err := weightFormat([]byte(config), quant) + + require.NoError(t, err) + assert.Equal(t, wantFormat, format) + }) + } + test("ModelOpt FP8 in hf_quant_config.json", modelOptFP8Config, modelOptFP8HFQuant, "fp8") + test("ModelOpt FP8 KV cache only", bf16Config, `{"quantization":{"quant_algo":"","kv_cache_quant_algo":"FP8"}}`, "fp8") + test("ModelOpt FP8 in config.json", `{"quantization_config":{"quant_method":"modelopt","quant_algo":"FP8"}}`, "", "fp8") + test("quant_method fp8", qwenFP8Config, "", "fp8") + test("compressed-tensors float 8-bit weights", compressedFP8Config, "", "fp8") + test("compressed-tensors FP8 KV cache", `{"quantization_config":{"quant_method":"compressed-tensors","kv_cache_scheme":{"type":"float","num_bits":8}}}`, "", "fp8") + test("compressed-tensors INT4", compressedInt4Config, "", "") + test("ModelOpt NVFP4", bf16Config, `{"quantization":{"quant_algo":"NVFP4"}}`, "") + test("unquantized", bf16Config, "", "") +} + +func TestUnsupportedWeights_RefusesOnlyTheFormatsAndModelTypesItNames(t *testing.T) { + refused := func(name, config, hfQuant string) { + t.Run(name, func(t *testing.T) { + err := checkWeights(macFP8Rule, "org/m", []byte(config), []byte(hfQuant)) + + require.ErrorContains(t, err, "org/m is an FP8 model: choose the model's BF16 version instead") + }) + } + allowed := func(name, config string, rule *UnsupportedWeights) { + t.Run(name, func(t *testing.T) { + assert.NoError(t, checkWeights(rule, "org/m", []byte(config), nil)) + }) + } + refused("the ModelOpt model QA loaded", modelOptFP8Config, modelOptFP8HFQuant) + refused("Qwen3 FP8", qwenFP8Config, "") + refused("compressed-tensors FP8", compressedFP8Config, "") + allowed("a model type the engine loads in FP8 anyway", qwen35FP8Config, macFP8Rule) + allowed("an unquantized model", bf16Config, macFP8Rule) + allowed("a platform without the rule", qwenFP8Config, nil) +} + +// vLLM-Metal refuses an FP8 KV cache and its mlx-lm loader has no FP8 path +// for most families, so only the Mac declares the rule. Linux, DGX Spark +// (linux/arm64) included, loads FP8 as before. +func TestBundledVLLMManifest_RefusesFP8OnlyOnTheMac(t *testing.T) { + reg := loadWithOverrides(t, t.TempDir()) + manifest, ok := reg.Get("vllm") + require.True(t, ok) + + mac := manifest.Platforms["darwin/arm64"].Runtime.UnsupportedWeights + require.NotNil(t, mac) + assert.Equal(t, []string{"fp8"}, mac.Formats) + assert.Nil(t, manifest.Platforms["linux/arm64"].Runtime.UnsupportedWeights, "DGX Spark loads FP8") + assert.Nil(t, manifest.Platforms["linux/amd64"].Runtime.UnsupportedWeights) +} + +func TestUnsupportedWeightsValidation(t *testing.T) { + test := func(name string, rule UnsupportedWeights, valid bool) { + t.Run(name, func(t *testing.T) { + p := Platform{Runtime: Runtime{Bin: "vllm", UnsupportedWeights: &rule}} + + err := p.validate("darwin/arm64") + + assert.Equal(t, valid, err == nil, "validate = %v", err) + }) + } + test("fp8 with a reason", UnsupportedWeights{Formats: []string{"fp8"}, Reason: "use BF16"}, true) + test("no reason", UnsupportedWeights{Formats: []string{"fp8"}}, false) + test("no formats", UnsupportedWeights{Reason: "use BF16"}, false) + test("a format it cannot read", UnsupportedWeights{Formats: []string{"int4"}, Reason: "use BF16"}, false) +} + +func TestWithoutUnsupportedWeights_HidesFP8RepositoriesFromATargetThatCannotLoadThem(t *testing.T) { + models := []CatalogModel{ + {ID: "nvidia/Llama-3.1-8B-Instruct-FP8"}, + {ID: "Qwen/Qwen3-8B", Tags: []string{"safetensors"}}, + {ID: "org/tagged-only", Tags: []string{"fp8"}}, + {ID: "org/fp8-in-the-middle-v2"}, + {ID: "org/notfp8model"}, + } + + assert.Equal(t, models, withoutUnsupportedWeights(models, nil), "a target without the rule sees everything") + var ids []string + for _, m := range withoutUnsupportedWeights(models, macFP8Rule) { + ids = append(ids, m.ID) + } + assert.Equal(t, []string{"Qwen/Qwen3-8B", "org/notfp8model"}, ids) +} + +// A Load of a model the platform cannot run is refused at the start, before +// vLLM loads it, and the engine is put back on the model it was serving. +func TestLaunchModel_KeepsThePreviousModelWhenThePlatformCannotLoadTheNewOne(t *testing.T) { + e := switchModelExecutor(t, "--model") + st, err := e.state("fake") + require.NoError(t, err) + modelsDir := t.TempDir() + repo := writeHFRepo(t, filepath.Join(modelsDir, "hub"), "nvidia/Llama-3.1-8B-Instruct-FP8", llmFiles(map[string]string{ + "config.json": modelOptFP8Config, hfQuantConfigFile: modelOptFP8HFQuant, "model.safetensors": "w", + })) + require.DirExists(t, repo) + st.mu.Lock() + st.modelsDir = modelsDir + st.plat.Runtime.UnsupportedWeights = macFP8Rule + // As the vLLM manifest does, so the engine's hub cache is modelsDir/hub. + st.plat.Runtime.Env = map[string]string{"HF_HOME": modelsDir} + st.mu.Unlock() + before := launchTextOf(t, e) + + _, err = loadModel(t, e, "nvidia/Llama-3.1-8B-Instruct-FP8") + + require.ErrorContains(t, err, "nvidia/Llama-3.1-8B-Instruct-FP8 is an FP8 model: choose the model's BF16 version instead") + require.ErrorContains(t, err, "the previous model was restored") + assert.Equal(t, before, launchTextOf(t, e)) + assert.True(t, mustStatus(t, e).Running, "the engine must be serving the previous model again") +} + +// A download the platform could never load is refused before it starts. +func TestPullHFModel_RefusesAModelThePlatformCannotLoad(t *testing.T) { + siblings := []hfSibling{ + {RFilename: "config.json", Size: 1}, + {RFilename: hfQuantConfigFile, Size: 1}, + {RFilename: "tokenizer.json", Size: 1}, + {RFilename: "model.safetensors", Size: 8_000_000_000}, + } + e, _, _ := hfPullFixture(t, siblings, http.StatusOK, "echo the download must not start >&2; exit 1") + hub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasPrefix(r.URL.Path, "/api/models/") { + assert.NoError(t, json.NewEncoder(w).Encode(map[string]any{"siblings": siblings})) + return + } + switch filepath.Base(r.URL.Path) { + case "config.json": + _, err := w.Write([]byte(modelOptFP8Config)) + assert.NoError(t, err) + case hfQuantConfigFile: + _, err := w.Write([]byte(modelOptFP8HFQuant)) + assert.NoError(t, err) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(hub.Close) + old := hfDefaultEndpoint + hfDefaultEndpoint = hub.URL + t.Cleanup(func() { hfDefaultEndpoint = old }) + st, err := e.state("vllm") + require.NoError(t, err) + st.mu.Lock() + st.plat.Runtime.UnsupportedWeights = macFP8Rule + st.mu.Unlock() + t.Setenv("HF_HUB_OFFLINE", "") + + _, err = e.PullModelStream(t.Context(), "vllm", "nvidia/Llama-3.1-8B-Instruct-FP8", nil) + + require.ErrorContains(t, err, "is an FP8 model: choose the model's BF16 version instead") + assert.NotContains(t, err.Error(), "the download must not start") +} + +// checkLaunchWeights reads the model a start serves at its revision, and for a +// model not in the cache, what the metadata fetch put there. +func TestCheckLaunchWeights_ReadsTheModelThisStartServes(t *testing.T) { + hub := t.TempDir() + writeHFRepo(t, hub, "org/cached-fp8", llmFiles(map[string]string{"config.json": qwenFP8Config, "model.safetensors": "w"})) + noFetch := func() (hfPendingModel, bool) { return hfPendingModel{}, false } + + require.ErrorContains(t, checkLaunchWeights(macFP8Rule, []string{"serve", "--model", "org/cached-fp8"}, hub, noFetch), "org/cached-fp8 is an FP8 model") + assert.NoError(t, checkLaunchWeights(macFP8Rule, []string{"serve", "--model", "org/unknown"}, hub, noFetch), "a model that cannot be read is left to vLLM") + assert.NoError(t, checkLaunchWeights(nil, []string{"serve", "--model", "org/cached-fp8"}, hub, noFetch)) + + pending := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(pending, "config.json"), []byte(compressedFP8Config), 0o644)) + fetch := func() (hfPendingModel, bool) { return hfPendingModel{snapshot: pending}, true } + assert.ErrorContains(t, checkLaunchWeights(macFP8Rule, []string{"serve", "--model", "org/not-yet-downloaded"}, hub, fetch), "org/not-yet-downloaded is an FP8 model") +} + +// engine:catalog applies the target's rule: a Mac is not offered FP8 vLLM +// models, and DGX Spark (linux/arm64) is offered them as before. +func TestRunCatalog_HidesFP8VLLMModelsOnlyFromAMac(t *testing.T) { + listing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.NoError(t, json.NewEncoder(w).Encode([]any{ + vllmRow("Qwen/Chat-8B", nil), + vllmRow("nvidia/Llama-3.1-8B-Instruct-FP8", nil), + })) + })) + t.Cleanup(listing.Close) + ex := NewExecutor(loadWithOverrides(t, t.TempDir()), NewReporter(nil), func(string, any) {}, t.TempDir()) + offered := func(name, platform, arch string, want []string) { + t.Run(name, func(t *testing.T) { + var out bytes.Buffer + m := NewManager(NewCodec(&out), ex, nil) + m.catalog.baseURL = listing.URL + id := json.RawMessage("1") + params := fmt.Sprintf(`{"engine":"vllm","platform":%q,"arch":%q}`, platform, arch) + + m.runCatalog(t.Context(), &Message{JSONRPC: "2.0", ID: &id, Method: "engine:catalog", Params: json.RawMessage(params)}) + + var reply struct { + Result CatalogResult `json:"result"` + } + require.NoError(t, json.Unmarshal(out.Bytes(), &reply), out.String()) + var ids []string + for _, model := range reply.Result.Models { + ids = append(ids, model.ID) + } + assert.ElementsMatch(t, want, ids) + }) + } + offered("a Mac", "darwin", "arm64", []string{"Qwen/Chat-8B"}) + offered("DGX Spark", "linux", "arm64", []string{"Qwen/Chat-8B", "nvidia/Llama-3.1-8B-Instruct-FP8"}) +} diff --git a/services/nvpair-job-scheduler/spec.md b/services/nvpair-job-scheduler/spec.md index cc1bbef56..1513e3a0e 100644 --- a/services/nvpair-job-scheduler/spec.md +++ b/services/nvpair-job-scheduler/spec.md @@ -55,7 +55,7 @@ capacity, model locality, latency, and affinity can grow later behind the same - **Simultaneous local burst**: before each auto-routed inference forward, the proxy atomically increments a local reservation for its chosen node. Concurrent requests therefore see one another even before workload feedback completes. -- **Cross-engine contention**: Ollama, LM Studio, and llama.cpp workloads all +- **Cross-engine contention**: Ollama, LM Studio, llama.cpp, and vLLM workloads all consume the destination node's execution resources, so any one of them changes every engine output. - **Periodic reconciliation**: the runtime-adjustable timer recomputes the same @@ -205,7 +205,7 @@ Requests (Broker → scheduler): | Method | Params | Result | |--------|--------|--------| -| `scheduler:get-status` | — | `{ interval_ms, engines: { ollama: EngineSchedule, lmstudio: EngineSchedule, llamacpp: EngineSchedule } }` | +| `scheduler:get-status` | — | `{ interval_ms, engines: { ollama: EngineSchedule, lmstudio: EngineSchedule, llamacpp: EngineSchedule, vllm: EngineSchedule } }` | | `scheduler:get-interval` | — | `{ interval_ms }` | | `scheduler:set-interval` | `{ interval_ms }` | `{ interval_ms }` — live; clamped to the floor (§7.5); effective next tick | | `scheduler:tick` | — | `{ ticked: true }` — force an immediate recompute + re-emit (tests/debug) | @@ -215,7 +215,7 @@ Requests (Broker → scheduler): `EngineSchedule`: ```json { - engine: string // "ollama" | "lmstudio" | "llamacpp" + engine: string // "ollama" | "lmstudio" | "llamacpp" | "vllm" emitted: [NodeRank] // last order emitted (empty if none yet) lastEmittedAt: number // epoch ms; 0 if never } @@ -289,8 +289,8 @@ Let `D` = current discovered-node `hostUuid`s: `scheduledOn` (not `originatedFrom`) is the key: we care where work runs, not where it came from. Engine remains part of workload identity and selects the downstream -proxy output, but it does not partition the load metric: Ollama, LM Studio, and -llama.cpp normally contend for the same node-level GPU, VRAM, CPU, and memory. +proxy output, but it does not partition the load metric: Ollama, LM Studio, +llama.cpp, and vLLM normally contend for the same node-level GPU, VRAM, CPU, and memory. Until resource affinity is observable, total node queue depth is the conservative signal. diff --git a/services/nvpair-node-scanner/README.md b/services/nvpair-node-scanner/README.md index ae2376e09..8cd2f30c3 100644 --- a/services/nvpair-node-scanner/README.md +++ b/services/nvpair-node-scanner/README.md @@ -11,7 +11,7 @@ Concretely, the daemon: - **Advertises** this node's one `_nvpair-node._tcp` record, built from a registry of the local services its parent (the broker) registers — `ni`/`ol`/`lm`/`er`/`wl`/`cl`/`em`/`ec` ports plus the node identity (`uuid=`, `cluster-uuid=` when clustered, `ip=`). One record per node covers every service. - **Browses** `_nvpair-node._tcp` for every node on the LAN (including itself) and maintains a queryable directory keyed by host UUID. -- **Enriches** each discovered node: GPU/CPU/memory inventory from its `ni` (node-info) port over plain HTTP, and the model list from its `em` (engine-manager) port (`GET /v1/models`) over cluster mTLS — a peer's model inventory is cluster data, so it is fetched only when this node holds a pin for the principal that peer advertises, while this node's own list is read over loopback in plaintext — the flat union, the per-engine breakdown (`modelsByEngine`), and the per-engine set of models loaded in memory (`loadedByEngine`), all carried onto the node's directory entry so remote cards can show loaded state. Each enrichment has a last-good cache so a transient fetch miss doesn't blank the node's card. Node-info enrichment can optionally be moved onto HTTPS (see the TLS flags below) — off by default and gated only on operator flags, never inferred per-node. +- **Enriches** each discovered node: GPU/CPU/memory inventory from its `ni` (node-info) port over plain HTTP, and the model list from its `em` (engine-manager) port (`GET /v1/models`) over cluster mTLS — a peer's model inventory is cluster data, so it is fetched only when this node holds a pin for the principal that peer advertises, while this node's own list is read over loopback in plaintext — the flat union, the per-engine breakdown (`modelsByEngine`), the per-engine set of models loaded in memory (`loadedByEngine`), and the downloaded inventory of an engine that advertises only its served model (`downloadedByEngine`, vLLM), all carried onto the node's directory entry so remote cards can show loaded state and keep a stopped vLLM's models. Each enrichment has a last-good cache so a transient fetch miss doesn't blank the node's card. Node-info enrichment can optionally be moved onto HTTPS (see the TLS flags below) — off by default and gated only on operator flags, never inferred per-node. - **Samples scheduling telemetry** from healthy nodes every two seconds with stable per-node jitter. Consecutive remote failures back off to 4, 8, 16, then 30 seconds and reset on success; local loopback sampling stays at the healthy diff --git a/services/nvpair-node-scanner/activity_test.go b/services/nvpair-node-scanner/activity_test.go index d804310f2..c9be6b8f2 100644 --- a/services/nvpair-node-scanner/activity_test.go +++ b/services/nvpair-node-scanner/activity_test.go @@ -132,9 +132,7 @@ func TestForget_ClearsActivity(t *testing.T) { d := activityDaemon() d.lastInfo = make(map[string]NodeInfoResponse) d.lastInfoAt = make(map[string]time.Time) - d.lastModels = make(map[string][]string) - d.lastModelsByEngine = make(map[string]map[string][]string) - d.lastLoadedByEngine = make(map[string]map[string][]string) + d.lastInventory = make(map[string]modelInventory) d.nodeInfoDown = make(map[string]bool) d.noteActivity("node", 0) diff --git a/services/nvpair-node-scanner/cluster_identity_refresh_test.go b/services/nvpair-node-scanner/cluster_identity_refresh_test.go index d29c6d136..00568b980 100644 --- a/services/nvpair-node-scanner/cluster_identity_refresh_test.go +++ b/services/nvpair-node-scanner/cluster_identity_refresh_test.go @@ -76,14 +76,12 @@ func startNodeInfoStub(t *testing.T, body string) (*nodeInfoStub, int) { // node is treated as a peer. func newIdentityRefreshTestDaemon() *daemon { return &daemon{ - reg: newRegistry("self-node-uuid", "", selfAddrs("127.0.0.1")), - dir: newDirectory(), - mesh: clustertrust.Open(""), - http: &http.Client{Timeout: nodeInfoFetchTimeout}, - lastInfo: make(map[string]NodeInfoResponse), - lastModels: make(map[string][]string), - lastModelsByEngine: make(map[string]map[string][]string), - lastLoadedByEngine: make(map[string]map[string][]string), + reg: newRegistry("self-node-uuid", "", selfAddrs("127.0.0.1")), + dir: newDirectory(), + mesh: clustertrust.Open(""), + http: &http.Client{Timeout: nodeInfoFetchTimeout}, + lastInfo: make(map[string]NodeInfoResponse), + lastInventory: make(map[string]modelInventory), } } diff --git a/services/nvpair-node-scanner/cluster_membership_test.go b/services/nvpair-node-scanner/cluster_membership_test.go index f951ac4f7..f17c1eb01 100644 --- a/services/nvpair-node-scanner/cluster_membership_test.go +++ b/services/nvpair-node-scanner/cluster_membership_test.go @@ -20,17 +20,15 @@ import ( func newMembershipTestDaemon(clusterDir string) *daemon { mesh := clustertrust.Open(clusterDir) return &daemon{ - mesh: mesh, - modelsClients: clustertrust.NewPeerClientPool(mesh, modelsFetchTimeout), - clusterDir: clusterDir, - baseDir: filepath.Dir(clusterDir), - reg: newRegistry("host-uuid", "", selfAddrs("127.0.0.1")), - dir: newDirectory(), - lastInfo: make(map[string]NodeInfoResponse), - lastInfoAt: make(map[string]time.Time), - lastModels: make(map[string][]string), - lastModelsByEngine: make(map[string]map[string][]string), - lastLoadedByEngine: make(map[string]map[string][]string), + mesh: mesh, + modelsClients: clustertrust.NewPeerClientPool(mesh, modelsFetchTimeout), + clusterDir: clusterDir, + baseDir: filepath.Dir(clusterDir), + reg: newRegistry("host-uuid", "", selfAddrs("127.0.0.1")), + dir: newDirectory(), + lastInfo: make(map[string]NodeInfoResponse), + lastInfoAt: make(map[string]time.Time), + lastInventory: make(map[string]modelInventory), // codec/responder nil: emit is nil-safe; reloadIdentity skips re-advertise. } } diff --git a/services/nvpair-node-scanner/connreuse_test.go b/services/nvpair-node-scanner/connreuse_test.go index bd96b00a7..3e8fb86c5 100644 --- a/services/nvpair-node-scanner/connreuse_test.go +++ b/services/nvpair-node-scanner/connreuse_test.go @@ -73,7 +73,8 @@ func TestModelsClient_ReusesPeerConnections(t *testing.T) { // pinned mTLS path instead of the plaintext self-fetch. const rounds = 5 for i := 0; i < rounds; i++ { - models, _, _, ok := d.fetchModels("localhost", port, "uuid-peer") + inventory, ok := d.fetchModels("localhost", port, "uuid-peer") + models := inventory.models require.True(t, ok, "round (%v, %v, %v)", i, ok, models) require.Len(t, models, 1, "round (%v, %v, %v)", i, ok, models) require.Equal(t, "m", models[0], "round (%v, %v, %v)", i, ok, models) diff --git a/services/nvpair-node-scanner/daemon.go b/services/nvpair-node-scanner/daemon.go index 42226db3d..cb700feba 100644 --- a/services/nvpair-node-scanner/daemon.go +++ b/services/nvpair-node-scanner/daemon.go @@ -91,24 +91,16 @@ type daemon struct { // ignores our own uuid — so this is the sole self-key writer lock. selfMu sync.Mutex - // lastInfo / lastModels cache each node's last successful enrichment (keyed by - // hostUuid) — GPU/CPU/memory from node-info and the model list from - // engine-manager — so a transient fetch failure reuses the prior values - // instead of blanking the node card. Same resilience the legacy scanner path - // had via its own lastInfo cache. - infoMu sync.Mutex - lastInfo map[string]NodeInfoResponse - lastInfoAt map[string]time.Time - lastModels map[string][]string - // lastModelsByEngine caches the per-engine attribution (engine name -> - // models) alongside the flat lastModels union, so a transient /v1/models - // miss reuses both rather than blanking the per-engine card. Keyed by - // hostUuid like the other enrichment caches. - lastModelsByEngine map[string]map[string][]string - // lastLoadedByEngine caches the per-engine loaded (in-memory) set from the - // same /v1/models fetch, so a transient miss reuses it rather than blanking - // a remote node's loaded state. Keyed by hostUuid. - lastLoadedByEngine map[string]map[string][]string + // lastInfo / lastInventory cache each node's last successful enrichment + // (keyed by hostUuid) — GPU/CPU/memory from node-info and the whole + // /v1/models answer from engine-manager (the flat union, the per-engine + // attribution, and the loaded and downloaded sets) — so a transient fetch + // failure reuses the prior values instead of blanking the node card. Same + // resilience the legacy scanner path had via its own lastInfo cache. + infoMu sync.Mutex + lastInfo map[string]NodeInfoResponse + lastInfoAt map[string]time.Time + lastInventory map[string]modelInventory // nodeInfoDown holds the nodes whose last node-info enrichment failed, so // the outage is reported when it starts and when it ends rather than once // per sweep. Guarded by infoMu with the caches it explains: the reason a @@ -350,24 +342,22 @@ func newDaemon(codec *Codec, mesh *clustertrust.Mesh, clusterDir string, tlsHTTP instance, _ := os.Hostname() d := &daemon{ - codec: codec, - mesh: mesh, - reg: reg, - instance: instance, - baseDir: baseDir, - clusterDir: clusterDir, - dir: newDirectory(), - http: &http.Client{Timeout: nodeInfoFetchTimeout, Transport: nodeInfoTransport(nil)}, - modelsHTTP: &http.Client{Timeout: modelsFetchTimeout}, - modelsClients: clustertrust.NewPeerClientPool(mesh, modelsFetchTimeout), - tlsHTTP: tlsHTTP, - lastInfo: make(map[string]NodeInfoResponse), - lastInfoAt: make(map[string]time.Time), - lastModels: make(map[string][]string), - lastModelsByEngine: make(map[string]map[string][]string), - lastLoadedByEngine: make(map[string]map[string][]string), - nodeInfoDown: make(map[string]bool), - lastActivityAt: make(map[string]time.Time), + codec: codec, + mesh: mesh, + reg: reg, + instance: instance, + baseDir: baseDir, + clusterDir: clusterDir, + dir: newDirectory(), + http: &http.Client{Timeout: nodeInfoFetchTimeout, Transport: nodeInfoTransport(nil)}, + modelsHTTP: &http.Client{Timeout: modelsFetchTimeout}, + modelsClients: clustertrust.NewPeerClientPool(mesh, modelsFetchTimeout), + tlsHTTP: tlsHTTP, + lastInfo: make(map[string]NodeInfoResponse), + lastInfoAt: make(map[string]time.Time), + lastInventory: make(map[string]modelInventory), + nodeInfoDown: make(map[string]bool), + lastActivityAt: make(map[string]time.Time), } // Anti-flap: before evicting a node that's missed the mDNS threshold, // confirm it still answers on node-info (or engine-manager as fallback). @@ -818,20 +808,12 @@ func (d *daemon) dropSelf(oldHostUUID, newHostUUID string) { if at, ok := d.lastInfoAt[oldHostUUID]; ok { d.lastInfoAt[newHostUUID] = at } - if models, ok := d.lastModels[oldHostUUID]; ok { - d.lastModels[newHostUUID] = models - } - if byEngine, ok := d.lastModelsByEngine[oldHostUUID]; ok { - d.lastModelsByEngine[newHostUUID] = byEngine - } - if loaded, ok := d.lastLoadedByEngine[oldHostUUID]; ok { - d.lastLoadedByEngine[newHostUUID] = loaded + if inventory, ok := d.lastInventory[oldHostUUID]; ok { + d.lastInventory[newHostUUID] = inventory } delete(d.lastInfo, oldHostUUID) delete(d.lastInfoAt, oldHostUUID) - delete(d.lastModels, oldHostUUID) - delete(d.lastModelsByEngine, oldHostUUID) - delete(d.lastLoadedByEngine, oldHostUUID) + delete(d.lastInventory, oldHostUUID) d.infoMu.Unlock() if existed { d.emit(noderec.NotifyNodeRemoved, noderec.DirectoryNode{HostUUID: oldHostUUID, Name: node.Name}) @@ -1334,52 +1316,42 @@ func (d *daemon) enrichModelsCandidates(node *noderec.DirectoryNode, hosts []str return } ask := func(host string) (modelInventory, bool) { - models, byEngine, loadedByEngine, fetched := d.fetchModels(host, em.Port, node.ClusterUUID) - if !fetched { - return modelInventory{}, false - } - return modelInventory{models: models, byEngine: byEngine, loadedByEngine: loadedByEngine}, true + return d.fetchModels(host, em.Port, node.ClusterUUID) } key := hostKey{hostUUID: node.HostUUID, service: noderec.ServiceEngineManager} if _, inventory, ok := askRemembered(&d.enrichHosts, key, hosts, ask); ok { d.infoMu.Lock() - if d.lastModels == nil { - d.lastModels = make(map[string][]string) - } - if d.lastModelsByEngine == nil { - d.lastModelsByEngine = make(map[string]map[string][]string) - } - if d.lastLoadedByEngine == nil { - d.lastLoadedByEngine = make(map[string]map[string][]string) - } - d.lastModels[node.HostUUID] = inventory.models - d.lastModelsByEngine[node.HostUUID] = inventory.byEngine - d.lastLoadedByEngine[node.HostUUID] = inventory.loadedByEngine + d.cacheInventory(node.HostUUID, inventory) d.infoMu.Unlock() - node.Models = inventory.models - node.ModelsByEngine = inventory.byEngine - node.LoadedByEngine = inventory.loadedByEngine + inventory.applyTo(node) return } d.infoMu.Lock() - cached, had := d.lastModels[node.HostUUID] - cachedByEngine := d.lastModelsByEngine[node.HostUUID] - cachedLoaded := d.lastLoadedByEngine[node.HostUUID] + cached, had := d.lastInventory[node.HostUUID] d.infoMu.Unlock() if had { - node.Models = cached - node.ModelsByEngine = cachedByEngine - node.LoadedByEngine = cachedLoaded + cached.applyTo(node) + } +} + +// cacheInventory records a node's last-good /v1/models answer. The caller holds +// infoMu. +func (d *daemon) cacheInventory(hostUUID string, inventory modelInventory) { + if d.lastInventory == nil { + d.lastInventory = make(map[string]modelInventory) } + d.lastInventory[hostUUID] = inventory } // fetchModels does one GET of a node's /v1/models, returning the // flat model-name union, the per-engine attribution (engine name -> models), -// the per-engine loaded (in-memory) set, and whether it succeeded. Shape: -// {"models":["name",...],"modelsByEngine":{"ollama":[...]},"loadedByEngine":{"ollama":[...]}}. -// modelsByEngine and loadedByEngine are additive (omitempty on the wire); a peer -// that doesn't send them yields nil maps, and the flat union still enriches the -// node. A present engine key with [] is a successful empty inventory. +// the per-engine loaded (in-memory) set, the per-engine downloaded inventory of +// an engine that advertises only what it serves, and whether it succeeded. Shape: +// {"models":["name",...],"modelsByEngine":{"ollama":[...]},"loadedByEngine":{"ollama":[...]}, +// "downloadedByEngine":{"vllm":[...]}}. The three maps are additive (omitempty on +// the wire); a peer that doesn't send them yields nil maps, and the flat union +// still enriches the node. A present engine key with [] is a successful empty +// inventory. // A peer's model inventory is cluster data, so a peer is fetched over cluster // mTLS keyed on the cluster principal it advertises (clusterUUID, from its // cluster-uuid= TXT): this node must hold a pin for that principal, or there is @@ -1387,32 +1359,38 @@ func (d *daemon) enrichModelsCandidates(node *noderec.DirectoryNode, hosts []str // dialed over loopback (both callers force that for self), which stays plain and // therefore keeps working when this node belongs to no cluster — a standalone // machine still has to show its own models. -func (d *daemon) fetchModels(ip string, port int, clusterUUID string) ([]string, map[string][]string, map[string][]string, bool) { +func (d *daemon) fetchModels(ip string, port int, clusterUUID string) (modelInventory, bool) { client, scheme, ok := d.modelsClient(ip, clusterUUID) if !ok { slog.Debug("model enrichment skipped: peer is not a pinned cluster member", "ip", ip, "clusterUuid", clusterUUID) - return nil, nil, nil, false + return modelInventory{}, false } url := scheme + "://" + net.JoinHostPort(ip, strconv.Itoa(port)) + "/v1/models" resp, err := client.Get(url) if err != nil { slog.Debug("model enrichment failed", "ip", ip, "url", url, "err", err) - return nil, nil, nil, false + return modelInventory{}, false } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { - return nil, nil, nil, false + return modelInventory{}, false } var body struct { - Models []string `json:"models"` - ModelsByEngine map[string][]string `json:"modelsByEngine"` - LoadedByEngine map[string][]string `json:"loadedByEngine"` + Models []string `json:"models"` + ModelsByEngine map[string][]string `json:"modelsByEngine"` + LoadedByEngine map[string][]string `json:"loadedByEngine"` + DownloadedByEngine map[string][]string `json:"downloadedByEngine"` } if err := json.NewDecoder(resp.Body).Decode(&body); err != nil { - return nil, nil, nil, false - } - return body.Models, body.ModelsByEngine, body.LoadedByEngine, true + return modelInventory{}, false + } + return modelInventory{ + models: body.Models, + byEngine: body.ModelsByEngine, + loadedByEngine: body.LoadedByEngine, + downloadedByEngine: body.DownloadedByEngine, + }, true } // modelsClient picks the transport for a model fetch: the plain client over @@ -1809,42 +1787,26 @@ func (d *daemon) refreshNodeModelsCandidates(hostUUID, guardIP string, hosts []s hosts = []string{loopbackHost} } ask := func(host string) (modelInventory, bool) { - models, byEngine, loadedByEngine, fetched := d.fetchModels(host, emPort, clusterUUID) - if !fetched { - return modelInventory{}, false - } - return modelInventory{models: models, byEngine: byEngine, loadedByEngine: loadedByEngine}, true + return d.fetchModels(host, emPort, clusterUUID) } key := hostKey{hostUUID: hostUUID, service: noderec.ServiceEngineManager} _, inventory, fetched := askRemembered(&d.enrichHosts, key, hosts, ask) if !fetched { return false } - models, byEngine, loadedByEngine := inventory.models, inventory.byEngine, inventory.loadedByEngine // Refresh the enrichment cache BEFORE touching the directory: the caches // (infoMu) and the directory (dir.mu) are guarded by different locks, so if // the directory were updated first, a concurrent onBrowse->enrichModelsAt whose - // own fetch fails would fall back to a still-stale lastModels and upsert the + // own fetch fails would fall back to a still-stale lastInventory and upsert the // old inventory, reverting this update. Writing the cache first means that // fallback reads the fresh value. The only side effect — caching a value for // a node applyModels then finds removed — is benign (it is the last-known-good // value the cache exists to preserve, and forget() clears it on eviction). d.infoMu.Lock() - if d.lastModels == nil { - d.lastModels = make(map[string][]string) - } - if d.lastModelsByEngine == nil { - d.lastModelsByEngine = make(map[string]map[string][]string) - } - if d.lastLoadedByEngine == nil { - d.lastLoadedByEngine = make(map[string]map[string][]string) - } - d.lastModels[hostUUID] = models - d.lastModelsByEngine[hostUUID] = byEngine - d.lastLoadedByEngine[hostUUID] = loadedByEngine + d.cacheInventory(hostUUID, inventory) d.infoMu.Unlock() - node, changed, valid := d.dir.applyModels(hostUUID, guardIP, emPort, models, byEngine, loadedByEngine) + node, changed, valid := d.dir.applyModels(hostUUID, guardIP, emPort, inventory) if !valid { return false } @@ -1861,9 +1823,7 @@ func (d *daemon) forget(hostUUID string) { d.infoMu.Lock() delete(d.lastInfo, hostUUID) delete(d.lastInfoAt, hostUUID) - delete(d.lastModels, hostUUID) - delete(d.lastModelsByEngine, hostUUID) - delete(d.lastLoadedByEngine, hostUUID) + delete(d.lastInventory, hostUUID) delete(d.nodeInfoDown, hostUUID) d.infoMu.Unlock() d.activityMu.Lock() diff --git a/services/nvpair-node-scanner/daemon_test.go b/services/nvpair-node-scanner/daemon_test.go index cfc2e8228..41919d4e9 100644 --- a/services/nvpair-node-scanner/daemon_test.go +++ b/services/nvpair-node-scanner/daemon_test.go @@ -38,17 +38,15 @@ func selfAddrs(ip string) []string { // nil-safe; publishSelf skips re-advertise when the responder is nil). func newSelfTestDaemon(hostUUID, ip string) *daemon { return &daemon{ - mesh: clustertrust.Open(""), - reg: newRegistry(hostUUID, "", selfAddrs(ip)), - instance: "myhost", - dir: newDirectory(), - http: &http.Client{Timeout: nodeInfoFetchTimeout}, - modelsHTTP: &http.Client{Timeout: modelsFetchTimeout}, - lastInfo: make(map[string]NodeInfoResponse), - lastInfoAt: make(map[string]time.Time), - lastModels: make(map[string][]string), - lastModelsByEngine: make(map[string]map[string][]string), - lastLoadedByEngine: make(map[string]map[string][]string), + mesh: clustertrust.Open(""), + reg: newRegistry(hostUUID, "", selfAddrs(ip)), + instance: "myhost", + dir: newDirectory(), + http: &http.Client{Timeout: nodeInfoFetchTimeout}, + modelsHTTP: &http.Client{Timeout: modelsFetchTimeout}, + lastInfo: make(map[string]NodeInfoResponse), + lastInfoAt: make(map[string]time.Time), + lastInventory: make(map[string]modelInventory), } } @@ -61,15 +59,13 @@ func newSelfTestDaemon(hostUUID, ip string) *daemon { func TestReloadIdentity_AdoptsClusterPrincipal(t *testing.T) { base := t.TempDir() d := &daemon{ - mesh: clustertrust.Open(""), // unclustered: no cluster dir -> no clusterUuid - reg: newRegistry("scanner-minted-X", "", selfAddrs("127.0.0.1")), - dir: newDirectory(), - baseDir: base, - lastInfo: make(map[string]NodeInfoResponse), - lastInfoAt: make(map[string]time.Time), - lastModels: make(map[string][]string), - lastModelsByEngine: make(map[string]map[string][]string), - lastLoadedByEngine: make(map[string]map[string][]string), + mesh: clustertrust.Open(""), // unclustered: no cluster dir -> no clusterUuid + reg: newRegistry("scanner-minted-X", "", selfAddrs("127.0.0.1")), + dir: newDirectory(), + baseDir: base, + lastInfo: make(map[string]NodeInfoResponse), + lastInfoAt: make(map[string]time.Time), + lastInventory: make(map[string]modelInventory), // codec/responder nil: emit is nil-safe; reloadIdentity skips re-advertise. } // Simulate this node already having been browsed under its scanner-minted uuid. diff --git a/services/nvpair-node-scanner/directory.go b/services/nvpair-node-scanner/directory.go index 70b7d04f3..0c5c39714 100644 --- a/services/nvpair-node-scanner/directory.go +++ b/services/nvpair-node-scanner/directory.go @@ -233,7 +233,7 @@ func (d *directory) snapshot(filter noderec.ServiceKey) []noderec.DirectoryNode // It confirms the node still exists AND still advertises the same em endpoint // (ip + port) the fetch targeted, so an in-flight result can never resurrect a // removed node or overwrite one that has since been re-addressed. It then -// compares the fetched flat, per-engine, and loaded inventories +// compares the fetched flat, per-engine, loaded, and downloaded inventories // order-insensitively with the stored entry and, only when any changed, updates // just the model fields in place — preserving every other field (identity, // services, node-info, trust) the browse path owns. @@ -241,7 +241,7 @@ func (d *directory) snapshot(filter noderec.ServiceKey) []noderec.DirectoryNode // Returns the (possibly updated) node, whether the inventory changed, and // whether the guarded apply was valid. ok == false means the result is stale // (node gone or re-addressed) and the caller must not cache or emit it. -func (d *directory) applyModels(hostUUID, ip string, emPort int, models []string, byEngine, loadedByEngine map[string][]string) (node noderec.DirectoryNode, changed, ok bool) { +func (d *directory) applyModels(hostUUID, ip string, emPort int, inventory modelInventory) (node noderec.DirectoryNode, changed, ok bool) { d.mu.Lock() defer d.mu.Unlock() n, present := d.nodes[hostUUID] @@ -252,12 +252,13 @@ func (d *directory) applyModels(hostUUID, ip string, emPort int, models []string if !has || n.IP != ip || em.Port != emPort { return n, false, false } - if sameStringSet(n.Models, models) && sameByEngine(n.ModelsByEngine, byEngine) && sameByEngine(n.LoadedByEngine, loadedByEngine) { + if sameStringSet(n.Models, inventory.models) && + sameByEngine(n.ModelsByEngine, inventory.byEngine) && + sameByEngine(n.LoadedByEngine, inventory.loadedByEngine) && + sameByEngine(n.DownloadedByEngine, inventory.downloadedByEngine) { return n, false, true } - n.Models = models - n.ModelsByEngine = byEngine - n.LoadedByEngine = loadedByEngine + inventory.applyTo(&n) d.nodes[hostUUID] = n return n, true, true } diff --git a/services/nvpair-node-scanner/enrich_hosts.go b/services/nvpair-node-scanner/enrich_hosts.go index 7997a192f..715328c5e 100644 --- a/services/nvpair-node-scanner/enrich_hosts.go +++ b/services/nvpair-node-scanner/enrich_hosts.go @@ -10,12 +10,22 @@ import ( "nvpair-shared/noderec" ) -// modelInventory is the three-part answer a node's engine manager gives, carried -// as one value so the concurrent ask can return it. +// modelInventory is a node engine manager's /v1/models answer, carried as one +// value so the concurrent ask can return it and the last-good cache keeps every +// part of it together. type modelInventory struct { - models []string - byEngine map[string][]string - loadedByEngine map[string][]string + models []string + byEngine map[string][]string + loadedByEngine map[string][]string + downloadedByEngine map[string][]string +} + +// applyTo copies the inventory onto a directory entry's model fields. +func (inv modelInventory) applyTo(node *noderec.DirectoryNode) { + node.Models = inv.models + node.ModelsByEngine = inv.byEngine + node.LoadedByEngine = inv.loadedByEngine + node.DownloadedByEngine = inv.downloadedByEngine } // hostKey identifies one node's endpoint for one service. Node-info and the engine diff --git a/services/nvpair-node-scanner/refresh_test.go b/services/nvpair-node-scanner/refresh_test.go index 78518fd23..9e00c8dc6 100644 --- a/services/nvpair-node-scanner/refresh_test.go +++ b/services/nvpair-node-scanner/refresh_test.go @@ -91,14 +91,12 @@ func portFromURL(t *testing.T, raw string) int { // TestRefreshSelfDialsLoopback). func newRefreshTestDaemon() *daemon { return &daemon{ - reg: newRegistry("self-node-uuid", "", selfAddrs("127.0.0.1")), - dir: newDirectory(), - modelsHTTP: &http.Client{Timeout: modelsFetchTimeout}, - lastInfo: make(map[string]NodeInfoResponse), - lastInfoAt: make(map[string]time.Time), - lastModels: make(map[string][]string), - lastModelsByEngine: make(map[string]map[string][]string), - lastLoadedByEngine: make(map[string]map[string][]string), + reg: newRegistry("self-node-uuid", "", selfAddrs("127.0.0.1")), + dir: newDirectory(), + modelsHTTP: &http.Client{Timeout: modelsFetchTimeout}, + lastInfo: make(map[string]NodeInfoResponse), + lastInfoAt: make(map[string]time.Time), + lastInventory: make(map[string]modelInventory), } } @@ -263,9 +261,9 @@ func TestRefreshFetchFailure_RetainsInventory(t *testing.T) { assert.Equal(t, []string{"keep-me"}, got.Models, "Models after failure") assert.Equal(t, map[string][]string{"ollama": {"keep-me"}}, got.ModelsByEngine, "ModelsByEngine after failure") d.infoMu.Lock() - cachedFlat := d.lastModels["peer-C"] - cachedByEngine := d.lastModelsByEngine["peer-C"] + cached := d.lastInventory["peer-C"] d.infoMu.Unlock() + cachedFlat, cachedByEngine := cached.models, cached.byEngine assert.Equal(t, []string{"keep-me"}, cachedFlat, "cached flat inventory") assert.Equal(t, map[string][]string{"ollama": {"keep-me"}}, cachedByEngine, "cached per-engine inventory") } @@ -352,11 +350,40 @@ func TestRefreshModels_RetainsLoadedModelsAfterFetchFailure(t *testing.T) { got, _ := d.dir.get("peer-N") assert.Equal(t, map[string][]string{"ollama": {"keep-me"}}, got.LoadedByEngine, "LoadedByEngine after failure") d.infoMu.Lock() - cachedLoaded := d.lastLoadedByEngine["peer-N"] + cachedLoaded := d.lastInventory["peer-N"].loadedByEngine d.infoMu.Unlock() assert.Equal(t, map[string][]string{"ollama": {"keep-me"}}, cachedLoaded, "cached loaded inventory") } +// A stopped vLLM peer advertises no model for routing but still reports what it +// has downloaded. That inventory reaches the directory, a change to it alone is +// reported (an Eject or a Delete on the peer), and a failed fetch keeps it. +func TestRefreshModels_CarriesTheDownloadedModelsOfAStoppedVLLMPeer(t *testing.T) { + d := newRefreshTestDaemon() + stub, port := startModelsStub(t, http.StatusOK, map[string]any{ + "models": []string{}, + "downloadedByEngine": map[string][]string{"vllm": {"Qwen/Qwen3.5-4B", "mistralai/Ministral-3-3B-Instruct-2512"}}, + }) + seedEMNode(d, "peer-V", "127.0.0.1", port) + require.True(t, d.refreshNodeModels("peer-V", "127.0.0.1", port, ""), "initial populate should report a change") + got, _ := d.dir.get("peer-V") + assert.Equal(t, map[string][]string{"vllm": {"Qwen/Qwen3.5-4B", "mistralai/Ministral-3-3B-Instruct-2512"}}, got.DownloadedByEngine) + assert.Empty(t, got.ModelsByEngine, "nothing advertised for routing") + + stub.set(http.StatusOK, map[string]any{ + "models": []string{}, + "downloadedByEngine": map[string][]string{"vllm": {"Qwen/Qwen3.5-4B"}}, + }) + require.True(t, d.refreshNodeModels("peer-V", "127.0.0.1", port, ""), "a downloaded-only change should report a change") + got, _ = d.dir.get("peer-V") + assert.Equal(t, map[string][]string{"vllm": {"Qwen/Qwen3.5-4B"}}, got.DownloadedByEngine, "after the peer deleted a model") + + stub.set(http.StatusInternalServerError, nil) + assert.False(t, d.refreshNodeModels("peer-V", "127.0.0.1", port, ""), "a failed fetch should not report a change") + got, _ = d.dir.get("peer-V") + assert.Equal(t, map[string][]string{"vllm": {"Qwen/Qwen3.5-4B"}}, got.DownloadedByEngine, "after a failed fetch") +} + // TestApplyModelsGuard covers requirement 6 at the directory layer: a completed // fetch is discarded (ok == false) when the node was removed or re-addressed // while it was in flight, so it can neither resurrect a removed node nor @@ -365,7 +392,7 @@ func TestApplyModels_RejectsRemovedNodeAndStaleAddress(t *testing.T) { d := newDirectory() // Removed mid-sweep: node absent -> not resurrected. - _, changed, ok := d.applyModels("ghost", "127.0.0.1", 14322, []string{"a"}, nil, nil) + _, changed, ok := d.applyModels("ghost", "127.0.0.1", 14322, modelInventory{models: []string{"a"}}) assert.False(t, ok, "applyModels on an absent node = (changed (%v, %v)", changed, ok) assert.False(t, changed, "applyModels on an absent node = (changed (%v, %v)", changed, ok) _, present := d.get("ghost") @@ -378,11 +405,11 @@ func TestApplyModels_RejectsRemovedNodeAndStaleAddress(t *testing.T) { Services: map[noderec.ServiceKey]noderec.ServiceStatus{noderec.ServiceEngineManager: {Port: 14322}}, Models: []string{"old"}, }) - _, changed, ok = d.applyModels("peer-E", "127.0.0.1", 14322, []string{"new"}, nil, nil) + _, changed, ok = d.applyModels("peer-E", "127.0.0.1", 14322, modelInventory{models: []string{"new"}}) assert.False(t, ok, "applyModels with a stale IP = (changed (%v, %v)", changed, ok) assert.False(t, changed, "applyModels with a stale IP = (changed (%v, %v)", changed, ok) // em port changed -> also discarded. - _, changed, ok = d.applyModels("peer-E", "10.0.0.9", 99999, []string{"new"}, nil, nil) + _, changed, ok = d.applyModels("peer-E", "10.0.0.9", 99999, modelInventory{models: []string{"new"}}) assert.False(t, ok, "applyModels with a stale em port = (changed (%v, %v)", changed, ok) assert.False(t, changed, "applyModels with a stale em port = (changed (%v, %v)", changed, ok) got, _ := d.get("peer-E") diff --git a/services/nvpair-node-scanner/spec.md b/services/nvpair-node-scanner/spec.md index 3a78cd3e4..2bee8abd0 100644 --- a/services/nvpair-node-scanner/spec.md +++ b/services/nvpair-node-scanner/spec.md @@ -151,7 +151,8 @@ It exists so that discovery has exactly one implementation per host. One record "memory": { "total_bytes": 68719476736, "used_bytes": 21474836480 }, "models": ["llama3:8b", "qwen:0.5b"], "modelsByEngine": { "ollama": ["llama3:8b"], "lmstudio": ["qwen:0.5b"] }, - "loadedByEngine": { "ollama": ["llama3:8b"], "lmstudio": [] }, + "loadedByEngine": { "ollama": ["llama3:8b"], "lmstudio": [], "vllm": [] }, + "downloadedByEngine": { "vllm": ["Qwen/Qwen3.5-4B"] }, "lastSeen": 1767225600 } ``` @@ -262,7 +263,7 @@ The browser (`nvpair-shared/discovery`) is a scan-and-diff state machine over `_ Two independent decorations, both keyed by `hostUuid` and both backed by a last-known-good cache that a transient failure falls back to. A node may advertise one without the other. - **Hardware** — `GET http://:/v1/node-info`, 3-second timeout, plain HTTP. The response's `GPUs`, `cpu`, and `memory` land on the entry. This stays plain **even for a clustered peer**: the daemon must never infer mTLS from a node-level `cluster-uuid=`. The only way it becomes HTTPS is the operator setting `--client-cert` / `--client-key` / `--ca-bundle`, which is global and never per node. -- **Models** — `GET /v1/models` on the `em` port, 6-second timeout, deliberately longer than engine-manager's own 5-second models sweep so a peer with several engines is not cut off mid-sweep. The body is `{"models": [...], "modelsByEngine": {...}, "loadedByEngine": {...}}`; the last two are additive, so a peer that omits them yields nil maps while the flat union still enriches the node. A present engine key with `[]` is a successful empty inventory; a missing key means that engine was not running or queryable. +- **Models** — `GET /v1/models` on the `em` port, 6-second timeout, deliberately longer than engine-manager's own 5-second models sweep so a peer with several engines is not cut off mid-sweep. The body is `{"models": [...], "modelsByEngine": {...}, "loadedByEngine": {...}, "downloadedByEngine": {...}}`; the last three are additive, so a peer that omits them yields nil maps while the flat union still enriches the node. A present engine key with `[]` is a successful empty inventory; a missing key means that engine was not running or queryable. `downloadedByEngine` is the inventory of an engine that advertises only its served model (vLLM), reported even while it is stopped. All four are cached together as the node's last-good answer. - **Model transport is the cluster gate.** Which models a machine holds is cluster data, so a **peer** is fetched only over mTLS pinned to the principal it advertises. No pin, no fetch — there is deliberately no plaintext fallback for a peer. **This node's own** engine-manager is dialed over loopback in plaintext in every membership state, which is what keeps a standalone machine's own model list working. - **Dial host**: a peer is dialed at each address it published, in its own order, until one answers for the expected `hostUuid`; an answer from a different host is skipped rather than attributed to the entry. The address that answered is remembered per node and per service (`ni` and `em` can be reachable at different addresses), and later rounds ask only it until it stops answering — otherwise a multi-homed peer whose leading address this host cannot reach would cost that address's full timeout on every sweep. This node is always dialed at `127.0.0.1`, for enrichment, the periodic refresh, and telemetry, so a first-run firewall block on inbound to its own LAN address cannot blank the local entry. - **Node-info connections**: the node-info client holds at most one connection per origin, with a 1-second connect and TLS-handshake budget. A caller for an origin waits its turn *before* its own request timeout starts, so one waiting behind a slow but healthy response does not expire before its request is sent. A response body is drained up to 64 KiB so the connection can be reused. A peer whose node-info stops answering logs one warning, and one info line when it answers again. @@ -302,7 +303,7 @@ The broker also carries three feeds that are not directory events. It caches eac ### 7.8 Versioning -`v=1` in TXT identifies the record schema. New service keys and new `DirectoryNode` fields are **additive** and require no version bump: an unknown TXT key round-trips as a service port, and an unknown JSON field is ignored by an older consumer while `omitempty` keeps it off the wire when unset. `modelsByEngine`, `loadedByEngine`, and `ips` were added this way. A change to the meaning of an existing key, to the identity keying, or to the SRV-port-is-non-authoritative rule is a breaking change and bumps `v`; the schema is a flag-day rollout across the bundle rather than a mixed-version compatibility surface, since every binary ships together. The JSON-RPC method and notification names are frozen; new methods are additive. The binary's `main.Version` is stamped from `services/versions.json` and reported in `ready` and by `--version`. +`v=1` in TXT identifies the record schema. New service keys and new `DirectoryNode` fields are **additive** and require no version bump: an unknown TXT key round-trips as a service port, and an unknown JSON field is ignored by an older consumer while `omitempty` keeps it off the wire when unset. `modelsByEngine`, `loadedByEngine`, `downloadedByEngine`, and `ips` were added this way. A change to the meaning of an existing key, to the identity keying, or to the SRV-port-is-non-authoritative rule is a breaking change and bumps `v`; the schema is a flag-day rollout across the bundle rather than a mixed-version compatibility surface, since every binary ships together. The JSON-RPC method and notification names are frozen; new methods are additive. The binary's `main.Version` is stamped from `services/versions.json` and reported in `ready` and by `--version`. ## 8. Dependencies diff --git a/services/nvpair-proxy/README.md b/services/nvpair-proxy/README.md index 06968b631..059268303 100644 --- a/services/nvpair-proxy/README.md +++ b/services/nvpair-proxy/README.md @@ -73,7 +73,7 @@ parameter, because one flag cannot carry multiple engines' plans. | Field | Default | Description | |-------|---------|-------------| -| `engine` | *(required)* | Which engine to front: `ollama`, `lmstudio`, or `llamacpp`. An unknown name is rejected with the accepted values. | +| `engine` | *(required)* | Which engine to front: `ollama`, `lmstudio`, `llamacpp`, or `vllm`. An unknown name is rejected with the accepted values. | | `port` | per engine, see below | HTTP listen port for request forwarding. Must be 1–65535, or omitted for the engine's standalone default. `0` means "the default" rather than "pick an ephemeral port", and any other out-of-range value is rejected, because the facade announces the requested port in its `ready` notification and the broker would be told `0`. | | `aliasAddresses` | *(empty)* | Optional secondary `host:port` values for the same routing handler, one per loopback family so `localhost` resolves either way. Only literal loopback addresses are accepted; the broker uses this for a safe inherited local `OLLAMA_HOST`, and the aliases are not advertised to peers. Accepted only for an engine with an inherited host variable — today Ollama alone — and rejected for any other. | | `ignorePersistedPort` | `false` | Use `port` even when a saved port exists (used by broker-managed startup) | @@ -83,10 +83,10 @@ parameter, because one flag cannot carry multiple engines' plans. Everything engine-specific is one entry in `engines.go`, plus the shared identity in `nvpair-shared/engines`. -| | `"engine":"ollama"` | `"engine":"lmstudio"` | `"engine":"llamacpp"` | -|---|---|---|---| -| Facade id — error-ID prefix and broker relay namespace | `ollama-proxy` | `lmstudio-proxy` | `llamacpp-proxy` | -| Discovery service key | `ol` | `lm` | `lc` | +| | `"engine":"ollama"` | `"engine":"lmstudio"` | `"engine":"llamacpp"` | `"engine":"vllm"` | +|---|---|---|---|---| +| Facade id — error-ID prefix and broker relay namespace | `ollama-proxy` | `lmstudio-proxy` | `llamacpp-proxy` | `vllm-proxy` | +| Discovery service key | `ol` | `lm` | `lc` | `vl` | The **log component, supervisor label, and TUI health crash key are not in that table**: they name the process (`nvpair-proxy`), not a facade, because one @@ -95,17 +95,17 @@ Facade-scoped log records carry an `engine` field instead. The supervisor label and the health crash key are matched against each other, so they move together — see `nvpair-shared/engines` and `spec.md` §9. -| | `"engine":"ollama"` | `"engine":"lmstudio"` | `"engine":"llamacpp"` | -|---|---|---|---| -| Engine's own client-facing port | 11434 | 1234 | 8080 | -| Where PAIR relocates the engine | 11435 | 1235 | 8081 | -| Standalone port, used when `port` is omitted | 11435 | 1234 | 8080 | -| Persisted-port file (declared, not derived) | `proxy-port.json` | `lmstudio-proxy-port.json` | `llamacpp-proxy-port.json` | -| Model-list routes | `GET /api/tags` (native), `GET /v1/models` (OpenAI) | `GET /v1/models` (OpenAI) | `GET /models`, `GET /v1/models` (OpenAI; both query upstream `/models`) | -| Inference routes | `/api/generate`, `/api/chat`, `/api/embeddings`, `/api/embed`, plus the OpenAI and Anthropic Messages sets below | OpenAI and Anthropic Messages sets below | OpenAI set below | -| Model naming | untagged means `:latest`, so `llama3` and `llama3:latest` are one model | identifiers compared byte for byte | identifiers compared byte for byte | +| | `"engine":"ollama"` | `"engine":"lmstudio"` | `"engine":"llamacpp"` | `"engine":"vllm"` | +|---|---|---|---|---| +| Engine's own client-facing port | 11434 | 1234 | 8080 | 8000 | +| Where PAIR relocates the engine | 11435 | 1235 | 8081 | 8001 | +| Standalone port, used when `port` is omitted | 11435 | 1234 | 8080 | 8000 | +| Persisted-port file (declared, not derived) | `proxy-port.json` | `lmstudio-proxy-port.json` | `llamacpp-proxy-port.json` | `vllm-proxy-port.json` | +| Model-list routes | `GET /api/tags` (native), `GET /v1/models` (OpenAI) | `GET /v1/models` (OpenAI) | `GET /models`, `GET /v1/models` (OpenAI; both query upstream `/models`) | `GET /v1/models` (OpenAI) | +| Inference routes | `/api/generate`, `/api/chat`, `/api/embeddings`, `/api/embed`, plus the OpenAI and Anthropic Messages sets below | OpenAI and Anthropic Messages sets below | OpenAI set below | OpenAI set below | +| Model naming | untagged means `:latest`, so `llama3` and `llama3:latest` are one model | identifiers compared byte for byte | identifiers compared byte for byte | identifiers compared byte for byte | -All three facades classify these OpenAI-compatible POST routes as inference: +All four facades classify these OpenAI-compatible POST routes as inference: - `/v1/chat/completions` — [Chat Completions](https://developers.openai.com/api/reference/resources/chat/subresources/completions/methods/create) - `/v1/completions` — [Completions](https://developers.openai.com/api/reference/resources/completions/methods/create) @@ -119,8 +119,8 @@ The route table is a **classifier, not an allowlist**. An unlisted path is forwarded verbatim, which is how `/api/show`, `/api/pull`, `/api/ps`, `/api/version` and `OPTIONS` preflights keep working. -The broker enables all three facades by default. Local llama.cpp-compatible -clients use `8080`, while the managed `llama-server` stays on `8081`. +The broker enables all four facades by default. Local llama.cpp-compatible +clients use `8080`, while the managed `llama-server` stays on `8081`. vLLM clients use `8000`, while the managed `vllm serve` stays on `8001`. `--proxy-engines` can restrict a standalone broker or TUI to a subset. ### HTTP Reverse Proxy diff --git a/services/nvpair-proxy/engines.go b/services/nvpair-proxy/engines.go index 52a7544a7..7a4c89d70 100644 --- a/services/nvpair-proxy/engines.go +++ b/services/nvpair-proxy/engines.go @@ -151,6 +151,13 @@ var llamaCPPBaseRoutes = []route{ {Path: "/v1/models", UpstreamPath: "/models", Role: roleModelListOpenAIGET}, } +// vllmBaseRoutes is the engine-specific surface that vLLM exposes before the +// shared OpenAI routes are added. vLLM serves one model per process, so the +// fleet inventory is what makes /v1/models useful across nodes. +var vllmBaseRoutes = []route{ + {Path: "/v1/models", Role: roleModelListOpenAIGET}, +} + // openAIInferenceRoutes is the OpenAI-compatible inference surface. var openAIInferenceRoutes = []route{ {Path: "/v1/chat/completions", Role: roleInferencePOST}, @@ -170,9 +177,11 @@ func buildProfiles() []engineProfile { ollama, _ := engines.ByName("ollama") lmstudio, _ := engines.ByName("lmstudio") llamacpp, _ := engines.ByName("llamacpp") + vllm, _ := engines.ByName("vllm") ollamaRoutes := slices.Concat(ollamaBaseRoutes, openAIInferenceRoutes, anthropicInferenceRoutes) lmStudioRoutes := slices.Concat(lmStudioBaseRoutes, openAIInferenceRoutes, anthropicInferenceRoutes) llamaCPPRoutes := slices.Concat(llamaCPPBaseRoutes, openAIInferenceRoutes) + vllmRoutes := slices.Concat(vllmBaseRoutes, openAIInferenceRoutes) return []engineProfile{ { @@ -200,6 +209,13 @@ func buildProfiles() []engineProfile { ModelNaming: exactID, ReservedPersistedPort: 8081, }, + { + Engine: vllm, + StandalonePort: 8000, + Routes: vllmRoutes, + ModelNaming: exactID, + ReservedPersistedPort: 8001, + }, } } diff --git a/services/nvpair-proxy/engines_test.go b/services/nvpair-proxy/engines_test.go index 5c855b7e5..e08f78d3e 100644 --- a/services/nvpair-proxy/engines_test.go +++ b/services/nvpair-proxy/engines_test.go @@ -58,6 +58,8 @@ func TestRoleFor_ClassifiesOnlyDeclaredRoutes(t *testing.T) { require.True(t, ok, "lmstudio profile missing") llamacpp, ok := profileFor("llamacpp") require.True(t, ok, "llamacpp profile missing") + vllm, ok := profileFor("vllm") + require.True(t, ok, "vllm profile missing") accept := func(name string, profile engineProfile, method, path string, wantRole routeRole) { t.Run(name, func(t *testing.T) { @@ -104,6 +106,16 @@ func TestRoleFor_ClassifiesOnlyDeclaredRoutes(t *testing.T) { reject("lmstudio has no show", lmstudio, http.MethodPost, "/api/show") reject("llamacpp models post passthrough", llamacpp, http.MethodPost, "/models") + + accept("vllm chat", vllm, http.MethodPost, "/v1/chat/completions", roleInferencePOST) + accept("vllm completions", vllm, http.MethodPost, "/v1/completions", roleInferencePOST) + accept("vllm embeddings", vllm, http.MethodPost, "/v1/embeddings", roleInferencePOST) + accept("vllm responses", vllm, http.MethodPost, "/v1/responses", roleInferencePOST) + accept("vllm list", vllm, http.MethodGet, "/v1/models", roleModelListOpenAIGET) + // vLLM's own /health and /tokenize are forwarded verbatim, not classified. + reject("vllm health passthrough", vllm, http.MethodGet, "/health") + reject("vllm no native routes", vllm, http.MethodPost, "/api/chat") + reject("vllm wrong method on list", vllm, http.MethodPost, "/v1/models") } func TestIsInferenceRequest_FollowsTheProfile(t *testing.T) { diff --git a/services/nvpair-proxy/spec.md b/services/nvpair-proxy/spec.md index f2fa47d3a..ece4a60f4 100644 --- a/services/nvpair-proxy/spec.md +++ b/services/nvpair-proxy/spec.md @@ -100,7 +100,7 @@ The process starts with **no engine and no listener**. The broker then sends one A flag cannot express this. The broker plans a different port for each engine — Ollama's managed facade wants `:11434`, LM Studio's wants `:1234`, and -llama.cpp's managed facade wants `:8080`; any may be absent so the child keeps +llama.cpp's managed facade wants `:8080`, and vLLM's wants `:8000`; any may be absent so the child keeps its own persisted port — and a single-valued flag carries only one plan. ### 3.1 Why one process @@ -109,7 +109,7 @@ Between scheduler snapshots a facade takes short-lived **reservations** for work it has dispatched but that the scheduler has not yet observed. Those live in the process, and every facade shares them. If each engine had its own proxy process, none could see the others' reservations, so simultaneous bursts across -Ollama, LM Studio, and llama.cpp could select the same node, each incorrectly +Ollama, LM Studio, llama.cpp, and vLLM could select the same node, each incorrectly believing it was idle. Sharing the map is the reason the engines share a process. @@ -164,7 +164,7 @@ For a model-bearing inference request: 1. Filter a request-local discovery snapshot to nodes whose per-engine inventory advertises the requested model. Ollama normalizes the implicit `:latest` tag; - LM Studio and llama.cpp ids match exactly. An empty owner set returns a local + LM Studio, llama.cpp, and vLLM ids match exactly. An empty owner set returns a local `502` without contacting an engine. 2. Order the eligible owners: explicit `node/select` pin, then the scheduler's priority list, then deterministic default ordering. @@ -664,12 +664,12 @@ untouched, the primary listener stays up, and a warning is reported. ## 8. Ports -| | Ollama | LM Studio | llama.cpp | -| --- | --- | --- | --- | -| Engine's own client-facing port | 11434 | 1234 | 8080 | -| Where PAIR relocates the engine | 11435 | 1235 | 8081 | -| Standalone port, when `port` is omitted | 11435 | 1234 | 8080 | -| Persisted-port file | `proxy-port.json` | `lmstudio-proxy-port.json` | `llamacpp-proxy-port.json` | +| | Ollama | LM Studio | llama.cpp | vLLM | +| --- | --- | --- | --- | --- | +| Engine's own client-facing port | 11434 | 1234 | 8080 | 8000 | +| Where PAIR relocates the engine | 11435 | 1235 | 8081 | 8001 | +| Standalone port, when `port` is omitted | 11435 | 1234 | 8080 | 8000 | +| Persisted-port file | `proxy-port.json` | `lmstudio-proxy-port.json` | `llamacpp-proxy-port.json` | `vllm-proxy-port.json` | A port chosen at runtime via `set-port` is persisted per engine and restored when that facade is enabled, taking precedence over the requested port, so the diff --git a/services/nvpair-proxy/twofacade_test.go b/services/nvpair-proxy/twofacade_test.go index b35e1bb38..f392f54eb 100644 --- a/services/nvpair-proxy/twofacade_test.go +++ b/services/nvpair-proxy/twofacade_test.go @@ -193,7 +193,7 @@ func TestAddressedRequest_ReachesOnlyItsOwnFacade(t *testing.T) { } // An engine with no facade resolves to nothing rather than to a sibling. - assert.Nil(t, p.facadeFor("vllm"), "unknown engine resolved to the") + assert.Nil(t, p.facadeFor("nosuchengine"), "unknown engine resolved to the") assert.Nil(t, p.facadeFor(""), "an unaddressed message resolved to the") } diff --git a/services/nvpair-tui/ui/catalog_test.go b/services/nvpair-tui/ui/catalog_test.go index c2f066667..9c4957ef7 100644 --- a/services/nvpair-tui/ui/catalog_test.go +++ b/services/nvpair-tui/ui/catalog_test.go @@ -325,9 +325,9 @@ func TestCatalog_IgnoresOtherEnginesReply(t *testing.T) { // TestCatalogEmptyCatalogIsDistinctFromNoMatch checks the two empty states read // differently, because the fixes are different. func TestCatalogEmptyCatalog_IsDistinctFromNoMatch(t *testing.T) { - b := newCatalogBrowser(nil, "vllm", "vLLM", "this-host", false) + b := newCatalogBrowser(nil, "nosuchengine", "No Such Engine", "this-host", false) b.SetSize(100, 24) - b.update(catalogLoadedMsg{engine: "vllm", models: nil}) + b.update(catalogLoadedMsg{engine: "nosuchengine", models: nil}) assert.Contains(t, b.View(), "No catalog available") } diff --git a/services/nvpair-tui/ui/engineswire.go b/services/nvpair-tui/ui/engineswire.go index 453afb755..49a21031b 100644 --- a/services/nvpair-tui/ui/engineswire.go +++ b/services/nvpair-tui/ui/engineswire.go @@ -5,6 +5,7 @@ package ui import ( "context" + "encoding/json" "errors" "fmt" @@ -21,8 +22,11 @@ type engineStatus struct { DisplayName string `json:"display_name"` Installed bool `json:"installed"` Running bool `json:"running"` - Healthy bool `json:"healthy"` - Port int `json:"port"` + // Starting is true while a start is launching the engine and waiting for + // it to become ready, which for vLLM can be minutes. + Starting bool `json:"starting"` + Healthy bool `json:"healthy"` + Port int `json:"port"` } func (e engineStatus) label() string { @@ -35,11 +39,15 @@ func (e engineStatus) label() string { // modelsResult mirrors nvpair-engine-manager's ModelsResult, the engine:models // reply and the engine:models-changed payload. LoadedByEngine names the models // currently resident in memory, which the engine-manager polls for and pushes — -// so a client never has to poll to know what is loaded. +// so a client never has to poll to know what is loaded. DownloadedByEngine is +// the inventory of an engine that advertises only the model it is serving +// (vLLM): ModelsByEngine then holds just that one, and this holds every model it +// can load, even while it is stopped. type modelsResult struct { - Models []string `json:"models"` - ModelsByEngine map[string][]string `json:"modelsByEngine"` - LoadedByEngine map[string][]string `json:"loadedByEngine"` + Models []string `json:"models"` + ModelsByEngine map[string][]string `json:"modelsByEngine"` + LoadedByEngine map[string][]string `json:"loadedByEngine"` + DownloadedByEngine map[string][]string `json:"downloadedByEngine"` } // engineOp is one lifecycle request and how to describe it to the operator. @@ -133,6 +141,21 @@ var modelWires = map[string]modelWire{ "delete": modelOnly("delete_model"), "pull": modelOnly("pull_model"), }, + // vLLM takes the model under "model" alone. It serves one model per + // process, so load restarts it on the model and unload stops it. + engines.NameVLLM: { + "load": modelOnly("load_model"), + "unload": modelOnly("unload_model"), + "delete": modelOnly("delete_model"), + "pull": modelOnly("pull_model"), + }, +} + +// pullsWhileStopped are the engines whose download does not go through their +// server: vLLM downloads into the Hugging Face cache with its own tool. It +// mirrors the desktop's modelOpsWhenStopped capability. +var pullsWhileStopped = map[string]bool{ + engines.NameVLLM: true, } // downloadExamples show what a download name looks like for each engine. The @@ -142,6 +165,7 @@ var downloadExamples = map[string]string{ engines.NameOllama: "llama3.2", engines.NameLMStudio: "lmstudio-community/Qwen3-8B-GGUF", engines.NameLlamaCPP: "ggml-org/gemma-3-1b-it-GGUF:Q4_K_M", + engines.NameVLLM: "Qwen/Qwen3.5-4B", } // downloadPrompt is the placeholder for the download-by-name field, with an @@ -193,6 +217,10 @@ type engineOpMsg struct { // running on the engine. Distinct from both success and failure: nothing // went wrong, and nothing has finished either. detached bool + // stopped marks a Load that a Stop or Eject ended before the model was + // served. The user asked for that, so it is not a failure, and the model + // did not load either. + stopped bool } // engineCmd issues a lifecycle request against a node. An empty node means this @@ -239,11 +267,26 @@ func modelCmd(client *rpc.Client, node, engine string, act modelAction, model st } params = local } - return call(client, method, params, func(_ *rpc.Message, err error) tea.Msg { - return classifyOpResult(what, engine, act.op, err) + return call(client, method, params, func(msg *rpc.Message, err error) tea.Msg { + return classifyModelOpResult(what, engine, act.op, msg, err) }) } +// classifyModelOpResult is classifyOpResult for a model operation, which can +// also have been stopped: vLLM's Load starts the engine, and a Stop or Eject +// during it ends the Load with "stopped" in the reply. +func classifyModelOpResult(what, engine, op string, msg *rpc.Message, err error) tea.Msg { + if err == nil && msg != nil && op == "load" { + var reply struct { + Stopped bool `json:"stopped"` + } + if json.Unmarshal(msg.Result, &reply) == nil && reply.Stopped { + return engineOpMsg{what: what, engine: engine, stopped: true} + } + } + return classifyOpResult(what, engine, op, err) +} + // longRunningOps are the operations whose real duration is set by how much data // has to move or how slow an engine is to become ready, not by the RPC. // diff --git a/services/nvpair-tui/ui/engineswire_test.go b/services/nvpair-tui/ui/engineswire_test.go index ab2128349..95a75ecfe 100644 --- a/services/nvpair-tui/ui/engineswire_test.go +++ b/services/nvpair-tui/ui/engineswire_test.go @@ -5,6 +5,7 @@ package ui import ( "context" + "encoding/json" "errors" "testing" @@ -12,6 +13,7 @@ import ( "github.com/stretchr/testify/require" "nvpair-shared/engines" + "nvpair-tui/rpc" ) // wireOf builds a model operation's envelope, failing the test if the engine @@ -191,7 +193,7 @@ func TestDownloadPrompt_UsesTheEnginesSpelling(t *testing.T) { assert.Contains(t, downloadExamples, e.Name, "every engine needs a download example") } assert.Contains(t, downloadPrompt(engines.NameLlamaCPP, "llama.cpp"), ":Q4_K_M", "llama.cpp prompt must show the quantization suffix") - assert.Equal(t, "model name for vLLM", downloadPrompt("vllm", "vLLM")) + assert.Equal(t, "model name for No Such Engine", downloadPrompt("nosuchengine", "No Such Engine")) } // TestEveryEngineHasAModelWire is what makes adding an engine a matter of @@ -205,7 +207,7 @@ func TestEveryEngine_HasAModelWire(t *testing.T) { require.NoError(t, err, "%s must have a local wire for %s", e.Name, op) } } - _, err := modelActionWire("vllm", "load", "m") + _, err := modelActionWire("nosuchengine", "load", "m") assert.Error(t, err, "an engine with no wire must not be given one") } @@ -248,3 +250,27 @@ func TestLifecycleRemoteCoverage_MatchesManager(t *testing.T) { assert.False(t, engineOps[name].localOnly, "%s has a remote variant and should not be local-only", name) } } + +// A vLLM Load that a Stop or Eject ended replies with stopped set. That is what +// the user asked for, so it must not read as a failure, nor as a model loaded. +func TestModelOpResult_ReportsAStoppedLoadAsStopped(t *testing.T) { + test := func(name, op, result string, want engineOpMsg) { + t.Run(name, func(t *testing.T) { + msg := classifyModelOpResult("load org/m", "vllm", op, &rpc.Message{Result: json.RawMessage(result)}, nil) + assert.Equal(t, want, msg) + }) + } + test("a stopped Load", "load", `{"model":"org/m","restarted":false,"stopped":true}`, + engineOpMsg{what: "load org/m", engine: "vllm", stopped: true}) + test("a Load that finished", "load", `{"model":"org/m","restarted":true}`, + engineOpMsg{what: "load org/m", engine: "vllm"}) + test("another engine's Load result", "load", `{"ok":true}`, + engineOpMsg{what: "load org/m", engine: "vllm"}) + test("only a Load can be stopped", "delete", `{"stopped":true}`, + engineOpMsg{what: "load org/m", engine: "vllm"}) + + d := localDetail() + d.update(engineOpMsg{what: "load org/m", engine: "vllm", stopped: true}) + assert.Contains(t, d.status.render(), "was stopped before it finished") + assert.NotContains(t, d.status.render(), "failed") +} diff --git a/services/nvpair-tui/ui/nodedetail.go b/services/nvpair-tui/ui/nodedetail.go index ac624a4d4..6501d2715 100644 --- a/services/nvpair-tui/ui/nodedetail.go +++ b/services/nvpair-tui/ui/nodedetail.go @@ -7,6 +7,7 @@ import ( "context" "errors" "fmt" + "slices" "sort" "strconv" "strings" @@ -228,9 +229,10 @@ func newNodeDetail(client *rpc.Client, node nodeRow) *nodeDetail { // A remote node's models are already in the discovery snapshot the broker // enriched, so the list is populated before any request completes. d.models = modelsResult{ - Models: node.models, - ModelsByEngine: node.modelsByEngine, - LoadedByEngine: node.loadedByEngine, + Models: node.models, + ModelsByEngine: node.modelsByEngine, + LoadedByEngine: node.loadedByEngine, + DownloadedByEngine: node.downloadedByEngine, } d.refreshModels() d.adoptProbedTelemetry() @@ -264,7 +266,7 @@ func detailEngineColumns(w int, remote bool) []table.Column { cols := []column{ flexCol("ENGINE", 10, 1), fixedCol("INSTALLED", 9), - fixedCol("RUNNING", 7), + fixedCol("RUNNING", 8), // fits "starting" fixedCol("HEALTHY", 7), fixedCol("ENGINE PORT", 11), } @@ -627,6 +629,8 @@ func (d *nodeDetail) update(msg tea.Msg) (tea.Cmd, bool) { d.status.error("%s failed: %s", what, msg.err) case msg.detached: d.status.info("%s is continuing in the background - watch for progress", what) + case msg.stopped: + d.status.info("%s was stopped before it finished", what) default: d.status.ok("%s requested", what) } @@ -813,9 +817,10 @@ func (d *nodeDetail) handleNotification(msg *rpc.Message) tea.Cmd { clear(d.settings) case "engine:models-changed": - // The manager polls each running engine's resident set and pushes this - // on any change — explicit load/unload, LM Studio's JIT auto-load, and - // idle eviction alike. Re-reading is cheaper than merging the delta. + // The manager polls each engine's model lists and pushes this on any + // change — explicit load/unload, LM Studio's JIT auto-load, idle + // eviction, and a pull or delete alike. Re-reading is cheaper than + // merging the delta. if !d.remote() { return d.modelsCmd() } @@ -840,10 +845,12 @@ func (d *nodeDetail) handleNotification(msg *rpc.Message) tea.Cmd { d.node.models = n.Models d.node.modelsByEngine = n.ModelsByEngine d.node.loadedByEngine = n.LoadedByEngine + d.node.downloadedByEngine = n.DownloadedByEngine d.models = modelsResult{ - Models: n.Models, - ModelsByEngine: n.ModelsByEngine, - LoadedByEngine: n.LoadedByEngine, + Models: n.Models, + ModelsByEngine: n.ModelsByEngine, + LoadedByEngine: n.LoadedByEngine, + DownloadedByEngine: n.DownloadedByEngine, } d.refreshModels() @@ -1605,10 +1612,13 @@ func (d *nodeDetail) uninstallEngine(engine *engineStatus) tea.Cmd { } // pullTargetEngine is the engine a download goes into: the one highlighted in -// the engine pane when it can serve, else the only running engine. Downloading -// requires a running engine, so a stopped one is never chosen silently. +// the engine pane when it can download, else the only running engine. Most +// engines download through their running server, so a stopped one is never +// chosen silently. An engine that downloads without its server (vLLM, into the +// Hugging Face cache) can be chosen stopped when it is highlighted, as the +// desktop app allows. func (d *nodeDetail) pullTargetEngine() string { - if e := d.selectedEngine(); e != nil && e.Running { + if e := d.selectedEngine(); e != nil && (e.Running || e.Installed && pullsWhileStopped[e.Engine]) { return e.Engine } for _, e := range d.engines { @@ -1642,7 +1652,11 @@ func (d *nodeDetail) refreshEngines() { if e.Port != 0 { port = strconv.Itoa(e.Port) } - row := table.Row{e.label(), yesNo(e.Installed), yesNo(e.Running), yesNo(e.Healthy), port} + running := yesNo(e.Running) + if e.Starting { + running = "starting" + } + row := table.Row{e.label(), yesNo(e.Installed), running, yesNo(e.Healthy), port} if d.proxy != nil { row = append(row, d.proxyPortCell(e.Engine)) } @@ -1681,10 +1695,15 @@ func (d *nodeDetail) refreshModels() { // where it already was. selected := d.selectedModelKey() - engines := make([]string, 0, len(d.models.ModelsByEngine)) + engines := make([]string, 0, len(d.models.ModelsByEngine)+len(d.models.DownloadedByEngine)) for name := range d.models.ModelsByEngine { engines = append(engines, name) } + for name := range d.models.DownloadedByEngine { + if _, listed := d.models.ModelsByEngine[name]; !listed { + engines = append(engines, name) + } + } sort.Strings(engines) d.modelRows = d.modelRows[:0] @@ -1696,7 +1715,7 @@ func (d *nodeDetail) refreshModels() { for _, m := range residents { loaded[m] = true } - models := append([]string(nil), d.models.ModelsByEngine[engine]...) + models := managedModels(d.models, engine) sort.Strings(models) for _, m := range models { d.modelRows = append(d.modelRows, @@ -1941,6 +1960,25 @@ func (d *nodeDetail) hardwareBudget() int { return clampWidth(d.height/2, 1) } +// managedModels is the list of models to manage for one engine. An engine that +// reports a downloaded inventory (vLLM, which advertises only the model it is +// serving) is listed from that inventory, which is read even while it is +// stopped, plus anything loaded that is not in it, such as a model served from +// a local path. Every other engine's advertised list is its inventory. +func managedModels(res modelsResult, engine string) []string { + downloaded, ok := res.DownloadedByEngine[engine] + if !ok { + return slices.Clone(res.ModelsByEngine[engine]) + } + models := slices.Clone(downloaded) + for _, m := range res.LoadedByEngine[engine] { + if !slices.Contains(models, m) { + models = append(models, m) + } + } + return models +} + // emptyModelsHint explains an empty list in terms of the thing to fix, since // "no models" has several quite different causes. // emptyEnginesHint explains an engine list that came back empty. @@ -2025,6 +2063,12 @@ func (d *nodeDetail) emptyModelsHint() string { // something needs starting. Whether an engine is running is exactly // what cannot be seen from outside the cluster. return " No models advertised by this node." + case len(d.models.DownloadedByEngine) > 0 && d.remote(): + // An engine read its downloaded models and found none, which is known + // with the engine stopped. One that could not be read has no entry. + return " No models downloaded on this node's engines." + case len(d.models.DownloadedByEngine) > 0: + return " No models installed. Press p to download one." case !d.anyEngineRunning(): return " No models reported - start an engine to see its models." case d.remote(): diff --git a/services/nvpair-tui/ui/nodedetail_test.go b/services/nvpair-tui/ui/nodedetail_test.go index 6a168c5ee..7aa65cc8e 100644 --- a/services/nvpair-tui/ui/nodedetail_test.go +++ b/services/nvpair-tui/ui/nodedetail_test.go @@ -105,6 +105,84 @@ func TestUnreportedLoadState_IsUnknown(t *testing.T) { } } +// modelRowStates maps each model row to its LOADED column. +func modelRowStates(d *nodeDetail) map[string]string { + got := map[string]string{} + for _, row := range d.modelTable.Rows() { + got[row[0]] = row[2] + } + return got +} + +// Ejecting a vLLM model stops the engine, so it advertises nothing. Its +// downloaded models must keep their rows, so they can still be loaded or deleted. +func TestRemoteDetail_KeepsTheDownloadedModelsOfAStoppedVLLMPeer(t *testing.T) { + d := remoteDetail() + + d.update(discoveryPush(t, availableNode{ + HostUUID: "peer", + Name: "peer-host", + LoadedByEngine: map[string][]string{"vllm": {}}, + DownloadedByEngine: map[string][]string{"vllm": {"Qwen/Qwen3.5-4B", "mistralai/Ministral-3-3B-Instruct-2512"}}, + })) + + assert.Equal(t, map[string]string{ + "Qwen/Qwen3.5-4B": "no", + "mistralai/Ministral-3-3B-Instruct-2512": "no", + }, modelRowStates(d)) +} + +// With --served-model-name, vLLM answers to the alias, but the model manager +// lists the cached repository, which the loaded list names. The repository row +// shows as loaded and the alias gets no row of its own. +func TestModelRows_MarkTheCachedModelLoadedWhenVLLMServesItUnderAnAlias(t *testing.T) { + d := localDetail() + d.models = modelsResult{ + Models: []string{"qwen-api"}, + ModelsByEngine: map[string][]string{"vllm": {"qwen-api"}}, + LoadedByEngine: map[string][]string{"vllm": {"Qwen/Qwen3.5-4B"}}, + DownloadedByEngine: map[string][]string{"vllm": {"Qwen/Qwen3.5-4B", "mistralai/Ministral-3-3B-Instruct-2512"}}, + } + d.refreshModels() + + assert.Equal(t, map[string]string{ + "Qwen/Qwen3.5-4B": "yes", + "mistralai/Ministral-3-3B-Instruct-2512": "no", + }, modelRowStates(d)) +} + +// A model vLLM serves from a local path is not in the cache, but it is loaded +// and still gets a row. +func TestModelRows_ListAServedModelThatIsNotInTheDownloadedList(t *testing.T) { + d := localDetail() + d.models = modelsResult{ + ModelsByEngine: map[string][]string{"vllm": {"/models/qwen"}}, + LoadedByEngine: map[string][]string{"vllm": {"/models/qwen"}}, + DownloadedByEngine: map[string][]string{"vllm": {"Qwen/Qwen3.5-4B"}}, + } + d.refreshModels() + + assert.Equal(t, map[string]string{"/models/qwen": "yes", "Qwen/Qwen3.5-4B": "no"}, modelRowStates(d)) +} + +// "Nothing downloaded" is said only when an engine read its inventory and found +// it empty. A peer whose inventory could not be read is not described that way. +func TestEmptyModelsHint_SaysNothingIsDownloadedOnlyWhenTheInventoryWasRead(t *testing.T) { + hintFor := func(downloaded map[string][]string) string { + d := remoteDetail() + d.node.membership = membershipMember + d.models = modelsResult{DownloadedByEngine: downloaded} + d.refreshModels() + return d.emptyModelsHint() + } + t.Run("inventory read and empty", func(t *testing.T) { + assert.Contains(t, hintFor(map[string][]string{"vllm": {}}), "No models downloaded") + }) + t.Run("inventory not read", func(t *testing.T) { + assert.NotContains(t, hintFor(nil), "downloaded") + }) +} + // TestLocalDetailIgnoresDiscoveryModels checks this machine keeps using the // authoritative engine:models RPC rather than the discovery summary. func TestLocalDetail_IgnoresDiscoveryModels(t *testing.T) { @@ -218,7 +296,7 @@ func TestProxyIndexForEngine_NormalizesNamesAndRejectsUnknownEngine(t *testing.T assert.Equal(t, 1, p.indexForEngine("lmstudio")) // Case and padding must not decide whether a port renders. assert.Equal(t, 0, p.indexForEngine(" OLLAMA "), "normalisation") - assert.Equal(t, -1, p.indexForEngine("vllm"), "unknown engine") + assert.Equal(t, -1, p.indexForEngine("nosuchengine"), "unknown engine") } // seedSettings puts a settings snapshot in the cache, as a fetch or a push @@ -1442,3 +1520,30 @@ func TestServiceTab_OmitsPortSettings(t *testing.T) { assert.NotContains(t, strings.ToLower(it.label), "port", "Service tab still offers %q; ports belong on the node", it.label) } } + +// vLLM downloads into the Hugging Face cache without its server, as the desktop +// app allows, so a stopped vLLM can take a download when it is highlighted. An +// engine that downloads through its server still needs to be running. +func TestPullTargetEngine_AcceptsAStoppedEngineThatDownloadsWithoutItsServer(t *testing.T) { + test := func(name string, highlighted engineStatus, want string) { + t.Run(name, func(t *testing.T) { + d := localDetail() + d.engines = []engineStatus{highlighted} + d.refreshEngines() + assert.Equal(t, want, d.pullTargetEngine()) + }) + } + test("a stopped vLLM", engineStatus{Engine: "vllm", Installed: true}, "vllm") + test("a vLLM that is not installed", engineStatus{Engine: "vllm"}, "") + test("a stopped Ollama", engineStatus{Engine: "ollama", Installed: true}, "") + test("a running Ollama", engineStatus{Engine: "ollama", Installed: true, Running: true}, "ollama") +} + +// A vLLM start can download and load for minutes. It is shown as starting, +// not as not running. +func TestEngineTable_ShowsAnEngineThatIsStarting(t *testing.T) { + d := localDetail() + d.engines = []engineStatus{{Engine: "vllm", DisplayName: "vLLM", Installed: true, Starting: true}} + d.refreshEngines() + assert.Contains(t, d.View(), "starting") +} diff --git a/services/nvpair-tui/ui/nodes_test.go b/services/nvpair-tui/ui/nodes_test.go index 834584827..9a152e3ad 100644 --- a/services/nvpair-tui/ui/nodes_test.go +++ b/services/nvpair-tui/ui/nodes_test.go @@ -28,7 +28,7 @@ type removalCall struct { func removalBroker(t *testing.T, fail string) (*rpc.Client, <-chan removalCall) { t.Helper() c1, c2 := net.Pipe() - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancel(t.Context()) t.Cleanup(func() { cancel() c1.Close() diff --git a/services/nvpair-tui/ui/nodesmodel.go b/services/nvpair-tui/ui/nodesmodel.go index a999ae6fe..91b4bf738 100644 --- a/services/nvpair-tui/ui/nodesmodel.go +++ b/services/nvpair-tui/ui/nodesmodel.go @@ -125,9 +125,10 @@ type nodeRow struct { nodeInfoTLS bool probedTelemetry *nodeTelemetry - models []string - modelsByEngine map[string][]string - loadedByEngine map[string][]string + models []string + modelsByEngine map[string][]string + loadedByEngine map[string][]string + downloadedByEngine map[string][]string } // modelCount is the number of distinct models the node advertises. @@ -175,6 +176,7 @@ func mergeNodes(in nodeFeeds) []nodeRow { row.models = d.Models row.modelsByEngine = d.ModelsByEngine row.loadedByEngine = d.LoadedByEngine + row.downloadedByEngine = d.DownloadedByEngine row.presence = discoveredPresence(d) switch { case d.Trusted: diff --git a/services/nvpair-tui/ui/nodeswire.go b/services/nvpair-tui/ui/nodeswire.go index 6721b91c8..a1604c5ff 100644 --- a/services/nvpair-tui/ui/nodeswire.go +++ b/services/nvpair-tui/ui/nodeswire.go @@ -43,11 +43,13 @@ type availableNode struct { // The broker enriches each discovered node with the model inventory it // learned from that peer's engine-manager, so a remote node's models need // no extra call: Models is the flat de-duplicated union, ModelsByEngine - // attributes each to the engine serving it, and LoadedByEngine names those - // currently resident in memory. - Models []string `json:"models"` - ModelsByEngine map[string][]string `json:"modelsByEngine"` - LoadedByEngine map[string][]string `json:"loadedByEngine"` + // attributes each to the engine serving it, LoadedByEngine names those + // currently resident in memory, and DownloadedByEngine lists what an engine + // that advertises only its served model (vLLM) has downloaded. + Models []string `json:"models"` + ModelsByEngine map[string][]string `json:"modelsByEngine"` + LoadedByEngine map[string][]string `json:"loadedByEngine"` + DownloadedByEngine map[string][]string `json:"downloadedByEngine"` } // clusterIdentity is this node's principal, from cluster:get-node-id. diff --git a/services/nvpair-tui/ui/proxystatus_test.go b/services/nvpair-tui/ui/proxystatus_test.go index a4e46c6aa..4f4a9ef48 100644 --- a/services/nvpair-tui/ui/proxystatus_test.go +++ b/services/nvpair-tui/ui/proxystatus_test.go @@ -109,7 +109,7 @@ func TestPortForEngine_DistinguishesDownFromUnknown(t *testing.T) { p := newProxyTracker() _, ready := p.portForEngine("ollama") assert.False(t, ready, "a proxy that has never reported reads as ready") - port, ready := p.portForEngine("vllm") + port, ready := p.portForEngine("nosuchengine") assert.False(t, ready) assert.Zero(t, port) } diff --git a/services/nvpair-ui-broker/ENGINE_SETTINGS.md b/services/nvpair-ui-broker/ENGINE_SETTINGS.md index 5d5f5728e..c3ebccf0b 100644 --- a/services/nvpair-ui-broker/ENGINE_SETTINGS.md +++ b/services/nvpair-ui-broker/ENGINE_SETTINGS.md @@ -16,7 +16,7 @@ propagation is not part of this API. | `engine:preview-settings` | `{engine, nodeId?, expectedRevision, settings, resolution?}` | Normalized settings, errors, conflict, restart/rebind summary | | `engine:apply-settings` | Preview request plus `requestId` | `{revision, phase}` acknowledgement | -`engine` is `ollama`, `lmstudio`, or `llamacpp`. `settings` contains all three +`engine` is `ollama`, `lmstudio`, `llamacpp`, or `vllm`. `settings` contains all three fields: `serverPort`, `proxyPort`, `launchText`. The last field contains arguments and leading environment assignments, without the executable or startup subcommand. The argument grammar is diff --git a/services/nvpair-ui-broker/README.md b/services/nvpair-ui-broker/README.md index bf0ac1a3c..c8da2d897 100644 --- a/services/nvpair-ui-broker/README.md +++ b/services/nvpair-ui-broker/README.md @@ -23,7 +23,7 @@ namespace: | --- | --- | --- | | `nvpair-node-scanner` | Discovery daemon: advertises this host's one `_nvpair-node._tcp` record and browses the LAN | `discovery:*` | | `nvpair-node-info` | Local GPU / CPU / memory inventory over HTTP at `/v1/node-info` | — (HTTP only) | -| `nvpair-proxy` | One process hosting an inference proxy and router facade per enabled engine | `ollama-proxy:*`, `lmstudio-proxy:*`, `llamacpp-proxy:*` | +| `nvpair-proxy` | One process hosting an inference proxy and router facade per enabled engine | `ollama-proxy:*`, `lmstudio-proxy:*`, `llamacpp-proxy:*`, `vllm-proxy:*` | | `nvpair-engine-manager` | Local engine and model control plane; also serves `GET /v1/models` to peers | `engine:*` | | `nvpair-cluster-manager` | Node identity, trusted-node store, PIN pairing | `cluster:*`, `nodes:*` | | `nvpair-workload-manager` | Cluster workload relay between this node and peers | `workloads:*` | @@ -39,9 +39,9 @@ lifecycle, and relay rules. Two responsibilities live in the broker itself rather than in a worker: -- **Engine advertising.** The broker polls local Ollama, LM Studio, and - llama.cpp every 5 s and registers each running engine's promoted facade port - (`ol` / `lm` / `lc`) with the discovery daemon, so they are carried in this +- **Engine advertising.** The broker polls local Ollama, LM Studio, + llama.cpp, and vLLM every 5 s and registers each running engine's promoted facade port + (`ol` / `lm` / `lc` / `vl`) with the discovery daemon, so they are carried in this host's single `_nvpair-node` record. The model list is not part of that record — it is served over HTTP by `nvpair-engine-manager` on the `em` service and fetched by a peer's daemon @@ -74,7 +74,7 @@ Bidirectional newline-delimited JSON-RPC 2.0 — same conventions as every other | `--scanner-path ` | `./nvpair-node-scanner[.exe]` in the CWD | Explicit path to the `nvpair-node-scanner` binary the broker should spawn | | `--node-info-path ` | `./nvpair-node-info[.exe]` in the CWD | Explicit path to the `nvpair-node-info` binary the broker should spawn. When omitted and no default sibling exists, the broker runs without the local inventory server (non-fatal); when set to an invalid path, the broker exits with an error | | `--proxy-path ` | `./nvpair-proxy[.exe]` in the CWD | Explicit path to the `nvpair-proxy` binary. One process fronts every engine: the broker spawns it once and then sends a `facade/enable` per entry in `--proxy-engines`. Same optional semantics as `--node-info-path`: an absent default sibling means no local proxies (non-fatal); an invalid explicit path exits with an error | -| `--proxy-engines ` | `ollama,lmstudio,llamacpp` | Which engines to front with a proxy. An unrecognized name exits with an error rather than being skipped, so a typo cannot look like it worked. An engine left out is not started **and not prepared** — the broker will not relocate an engine whose facade nothing is going to claim | +| `--proxy-engines ` | `ollama,lmstudio,llamacpp,vllm` | Which engines to front with a proxy. An unrecognized name exits with an error rather than being skipped, so a typo cannot look like it worked. An engine left out is not started **and not prepared** — the broker will not relocate an engine whose facade nothing is going to claim | | `--workload-manager-path ` | `./nvpair-workload-manager[.exe]` in the CWD | Explicit path to the `nvpair-workload-manager` binary the broker spawns for the cluster workload relay. Same optional semantics as `--node-info-path`: an absent default sibling means no workload relay (non-fatal); an invalid explicit path exits with an error | | `--errors-path ` | `./nvpair-errors[.exe]` in the CWD | Explicit path to the `nvpair-errors` binary the broker spawns (with `--peer-sync`) for the service-error pipeline. Same optional semantics as `--node-info-path`: an absent default sibling means the error pipeline is disabled — producers' errors are dropped (non-fatal); an invalid explicit path exits with an error | | `--engine-manager-path ` | `./nvpair-engine-manager[.exe]` in the CWD | Explicit path to the `nvpair-engine-manager` binary the broker spawns for engine management. Same optional semantics as `--node-info-path` | @@ -97,7 +97,7 @@ On startup — **before** emitting `app:ready` — the broker spawns the scanner **`nvpair-node-info`** is spawned next. It's a server, not an event source: it stands up the local `/v1/node-info` HTTP endpoint (GPU/CPU/memory inventory). It does not advertise itself — the broker registers its `ni` port with the scanner daemon, which carries it in the node record, and a peer's daemon fetches `/v1/node-info` over plain HTTP to enrich the node. The broker doesn't read anything back from node-info's stdout (drained and discarded). Spawning it is **optional**: if the binary can't be resolved (and no `--node-info-path` override was given) the broker logs a warning and continues serving discovery without it. **Engine advertising.** The broker runs an internal 5 s poll loop against local -Ollama (`GET /`), LM Studio (`GET /v1/models`), and llama.cpp (`GET /health`) +Ollama (`GET /`), LM Studio (`GET /v1/models`), llama.cpp (`GET /health`), and vLLM (`GET /health`) at their configured backend ports and reconciles this node's engine registration with the scanner daemon: @@ -122,6 +122,10 @@ withdrawn or answered with an error on its own, leaving the others serving. Ollama's standalone default is `:11435`; with managed port ownership enabled (the default), the broker starts settings and engine-manager first, claims `:11434` with that facade, and only then moves a stopped default-port Ollama backend to a free port. Custom backend ports are preserved. When the inherited `OLLAMA_HOST` names a distinct local plaintext port, the broker also gives the facade that normalized loopback-only alias so clients already using the variable enter the same routing path; `localhost` reserves both canonical loopback families atomically, while remote and HTTPS targets are ignored. The alias port is reserved against every configured engine, local or remote engine start override, every facade's control plane, and every managed engine's backend port plan, so a backend that has to move can never land on the alias. A running Ollama or unknown owner on either requested port is never stopped or moved: the primary uses a safe fallback when needed, and an occupied alias remains with its owner while the broker reports a warning. +vLLM is a managed engine handled exactly like llama.cpp below, with the +`vllm serve` process on `:8001` and its default-enabled facade on `:8000` or a +safe fallback. + llama.cpp is a managed engine: engine-manager owns its lifecycle and starts it on its configured loopback port, `:8081` by default. The broker places its default-enabled facade on `:8080` or a safe fallback, without automatically @@ -160,7 +164,7 @@ Two classes of proxy notification are **not** re-emitted under any telemetry, and discovery snapshot. It smooths fresh utilization into pressure 0–3, uses neutral pressure 1 for invalid/missing/older-than-10-second samples, and orders by `pending + gpuPressure`, then pressure, then stable UUID. Load is -node-wide across Ollama, LM Studio, and llama.cpp because they normally contend +node-wide across Ollama, LM Studio, llama.cpp, and vLLM because they normally contend for the same resources. Each engine-specific `schedule:priority` carries `{engine,nodes,ranks}` and refreshes when order, pending counts, or pressure changes. The broker caches, generation-orders, and replays the full @@ -284,7 +288,7 @@ ServiceError = { id, message, timestamp, nodeId?, severity?, action?, engineType #### `engine:` -**Opt-in** (`engine:subscribe`). The notifications `nvpair-engine-manager` emits — `engine:ready`, `engine:state-changed` (full `EngineStatus`), `engine:models-changed` (`{engine, models}`, pushed when an engine's set of models loaded in memory changes), `engine:install-progress`, `engine:pull-progress` (live progress for a local model pull), `engine:remote-progress` (live progress for a remote install/pull) — re-emitted verbatim (their method names are already `engine:`-prefixed). Off by default. +**Opt-in** (`engine:subscribe`). The notifications `nvpair-engine-manager` emits — `engine:ready`, `engine:state-changed` (full `EngineStatus`), `engine:models-changed` (`{engine, models}`, pushed when an engine's served, loaded-in-memory or downloaded model list changes), `engine:install-progress`, `engine:pull-progress` (live progress for a local model pull), `engine:remote-progress` (live progress for a remote install/pull) — re-emitted verbatim (their method names are already `engine:`-prefixed). Off by default. #### `connection/cluster-identity` / `connection/cluster-auto-sync` @@ -336,6 +340,7 @@ Field-by-field: | `models` | string[] | the node's available model names, enriched by the daemon from the node's engine-manager `em` endpoint (`GET /v1/models`). Omitted when the node advertises no engine-manager or no engine is running. | | `modelsByEngine` | object | the same models attributed to the engine that serves each (keyed by engine name, e.g. `ollama` / `lmstudio`). Additive alongside the flat `models` union. A present engine key with `[]` means its inventory was successfully queried and is empty; a missing key means it was not running/queryable. Omitted when no engine inventory was successfully reported. | | `loadedByEngine` | object | the models currently **loaded in memory** for each engine (normally a subset of `modelsByEngine`), keyed by engine name. Enriched from the peer's `loadedByEngine`. An engine key with an empty list means "running, nothing loaded"; a missing key means loaded state wasn't reported. Omitted when no engine reports loaded state. | +| `downloadedByEngine` | object | every model an engine can load, for an engine that advertises only the model it is serving (vLLM, whose `modelsByEngine` entry is just that served name). Reported even while the engine is stopped, so a remote model manager keeps its Load and Delete rows; its `loadedByEngine` entry names models the same way. An engine key with an empty list means the inventory was read and is empty; a missing key means it was not read. Omitted when no engine reports one. | A node that the scanner reports as removed is deleted from the snapshot; `lastSeen` is not preserved for removed nodes. The response is wrapped in `{nodes: [...]}` (rather than a bare array) so we can grow summary fields later without breaking clients. An empty list is a normal early-startup state, not an error. @@ -435,7 +440,7 @@ Two relay-specific error cases: - If no proxy is being supervised (or it has exited), the broker replies with error `-32000` `"ollama-proxy not available"`. - `ollama-proxy:shutdown` is **refused** with error `-32601` — the broker owns the proxy's lifecycle, so a client can't terminate it independently. Shut the broker down instead (which tears the proxy down with it). -#### `ollama-proxy:set-port` / `lmstudio-proxy:set-port` / `llamacpp-proxy:set-port` +#### `ollama-proxy:set-port` / `lmstudio-proxy:set-port` / `llamacpp-proxy:set-port` / `vllm-proxy:set-port` **Intercepted, not relayed verbatim.** For every engine profile, a port-only caller (none in tree; `nvpair-tui` uses `engine:apply-settings`) gets to move a single port without rendering the whole launch settings form, but the change still runs through the same authoritative settings operation the desktop editor uses, so a port set from the terminal cannot diverge from one set from the UI. The namespace selects the engine. The broker reads that engine's current settings, substitutes the requested proxy port, and applies the result while preserving the server port and launch arguments. Accepted settings are saved in the settings journal. @@ -452,7 +457,7 @@ Automatic conflict resolution still exists, but only for a port the user did not The LM Studio counterpart of the `ollama-proxy:*` surface runs the supervised `lmstudio-proxy` on compatibility port `:1234` and tracks the managed LM Studio backend on `:1235`. With managed port ownership enabled (the default), the broker identifies and moves an existing LM Studio server through engine-manager before allowing the proxy to claim `1234`; unknown owners are left untouched and force a warned proxy fallback. Disabling managed ownership preserves explicit custom backend and proxy ports. `lmstudio-proxy:get-status` reports the actual bound port; `lmstudio-proxy:subscribe` / `lmstudio-proxy:unsubscribe` opt into / out of its `lmstudio-proxy:` stream; `lmstudio-proxy:set-port` uses the settings operation described above; and other `lmstudio-proxy:` requests are relayed with the prefix translated (`nodes/list`, `node/select`, `node/add-manual`, `node/remove-manual`, ...). `lmstudio-proxy:shutdown` is refused because the broker owns lifecycle ordering. Workload and error events feed the shared streams exactly as Ollama's do. -The same broker-local status, subscription, and port-setting methods apply to `llamacpp-proxy:*`. llama.cpp port changes use the same validation, journal, and application path; remaining methods use the generic facade relay, with `llamacpp-proxy:shutdown` refused. +The same broker-local status, subscription, and port-setting methods apply to `llamacpp-proxy:*` and `vllm-proxy:*`. llama.cpp port changes use the same validation, journal, and application path; remaining methods use the generic facade relay, with `llamacpp-proxy:shutdown` refused. #### `ollama-proxy:subscribe` @@ -627,8 +632,8 @@ Attach to a pre-existing endpoint: ## What this version intentionally does NOT do (yet) - **Engine-advertise control surface.** Engine registration is auto-driven only: - the broker tracks local Ollama, LM Studio, and llama.cpp and registers `ol`, - `lm`, or `lc` with the daemon while each is up. There's no manual-advertise + the broker tracks local Ollama, LM Studio, llama.cpp, and vLLM and registers `ol`, + `lm`, `lc`, or `vl` with the daemon while each is up. There's no manual-advertise RPC (custom service, port, name, or TXT), and no way to advertise anything other than the detected engines. - **node-info control surface.** node-info is spawned and torn down with the broker, and the broker pushes it only two things over stdin: the log level, and this node's cluster principal (`nodeinfo:set-cluster-identity`, sent on spawn and on every membership or pin-set change, because node-info holds no cluster dir and so cannot read membership itself). Otherwise it's hands-off: the broker registers its port with the daemon (which enriches over plain HTTP) but doesn't pass through TLS material (`--cert` / `--key` / `--client-ca`) or a custom `--port`, and exposes no RPC to query or reconfigure it. It runs with its own defaults plus those two pushes. diff --git a/services/nvpair-ui-broker/broker.go b/services/nvpair-ui-broker/broker.go index 14e75c9ad..5464d2622 100644 --- a/services/nvpair-ui-broker/broker.go +++ b/services/nvpair-ui-broker/broker.go @@ -114,6 +114,14 @@ type AvailableNode struct { // consumer show which of a remote node's models are loaded. Omitted // when no engine reports loaded state. LoadedByEngine map[string][]string `json:"loadedByEngine,omitempty"` + // DownloadedByEngine lists every model an engine can load when that engine + // advertises only the one it is serving (vLLM), so ModelsByEngine holds just + // that one. It is reported even while the engine is stopped, so a remote + // node's model manager keeps its Load and Delete rows, and its LoadedByEngine + // entry names models the same way. An engine key with an empty list means the + // inventory was read and is empty; a missing key means it was not read. + // Omitted when no engine reports one. + DownloadedByEngine map[string][]string `json:"downloadedByEngine,omitempty"` } // GetNodesResult is the response to "discovery:get-nodes". Wrapped in an @@ -3423,7 +3431,20 @@ func (b *Broker) relayToEngineNow(msg *Message) { } requestedEngine, requestedPort, isEnginePortAssignment := enginePortAssignmentRequest(msg.Method, msg.Params) - if isEnginePortAssignment || msg.Method == "engine:stop" || msg.Method == "engine:start" || msg.Method == "engine:restart" { + switch { + case msg.Method == "engine:stop": + // A Stop must reach the engine manager while a Start is still waiting for + // readiness, which can take minutes, because cancelling that Start is + // what it is for. The Start holds the lock for that whole time, so a + // Stop that waited for it would only arrive after the Start had finished + // on its own. A Stop changes no launch settings, so it takes the lock + // only afterwards, to reconcile, once its reply is sent. + defer func() { + b.engineConfigMu.Lock() + defer b.engineConfigMu.Unlock() + b.reconcileLegacySettingsLocked() + }() + case isEnginePortAssignment || msg.Method == "engine:start" || msg.Method == "engine:restart": b.engineConfigMu.Lock() defer b.engineConfigMu.Unlock() defer b.reconcileLegacySettingsLocked() diff --git a/services/nvpair-ui-broker/broker_lifecycle_test.go b/services/nvpair-ui-broker/broker_lifecycle_test.go index 084e608bf..51d9e73c6 100644 --- a/services/nvpair-ui-broker/broker_lifecycle_test.go +++ b/services/nvpair-ui-broker/broker_lifecycle_test.go @@ -5,6 +5,7 @@ package main import ( "context" + "encoding/json" "net" "path/filepath" "sync" @@ -255,3 +256,68 @@ func TestUnbindableLMStudioFacade_FinishesWithoutRestartingTheProcess(t *testing assert.Same(t, proxy, b.getLMStudioProxy(), "a facade port failure replaced the shared proxy handle") assert.Same(t, proxy, b.getProxy(), "a facade port failure disturbed the other engine's handle on the same process") } + +// A Stop has to reach the engine manager while a Start is still waiting for +// readiness: cancelling that Start is what it is for. It used to wait for the +// lock the Start holds, so it arrived only after the Start finished by itself. +func TestRelayToEngine_DeliversAStopWhileAStartIsWaitingForReadiness(t *testing.T) { + worker, engineManager := newTestRPCWorkerPipe(t) + brokerConn, clientConn := net.Pipe() + t.Cleanup(func() { + _ = brokerConn.Close() + _ = clientConn.Close() + }) + b := &Broker{codec: NewCodec(brokerConn)} + b.setEngineMgr(worker) + replies := make(chan *Message, 2) + go func() { + client := NewCodec(clientConn) + for { + msg, err := client.Read() + if err != nil { + return + } + replies <- msg + } + }() + + stopArrived := make(chan struct{}) + go func() { + var start *Message + for { + request, err := engineManager.Read() + if err != nil { + return + } + switch request.Method { + case "engine:start": + start = request // held: the engine is still loading + case "engine:stop": + close(stopArrived) + _ = engineManager.Respond(request.ID, map[string]any{"engine": "vllm", "running": false}) + if start != nil { + _ = engineManager.RespondError(start.ID, -32000, "start cancelled") + } + return + } + } + }() + + startID, stopID := json.RawMessage(`1`), json.RawMessage(`2`) + b.relayToEngine(&Message{JSONRPC: "2.0", ID: &startID, Method: "engine:start", Params: json.RawMessage(`{"engine":"vllm"}`)}) + time.Sleep(100 * time.Millisecond) // the Start is in flight and holds the lock + b.relayToEngine(&Message{JSONRPC: "2.0", ID: &stopID, Method: "engine:stop", Params: json.RawMessage(`{"engine":"vllm"}`)}) + + select { + case <-stopArrived: + case <-time.After(5 * time.Second): + require.FailNow(t, "the Stop did not reach the engine manager while the Start was in flight") + } + for range 2 { + select { + case <-replies: + case <-time.After(5 * time.Second): + require.FailNow(t, "the Start and the Stop were not both answered") + } + } +} diff --git a/services/nvpair-ui-broker/discovery.go b/services/nvpair-ui-broker/discovery.go index 8f2165336..09a16b592 100644 --- a/services/nvpair-ui-broker/discovery.go +++ b/services/nvpair-ui-broker/discovery.go @@ -71,6 +71,11 @@ type EnrichedNode struct { // peer's engine-manager /v1/models loadedByEngine field. Carried through to // AvailableNode so a remote node's cards can reflect loaded state. LoadedByEngine map[string][]string `json:"loadedByEngine,omitempty"` + // DownloadedByEngine is the downloaded inventory of an engine that advertises + // only the model it is serving (vLLM), enriched by the daemon from the peer's + // engine-manager /v1/models downloadedByEngine field. Carried through to + // AvailableNode so a remote node's model manager keeps its rows. + DownloadedByEngine map[string][]string `json:"downloadedByEngine,omitempty"` } // directoryToEnriched projects the promoted daemon's DirectoryNode onto the @@ -95,20 +100,21 @@ func directoryToEnriched(n noderec.DirectoryNode) EnrichedNode { port = ni.Port } return EnrichedNode{ - ID: n.Name, - HostUUID: n.HostUUID, - Host: n.Name, - Port: port, - Addresses: addrs, - TXT: txt, - GPUs: n.GPUs, - CPU: n.CPU, - Memory: n.Memory, - Trusted: n.Trusted, - Clustered: n.Clustered(), - Models: n.Models, - ModelsByEngine: n.ModelsByEngine, - LoadedByEngine: n.LoadedByEngine, + ID: n.Name, + HostUUID: n.HostUUID, + Host: n.Name, + Port: port, + Addresses: addrs, + TXT: txt, + GPUs: n.GPUs, + CPU: n.CPU, + Memory: n.Memory, + Trusted: n.Trusted, + Clustered: n.Clustered(), + Models: n.Models, + ModelsByEngine: n.ModelsByEngine, + LoadedByEngine: n.LoadedByEngine, + DownloadedByEngine: n.DownloadedByEngine, } } @@ -311,18 +317,19 @@ func (sn storedNode) toAvailable() AvailableNode { candidates = nil } return AvailableNode{ - ID: n.ID, - Name: n.ID, - HostUUID: n.HostUUID, - IPAddress: primary, - IPAddresses: candidates, - Port: n.Port, - LastSeen: sn.lastSeen.Unix(), - Trusted: n.Trusted, - Clustered: n.Clustered, - Models: n.Models, - ModelsByEngine: n.ModelsByEngine, - LoadedByEngine: n.LoadedByEngine, + ID: n.ID, + Name: n.ID, + HostUUID: n.HostUUID, + IPAddress: primary, + IPAddresses: candidates, + Port: n.Port, + LastSeen: sn.lastSeen.Unix(), + Trusted: n.Trusted, + Clustered: n.Clustered, + Models: n.Models, + ModelsByEngine: n.ModelsByEngine, + LoadedByEngine: n.LoadedByEngine, + DownloadedByEngine: n.DownloadedByEngine, } } diff --git a/services/nvpair-ui-broker/discovery_test.go b/services/nvpair-ui-broker/discovery_test.go index 39362b3c9..810052c75 100644 --- a/services/nvpair-ui-broker/discovery_test.go +++ b/services/nvpair-ui-broker/discovery_test.go @@ -102,6 +102,28 @@ func TestDiscoveryStoreScannerProjection_Wins(t *testing.T) { require.Equal(t, "manual-host", snap[0].ID, "manual projection should surface after scanner removal (%v)", snap) } +// A peer's model lists reach the desktop's AvailableNode unchanged, including +// the downloaded inventory of a stopped vLLM, so a remote model manager keeps +// its rows after an Eject. +func TestDiscoveryStore_CarriesEveryModelListFromTheScanner(t *testing.T) { + s := newDiscoveryStore() + s.Upsert(directoryToEnriched(noderec.DirectoryNode{ + HostUUID: "peer-uuid", + Name: "peer", + Models: []string{"qwen-api"}, + ModelsByEngine: map[string][]string{"vllm": {"qwen-api"}}, + LoadedByEngine: map[string][]string{"vllm": {"Qwen/Qwen3.5-4B"}}, + DownloadedByEngine: map[string][]string{"vllm": {"Qwen/Qwen3.5-4B", "mistralai/Ministral-3-3B-Instruct-2512"}}, + }), sourceScanner) + + snap := s.Snapshot() + require.Len(t, snap, 1) + assert.Equal(t, []string{"qwen-api"}, snap[0].Models) + assert.Equal(t, map[string][]string{"vllm": {"qwen-api"}}, snap[0].ModelsByEngine) + assert.Equal(t, map[string][]string{"vllm": {"Qwen/Qwen3.5-4B"}}, snap[0].LoadedByEngine) + assert.Equal(t, map[string][]string{"vllm": {"Qwen/Qwen3.5-4B", "mistralai/Ministral-3-3B-Instruct-2512"}}, snap[0].DownloadedByEngine) +} + func newManualTestBroker() *Broker { return &Broker{ manualNodeKeys: make(map[string]string), diff --git a/services/nvpair-ui-broker/engineproxy.go b/services/nvpair-ui-broker/engineproxy.go index 43a5ff91b..0b1e102b8 100644 --- a/services/nvpair-ui-broker/engineproxy.go +++ b/services/nvpair-ui-broker/engineproxy.go @@ -165,6 +165,7 @@ func buildEngineProxyProfiles() []engineProxyProfile { // llama.cpp's managed process is stopped directly by engine-manager. "lmstudio": {Ownership: managedEngine, HealthProbePath: "/v1/models"}, "llamacpp": {Ownership: managedEngine, HealthProbePath: "/health"}, + "vllm": {Ownership: managedEngine, HealthProbePath: "/health"}, } out := make([]engineProxyProfile, 0, len(engines.All())) for _, e := range engines.All() { diff --git a/services/nvpair-ui-broker/engineproxy_test.go b/services/nvpair-ui-broker/engineproxy_test.go index 78c372372..665fc44dd 100644 --- a/services/nvpair-ui-broker/engineproxy_test.go +++ b/services/nvpair-ui-broker/engineproxy_test.go @@ -127,13 +127,13 @@ func TestParseProxyEngines_ValidatesSelectedEngineNames(t *testing.T) { require.Equal(t, append([]string{}, want...), append([]string{}, got...)) }) } - test("default set", "ollama,lmstudio,llamacpp", []string{"ollama", "lmstudio", "llamacpp"}, false) + test("default set", "ollama,lmstudio,llamacpp,vllm", []string{"ollama", "lmstudio", "llamacpp", "vllm"}, false) test("single engine", "lmstudio", []string{"lmstudio"}, false) test("whitespace and blanks are tolerated", " ollama , , lmstudio ", []string{"ollama", "lmstudio"}, false) test("duplicates collapse", "ollama,ollama", []string{"ollama"}, false) test("empty selects nothing", "", nil, false) // Silently fronting the engines it did recognize would look like the flag worked. - test("unknown engine fails", "ollama,vllm", nil, true) + test("unknown engine fails", "ollama,nosuchengine", nil, true) } // An engine left out of --proxy-engines is not started, and neither is any diff --git a/services/nvpair-ui-broker/enginesettings_test.go b/services/nvpair-ui-broker/enginesettings_test.go index 3e86c5226..f01dc0a10 100644 --- a/services/nvpair-ui-broker/enginesettings_test.go +++ b/services/nvpair-ui-broker/enginesettings_test.go @@ -61,7 +61,7 @@ func newSettingsHarness(t *testing.T) *settingsHarness { func newSettingsHarnessForEngine(t *testing.T, engine string) *settingsHarness { t.Helper() - ports := make([]int, 4) + ports := make([]int, 5) releases := make([]func(), len(ports)) for i := range ports { ports[i], releases[i] = reserveSettingsPort(t) @@ -77,6 +77,7 @@ func newSettingsHarnessForEngine(t *testing.T, engine string) *settingsHarness { "ollama": {ready: true, port: ports[1]}, "lmstudio": {ready: true, port: ports[2]}, "llamacpp": {ready: true, port: ports[3]}, + "vllm": {ready: true, port: ports[4]}, }} for _, profile := range engineProxyProfiles { h.b.setEngineProxyHandle(profile, proxy) diff --git a/services/nvpair-ui-broker/proxyaddressing_test.go b/services/nvpair-ui-broker/proxyaddressing_test.go index 436b40b79..b907a5786 100644 --- a/services/nvpair-ui-broker/proxyaddressing_test.go +++ b/services/nvpair-ui-broker/proxyaddressing_test.go @@ -197,7 +197,7 @@ func TestReadiness_IsTrackedPerEngine(t *testing.T) { // An engine that never announced itself is not ready, rather than // inheriting a sibling's port. - ready, port := p.Status("vllm") + ready, port := p.Status("nosuchengine") assert.False(t, ready, "an unannounced engine reported ready") assert.Equal(t, 0, port, "an unannounced engine reported a port") } diff --git a/services/nvpair-ui-broker/spec.md b/services/nvpair-ui-broker/spec.md index 0661a0d2c..4a7f6494e 100644 --- a/services/nvpair-ui-broker/spec.md +++ b/services/nvpair-ui-broker/spec.md @@ -12,12 +12,12 @@ The single front door every PAIR client talks to, and the supervisor of every ot **In scope** - Process supervision: spawn each worker as a child with stdio JSON-RPC pipes, watch it, restart it on unexpected exit under a bounded backoff policy, and stop it in a defined order on shutdown within a bounded teardown budget. - Startup ordering and the `app:ready` gate — the one signal a client waits on before assuming the node's service tree is live. -- Namespace relay: forward `ollama-proxy:*`, `lmstudio-proxy:*`, and `llamacpp-proxy:*` to that engine's `nvpair-proxy` facade, and `engine:*`, `settings/*`, `cluster:*`, `nodes:*`, and the manual-node methods to their owning worker verbatim, and map each worker's result or JSON-RPC error straight back. +- Namespace relay: forward `ollama-proxy:*`, `lmstudio-proxy:*`, `llamacpp-proxy:*`, and `vllm-proxy:*` to that engine's `nvpair-proxy` facade, and `engine:*`, `settings/*`, `cluster:*`, `nodes:*`, and the manual-node methods to their owning worker verbatim, and map each worker's result or JSON-RPC error straight back. - Engine settings operations: `engine:get-settings`, `engine:preview-settings`, and `engine:apply-settings` for this node or a pinned peer, revision-checked and journaled, with every port setter routed through the same operation (§7.10). - Notification projection: consume every worker's stdout notifications, route the error-pipeline ones into `nvpair-errors`, translate or re-emit the rest as the client-facing push stream, and drop what is internal. -- Opt-in subscription state per domain (`discovery:`, `ollama-proxy:`, `lmstudio-proxy:`, `llamacpp-proxy:`, `engine:`, `workloads:`) so a client only receives the streams it asked for. +- Opt-in subscription state per domain (`discovery:`, `ollama-proxy:`, `lmstudio-proxy:`, `llamacpp-proxy:`, `vllm-proxy:`, `engine:`, `workloads:`) so a client only receives the streams it asked for. - The node's discovery service registration: registering each worker's LAN service key and port with the scanner daemon on behalf of the worker, and replaying the registration set across every scanner respawn. -- Proxy/engine port reconciliation for the Ollama and LM Studio compatibility ports, driven by the `settings/get-force-ports` policy and live `engine:status`, before each facade is enabled. llama.cpp has no compatibility-port takeover; its facade is placed on its default port or a safe fallback without moving the engine (§7.7). +- Proxy/engine port reconciliation for the Ollama and LM Studio compatibility ports, driven by the `settings/get-force-ports` policy and live `engine:status`, before each facade is enabled. llama.cpp and vLLM have no compatibility-port takeover; each facade is placed on its default port or a safe fallback without moving the engine (§7.7). - Scheduling telemetry: a source-aware cache of compact GPU telemetry fed by the scanner and by `nvpair-manual-nodes`, aged locally and relayed to `nvpair-job-scheduler`, and replayed when the scheduler restarts. - Discovery evidence relay: forwarding each proxy's peer-activity reports and `nvpair-node-info`'s peer-observed addresses to `nvpair-node-scanner`. - Merging manual nodes into the same discovery snapshot the scanner feeds, so one node list answers `discovery:get-nodes`. @@ -91,7 +91,7 @@ The single front door every PAIR client talks to, and the supervisor of every ot - Maintain a merged discovery store fed by the scanner and by `nvpair-manual-nodes`, keyed by each node's operational identity. - Maintain the workload store from local-origin proxy events and peer-origin workload-manager events, persist its terminal history, load it on startup, and periodically retire remote-origin records whose origin has gone silent about them. - Register this node's LAN service keys with the scanner daemon and replay them on every scanner spawn. -- Reconcile the Ollama and LM Studio compatibility ports against the `force_ports` policy and live engine status before enabling each facade, and place the llama.cpp facade on its default port or a safe fallback with no compatibility-port gate. +- Reconcile the Ollama and LM Studio compatibility ports against the `force_ports` policy and live engine status before enabling each facade, and place the llama.cpp and vLLM facades on their default ports or a safe fallback with no compatibility-port gate. - Cache compact GPU telemetry from the scanner and `nvpair-manual-nodes` with the scanner taking precedence, relay it to `nvpair-job-scheduler`, and replay the cache when the scheduler respawns. - Forward a proxy's `node/activity` to the scanner as `discovery:node-activity`, and `nvpair-node-info`'s `nodeinfo:observed-addresses` to the scanner as `discovery:set-observed-addresses`. - Restore the persisted cluster id into `nvpair-cluster-manager` before `app:ready`, and persist a later `cluster:identity-changed` back into `nvpair-node-settings`. @@ -138,6 +138,7 @@ The single front door every PAIR client talks to, and the supervisor of every ot models?: string[] modelsByEngine?: { [engine: string]: string[] } loadedByEngine?: { [engine: string]: string[] } + downloadedByEngine?: { [engine: string]: string[] } } ``` @@ -146,7 +147,7 @@ The single front door every PAIR client talks to, and the supervisor of every ot { subscribed: boolean } ``` -`ProxyStatusResult` — the response to `ollama-proxy:get-status`, `lmstudio-proxy:get-status`, and `llamacpp-proxy:get-status`, answered per engine from the broker's captured facade state with no round-trip to the proxy. It is per engine because the facades bind different ports. The zero value is a valid "not available" answer, so none of these methods errors: +`ProxyStatusResult` — the response to `ollama-proxy:get-status`, `lmstudio-proxy:get-status`, `llamacpp-proxy:get-status`, and `vllm-proxy:get-status`, answered per engine from the broker's captured facade state with no round-trip to the proxy. It is per engine because the facades bind different ports. The zero value is a valid "not available" answer, so none of these methods errors: ```json { ready: boolean, port: number } ``` @@ -193,10 +194,10 @@ Requests the broker answers itself (client → broker): | `ollama-proxy:subscribe` | — | `SubscriptionResult{subscribed:true}`, then a baseline `ollama-proxy:ready` if that facade has come up | | `ollama-proxy:unsubscribe` | — | `SubscriptionResult{subscribed:false}` | | `ollama-proxy:set-port` | `{ port }` | `{ port }`; runs `engine:apply-settings` with the new proxy port; `-32000` when the port is in use or the change fails | -| `lmstudio-proxy:get-status`, `llamacpp-proxy:get-status` | — | `ProxyStatusResult` | -| `lmstudio-proxy:subscribe`, `llamacpp-proxy:subscribe` | — | `SubscriptionResult{subscribed:true}`, then a baseline `-proxy:ready` if that facade has come up | -| `lmstudio-proxy:unsubscribe`, `llamacpp-proxy:unsubscribe` | — | `SubscriptionResult{subscribed:false}` | -| `lmstudio-proxy:set-port`, `llamacpp-proxy:set-port` | `{ port }` | same as `ollama-proxy:set-port`, for that engine | +| `lmstudio-proxy:get-status`, `llamacpp-proxy:get-status`, `vllm-proxy:get-status` | — | `ProxyStatusResult` | +| `lmstudio-proxy:subscribe`, `llamacpp-proxy:subscribe`, `vllm-proxy:subscribe` | — | `SubscriptionResult{subscribed:true}`, then a baseline `-proxy:ready` if that facade has come up | +| `lmstudio-proxy:unsubscribe`, `llamacpp-proxy:unsubscribe`, `vllm-proxy:unsubscribe` | — | `SubscriptionResult{subscribed:false}` | +| `lmstudio-proxy:set-port`, `llamacpp-proxy:set-port`, `vllm-proxy:set-port` | `{ port }` | same as `ollama-proxy:set-port`, for that engine | | `engine:get-settings` | `{ engine, nodeId? }` | the engine's full settings snapshot, local or from a pinned peer | | `engine:preview-settings` | `{ engine, nodeId?, expectedRevision, settings, resolution? }` | normalized settings, validation errors, conflicts, and a restart/rebind summary; changes nothing | | `engine:apply-settings` | the preview request plus `requestId` | `{ revision, phase }` receipt; progress arrives as `engine:settings-changed` | @@ -213,7 +214,7 @@ Requests the broker answers itself (client → broker): | `engine:set-reserved-port` | — | `-32601` — private engine-manager method, refused at the boundary | | `internal:set-reserved-port` | — | `-32601` — private engine-manager method, refused at the boundary | | `shutdown` | — | `null`, then the ordered teardown begins | -| `ollama-proxy:shutdown`, `lmstudio-proxy:shutdown`, `llamacpp-proxy:shutdown` | — | `-32601` — the broker owns the proxy lifecycle and will not let a client kill it | +| `ollama-proxy:shutdown`, `lmstudio-proxy:shutdown`, `llamacpp-proxy:shutdown`, `vllm-proxy:shutdown` | — | `-32601` — the broker owns the proxy lifecycle and will not let a client kill it | The facade rows are one profile-driven handler applied to every engine in `nvpair-shared/engines`, so each `-proxy:` namespace has the same broker-local surface. @@ -221,7 +222,7 @@ Anything else is matched by namespace prefix and relayed verbatim: | Namespace | Owning worker | Relay shape | |-----------|---------------|-------------| -| `ollama-proxy:*`, `lmstudio-proxy:*`, `llamacpp-proxy:*` | that engine's `nvpair-proxy` facade | bounded synchronous `Call`; the client prefix is replaced by the engine address (`ollama-proxy:nodes/list` → `ollama:nodes/list`) | +| `ollama-proxy:*`, `lmstudio-proxy:*`, `llamacpp-proxy:*`, `vllm-proxy:*` | that engine's `nvpair-proxy` facade | bounded synchronous `Call`; the client prefix is replaced by the engine address (`ollama-proxy:nodes/list` → `ollama:nodes/list`) | | `engine:*` | `nvpair-engine-manager` | asynchronous, **no broker-imposed timeout**, method verbatim | | `settings/*` | `nvpair-node-settings` | bounded synchronous `Call`, method verbatim | | `cluster:*`, `nodes:*` | `nvpair-cluster-manager` | bounded synchronous `Call`, method verbatim | @@ -235,7 +236,7 @@ Notifications the broker emits to the client, and what produces each: |---------------------|------|-------------------------| | `app:ready{version}` | always | the broker itself, once, after the worker tree is started | | `discovery:nodes-changed` | `discovery:subscribe` | the merged discovery store (scanner + manual nodes); payload is a bare `AvailableNode[]` | -| `ollama-proxy:`, `lmstudio-proxy:`, `llamacpp-proxy:` (notably `…:ready`) | that namespace's `…:subscribe` | the facade's engine-addressed notification, with the engine address replaced by the client prefix | +| `ollama-proxy:`, `lmstudio-proxy:`, `llamacpp-proxy:`, `vllm-proxy:` (notably `…:ready`) | that namespace's `…:subscribe` | the facade's engine-addressed notification, with the engine address replaced by the client prefix | | `engine:ready`, `engine:state-changed`, `engine:models-changed`, `engine:install-progress`, `engine:pull-progress`, `engine:remote-progress` | `engine:subscribe` | `nvpair-engine-manager`, verbatim (already `engine:`-prefixed) | | `engine:settings-changed` | always | the broker itself, a full settings snapshot per engine on every settings change | | `workloads:upsert`, `workloads:remove` | `workloads:subscribe` | `nvpair-workload-manager` (peer-origin) and `nvpair-proxy` (local-origin), after the store accepts the transition | @@ -291,7 +292,7 @@ Each worker's reader goroutine invokes a forwarder. Every forwarder calls `dispa | `forwardEngineNotification` | `nvpair-engine-manager` | `engine:settings-request` / `engine:settings-cancel` carry a paired peer's settings operation into the broker (§7.10); `engine:ready` also triggers LM Studio proxy reconciliation and re-sends the settings projection; `discovery:subscribe` is claimed as a worker relay subscription; everything else is re-emitted verbatim to `engine:subscribe`'d clients | | `forwardProxyProcessNotification` | `nvpair-proxy` | an engine-addressed notification goes to that engine's handler, with that engine's spawn generation for Ollama and LM Studio; an unaddressed one is process-scoped and handled once — `workload:*` is stamped and routed to the workload manager, and `node/activity` is forwarded to the scanner as `discovery:node-activity` | | `forwardProxyNotificationForGeneration` / `forwardLMStudioProxyNotificationForGeneration` | one facade | stale generations dropped; a `bind-failed` on the managed facade port blocks the facade, otherwise selects a fallback port unless the user set the port explicitly; `ready` triggers port reconciliation; the remainder is emitted as `-proxy:` to that namespace's subscribers | -| `forwardDefaultEngineProxyNotification` | the llama.cpp facade | no port reconciliation and no generation check: error-pipeline frames are claimed, and the remainder is emitted as `llamacpp-proxy:` to that namespace's subscribers | +| `forwardDefaultEngineProxyNotification` | the llama.cpp and vLLM facades | no port reconciliation and no generation check: error-pipeline frames are claimed, and the remainder is emitted as `llamacpp-proxy:` or `vllm-proxy:` to that namespace's subscribers | | `forwardManualNodesNotification` | `nvpair-manual-nodes` | `node/discovered` and `node/updated` upsert into the discovery store, cache the node's GPU telemetry as the manual source, and bridge it into the Ollama or LM Studio facade whose engine is reachable (manual nodes have no llama.cpp leg); `node/removed` removes; `ready` is logged; nothing reaches the client directly — manual nodes surface inside `discovery:nodes-changed` | | `forwardSettingsNotification` | `nvpair-node-settings` | only `connection/cluster-identity` and `connection/cluster-auto-sync` are re-emitted, verbatim and unconditionally; everything else is debug-logged and dropped | | `forwardClusterManagerNotification` | `nvpair-cluster-manager` | every notification relayed verbatim and unconditionally; `cluster:identity-changed` additionally persists the identity into settings and reconverges the node's advertised identity; `cluster:trust-changed` additionally re-derives the scanner's peer-trust annotation and re-publishes the principal to `nvpair-node-info` | @@ -330,6 +331,8 @@ The broker owns the ordering between the compatibility ports the ecosystem expec llama.cpp has **no compatibility-port takeover**. `nvpair-engine-manager` owns its process and starts it on its configured loopback port (`:8081` by default); the broker records that port without a policy read, a startup gate, or any relocation of the engine. It requests the facade on `:8080`, or on the port already chosen for this broker lifetime, and on a bind failure retries once on the next free port that avoids the failed port, the engine's default port, the inherited `OLLAMA_HOST` alias, and every other facade's planned, backend, and persisted ports. A facade port set explicitly through the settings journal is never replaced by a fallback; its bind failure is logged and the facade stays down. llama.cpp never reports a port-ownership block. Its only `llamacpp-proxy:port-ownership-blocked` entry is the warning every engine gets when the `nvpair-proxy` binary does not resolve ("NVPAIR is not fronting llama.cpp on port 8080: proxy binary not resolved."). +vLLM is placed the same way: `nvpair-engine-manager` starts `vllm serve` on its configured loopback port (`:8001` by default), and the broker requests the vLLM facade on `:8000` or a safe fallback, with no compatibility-port gate and no port-ownership block. + A port the user chooses is never silently moved. `engine:set-port` and every `-proxy:set-port` are intercepted and run through `engine:apply-settings` (§7.10), which validates both of the engine's ports against registered services, aliases, other engines and facades, and occupied listeners. A port already in use is refused with `-32000`, and a collision with the `OLLAMA_HOST` alias is refused with a message naming it. Automatic steering remains only for a port the user did not just choose: when a facade announces a restored port a running engine has since taken, the broker moves the facade to a free port and surfaces the sticky warning `ollama-proxy:port-bumped`, never changing the engine's port. An `engine:*` request that depends on a facade port still being reconciled is parked on the corresponding gate and either replayed when the gate opens or answered with a retry-me error after 5 seconds. ### 7.8 Discovery registration and relay @@ -441,7 +444,7 @@ The broker is a single Go binary that runs as a child of its client and as the p | `--engine-manager-path` | `./nvpair-engine-manager` | engine lifecycle and model operations | | `--node-info-path` | `./nvpair-node-info` | local hardware advertisement | | `--proxy-path` | `./nvpair-proxy` | the one proxy process that hosts every engine's facade | -| `--proxy-engines` | every engine in `nvpair-shared/engines` (today `ollama,lmstudio,llamacpp`) | which engines to enable a facade for; an unrecognized name is fatal, and an engine left out is neither enabled nor prepared | +| `--proxy-engines` | every engine in `nvpair-shared/engines` (today `ollama,lmstudio,llamacpp,vllm`) | which engines to enable a facade for; an unrecognized name is fatal, and an engine left out is neither enabled nor prepared | | `--workload-manager-path` | `./nvpair-workload-manager` | cluster workload relay | | `--manual-nodes-path` | `./nvpair-manual-nodes` | user-entered nodes | | `--cluster-manager-path` | `./nvpair-cluster-manager` | identity, pairing, membership | diff --git a/services/readme.md b/services/readme.md index a880ed572..fab185fc8 100644 --- a/services/readme.md +++ b/services/readme.md @@ -12,7 +12,7 @@ node offers and where to reach them. What a discovered node can actually serve is a separate question, answered after discovery. A node may be running [Ollama](https://ollama.com/), LM Studio, -llama.cpp, any combination, or none, and its model inventory is fetched over +llama.cpp, vLLM, any combination, or none, and its model inventory is fetched over HTTP from its engine-manager rather than crammed into mDNS TXT records, which are too small to carry it. @@ -28,11 +28,13 @@ can serve that loopback-only address from the same router when it is free — `localhost` is claimed on IPv4 and IPv6 together, and remote or HTTPS targets are never intercepted. -The llama.cpp facade is enabled by default alongside Ollama and LM Studio. It -exposes OpenAI-compatible traffic on `http://localhost:8080` while the managed -router runs on `8081`; a safe fallback is reported when the facade port is -occupied. The desktop and TUI expose its managed lifecycle, downloads, -inventory, and load/unload actions. `--proxy-engines` can still restrict the +The llama.cpp and vLLM facades are enabled by default alongside Ollama and LM Studio. +The llama.cpp facade exposes OpenAI-compatible traffic on `http://localhost:8080` +while the managed router runs on `8081`. The vLLM facade (the engine itself runs on +Linux and Apple Silicon macOS) uses `http://localhost:8000` with its managed engine on `8001`. A safe fallback is reported when a facade port is +occupied. The desktop and TUI expose each engine's managed lifecycle and +inventory, plus llama.cpp's downloads and load/unload actions; vLLM serves one +model per process, so loading another downloaded model restarts it on that model. `--proxy-engines` can still restrict the facades a standalone broker or TUI starts. PAIR ships a graphical UI alongside these services. The UI launches @@ -70,7 +72,7 @@ configuration. The broker feeds every accepted local or peer workload transition plus compact GPU telemetry to the scheduler. Queued and running work is counted by destination -node across Ollama, LM Studio, and llama.cpp together. Fresh maximum-GPU +node across Ollama, LM Studio, llama.cpp, and vLLM together. Fresh maximum-GPU utilization is smoothed into pressure 0–3; missing or stale telemetry is neutral. Rankings use `pending + gpuPressure`, and each facade adds local reservations before choosing, so bursts spread without waiting for workload diff --git a/services/shared/engines/engines.go b/services/shared/engines/engines.go index ed5a7d9a1..3db23e56e 100644 --- a/services/shared/engines/engines.go +++ b/services/shared/engines/engines.go @@ -131,6 +131,7 @@ const ( NameOllama = "ollama" NameLMStudio = "lmstudio" NameLlamaCPP = "llamacpp" + NameVLLM = "vllm" ) // all is the ordered engine set. Ollama is first; see the package comment. @@ -159,6 +160,14 @@ var all = []Engine{ EnginePortBase: 8081, PortFile: "llamacpp-proxy-port.json", }, + { + Name: NameVLLM, + DisplayName: "vLLM", + DiscoveryService: noderec.ServiceVLLM, + FacadePort: 8000, + EnginePortBase: 8001, + PortFile: "vllm-proxy-port.json", + }, } // All returns the engine set in preparation order. The result is a copy, so a diff --git a/services/shared/engines/engines_test.go b/services/shared/engines/engines_test.go index 38994adce..1de832f86 100644 --- a/services/shared/engines/engines_test.go +++ b/services/shared/engines/engines_test.go @@ -23,7 +23,7 @@ func TestOllama_IsPreparedFirst(t *testing.T) { } func TestEngineNames_ReturnsCanonicalOrder(t *testing.T) { - assert.Equal(t, []string{"ollama", "lmstudio", "llamacpp"}, Names()) + assert.Equal(t, []string{"ollama", "lmstudio", "llamacpp", "vllm"}, Names()) } // There are two proxy identities: ComponentName per facade, ProxyComponent per @@ -112,7 +112,7 @@ func TestEngineLookup_ResolvesDeclaredNamesAndKeys(t *testing.T) { assert.Equal(t, e.ComponentName(), byName.ComponentName()) } - _, ok := ByName("vllm") + _, ok := ByName("nosuchengine") assert.False(t, ok, "ByName should report ok=false for an unknown engine") } diff --git a/services/shared/noderec/noderec.go b/services/shared/noderec/noderec.go index f44a263d6..983fe81ee 100644 --- a/services/shared/noderec/noderec.go +++ b/services/shared/noderec/noderec.go @@ -10,7 +10,7 @@ // whose TXT map carries a schema version, the node's identity, its LAN address, // and one compact key per local service port, e.g.: // -// v=1;uuid=;cluster-uuid=;ip=192.168.1.10;ni=14318;ol=11434;lm=1234;lc=8080;er=14319;wl=14320;cl=14321;em=14322 +// v=1;uuid=;cluster-uuid=;ip=192.168.1.10;ni=14318;ol=11434;lm=1234;lc=8080;vl=8000;er=14319;wl=14320;cl=14321;em=14322 // // Design decisions this package encodes: // - SRV port is a fixed, NON-authoritative constant; consumers ignore it and @@ -88,6 +88,7 @@ const ( ServiceOllama ServiceKey = "ol" ServiceLMStudio ServiceKey = "lm" ServiceLlamaCPP ServiceKey = "lc" + ServiceVLLM ServiceKey = "vl" ServiceErrors ServiceKey = "er" ServiceWorkload ServiceKey = "wl" ServiceCluster ServiceKey = "cl" @@ -105,7 +106,7 @@ const ( // serviceKeyOrder is the deterministic emit order for service ports in TXT. var serviceKeyOrder = []ServiceKey{ - ServiceNodeInfo, ServiceOllama, ServiceLMStudio, ServiceLlamaCPP, + ServiceNodeInfo, ServiceOllama, ServiceLMStudio, ServiceLlamaCPP, ServiceVLLM, ServiceErrors, ServiceWorkload, ServiceCluster, ServiceEngineManager, ServiceEngineControl, } @@ -515,7 +516,16 @@ type DirectoryNode struct { // empty list means "running, nothing loaded"; a missing key means the peer // didn't report loaded state for it. Omitted when the peer reports none. LoadedByEngine map[string][]string `json:"loadedByEngine,omitempty"` - LastSeen int64 `json:"lastSeen"` // Unix seconds + // DownloadedByEngine is the model inventory of an engine that advertises + // only the model it is serving (vLLM), so ModelsByEngine holds just that one + // for routing. It lists every model the engine can load, read even while the + // engine is stopped, so a model manager keeps its Load and Delete rows. Its + // LoadedByEngine entry names models the same way. An engine key with an empty + // list means the inventory was read and is empty; a missing key means it was + // not read. Enriched from engine-manager's em /v1/models downloadedByEngine + // field; omitted when no engine reports one. + DownloadedByEngine map[string][]string `json:"downloadedByEngine,omitempty"` + LastSeen int64 `json:"lastSeen"` // Unix seconds } // Clustered reports whether the node advertises a cluster identity. diff --git a/services/tests/vllm_interop_test.go b/services/tests/vllm_interop_test.go new file mode 100644 index 000000000..6b2112f66 --- /dev/null +++ b/services/tests/vllm_interop_test.go @@ -0,0 +1,150 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package tests + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestBrokerDefaults_StartTheVLLMProxy(t *testing.T) { + stdin, msgs, stderr, cleanup := startBrokerWith(t, "--proxy-path", proxyBin) + t.Cleanup(cleanup) + go func() { + for range stderr { + } + }() + + waitForMethod(t, msgs, "app:ready", 10*time.Second) + assert.Positive(t, waitEngineProxyReady(t, "vllm-proxy", stdin, msgs, 15*time.Second), "the default vLLM proxy must be a listening facade") +} + +func TestBrokerVLLMProxySetPort_RejectsInvalidPorts(t *testing.T) { + stdin, msgs, stderr, cleanup := startBrokerWith(t, + "--proxy-path", proxyBin, "--proxy-engines", "vllm", + ) + t.Cleanup(cleanup) + go func() { + for range stderr { + } + }() + waitForMethod(t, msgs, "app:ready", 10*time.Second) + nextID := 7400 + reject := func(name, params string) { + t.Run(name, func(t *testing.T) { + id := nextID + nextID++ + _, err := fmt.Fprintf(stdin, `{"jsonrpc":"2.0","id":%d,"method":"vllm-proxy:set-port","params":%s}`+"\n", id, params) + require.NoError(t, err, "write proxy port request") + + response := waitForResponse(t, msgs, 10*time.Second) + + require.NotNil(t, response.ID) + require.Equal(t, fmt.Sprint(id), string(*response.ID), "response ID") + require.NotNil(t, response.Error, "want the broker's invalid-port rejection") + assert.Equal(t, -32602, response.Error.Code) + assert.Equal(t, "port must be between 1 and 65535", response.Error.Message) + }) + } + reject("missing port", `{}`) + reject("zero port", `{"port":0}`) + reject("negative port", `{"port":-1}`) + reject("oversized port", `{"port":65536}`) +} + +// vLLM serves one model per process under its Hugging Face id, so a facade in +// front of it must route only the exact id a node advertises. +func TestVLLMFacade_ServesTheModelListAndRoutesOnlyExactModelIDs(t *testing.T) { + const model = "Qwen/Qwen2.5-0.5B-Instruct" + var modelListHits atomic.Int32 + var inferenceHits atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == http.MethodGet && r.URL.Path == "/v1/models": + modelListHits.Add(1) + assert.NoError(t, json.NewEncoder(w).Encode(map[string]any{ + "object": "list", + "data": []map[string]string{{"id": model}}, + })) + case r.Method == http.MethodPost && r.URL.Path == "/v1/chat/completions": + inferenceHits.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + _, _ = io.WriteString(w, `{"choices":[]}`) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(upstream.Close) + + stdin, msgs, stderr, cleanup := startBrokerWith(t, + "--proxy-path", proxyBin, "--proxy-engines", "vllm", + ) + t.Cleanup(cleanup) + go func() { + for range stderr { + } + }() + + waitForMethod(t, msgs, "app:ready", 10*time.Second) + proxyPort := waitEngineProxyReady(t, "vllm-proxy", stdin, msgs, 15*time.Second) + callBrokerRPC(t, stdin, msgs, 7500, "vllm-proxy:node/add-manual", map[string]any{ + "id": "vllm-owner", + "host": "127.0.0.1", + "port": portOfURL(t, upstream.URL), + "addresses": []string{"127.0.0.1"}, + "models": []string{model}, + }) + + client := &http.Client{Timeout: 5 * time.Second} + t.Cleanup(client.CloseIdleConnections) + + t.Run("serves the OpenAI model list", func(t *testing.T) { + response, err := client.Get(fmt.Sprintf("http://127.0.0.1:%d/v1/models", proxyPort)) + require.NoError(t, err, "get model list") + defer response.Body.Close() + var list struct { + Object string `json:"object"` + Data []struct { + ID string `json:"id"` + } `json:"data"` + } + require.NoError(t, json.NewDecoder(response.Body).Decode(&list), "decode model list") + var ids []string + for _, item := range list.Data { + ids = append(ids, item.ID) + } + assert.Equal(t, http.StatusOK, response.StatusCode) + assert.Equal(t, "list", list.Object) + assert.Contains(t, ids, model) + assert.NotZero(t, modelListHits.Load(), "the list must come from the upstream engine") + }) + + post := func(t *testing.T, requestedModel string) int { + t.Helper() + endpoint := fmt.Sprintf("http://127.0.0.1:%d/v1/chat/completions", proxyPort) + response, err := client.Post(endpoint, "application/json", + bytes.NewBufferString(fmt.Sprintf(`{"model":%q,"messages":[]}`, requestedModel))) + require.NoError(t, err, "post inference") + _, err = io.Copy(io.Discard, response.Body) + require.NoError(t, err, "read inference response") + require.NoError(t, response.Body.Close(), "close inference response") + return response.StatusCode + } + t.Run("routes only the exact advertised model id", func(t *testing.T) { + require.Equal(t, http.StatusOK, post(t, model), "the exact model id") + require.Equal(t, http.StatusBadGateway, post(t, "qwen/qwen2.5-0.5b-instruct"), "a case-changed model id") + assert.Equal(t, int32(1), inferenceHits.Load(), "only the exact match reaches the upstream") + }) +}