diff --git a/desktop/docs/services-api.md b/desktop/docs/services-api.md index c972f764d..4bf199ba3 100644 --- a/desktop/docs/services-api.md +++ b/desktop/docs/services-api.md @@ -182,6 +182,7 @@ | `engine:install` | request (we call) | ✅ yes | | `engine:install-tool` | request (we call) | ✅ yes | | `engine:logs` | request (we call) | ⚠️ not called | +| `engine:model-catalog` | request (we call) | ✅ yes | | `engine:models` | request (we call) | ✅ yes | | `engine:prepare-shutdown` | request (we call) | ✅ yes | | `engine:preview-launch` | request (we call) | ⚠️ not called | @@ -395,6 +396,7 @@ Client surface: requests are sent to the broker; notifications are consumed from | `engine:get-load-profile` | request (we call) | ✅ yes | | `engine:get-settings` | request (we call) | ✅ yes | | `engine:install` | request (we call) | ✅ yes | +| `engine:model-catalog` | request (we call) | ✅ yes | | `engine:preview-load-profile` | request (we call) | ✅ yes | | `engine:preview-settings` | request (we call) | ✅ yes | | `engine:remote-cancel-pull` | request (we call) | ✅ yes | diff --git a/desktop/src/electron/service-bridge/json-rpc-subprocess.ts b/desktop/src/electron/service-bridge/json-rpc-subprocess.ts index 0a0c443f0..b212445fb 100644 --- a/desktop/src/electron/service-bridge/json-rpc-subprocess.ts +++ b/desktop/src/electron/service-bridge/json-rpc-subprocess.ts @@ -62,6 +62,9 @@ export interface JsonRpcInboundRequest { id: JsonRpcId } +// Go-owned llama.cpp model hub contract; screens are added separately. +type LlamaCatalogParams = { query?: string } + interface PendingCall { resolve: (value: JsonValue | undefined) => void reject: (error: Error) => void @@ -174,6 +177,16 @@ export class JsonRpcSubprocess extends EventEmitter { }) } + call( + method: 'engine:model-catalog', + params?: LlamaCatalogParams, + timeoutMs?: number | null + ): Promise + call( + method: string, + params?: JsonValue, + timeoutMs?: number | null + ): Promise async call( method: string, params?: JsonValue, diff --git a/docs/engine-lifecycle.mdx b/docs/engine-lifecycle.mdx index d14ebe00e..912ec0f30 100644 --- a/docs/engine-lifecycle.mdx +++ b/docs/engine-lifecycle.mdx @@ -11,8 +11,8 @@ those engines for you, or work with a copy you already have. A macOS arm64 **llama.cpp engine** also appears in Engine settings: it can be installed, started, stopped and configured, and is served by its own routing proxy on **1237** (engine listener **127.0.0.1:8082**). Select a regular GGUF inside PAIR's -llama.cpp model cache with allowlisted `--model`/`-m`; model hub/download/delete -are not available yet. Its official b12171 macOS build is **not notarized**, like PAIR itself. +llama.cpp model cache with allowlisted `--model`/`-m`; the Go service now +provides catalog/download/cancel/delete RPCs, with desktop/TUI hub screens deferred. Its official b12171 macOS build is **not notarized**, like PAIR itself. This page explains what you can expect when you install, start, stop, update, or remove an engine. @@ -72,11 +72,15 @@ Refer also to [Getting started](getting-started.mdx). Install llama.cpp from its node row (not first-run automatic selection). The managed listener defaults to `127.0.0.1:8082`; Settings exposes the server port -and launch arguments/environment but no proxy port. PAIR protects all listener +and proxy ports and launch arguments/environment. PAIR protects all listener controls, including Unix-socket host values and port-changing presets. Readiness -requires the `/props` router identity, not merely HTTP success. Non-loopback +requires the `/props` router identity (or the inherited model descriptor in +single-model mode), not merely HTTP success. Non-loopback binding still requires the existing explicit `engine_allow_lan_bind` permission. -There is no model hub, model download/delete, load profile or routed endpoint yet. +The Go model hub supplies a bounded public Q4_K_M catalog and PAIR-owned +verified downloads, asynchronous cancel and manifest-owned delete. See the +[model hub contract](../services/nvpair-engine-manager/LLAMACPP.md#go-model-hub-and-owned-cache). +The dedicated proxy routes inference; model hub screens and load profiles remain deferred. Installation verifies the pinned SHA-256, bounded contained extraction, expected server/libraries/Metal layout, absence of quarantine and existing signature diff --git a/docs/release-readiness.mdx b/docs/release-readiness.mdx index b459453c8..382dc6f4e 100644 --- a/docs/release-readiness.mdx +++ b/docs/release-readiness.mdx @@ -16,7 +16,9 @@ The current source adds `llamacpp-proxy` (default port 1237, engine 8082, service `lc`): **13 supervised workers / 15 executables**. The broker, desktop and TUI inventories, build/install lists and contract documentation include it. Native b12171 Messages passthrough and exact-ID model routing are supported; -model hub/download/delete remain deferred. Historical packaged-app evidence +the Go model hub provides catalog/download/cancel/delete; desktop/TUI hub +screens remain deferred. The new hub uses fixture/local-server tests, not live +model downloads or packaged certification. Historical packaged-app evidence below predates this worker and does not certify its packaged or multi-node inference behavior. Validate those before release. diff --git a/services/nvpair-engine-manager/LLAMACPP.md b/services/nvpair-engine-manager/LLAMACPP.md index 5a8682d6d..c0ddae1be 100644 --- a/services/nvpair-engine-manager/LLAMACPP.md +++ b/services/nvpair-engine-manager/LLAMACPP.md @@ -9,8 +9,8 @@ PAIR manages the official **b12171 macOS arm64** build. Other platforms are unsupported. PAIR installs, starts, stops and configures the engine and serves it through the dedicated `llamacpp-proxy` worker, default port **1237**. Routing uses exact model IDs from the engine's `/v1/models` inventory, with separate -`llamacpp` attribution. There is no model hub, download/delete adapter or load -profile yet. +`llamacpp` attribution. The Go model hub supports catalog, downloads, cancel +and delete over JSON-RPC; desktop/TUI hub screens and load profiles are deferred. The macOS build is **not notarized**, like this fork's PAIR application. The release is checksum-pinned, not publisher-trusted by that checksum. Its existing @@ -46,7 +46,112 @@ inherited descriptor's `model_path` from `tools/server/server-context.cpp` A different service returning HTTP 200 is not an engine match. Foreground-process shutdown uses the existing process lifecycle, not command-mode `wait_for_exit`. The model cache is `LLAMA_CACHE={install_dir}-models`, a PAIR-managed sibling -that survives engine uninstall. Nothing in this PR downloads models. +that survives engine uninstall. PAIR downloads models itself; llama-server is +kept offline and none of its download launch options are enabled. + +## Go model hub and owned cache + +`engine:model-catalog {query?: string}` returns `{models: [{id, file, size?, +sha256?}]}`. Empty query reads ggml-org, bartowski and unsloth, up to 50 repos +each. Public search admits at most 50 repos and 100 query characters; at most +20 search queries are cached for six hours. The aggregate metadata deadline is +20 seconds (broker 25 seconds). Only public, ungated text-generation GGUF repos +with exactly one single-file Q4_K_M candidate are listed. IDs are +`owner/repo:Q4_K_M`. All path components use bounded strict ASCII validation. + +Existing `engine:action` supports llama.cpp `pull_model {model: id}` with +`opId`, `list_downloaded` returning `{models: [{id, file, size, sha256, +finalHost}]}`, and `delete_model {model: id}`. Inventory combines the ownership +manifest with a bounded (4096 entries) no-follow regular-GGUF scan; manually +supplied GGUFs have their filename as ID and empty checksum/host and are never +deletable by PAIR. Recorded files must match their descriptor-derived file identity, size and valid +single-file GGUF metadata to appear. The live `list_models` remains the +llama-server served inventory, not this installed list. Installed scans check a +five-second context budget between bounded per-file metadata reads (broker ten +seconds; future clients fifteen seconds). Filesystem syscall latency itself +cannot be preempted by a Go context. + +PAIR resolves metadata again before transfer, requiring LFS SHA-256 and size +and a pinned repository commit. Limit: 64 GiB per file plus 256 MiB free-space +reserve. During streaming, fstatfs on the open transfer descriptor rechecks +space every 256 MiB or five seconds, whichever occurs first at a streaming +boundary (a blocked read writes no additional bytes). Remaining free space must +cover the unreceived file bytes plus the 256 MiB safety margin. A drop below that +threshold aborts with "not enough disk space" through the existing progress/error +state and identity-checked cleanup of this operation's temp. Transfers have a +30-minute operation bound, 10-second dial/TLS/header +bounds, a 64 KiB response-header cap and a 30-second idle-read bound. Redirects +must stay HTTPS, cannot contain credentials and are limited to five; the final +host is recorded. Bodies and transfer URLs are never logged. An exclusive +random temporary file is created inside the pinned cache descriptor; exact +size, SHA-256 and single-file GGUF metadata are checked before fsync and an +atomic no-replace rename. The small ownership manifest is fsynced and atomically +replaced. Existing destinations are never overwritten. A crash between file +commit and manifest commit can leave an unrecorded file, never a deletable +owned entry; abandoned crash temp files are not automatically reclaimed. + +The cache is in the user's own app data: entry swaps require another process +running as the same user, not a different unprivileged user. Staging uses +`openat(cacheFd, randomName, O_CREAT|O_EXCL|O_NOFOLLOW|O_RDWR|O_CLOEXEC, 0600)` +through the no-follow directory walk and records dev/ino immediately with fstat. +After hashing and fsync through that descriptor, commit uses renameatx_np with +RENAME_EXCL on macOS (renameat2 RENAME_NOREPLACE on Linux), then fstatat with +AT_SYMLINK_NOFOLLOW compares the final entry's dev/ino and size to the verified +file. A mismatch fails with no ownership record and leaves the unrelated entry +untouched. Only a match permits directory fsync and manifest commit; the manifest +records dev/ino (`fileIdentity`), size and SHA-256. Generated 64-hex `.gguf` names +are reserved case-insensitively, matching default macOS filesystem semantics: +unrecorded commits cannot become selectable manual models, and case aliases +cannot bypass recorded-model identity checks. +Selection and the actual opened launch descriptor must match a recorded model's +dev/ino/size; replaced files are refused without re-hashing up to 64 GiB. + +macOS provides no rename-by-descriptor or unlink-by-descriptor. A same-user +process can still change an entry after the last identity check; the residual +check-to-use window is the minimum the OS allows, not an atomicity guarantee. +Launch validates the inherited descriptor, so a later pathname replacement +cannot redirect llama-server. These checks detect entry replacement, not writes +in place by a process already authorized as the same user. + +`engine:cancel-pull` accepts asynchronously; only the terminal progress update +confirms transfer release and removal of that operation's own temp file. +The shared pull tracker reports `cancelled` only when that operation's user +cancel was requested and the returned error wraps `context.Canceled`. Metadata +requests and streaming reads preserve that cause, including transports whose +body read errors lose it when cancellation tears down the connection. The pull +then returns `{"status":"cancelled"}` without a `pull-failed` engine error. +Unrelated failures (including checksum/size mismatch, HTTP status, disk space, +identity checks and manifest writes) remain errors regardless of cancel timing, +even when cancellation arrives during teardown. Independent cleanup failures +also remain errors. +Completion is fenced against cancel acceptance after the post-rename identity +check and successful manifest write/fsync. A cancel accepted before that fence +rolls back any newly written manifest entry and removes the operation's +identity-matched published file as well as its temp. After the fence, cancellation +is no longer accepted: the pull reports success and the model remains installed. +No filesystem or network I/O is performed while holding the cancellation lock. +Before temp cleanup or completed-model deletion, fstatat with +AT_SYMLINK_NOFOLLOW checks dev/ino (and the recorded size for a completed model) +immediately before unlinkat. A changed entry is left untouched and reports +"cleanup skipped" through the terminal error state, not successful cancellation. +The original temp descriptor stays open until cleanup finishes. The same check +protects manifest staging cleanup and rollback of an unrecorded model. The +remaining check-to-unlink window cannot be eliminated on macOS; no unlink-by-fd +primitive exists. Unlike Ollama/LM Studio in #50, only this operation's +identity-matched temp is removed; shared or substituted files are not knowingly +removed. Concurrent +cache mutations are refused (including duplicate pulls), not queued. #50's +operation-ID fencing, broker worker-loss interruption and epoch recovery apply +unchanged. Verified-peer downloads/cancel/delete use the existing authenticated +remote methods and streams. Start is refused during a cache mutation; deletion +conservatively requires the engine to be stopped and no lifecycle op in flight. +Every cache ancestor is opened with openat/no-follow from `/`; mutations stay +relative to that descriptor even after a parent swap. Delete requires a matching +ownership manifest record, file identity and regular file, and never recursively removes data. + +This intentionally differs from NVIDIA/Personal-AI-Router@6d2d67f: PAIR does not +ask llama-server's router-mode API to download. All checks and #50 tracking stay +inside Go. No real model was downloaded during development. ## Launch allowlist (b12171) @@ -82,9 +187,9 @@ reopens the path with `llama_file` for reads or `mmap` on that descriptor. macOS sibling paths, so split filenames and `split.*` metadata (including renamed shards) are explicitly rejected at spawn. Metadata inspection is bounded to 64 MiB; malformed or unsupported GGUF metadata fails closed. Multi-file models -and mmproj are not supported in this PR. The user supplies a single-file model -for now; PR 3 adds the hub that fills the cache. Omitting `--model` retains router -mode. +and mmproj are not supported. The hub fills the cache with verified single-file +models. `--model` accepts the installed catalog ID, cache filename or contained +path. Omitting `--model` retains router mode. Other allowed value options (documented short aliases included): `--ctx-size/-c`, `--batch-size/-b`, `--ubatch-size/-ub` (context 0–1048576, batches 1–1048576), diff --git a/services/nvpair-engine-manager/README.md b/services/nvpair-engine-manager/README.md index 50db31fec..2505b0707 100644 --- a/services/nvpair-engine-manager/README.md +++ b/services/nvpair-engine-manager/README.md @@ -14,7 +14,18 @@ is a JSON manifest, not code. [llama.cpp](LLAMACPP.md) is supported on macOS arm64, pinned to b12171. Its macOS build is not notarized, like PAIR itself. The dedicated `llamacpp-proxy` serves it on port 1237. Select a cached GGUF with the allowlisted `--model`/`-m`; -`list_models` reads its OpenAI inventory. Hub/download/delete remain deferred. +`list_models` reads its OpenAI inventory. `engine:model-catalog` accepts +`{query?: string}` and returns `{models: [{id, file, size?, sha256?}]}` from +public, ungated, text-generation GGUF repositories. IDs are `owner/repo:Q4_K_M`; +only one unsplit Q4_K_M file is admitted. Empty query reads ggml-org, bartowski +and unsloth (50 repositories each). Public search is capped at 50 repositories, +100 query characters, 20 cached queries and a six-hour TTL. Requests have a +20-second aggregate timeout (broker: 25 seconds); metadata bodies are bounded +to 4 MiB. Existing `engine:action` supports llama.cpp `pull_model`, +`list_downloaded` and `delete_model`, including existing operation-ID progress, +asynchronous cancel and verified-peer remote management. See [the owned-cache +contract](LLAMACPP.md#go-model-hub-and-owned-cache) for checks, inventory and +cleanup. Desktop/TUI catalog screens are deferred to PR 4. The bundled manifests under `manifests/` are the working reference for manifest authoring. [Pinned archive installs](MANIFEST.md#pinned-macos-archive-installs) @@ -762,8 +773,10 @@ its own. Pulls are attributed by both model and operation ID, locally and remotely. The pull accepts an optional caller-generated `opId`; cancel requires that exact ID plus engine and model. A stale ID cannot cancel a replacement. Concurrent -different models are independent; a duplicate model in flight is refused. -No PAIR download queue or cache-file cleanup is introduced. +different models are independent for engine-managed caches; a duplicate model +in flight is refused. llama.cpp cache mutations are conservatively serialized +by rejecting concurrent mutations, not by queuing them. No PAIR download queue +is introduced. Cancel returns once accepted, **not** once stopped. Clients keep **Canceling** until the original pull reports `cancelled`, `done`, `interrupted`, or `error`. A successful @@ -794,8 +807,12 @@ progress subscription dropped its last frame. cancellation while retaining completed files for retry. No admin-task/session integration is added here. -**PAIR never inspects or deletes cache files when cancelling.** Partial data is -left for the engine to manage; the desktop briefly says “Partial data is kept so +- **llama.cpp:** PAIR streams and verifies the download itself, and removes only + its exclusive operation-owned temp file before reporting cancellation. A + cleanup failure reports an error. See the [owned-cache contract](LLAMACPP.md#go-model-hub-and-owned-cache). + +**PAIR never inspects or deletes shared engine cache files when cancelling.** +Ollama/LM Studio partial data is left for the engine to manage; the desktop briefly says “Partial data is kept so the download can resume.” This intentionally diverges from NVIDIA/Personal-AI-Router#111's snapshot/stillness cleanup: a shared engine cache does not establish file ownership, even when a partial appears during our pull. @@ -841,8 +858,9 @@ and launcher exit; survivors produce an error (also logged on shutdown). Signal/kill helpers each get at most 500 ms, and group/job verification shares the final 2 s reap deadline: at most 3.5 s per CLI shutdown sequence. Inherited-stdio waiting is capped at 1 s. The aggregate download join is 4 s, -regardless of transfer count. Joins abandon only after the fixed deadline; no -cache data is inspected or removed. +regardless of transfer count. Joins abandon only after the fixed deadline; +shared engine cache data is never inspected or removed. llama.cpp removes only +its own transfer temp file. The outward budgets are download stop sequence **3.5 s < aggregate join 4 s < broker prepare-shutdown 5 s < engine-manager worker EOF grace 6 s** (see diff --git a/services/nvpair-engine-manager/actions.go b/services/nvpair-engine-manager/actions.go index d127c7a4c..7b20c517f 100644 --- a/services/nvpair-engine-manager/actions.go +++ b/services/nvpair-engine-manager/actions.go @@ -27,6 +27,16 @@ func (e *Executor) Action(ctx context.Context, engine, action string, params jso if err != nil { return nil, err } + if engine == "llamacpp" { + switch action { + case "pull_model": + return e.PullModelStream(ctx, engine, modelFromParams(params), params) + case "list_downloaded": + return e.listLlamaModels(ctx) + case "delete_model": + return e.deleteLlamaModel(ctx, modelFromParams(params)) + } + } act, ok := st.manifest.Actions[action] if !ok { return nil, fmt.Errorf("engine %q has no action %q", engine, action) diff --git a/services/nvpair-engine-manager/capabilities.go b/services/nvpair-engine-manager/capabilities.go index bac3e8b34..9237325b7 100644 --- a/services/nvpair-engine-manager/capabilities.go +++ b/services/nvpair-engine-manager/capabilities.go @@ -49,6 +49,9 @@ func capabilitiesForPlatform(m *Manifest, p Platform, goos string, adopted bool, _, list := m.Actions["list_models"] _, pull := m.Actions["pull_model"] deleteAction, deleteModel := m.Actions["delete_model"] + if m.Engine == "llamacpp" { + list, pull, deleteModel = true, true, true + } uninstall := p.Uninstall != nil && (len(p.Uninstall.Run) > 0 || len(p.Uninstall.Remove) > 0) && !adopted if uninstall && p.Runtime.modeOrDefault() == "command" { // Uninstall declines a command-mode engine PAIR has no record of diff --git a/services/nvpair-engine-manager/download_shutdown.go b/services/nvpair-engine-manager/download_shutdown.go index 451263b6f..eec751a35 100644 --- a/services/nvpair-engine-manager/download_shutdown.go +++ b/services/nvpair-engine-manager/download_shutdown.go @@ -11,7 +11,8 @@ import ( // Downloads stop concurrently with engines. The CLI shutdown sequence has at // most 500ms signal helper + 500ms grace + 500ms kill helper + 2s reap (3.5s). // Its 4s aggregate join is below the broker's 6s engine-manager EOF grace; -// StopAll runs alongside it, not afterward. No cache files are touched. +// StopAll runs alongside it, not afterward. Shared engine caches are untouched; +// the Go llama.cpp adapter removes only its own transfer temp file. const downloadShutdownBudget = 4 * time.Second const downloadShutdownSignalBudget = 500 * time.Millisecond const downloadShutdownGrace = 500 * time.Millisecond diff --git a/services/nvpair-engine-manager/executor.go b/services/nvpair-engine-manager/executor.go index 6f1b746ee..0084022ec 100644 --- a/services/nvpair-engine-manager/executor.go +++ b/services/nvpair-engine-manager/executor.go @@ -42,10 +42,11 @@ type EngineStatus struct { // engineState is the per-engine runtime state. type engineState struct { - manifest *Manifest - plat *Platform - logs *logBuffer - installDir string + manifest *Manifest + plat *Platform + logs *logBuffer + installDir string + llamaCacheBusy bool // reservation only; no mutex held across I/O // modelsDir is the resolved {models_dir}: the engine's model store, which no // removal this service performs may delete. Empty when the manifest declares // none, in which case uninstall has nothing to preserve. @@ -120,6 +121,9 @@ func (st *engineState) unlockOperation() { // layer runs the long ones (install, start) in goroutines so the read // loop stays responsive. type Executor struct { + llamaHub *llamaHub + llamaFreeSpace func(*os.File) (uint64, error) + llamaCommit func(*os.File, string, string) error loadProfilesMu sync.Mutex // disk state only; never held across engine calls or notifications settingsHub settings.Hub settingsParent func(context.Context, string, settings.Request, string) (json.RawMessage, error) @@ -182,6 +186,9 @@ type Executor struct { func NewExecutor(reg *Registry, reporter *Reporter, emit func(string, any), baseDir string) *Executor { return &Executor{ + llamaHub: newLlamaHub(), + llamaFreeSpace: llamaCacheFree, + llamaCommit: llamaCacheCommit, reg: reg, tools: defaultToolEnv(), reporter: reporter, diff --git a/services/nvpair-engine-manager/lifecycle.go b/services/nvpair-engine-manager/lifecycle.go index 8bca452f1..276e9de4b 100644 --- a/services/nvpair-engine-manager/lifecycle.go +++ b/services/nvpair-engine-manager/lifecycle.go @@ -118,6 +118,11 @@ func (e *Executor) doStart(ctx context.Context, st *engineState, engine string, } ctx, cancel := context.WithCancel(ctx) st.mu.Lock() + if st.llamaCacheBusy { + st.mu.Unlock() + cancel() + return fmt.Errorf("llama.cpp model cache busy") + } st.startCancel = cancel st.mu.Unlock() defer func() { diff --git a/services/nvpair-engine-manager/llamacpp_cache_unix.go b/services/nvpair-engine-manager/llamacpp_cache_unix.go new file mode 100644 index 000000000..7040acc60 --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_cache_unix.go @@ -0,0 +1,129 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "fmt" + "os" + "path/filepath" + "runtime" + "strings" + + "golang.org/x/sys/unix" +) + +// Pin the canonical directory using the same component-by-component no-follow +// walk as model launch. All later mutations are relative to this descriptor. +func openLlamaCache(root string, create bool) (*os.File, error) { + root, err := filepath.Abs(root) + if err != nil { + return nil, err + } + // Only canonical system aliases are normalized. Never resolve arbitrary + // parent links: that would bless a swapped ancestor before the walk. + if runtime.GOOS == "darwin" { + for _, alias := range []string{"/tmp", "/var"} { + if root == alias || strings.HasPrefix(root, alias+"/") { + root = "/private" + root + break + } + } + } + parent := filepath.Dir(root) + dir, err := openLlamaCacheAt(parent) + if err != nil { + return nil, err + } + defer dir.Close() + name := filepath.Base(root) + if create { + if err := unix.Mkdirat(int(dir.Fd()), name, 0700); err != nil && err != unix.EEXIST { + return nil, err + } + } + fd, err := unix.Openat(int(dir.Fd()), name, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0) + if err != nil { + return nil, err + } + return os.NewFile(uintptr(fd), filepath.Join(parent, name)), nil +} +func openLlamaCacheAt(root string) (*os.File, error) { + if !filepath.IsAbs(root) || filepath.Clean(root) != root { + return nil, fmt.Errorf("invalid cache root") + } + flags := unix.O_RDONLY | unix.O_DIRECTORY | unix.O_NOFOLLOW | unix.O_CLOEXEC + fd, err := unix.Open("/", flags, 0) + if err != nil { + return nil, err + } + for _, part := range strings.Split(strings.TrimPrefix(root, "/"), "/") { + if part == "" { + continue + } + next, err := unix.Openat(fd, part, flags, 0) + _ = unix.Close(fd) + if err != nil { + return nil, fmt.Errorf("cache directory changed or unavailable") + } + fd = next + } + return os.NewFile(uintptr(fd), root), nil +} +func llamaCacheOpen(dir *os.File, name string, create bool) (*os.File, error) { + if name == "" || strings.ContainsAny(name, "/\\") || strings.Contains(name, "..") { + return nil, fmt.Errorf("invalid cache file") + } + flags := unix.O_RDONLY | unix.O_NOFOLLOW | unix.O_CLOEXEC | unix.O_NONBLOCK + if create { + flags = unix.O_RDWR | unix.O_CREAT | unix.O_EXCL | unix.O_NOFOLLOW | unix.O_CLOEXEC + } + fd, err := unix.Openat(int(dir.Fd()), name, flags, 0600) + if err != nil { + return nil, err + } + f := os.NewFile(uintptr(fd), name) + st, err := f.Stat() + if err != nil || !st.Mode().IsRegular() { + _ = f.Close() + return nil, fmt.Errorf("cache entry is not a regular file") + } + return f, nil +} +func llamaCacheRemove(dir *os.File, name string) error { return unix.Unlinkat(int(dir.Fd()), name, 0) } +func llamaCacheRename(dir *os.File, from, to string) error { + return unix.Renameat(int(dir.Fd()), from, int(dir.Fd()), to) +} +func llamaCacheIdentity(f *os.File) string { + info, err := f.Stat() + if err != nil { + return "" + } + var st unix.Stat_t + if err := unix.Fstat(int(f.Fd()), &st); err != nil || !info.Mode().IsRegular() { + return "" + } + return fmt.Sprintf("%d:%d", st.Dev, st.Ino) +} + +// The cache belongs to the user; swaps require another same-user process. +// fstatat never follows a substituted link. macOS has no rename/unlink by fd: +// callers must check at the last possible point and fail closed on mismatch. +func llamaCacheMatch(dir *os.File, name, identity string, size int64) error { + var st unix.Stat_t + if err := unix.Fstatat(int(dir.Fd()), name, &st, unix.AT_SYMLINK_NOFOLLOW); err != nil { + return err + } + if identity == "" || fmt.Sprintf("%d:%d", st.Dev, st.Ino) != identity || st.Mode&unix.S_IFMT != unix.S_IFREG || (size >= 0 && st.Size != size) { + return fmt.Errorf("model cache entry changed; refusing unverified file") + } + return nil +} + +func llamaCacheFree(dir *os.File) (uint64, error) { + var st unix.Statfs_t + err := unix.Fstatfs(int(dir.Fd()), &st) + return uint64(st.Bavail) * uint64(st.Bsize), err +} diff --git a/services/nvpair-engine-manager/llamacpp_cache_unix_test.go b/services/nvpair-engine-manager/llamacpp_cache_unix_test.go new file mode 100644 index 000000000..2db66f94a --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_cache_unix_test.go @@ -0,0 +1,156 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +func TestLlamaDeleteParentSwapCannotEscape(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, nil) + if err := pullLlamaTest(ex, "download"); err != nil { + t.Fatal(err) + } + parent := filepath.Dir(cache) + outside := t.TempDir() + outsideCache := filepath.Join(outside, "llamacpp-models") + if err := os.Mkdir(outsideCache, 0700); err != nil { + t.Fatal(err) + } + for _, name := range []string{llamaManifestName, llamaOwnedFilename(llamaTestID)} { + data, err := os.ReadFile(filepath.Join(cache, name)) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(outsideCache, name), data, 0600); err != nil { + t.Fatal(err) + } + } + if err := os.Rename(parent, parent+".saved"); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = os.Remove(parent); _ = os.Rename(parent+".saved", parent) }) + if err := os.Symlink(outside, parent); err != nil { + t.Fatal(err) + } + if _, err := ex.ModelDelete(context.Background(), "llamacpp", llamaTestID); err == nil { + t.Fatal("swapped parent was accepted") + } + if _, err := os.Stat(filepath.Join(outsideCache, llamaOwnedFilename(llamaTestID))); err != nil { + t.Fatal("outside file changed") + } +} +func TestLlamaDeleteRefusesReplacedRecordedFile(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, nil) + if err := pullLlamaTest(ex, "original"); err != nil { + t.Fatal(err) + } + path := filepath.Join(cache, llamaOwnedFilename(llamaTestID)) + if err := os.Rename(path, path+".saved"); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, llamaTestBytes(), 0600); err != nil { + t.Fatal(err) + } + if _, err := ex.ModelDelete(context.Background(), "llamacpp", llamaTestID); err == nil { + t.Fatal("deleted replacement that PAIR did not record") + } + if _, err := os.Stat(path); err != nil { + t.Fatal("replacement disappeared") + } +} + +func TestLlamaOwnershipManifestEntryBound(t *testing.T) { + dir, err := openLlamaCache(filepath.Join(t.TempDir(), "cache"), true) + if err != nil { + t.Fatal(err) + } + defer dir.Close() + if err := writeLlamaManifest(dir, llamaCacheManifest{Models: make([]llamaCacheRecord, 4097)}); err == nil { + t.Fatal("wrote a manifest the reader must reject") + } + if _, err := llamaCacheOpen(dir, llamaManifestName, false); !os.IsNotExist(err) { + t.Fatal("oversized manifest was committed") + } +} + +func TestLlamaCommittedDescriptorDoesNotClaimReplacement(t *testing.T) { + cache := filepath.Join(t.TempDir(), "cache") + dir, err := openLlamaCache(cache, true) + if err != nil { + t.Fatal(err) + } + defer dir.Close() + verified, err := llamaCacheOpen(dir, "temp", true) + if err != nil { + t.Fatal(err) + } + defer verified.Close() + if _, err := verified.Write(llamaTestBytes()); err != nil { + t.Fatal(err) + } + if err := llamaCacheCommit(dir, "temp", "model.gguf"); err != nil { + t.Fatal(err) + } + original := llamaCacheIdentity(verified) + if original == "" { + t.Fatal("missing committed identity") + } + if err := os.Rename(filepath.Join(cache, "model.gguf"), filepath.Join(cache, "saved.gguf")); err != nil { + t.Fatal(err) + } + replacement, err := llamaCacheOpen(dir, "model.gguf", true) + if err != nil { + t.Fatal(err) + } + defer replacement.Close() + if llamaCacheIdentity(verified) == llamaCacheIdentity(replacement) { + t.Fatal("replacement acquired verified descriptor identity") + } +} + +func TestLlamaCachePinnedDirectoryAndNoClobber(t *testing.T) { + root := t.TempDir() + cache := filepath.Join(root, "models") + dir, err := openLlamaCache(cache, true) + if err != nil { + t.Fatal(err) + } + defer dir.Close() + if err := os.Rename(cache, cache+".saved"); err != nil { + t.Fatal(err) + } + outside := t.TempDir() + _ = os.WriteFile(filepath.Join(outside, "model.gguf"), []byte("outside"), 0600) + if err := os.Symlink(outside, cache); err != nil { + t.Fatal(err) + } + file, err := llamaCacheOpen(dir, "model.gguf", true) + if err != nil { + t.Fatal(err) + } + file.Close() + if err := llamaCacheRemove(dir, "model.gguf"); err != nil { + t.Fatal(err) + } + data, _ := os.ReadFile(filepath.Join(outside, "model.gguf")) + if string(data) != "outside" { + t.Fatal("escaped pinned directory") + } + for _, name := range []string{"temp", "existing.gguf"} { + f, err := llamaCacheOpen(dir, name, true) + if err != nil { + t.Fatal(err) + } + f.Close() + } + if err := llamaCacheCommit(dir, "temp", "existing.gguf"); err == nil { + t.Fatal("overwrote an existing file") + } +} diff --git a/services/nvpair-engine-manager/llamacpp_cache_windows.go b/services/nvpair-engine-manager/llamacpp_cache_windows.go new file mode 100644 index 000000000..3e8217b3e --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_cache_windows.go @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "fmt" + "os" +) + +func llamaCacheUnsupported() error { + return fmt.Errorf("llama.cpp model cache operations are unsupported on Windows") +} +func openLlamaCache(string, bool) (*os.File, error) { return nil, llamaCacheUnsupported() } +func llamaCacheOpen(*os.File, string, bool) (*os.File, error) { return nil, llamaCacheUnsupported() } +func llamaCacheRemove(*os.File, string) error { return llamaCacheUnsupported() } +func llamaCacheRename(*os.File, string, string) error { return llamaCacheUnsupported() } +func llamaCacheCommit(*os.File, string, string) error { return llamaCacheUnsupported() } +func llamaCacheIdentity(*os.File) string { return "" } +func llamaCacheMatch(*os.File, string, string, int64) error { return llamaCacheUnsupported() } +func llamaCacheFree(*os.File) (uint64, error) { return 0, llamaCacheUnsupported() } diff --git a/services/nvpair-engine-manager/llamacpp_cancel_test.go b/services/nvpair-engine-manager/llamacpp_cancel_test.go new file mode 100644 index 000000000..d4f7adb6a --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_cancel_test.go @@ -0,0 +1,275 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +// These exercise the public pull result as well as progress and the errors +// pipeline: merely suppressing the terminal error notification is not enough. +func TestLlamaDownloadCancellationPhases(t *testing.T) { + for _, phase := range []string{"metadata", "free-space", "temp-creation", "streaming", "verification", "rename", "post-rename", "manifest"} { + t.Run(phase, func(t *testing.T) { + started := make(chan struct{}) + ex, cache := llamaDownloadExecutor(t, nil, nil) + cancel := func() { + if !ex.CancelModelPull("llamacpp", llamaTestID, "cancel-phase") { + t.Error("cancel refused") + } + } + switch phase { + case "metadata": + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + close(started) + <-r.Context().Done() + })) + t.Cleanup(server.Close) + ex.llamaHub.endpoint = server.URL + ex.llamaHub.client.Transport = server.Client().Transport + case "free-space", "temp-creation": + ex.llamaFreeSpace = func(*os.File) (uint64, error) { + cancel() + if phase == "free-space" { + return 0, fmt.Errorf("space check interrupted: %w", context.Canceled) + } + // Deny staging creation after cancellation; no temp exists yet. + if err := os.Chmod(cache, 0500); err != nil { + t.Error(err) + } + t.Cleanup(func() { _ = os.Chmod(cache, 0700) }) + return 1 << 40, nil + } + case "streaming", "verification": + ex, cache = llamaDownloadExecutor(t, func(w http.ResponseWriter, r *http.Request) { + if phase == "streaming" { + _, _ = w.Write(llamaTestBytes()[:8]) + w.(http.Flusher).Flush() + close(started) + <-r.Context().Done() + return + } + cancel() + _, _ = w.Write(llamaTestBytes()) + }, func(repo *llamaHubRepo) { + if phase == "verification" { + repo.Siblings[0].LFS.SHA256 = strings.Repeat("b", 64) + } + }) + case "rename", "post-rename", "manifest": + ex.llamaCommit = func(dir *os.File, from, to string) error { + cancel() + if phase == "rename" { + return fmt.Errorf("rename interrupted: %w", context.Canceled) + } + if err := llamaCacheCommit(dir, from, to); err != nil { + return err + } + if phase == "manifest" { + return os.Mkdir(filepath.Join(cache, llamaManifestName), 0700) + } + return nil + } + } + events, unsub := ex.progress.subscribe("llamacpp") + defer unsub() + type outcome struct { + result json.RawMessage + err error + } + done := make(chan outcome, 1) + go func() { + result, err := ex.PullModelStream(context.Background(), "llamacpp", llamaTestID, nil, "cancel-phase") + done <- outcome{result, err} + }() + if phase == "metadata" || phase == "streaming" { + select { + case <-started: + cancel() + case <-time.After(3 * time.Second): + t.Fatal("request did not start") + } + } + select { + case got := <-done: + if got.err != nil || string(got.result) != `{"status":"cancelled"}` { + t.Errorf("pull result = %s, %v; want cancelled success response", got.result, got.err) + } + case <-time.After(3 * time.Second): + t.Fatal("pull did not settle") + } + terminal := "" + for len(events) > 0 { + ev := <-events + if ev.Stage == "cancelled" || ev.Stage == "error" || ev.Stage == "done" { + terminal = ev.Stage + } + } + if terminal != "cancelled" { + t.Errorf("terminal stage = %q", terminal) + } + if errors := ex.reporter.snapshot(); len(errors) != 0 { + t.Errorf("cancel reported engine failure: %+v", errors) + } + entries, err := os.ReadDir(cache) + if err != nil { + t.Fatal(err) + } + for _, entry := range entries { + if strings.HasSuffix(entry.Name(), ".tmp") || strings.HasSuffix(entry.Name(), ".gguf") { + t.Errorf("operation file leaked: %s", entry.Name()) + } + } + if phase == "post-rename" { + dir, err := openLlamaCache(cache, false) + if err != nil { + t.Fatal(err) + } + defer dir.Close() + manifest, err := readLlamaManifest(dir) + if err != nil || len(manifest.Models) != 0 { + t.Errorf("cancelled model recorded: %+v %v", manifest, err) + } + } + }) + } +} + +func TestLlamaDownloadFailureBeforeCancelRemainsError(t *testing.T) { + ex, _ := llamaDownloadExecutor(t, nil, func(repo *llamaHubRepo) { + repo.Siblings[0].LFS.SHA256 = strings.Repeat("b", 64) + }) + _, err := ex.trackedPull(context.Background(), "llamacpp", llamaTestID, "failed", true, func(ctx context.Context) (json.RawMessage, error) { + result, err := ex.pullLlamaModel(ctx, llamaTestID) + if ex.CancelModelPull("llamacpp", llamaTestID, "failed") { + t.Error("cancel accepted after primary failure") + } + return result, err + }) + if err == nil || !strings.Contains(err.Error(), "checksum mismatch") { + t.Fatalf("primary failure lost: %v", err) + } + if errors := ex.reporter.snapshot(); len(errors) != 1 || !strings.Contains(errors[0].ID, "pull-failed:llamacpp:") { + t.Fatalf("primary failure not reported: %+v", errors) + } +} + +func TestLlamaDownloadCancelAfterCommitSucceeds(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, nil) + events, unsub := ex.progress.subscribe("llamacpp") + defer unsub() + result, err := ex.trackedPull(context.Background(), "llamacpp", llamaTestID, "committed", true, func(ctx context.Context) (json.RawMessage, error) { + result, err := ex.pullLlamaModel(ctx, llamaTestID) + if err != nil { + return result, err + } + // The manifest is committed, but trackedPull has not published done yet. + if ex.CancelModelPull("llamacpp", llamaTestID, "committed") { + t.Error("cancel accepted after completion boundary") + } + return result, nil + }) + if err != nil || !strings.Contains(string(result), llamaTestID) { + t.Fatalf("completed pull failed: %s %v", result, err) + } + if _, err := os.Stat(filepath.Join(cache, llamaOwnedFilename(llamaTestID))); err != nil { + t.Fatal(err) + } + models, err := ex.listLlamaModels(context.Background()) + if err != nil || !strings.Contains(string(models), llamaTestID) { + t.Fatalf("committed model missing: %s %v", models, err) + } + completed := false + for len(events) > 0 { + ev := <-events + if ev.Stage == "cancelled" || ev.Stage == "error" { + t.Fatalf("completed download changed outcome: %+v", ev) + } + if ev.Stage == "done" { + completed = true + } + } + if !completed { + t.Fatal("completed download did not report success") + } + if errors := ex.reporter.snapshot(); len(errors) != 0 { + t.Fatalf("completed download reported failure: %+v", errors) + } +} + +type gatedCloseBody struct { + io.ReadCloser + closing chan<- struct{} + release <-chan struct{} + once sync.Once +} + +func (b *gatedCloseBody) Close() error { + b.once.Do(func() { close(b.closing) }) + <-b.release + return b.ReadCloser.Close() +} + +type gatedCloseTransport struct { + http.RoundTripper + closing chan<- struct{} + release <-chan struct{} +} + +func (g gatedCloseTransport) RoundTrip(r *http.Request) (*http.Response, error) { + resp, err := g.RoundTripper.RoundTrip(r) + if err == nil && !strings.HasPrefix(r.URL.Path, "/api/models/") { + resp.Body = &gatedCloseBody{ReadCloser: resp.Body, closing: g.closing, release: g.release} + } + return resp, err +} + +// A failure observed before the response body closes stays a failure even if +// the user cancels while Close is still running. +func TestLlamaDownloadCancelDuringBodyCloseKeepsFailure(t *testing.T) { + ex, _ := llamaDownloadExecutor(t, nil, func(repo *llamaHubRepo) { + repo.Siblings[0].LFS.SHA256 = strings.Repeat("b", 64) + }) + closing, release := make(chan struct{}), make(chan struct{}) + ex.llamaHub.client.Transport = gatedCloseTransport{RoundTripper: ex.llamaHub.client.Transport, closing: closing, release: release} + done := make(chan error, 1) + go func() { + _, err := ex.trackedPull(context.Background(), "llamacpp", llamaTestID, "close-race", true, func(ctx context.Context) (json.RawMessage, error) { + return ex.pullLlamaModel(ctx, llamaTestID) + }) + done <- err + }() + select { + case <-closing: + case <-time.After(10 * time.Second): + t.Fatal("download body was never closed") + } + _ = ex.CancelModelPull("llamacpp", llamaTestID, "close-race") + close(release) + var err error + select { + case err = <-done: + case <-time.After(10 * time.Second): + t.Fatal("pull did not finish") + } + if err == nil || !strings.Contains(err.Error(), "checksum mismatch") { + t.Fatalf("observed failure relabelled: %v", err) + } + if errors := ex.reporter.snapshot(); len(errors) != 1 || !strings.Contains(errors[0].ID, "pull-failed:llamacpp:") { + t.Fatalf("primary failure not reported: %+v", errors) + } +} diff --git a/services/nvpair-engine-manager/llamacpp_catalog.go b/services/nvpair-engine-manager/llamacpp_catalog.go new file mode 100644 index 000000000..3be1863f7 --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_catalog.go @@ -0,0 +1,248 @@ +// 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" + "net/http" + "net/url" + "regexp" + "slices" + "strings" + "sync" + "time" + "unicode/utf8" +) + +const llamaCatalogTTL = 6 * time.Hour +const llamaCatalogTimeout = 20 * time.Second +const llamaMaxModelSize int64 = 64 << 30 + +var llamaComponent = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.-]{0,95}$`) +var llamaFilename = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.-]{0,199}\.gguf$`) +var llamaSHA = regexp.MustCompile(`^[a-f0-9]{64}$`) + +// LlamaCatalogEntry is one single-file Q4_K_M model. Metadata is resolved again +// immediately before transfer; a catalog row is never authority for a download. +type LlamaCatalogEntry struct { + ID string `json:"id"` + File string `json:"file"` + Size int64 `json:"size,omitempty"` + SHA256 string `json:"sha256,omitempty"` +} +type llamaHubRepo struct { + ID string `json:"id"` + SHA string `json:"sha"` + Private *bool `json:"private"` + Gated json.RawMessage `json:"gated"` + Pipeline string `json:"pipeline_tag"` + Tags []string `json:"tags"` + Siblings []struct { + Name string `json:"rfilename"` + Size int64 `json:"size"` + LFS *struct { + SHA256 string `json:"sha256"` + OID string `json:"oid"` + Size int64 `json:"size"` + } `json:"lfs"` + } `json:"siblings"` +} +type llamaCatalogCache struct { + at time.Time + rows []LlamaCatalogEntry +} +type llamaHub struct { + client *http.Client + endpoint string + now func() time.Time + mu sync.Mutex + cache map[string]llamaCatalogCache +} + +func newLlamaHub() *llamaHub { + tr := &http.Transport{Proxy: http.ProxyFromEnvironment, DialContext: llamaDial, TLSHandshakeTimeout: 10 * time.Second, ResponseHeaderTimeout: 10 * time.Second, MaxResponseHeaderBytes: 64 << 10, IdleConnTimeout: 30 * time.Second} + return &llamaHub{endpoint: "https://huggingface.co", now: time.Now, cache: make(map[string]llamaCatalogCache), client: &http.Client{Transport: tr, CheckRedirect: func(req *http.Request, via []*http.Request) error { + if req.URL.Scheme != "https" || req.URL.User != nil || len(via) >= 5 { + return fmt.Errorf("unsafe model redirect") + } + return nil + }}} +} +func validLlamaComponent(s string) bool { + return llamaComponent.MatchString(s) && !strings.Contains(s, "..") +} +func validLlamaRepo(s string) bool { + p := strings.Split(s, "/") + return len(p) == 2 && validLlamaComponent(p[0]) && validLlamaComponent(p[1]) +} +func validLlamaFile(s string) bool { + return llamaFilename.MatchString(s) && !strings.Contains(s, "..") && !llamaSplitName.MatchString(s) +} +func llamaRepoID(id string) (string, error) { + repo, ok := strings.CutSuffix(id, ":Q4_K_M") + if !ok || !validLlamaRepo(repo) { + return "", fmt.Errorf("invalid llama.cpp model ID") + } + return repo, nil +} +func (r llamaHubRepo) entry(requireLFS bool) (LlamaCatalogEntry, bool) { + if !validLlamaRepo(r.ID) || r.Private == nil || *r.Private || string(r.Gated) != "false" || r.Pipeline != "text-generation" || !slices.Contains(r.Tags, "gguf") { + return LlamaCatalogEntry{}, false + } + var entries []LlamaCatalogEntry + for _, f := range r.Siblings { + if !validLlamaFile(f.Name) || !strings.HasSuffix(strings.ToUpper(f.Name), "Q4_K_M.GGUF") { + continue + } + row := LlamaCatalogEntry{ID: r.ID + ":Q4_K_M", File: f.Name} + if f.LFS != nil { + row.Size = f.LFS.Size + row.SHA256 = f.LFS.SHA256 + if row.SHA256 == "" { + row.SHA256 = strings.TrimPrefix(f.LFS.OID, "sha256:") + } + if row.Size <= 0 || row.Size > llamaMaxModelSize || !llamaSHA.MatchString(row.SHA256) || f.Size != row.Size { + continue + } + } else if requireLFS { + continue + } + entries = append(entries, row) + } + if len(entries) != 1 { + return LlamaCatalogEntry{}, false + } + return entries[0], true +} + +// Preserve cancellation and deadline causes without exposing request URLs or +// remote response contents in user-facing errors or the engine error pipeline. +func llamaRequestError(message string, err error) error { + for _, cause := range []error{context.Canceled, context.DeadlineExceeded} { + if errors.Is(err, cause) { + return fmt.Errorf("%s: %w", message, cause) + } + } + return errors.New(message) +} + +func llamaReadError(ctx context.Context, message string, err error) error { + // Some transports return a plain read error when cancellation tears down + // the body. Only at that read boundary may the operation context supply + // the missing cause; never use this for integrity or filesystem failures. + if ctx.Err() != nil { + p, tracked := ctx.Value(activePullKey{}).(*activePull) + if !tracked || p.requested.Load() || errors.Is(ctx.Err(), context.DeadlineExceeded) { + err = ctx.Err() + } + } + return llamaRequestError(message, err) +} + +func (h *llamaHub) get(ctx context.Context, path string, out any) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, h.endpoint+path, nil) + if err != nil { + return llamaRequestError("invalid catalog request", err) + } + resp, err := h.client.Do(req) + if err != nil { + return llamaRequestError("model metadata request failed", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("model metadata returned HTTP %d", resp.StatusCode) + } + data, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20+1)) + if err != nil { + return llamaReadError(ctx, "model metadata could not be read", err) + } + if len(data) > 4<<20 { + return fmt.Errorf("model metadata exceeds limit") + } + if json.Unmarshal(data, out) != nil { + return fmt.Errorf("invalid model metadata") + } + return nil +} +func (h *llamaHub) catalog(ctx context.Context, query string) ([]LlamaCatalogEntry, error) { + if !utf8.ValidString(query) || utf8.RuneCountInString(query) > 100 || strings.IndexFunc(query, func(r rune) bool { return r < 32 || r == 127 }) >= 0 { + return nil, fmt.Errorf("query must be at most 100 characters without control characters") + } + key := strings.TrimSpace(query) + h.mu.Lock() + cached, ok := h.cache[key] + h.mu.Unlock() + if ok && h.now().Sub(cached.at) < llamaCatalogTTL { + return slices.Clone(cached.rows), nil + } + ctx, cancel := context.WithTimeout(ctx, llamaCatalogTimeout) + defer cancel() + publishers := []string{"ggml-org", "bartowski", "unsloth"} + if key != "" { + publishers = []string{""} + } + rows := make([]LlamaCatalogEntry, 0) + for _, publisher := range publishers { + q := url.Values{"filter": {"gguf"}, "pipeline_tag": {"text-generation"}, "limit": {"50"}, "full": {"true"}} + if key != "" { + q.Set("search", key) + } else { + q.Set("author", publisher) + } + var repos []llamaHubRepo + if err := h.get(ctx, "/api/models?"+q.Encode(), &repos); err != nil { + return nil, err + } + for i, r := range repos { + if i >= 50 { + break + } + if publisher != "" && !strings.HasPrefix(r.ID, publisher+"/") { + continue + } + if row, ok := r.entry(false); ok { + rows = append(rows, row) + } + } + } + h.mu.Lock() + queries := len(h.cache) + if _, ok := h.cache[""]; ok { + queries-- + } + if _, exists := h.cache[key]; key != "" && !exists && queries >= 20 { + var oldest string + var at time.Time + for k, v := range h.cache { + if k != "" && (at.IsZero() || v.at.Before(at)) { + oldest, at = k, v.at + } + } + delete(h.cache, oldest) + } + h.cache[key] = llamaCatalogCache{at: h.now(), rows: slices.Clone(rows)} + h.mu.Unlock() + return rows, nil +} +func (h *llamaHub) resolve(ctx context.Context, id string) (LlamaCatalogEntry, string, error) { + repo, err := llamaRepoID(id) + if err != nil { + return LlamaCatalogEntry{}, "", err + } + ctx, cancel := context.WithTimeout(ctx, llamaCatalogTimeout) + defer cancel() + var r llamaHubRepo + if err := h.get(ctx, "/api/models/"+repo+"?blobs=true", &r); err != nil { + return LlamaCatalogEntry{}, "", err + } + row, ok := r.entry(true) + if !ok || row.ID != id || !regexp.MustCompile(`^[a-f0-9]{40}$`).MatchString(r.SHA) { + return LlamaCatalogEntry{}, "", fmt.Errorf("model is not a public single-file Q4_K_M generative GGUF with verified LFS metadata") + } + return row, h.endpoint + "/" + repo + "/resolve/" + r.SHA + "/" + url.PathEscape(row.File), nil +} diff --git a/services/nvpair-engine-manager/llamacpp_catalog_test.go b/services/nvpair-engine-manager/llamacpp_catalog_test.go new file mode 100644 index 000000000..c584cfe80 --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_catalog_test.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 ( + "context" + "encoding/json" + "fmt" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// Shape and names verified with read-only Hugging Face API metadata, 2026-11-07. +const llamaRepoFixture = `{"id":"bartowski/Llama-3.2-1B-Instruct-GGUF","sha":"0123456789012345678901234567890123456789","private":false,"gated":false,"pipeline_tag":"text-generation","tags":["gguf"],"siblings":[{"rfilename":"Llama-3.2-1B-Instruct-Q4_K_M.gguf","size":24,"lfs":{"sha256":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","size":24}}]}` + +func fixtureLlamaRepo(t *testing.T) llamaHubRepo { + t.Helper() + var r llamaHubRepo + if err := json.Unmarshal([]byte(llamaRepoFixture), &r); err != nil { + t.Fatal(err) + } + return r +} +func TestLlamaCatalogValidation(t *testing.T) { + r := fixtureLlamaRepo(t) + if _, ok := r.entry(true); !ok { + t.Fatal("valid fixture rejected") + } + for _, id := range []string{"../outside/repo", "a/..", "a/b/c", "a/" + strings.Repeat("x", 100), "a/b\n", "a/b\\c"} { + r := fixtureLlamaRepo(t) + r.ID = id + if _, ok := r.entry(true); ok { + t.Fatalf("accepted %q", id) + } + } + for _, name := range []string{"../Q4_K_M.gguf", "sub/Q4_K_M.gguf", "x..Q4_K_M.gguf", strings.Repeat("x", 210) + "Q4_K_M.gguf", "x-Q4_K_M-00001-of-00002.gguf", "x-Q4_K_M\n.gguf"} { + r := fixtureLlamaRepo(t) + r.Siblings[0].Name = name + if _, ok := r.entry(true); ok { + t.Fatalf("accepted %q", name) + } + } + for _, mutate := range []func(*llamaHubRepo){func(r *llamaHubRepo) { r.Gated = json.RawMessage(`"auto"`) }, func(r *llamaHubRepo) { r.Private = nil }, func(r *llamaHubRepo) { r.Pipeline = "feature-extraction" }, func(r *llamaHubRepo) { r.Siblings[0].LFS.Size = llamaMaxModelSize + 1 }, func(r *llamaHubRepo) { r.Siblings = append(r.Siblings, r.Siblings[0]) }} { + r := fixtureLlamaRepo(t) + mutate(&r) + if _, ok := r.entry(true); ok { + t.Fatal("hostile fixture accepted") + } + } +} +func TestLlamaCatalogCuratedAndRPC(t *testing.T) { + publishers := map[string]bool{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + author := r.URL.Query().Get("author") + publishers[author] = true + rows := make([]llamaHubRepo, 60) + for i := range rows { + rows[i] = fixtureLlamaRepo(t) + rows[i].ID = fmt.Sprintf("%s/model%d", author, i) + } + _ = json.NewEncoder(w).Encode(rows) + })) + defer srv.Close() + ex := NewExecutor(NewRegistry(), NewReporter(nil), nil, t.TempDir()) + ex.llamaHub.endpoint = srv.URL + server, client := net.Pipe() + defer server.Close() + defer client.Close() + _ = client.SetDeadline(time.Now().Add(3 * time.Second)) + m := NewManager(NewCodec(server), ex, nil) + defer m.remoteHTTP.CloseIdle() + defer m.readyHTTP.CloseIdle() + defer m.settingsHTTP.CloseIdle() + id := json.RawMessage("1") + m.handleMessage(context.Background(), &Message{JSONRPC: "2.0", ID: &id, Method: "engine:model-catalog", Params: json.RawMessage(`{}`)}) + var response Message + if err := json.NewDecoder(client).Decode(&response); err != nil { + t.Fatal(err) + } + var result struct { + Models []LlamaCatalogEntry `json:"models"` + } + if err := json.Unmarshal(response.Result, &result); err != nil { + t.Fatal(err) + } + if response.Error != nil || len(result.Models) != 150 || !publishers["ggml-org"] || !publishers["bartowski"] || !publishers["unsloth"] { + t.Fatalf("curated RPC result: %d %v", len(result.Models), response.Error) + } +} + +func TestLlamaCatalogBoundsCacheTTL(t *testing.T) { + calls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + if r.URL.Query().Get("limit") != "50" { + t.Error("missing API bound") + } + rows := make([]llamaHubRepo, 70) + for i := range rows { + rows[i] = fixtureLlamaRepo(t) + } + _ = json.NewEncoder(w).Encode(rows) + })) + defer srv.Close() + h := newLlamaHub() + h.endpoint = srv.URL + now := time.Now() + h.now = func() time.Time { return now } + rows, err := h.catalog(context.Background(), "llama") + if err != nil || len(rows) != 50 { + t.Fatalf("rows=%d err=%v", len(rows), err) + } + rows[0].ID = "mutated" + rows, err = h.catalog(context.Background(), "llama") + if err != nil || calls != 1 || rows[0].ID == "mutated" { + t.Fatal("cache not isolated") + } + now = now.Add(llamaCatalogTTL) + _, _ = h.catalog(context.Background(), "llama") + if calls != 2 { + t.Fatal("TTL ignored") + } + for i := 0; i < 25; i++ { + _, _ = h.catalog(context.Background(), fmt.Sprintf("q%d", i)) + } + if len(h.cache) > 20 { + t.Fatalf("unbounded cache: %d", len(h.cache)) + } + before := calls + for _, q := range []string{strings.Repeat("x", 101), "x\n"} { + if _, err := h.catalog(context.Background(), q); err == nil { + t.Fatal("query accepted") + } + } + if calls != before { + t.Fatal("invalid query reached network") + } +} diff --git a/services/nvpair-engine-manager/llamacpp_cause_test.go b/services/nvpair-engine-manager/llamacpp_cause_test.go new file mode 100644 index 000000000..7e306a264 --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_cause_test.go @@ -0,0 +1,251 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" +) + +type llamaCauseTransport func(*http.Request) (*http.Response, error) + +func (f llamaCauseTransport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +type llamaCauseBody struct { + read func([]byte) (int, error) + close func() error +} + +func (b llamaCauseBody) Read(p []byte) (int, error) { return b.read(p) } +func (b llamaCauseBody) Close() error { return b.close() } + +// Losing the metadata cause used to make the live metadata cancel report a +// pull-failed engine error once order-based classification was removed. +func TestLlamaMetadataPreservesCancellationCause(t *testing.T) { + for _, phase := range []string{"request", "body-wrapped", "body-unwrapped"} { + t.Run(phase, func(t *testing.T) { + h := newLlamaHub() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + h.client.Transport = llamaCauseTransport(func(r *http.Request) (*http.Response, error) { + if phase == "request" { + cancel() + return nil, fmt.Errorf("transport stopped: %w", r.Context().Err()) + } + body := llamaCauseBody{read: func([]byte) (int, error) { + cancel() + if phase == "body-wrapped" { + return 0, fmt.Errorf("read stopped: %w", r.Context().Err()) + } + return 0, io.ErrUnexpectedEOF + }, close: func() error { return nil }} + return &http.Response{StatusCode: 200, Body: body, Request: r}, nil + }) + var repo llamaHubRepo + if err := h.get(ctx, "/api/models/test", &repo); !errors.Is(err, context.Canceled) { + t.Fatalf("cancellation cause lost: %v", err) + } + }) + } +} + +// The same accepted cancel must not relabel a real filesystem failure, but a +// filesystem helper returning the cancellation cause must still be cancelled. +func TestLlamaDownloadClassifiesByCause(t *testing.T) { + for _, phase := range []string{"free-space", "commit"} { + for _, cancelled := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/cancellation=%v", phase, cancelled), func(t *testing.T) { + ex, _ := llamaDownloadExecutor(t, nil, nil) + cause := errors.New("filesystem failed") + if cancelled { + cause = context.Canceled + } + fail := func() error { + if !ex.CancelModelPull("llamacpp", llamaTestID, "cause") { + t.Error("cancel refused") + } + return fmt.Errorf("phase failed: %w", cause) + } + if phase == "free-space" { + ex.llamaFreeSpace = func(*os.File) (uint64, error) { return 0, fail() } + } else { + ex.llamaCommit = func(*os.File, string, string) error { return fail() } + } + result, err := ex.PullModelStream(context.Background(), "llamacpp", llamaTestID, nil, "cause") + if cancelled { + if err != nil || string(result) != `{"status":"cancelled"}` { + t.Fatalf("want cancelled, got %s %v", result, err) + } + if len(ex.reporter.snapshot()) != 0 { + t.Fatal("cancellation reported as failure") + } + } else { + if err == nil || !strings.Contains(err.Error(), "filesystem failed") { + t.Fatalf("real failure lost: %s %v", result, err) + } + if len(ex.reporter.snapshot()) != 1 { + t.Fatal("failure not reported") + } + } + }) + } + } +} + +// A real blocked HTTP client request must wrap context.Canceled; plain body +// read errors need the operation context as their cause at the read boundary. +func TestLlamaDownloadCancellationCausePhases(t *testing.T) { + for _, phase := range []string{"metadata-blocked", "request-blocked", "streaming-wrapped", "streaming-unwrapped", "after-commit"} { + t.Run(phase, func(t *testing.T) { + ex, _ := llamaDownloadExecutor(t, nil, nil) + cancel := func() bool { return ex.CancelModelPull("llamacpp", llamaTestID, "cause-phase") } + if strings.HasSuffix(phase, "blocked") { + started := make(chan struct{}) + if phase == "metadata-blocked" { + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + close(started) + <-r.Context().Done() + })) + t.Cleanup(server.Close) + ex.llamaHub.endpoint = server.URL + ex.llamaHub.client.Transport = server.Client().Transport + } else { + ex, _ = llamaDownloadExecutor(t, func(w http.ResponseWriter, r *http.Request) { + close(started) + <-r.Context().Done() + }, nil) + } + done := make(chan struct{}) + defer close(done) + go func() { + select { + case <-started: + if !cancel() { + t.Error("blocked request cancel refused") + } + case <-done: + } + }() + } else { + transport := ex.llamaHub.client.Transport + ex.llamaHub.client.Transport = llamaCauseTransport(func(r *http.Request) (*http.Response, error) { + resp, err := transport.RoundTrip(r) + if err != nil || strings.HasPrefix(r.URL.Path, "/api/models/") { + return resp, err + } + body := resp.Body + if phase == "after-commit" { + resp.Body = llamaCauseBody{read: body.Read, close: func() error { + if cancel() { + t.Error("cancel accepted after manifest commit") + } + return body.Close() + }} + } else { + resp.Body = llamaCauseBody{read: func([]byte) (int, error) { + if !cancel() { + t.Error("streaming cancel refused") + } + if phase == "streaming-wrapped" { + return 0, fmt.Errorf("transport read: %w", r.Context().Err()) + } + return 0, io.ErrUnexpectedEOF + }, close: body.Close} + } + return resp, nil + }) + } + events, unsub := ex.progress.subscribe("llamacpp") + defer unsub() + ctx, stop := context.WithTimeout(context.Background(), 5*time.Second) + defer stop() + result, err := ex.PullModelStream(ctx, "llamacpp", llamaTestID, nil, "cause-phase") + want := "cancelled" + if phase == "after-commit" { + want = "done" + } + if err != nil { + t.Fatalf("pull failed: %v", err) + } + if want == "cancelled" && string(result) != `{"status":"cancelled"}` { + t.Fatalf("cancellation lost: %s", result) + } + if want == "done" && !strings.Contains(string(result), llamaTestID) { + t.Fatalf("commit lost: %s", result) + } + terminal := "" + for len(events) > 0 { + ev := <-events + if ev.Stage == "done" || ev.Stage == "cancelled" || ev.Stage == "error" { + terminal = ev.Stage + } + } + if terminal != want { + t.Fatalf("terminal = %q, want %q", terminal, want) + } + if len(ex.reporter.snapshot()) != 0 { + t.Fatal("cancellation or success reported a pull failure") + } + }) + } +} + +func TestLlamaDownloadChecksumCancelDuringTeardown(t *testing.T) { + for _, point := range []string{"body-close", "cache-close", "manifest-read"} { + t.Run(point, func(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, func(repo *llamaHubRepo) { + repo.Siblings[0].LFS.SHA256 = strings.Repeat("b", 64) + }) + transport := ex.llamaHub.client.Transport + ex.llamaHub.client.Transport = llamaCauseTransport(func(r *http.Request) (*http.Response, error) { + resp, err := transport.RoundTrip(r) + if err != nil || strings.HasPrefix(r.URL.Path, "/api/models/") { + return resp, err + } + body := resp.Body + resp.Body = llamaCauseBody{read: body.Read, close: func() error { + // These nested teardown operations all run before pullLlamaModel + // returns. Cancellation must not depend on a helper's Close order. + if point != "body-close" { + dir, err := openLlamaCache(cache, false) + if err != nil { + t.Error(err) + return body.Close() + } + if !ex.CancelModelPull("llamacpp", llamaTestID, "teardown") { + t.Error("cancel refused during nested teardown") + } + if point == "manifest-read" { + _, err = readLlamaManifest(dir) + if err != nil { + t.Error(err) + } + } + _ = dir.Close() + } else if !ex.CancelModelPull("llamacpp", llamaTestID, "teardown") { + t.Error("cancel refused during body teardown") + } + return body.Close() + }} + return resp, nil + }) + if err := pullLlamaTest(ex, "teardown"); err == nil || !strings.Contains(err.Error(), "checksum mismatch") { + t.Fatalf("checksum failure lost: %v", err) + } + if len(ex.reporter.snapshot()) != 1 { + t.Fatal("checksum failure not reported") + } + }) + } +} diff --git a/services/nvpair-engine-manager/llamacpp_cleanup_test.go b/services/nvpair-engine-manager/llamacpp_cleanup_test.go new file mode 100644 index 000000000..57da318da --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_cleanup_test.go @@ -0,0 +1,116 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "context" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func TestLlamaDownloadCleanupLeavesSubstitutedTemp(t *testing.T) { + for _, cancelPull := range []bool{false, true} { + name := "failure" + if cancelPull { + name = "cancel" + } + t.Run(name, func(t *testing.T) { + started := make(chan error, 1) + var cache, substituted string + ex, path := llamaDownloadExecutor(t, func(w http.ResponseWriter, r *http.Request) { + files, err := os.ReadDir(cache) + if err != nil { + started <- err + w.WriteHeader(500) + return + } + for _, f := range files { + if strings.HasPrefix(f.Name(), ".pair-") && strings.HasSuffix(f.Name(), ".tmp") { + substituted = filepath.Join(cache, f.Name()) + if err = os.Rename(substituted, substituted+".saved"); err == nil { + err = os.WriteFile(substituted, []byte("unrelated"), 0600) + } + break + } + } + if substituted == "" && err == nil { + err = os.ErrNotExist + } + started <- err + if cancelPull { + <-r.Context().Done() + } else { + w.WriteHeader(500) + } + }, nil) + cache = path + events, unsub := ex.progress.subscribe("llamacpp") + defer unsub() + done := make(chan error, 1) + go func() { done <- pullLlamaTest(ex, "cleanup-swap") }() + select { + case err := <-started: + if err != nil { + t.Fatal(err) + } + case <-time.After(3 * time.Second): + t.Fatal("transfer not started") + } + if cancelPull && !ex.CancelModelPull("llamacpp", llamaTestID, "cleanup-swap") { + t.Fatal("cancel refused") + } + select { + case err := <-done: + if err == nil || !strings.Contains(err.Error(), "cleanup skipped") { + t.Errorf("substituted cleanup not reported: %v", err) + } + case <-time.After(3 * time.Second): + t.Fatal("transfer did not settle") + } + data, err := os.ReadFile(substituted) + if err != nil || string(data) != "unrelated" { + t.Fatalf("substituted file changed: %q %v", data, err) + } + failed := false + for len(events) > 0 { + ev := <-events + if ev.Stage == "cancelled" || ev.Stage == "done" { + t.Fatalf("claimed success: %+v", ev) + } + if ev.Stage == "error" { + failed = true + } + } + if !failed { + t.Fatal("cleanup failure not emitted through progress") + } + }) + } +} + +func TestLlamaDeleteLeavesReplacedCompletedModel(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, nil) + if err := pullLlamaTest(ex, "delete-swap"); err != nil { + t.Fatal(err) + } + path := filepath.Join(cache, llamaOwnedFilename(llamaTestID)) + if err := os.Rename(path, path+".saved"); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, llamaTestBytes(), 0600); err != nil { + t.Fatal(err) + } + if _, err := ex.deleteLlamaModel(context.Background(), llamaTestID); err == nil || !strings.Contains(err.Error(), "changed") { + t.Fatalf("replacement accepted: %v", err) + } + if _, err := os.Stat(path); err != nil { + t.Fatalf("replacement deleted: %v", err) + } +} diff --git a/services/nvpair-engine-manager/llamacpp_commit_darwin.go b/services/nvpair-engine-manager/llamacpp_commit_darwin.go new file mode 100644 index 000000000..b7fc95a59 --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_commit_darwin.go @@ -0,0 +1,14 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "os" + + "golang.org/x/sys/unix" +) + +func llamaCacheCommit(dir *os.File, from, to string) error { + return unix.RenameatxNp(int(dir.Fd()), from, int(dir.Fd()), to, unix.RENAME_EXCL) +} diff --git a/services/nvpair-engine-manager/llamacpp_commit_linux.go b/services/nvpair-engine-manager/llamacpp_commit_linux.go new file mode 100644 index 000000000..daeec2875 --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_commit_linux.go @@ -0,0 +1,14 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "os" + + "golang.org/x/sys/unix" +) + +func llamaCacheCommit(dir *os.File, from, to string) error { + return unix.Renameat2(int(dir.Fd()), from, int(dir.Fd()), to, unix.RENAME_NOREPLACE) +} diff --git a/services/nvpair-engine-manager/llamacpp_download.go b/services/nvpair-engine-manager/llamacpp_download.go new file mode 100644 index 000000000..22454aa60 --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_download.go @@ -0,0 +1,470 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "os" + "strings" + "time" +) + +const llamaManifestName = ".pair-models.json" + +type llamaOwnedModel struct { + ID string `json:"id"` + File string `json:"file"` + Size int64 `json:"size"` + SHA256 string `json:"sha256"` + FinalHost string `json:"finalHost"` +} +type llamaCacheRecord struct { + llamaOwnedModel + Identity string `json:"fileIdentity"` +} +type llamaCacheManifest struct { + Models []llamaCacheRecord `json:"models"` +} + +func llamaOwnedFilename(id string) string { + sum := sha256.Sum256([]byte(id)) + return hex.EncodeToString(sum[:]) + ".gguf" +} +func readLlamaManifest(dir *os.File) (llamaCacheManifest, error) { + out := llamaCacheManifest{Models: []llamaCacheRecord{}} + f, err := llamaCacheOpen(dir, llamaManifestName, false) + if os.IsNotExist(err) { + return out, nil + } + if err != nil { + return out, err + } + defer f.Close() + data, err := io.ReadAll(io.LimitReader(f, (1<<20)+1)) + if err != nil { + return out, fmt.Errorf("read model ownership manifest: %w", err) + } + if len(data) > 1<<20 { + return out, fmt.Errorf("invalid model ownership manifest") + } + if json.Unmarshal(data, &out) != nil || len(out.Models) > 4096 { + return out, fmt.Errorf("invalid model ownership manifest") + } + seen := map[string]bool{} + for _, m := range out.Models { + if _, err := llamaRepoID(m.ID); err != nil || m.File != llamaOwnedFilename(m.ID) || m.Size <= 0 || m.Size > llamaMaxModelSize || !llamaSHA.MatchString(m.SHA256) || seen[m.ID] || m.Identity == "" { + return out, fmt.Errorf("invalid model ownership record") + } + seen[m.ID] = true + } + return out, nil +} +func llamaTempName() string { + var bytes [16]byte + if _, err := rand.Read(bytes[:]); err != nil { + panic("system random unavailable") + } + return ".pair-" + hex.EncodeToString(bytes[:]) + ".tmp" +} +func writeLlamaManifest(dir *os.File, m llamaCacheManifest) (retErr error) { + if len(m.Models) > 4096 { + return fmt.Errorf("model ownership manifest full") + } + data, err := json.Marshal(m) + if err != nil || len(data) > 1<<20 { + return fmt.Errorf("model ownership manifest full") + } + name := llamaTempName() + f, err := llamaCacheOpen(dir, name, true) + if err != nil { + return err + } + identity := llamaCacheIdentity(f) + defer func() { + if err := llamaCacheRemoveOwned(dir, name, identity, -1); err != nil && !os.IsNotExist(err) { + retErr = fmt.Errorf("manifest temporary file could not be removed: %w", err) + } + _ = f.Close() + }() + if _, err = f.Write(data); err != nil { + return err + } + if err = f.Sync(); err != nil { + return err + } + if err = llamaCacheRename(dir, name, llamaManifestName); err != nil { + return err + } + if err := llamaCacheMatch(dir, llamaManifestName, identity, int64(len(data))); err != nil { + return fmt.Errorf("model ownership manifest changed during commit: %w", err) + } + return dir.Sync() +} + +// A reservation is a state flag, not a lock held across filesystem/network I/O. +// Reject concurrent cache mutations rather than queuing an unbounded wait. +func (e *Executor) reserveLlamaCache(deleting bool) (*engineState, func(), error) { + st, err := e.state("llamacpp") + if err != nil { + return nil, nil, err + } + st.mu.Lock() + if st.llamaCacheBusy || (deleting && (st.running || st.opDone != nil)) { + st.mu.Unlock() + return nil, nil, fmt.Errorf("llama.cpp model cache busy or engine in use; stop the engine before deleting") + } + st.llamaCacheBusy = true + st.mu.Unlock() + return st, func() { st.mu.Lock(); st.llamaCacheBusy = false; st.mu.Unlock() }, nil +} +func (e *Executor) pullLlamaModel(ctx context.Context, id string) (result json.RawMessage, retErr error) { + // trackedPull classifies by the returned cause, never by cancel timing. + // Disable further cancels on every exit, after releasing operation resources. + defer e.finishPull(ctx) + if _, err := llamaRepoID(id); err != nil { + return nil, err + } + st, release, err := e.reserveLlamaCache(false) + if err != nil { + return nil, err + } + defer release() + dir, err := openLlamaCache(st.installDir+"-models", true) + if err != nil { + return nil, fmt.Errorf("open model cache: %w", err) + } + defer dir.Close() + manifest, err := readLlamaManifest(dir) + if err != nil { + return nil, err + } + if len(manifest.Models) >= 4096 { + return nil, fmt.Errorf("model ownership manifest full") + } + for _, m := range manifest.Models { + if m.ID == id { + return nil, fmt.Errorf("model already recorded; delete it before downloading again") + } + } + row, downloadURL, err := e.llamaHub.resolve(ctx, id) + if err != nil { + return nil, err + } + free, err := e.llamaFreeSpace(dir) + if err != nil { + return nil, fmt.Errorf("check model disk space: %w", err) + } + if free < uint64(row.Size)+(256<<20) { + return nil, fmt.Errorf("not enough disk space: free space must cover the model and 256 MiB reserve") + } + if err := ctx.Err(); err != nil { + return nil, err + } + final := llamaOwnedFilename(id) + if f, err := llamaCacheOpen(dir, final, false); err == nil { + _ = f.Close() + return nil, fmt.Errorf("destination already exists; unrecorded files are never replaced") + } else if !os.IsNotExist(err) { + return nil, err + } + temp := llamaTempName() + f, err := llamaCacheOpen(dir, temp, true) + if err != nil { + return nil, fmt.Errorf("create temporary model: %w", err) + } + // Capture dev/ino immediately, before network I/O or writes. + identity := llamaCacheIdentity(f) + if identity == "" { + _ = f.Close() + return nil, fmt.Errorf("cannot identify temporary model") + } + published, manifestWritten, completed := false, false, false + originalCount := len(manifest.Models) + defer func() { + if !completed && published { + if manifestWritten { + manifest.Models = manifest.Models[:originalCount] + if err := writeLlamaManifest(dir, manifest); err != nil { + result, retErr = nil, fmt.Errorf("model ownership rollback failed: %w", err) + } + } + if err := llamaCacheRemoveOwned(dir, final, identity, row.Size); err != nil && !os.IsNotExist(err) { + result, retErr = nil, fmt.Errorf("%v; model cleanup: %w", retErr, err) + } + } + if err := llamaCacheRemoveOwned(dir, temp, identity, -1); err != nil && !os.IsNotExist(err) { + result = nil + retErr = fmt.Errorf("model transfer stopped but its temporary file could not be removed: %w", err) + } + _ = f.Close() + }() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil) + if err != nil { + return nil, llamaRequestError("invalid model download request", err) + } + resp, err := e.llamaHub.client.Do(req) + if err != nil { + return nil, llamaRequestError("model transfer failed", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("model download returned HTTP %d", resp.StatusCode) + } + if resp.ContentLength >= 0 && resp.ContentLength != row.Size { + return nil, fmt.Errorf("model size mismatch") + } + hash := sha256.New() + reader := io.LimitReader(resp.Body, row.Size+1) + buf := make([]byte, 128<<10) + var total, spaceCheckedAt int64 + lastSpaceCheck := time.Now() + last := -1 + for { + n, readErr := reader.Read(buf) + if n > 0 { + total += int64(n) + if total > row.Size { + return nil, fmt.Errorf("model exceeds expected size") + } + if _, err := f.Write(buf[:n]); err != nil { + return nil, err + } + _, _ = hash.Write(buf[:n]) + // Recheck the pinned transfer filesystem at streaming boundaries: + // at most 256 MiB or five seconds of active streaming between checks. + // No state lock is held during fstatfs or network/file waits. + if total-spaceCheckedAt >= 256<<20 || time.Since(lastSpaceCheck) >= 5*time.Second { + free, err := e.llamaFreeSpace(f) + if err != nil { + return nil, fmt.Errorf("check remaining model disk space: %w", err) + } + if free < uint64(row.Size-total)+(256<<20) { + return nil, fmt.Errorf("not enough disk space for remaining model data and 256 MiB reserve") + } + spaceCheckedAt, lastSpaceCheck = total, time.Now() + } + pct := int(total * 100 / row.Size) + if pct != last { + last = pct + e.emitPullUpdate(ctx, ProgressEvent{Stage: "downloading", Percent: pct}) + } + } + if readErr == io.EOF { + break + } + if readErr != nil { + return nil, llamaReadError(ctx, "model transfer interrupted", readErr) + } + if err := ctx.Err(); err != nil { + return nil, err + } + } + if total != row.Size { + return nil, fmt.Errorf("model size mismatch") + } + if hex.EncodeToString(hash.Sum(nil)) != row.SHA256 { + return nil, fmt.Errorf("model checksum mismatch") + } + if err := rejectSplitGGUF(f); err != nil { + return nil, err + } + if err := f.Sync(); err != nil { + return nil, err + } + // Keep the verified descriptor open through rename and ownership commit. + // A replacement at the final pathname must never acquire this claim. + if ctx.Err() != nil { + return nil, ctx.Err() + } + if err := e.llamaCommit(dir, temp, final); err != nil { + return nil, fmt.Errorf("commit model: %w", err) + } + published = true + // Rename is pathname-based, even with RENAME_EXCL/NOREPLACE. Check the + // published entry against the verified fd; never claim or remove a mismatch. + if err := llamaCacheMatch(dir, final, identity, row.Size); err != nil { + return nil, fmt.Errorf("committed model changed; download not recorded: %w", err) + } + if err := dir.Sync(); err != nil { + return nil, err + } + record := llamaOwnedModel{ID: id, File: final, Size: row.Size, SHA256: row.SHA256, FinalHost: resp.Request.URL.Hostname()} + manifest.Models = append(manifest.Models, llamaCacheRecord{llamaOwnedModel: record, Identity: identity}) + if err := ctx.Err(); err != nil { + return nil, err + } + if err := writeLlamaManifest(dir, manifest); err != nil { + return nil, fmt.Errorf("write model ownership manifest: %w", err) + } + manifestWritten = true + // Completion is fenced against cancel acceptance only after the renamed + // file's identity matches and the manifest is written. An earlier accepted + // cancel rolls both back; a later request cannot undo a completed download. + if err := ctx.Err(); err != nil { + return nil, err + } + if e.finishPull(ctx) { + return nil, context.Canceled + } + completed = true + return json.Marshal(record) +} +func (e *Executor) listLlamaModels(ctx context.Context) (json.RawMessage, error) { + ctx, cancel := context.WithTimeout(ctx, modelsTimeout) + defer cancel() + if err := ctx.Err(); err != nil { + return nil, err + } + st, err := e.state("llamacpp") + if err != nil { + return nil, err + } + dir, err := openLlamaCache(st.installDir+"-models", false) + if os.IsNotExist(err) { + return json.RawMessage(`{"models":[]}`), nil + } + if err != nil { + return nil, err + } + defer dir.Close() + manifest, err := readLlamaManifest(dir) + if err != nil { + return nil, err + } + models := []llamaOwnedModel{} + // Verify the manifest against a bounded no-follow scan of the flat cache. + // Manually supplied single-file GGUFs remain selectable, but never deletable. + entries, err := dir.ReadDir(4097) + if err != nil && err != io.EOF { + return nil, err + } + if len(entries) > 4096 { + return nil, fmt.Errorf("model cache has too many entries") + } + recorded := map[string]bool{} + for _, m := range manifest.Models { + if err := ctx.Err(); err != nil { + return nil, err + } + recorded[m.File] = true + f, err := llamaCacheOpen(dir, m.File, false) + if err != nil { + continue + } + info, err := f.Stat() + valid := err == nil && info.Size() == m.Size && llamaCacheIdentity(f) == m.Identity && rejectSplitGGUF(f) == nil + _ = f.Close() + if valid { + models = append(models, m.llamaOwnedModel) + } + } + for _, entry := range entries { + if err := ctx.Err(); err != nil { + return nil, err + } + name := entry.Name() + if recorded[name] || llamaManagedFilename(name) || !validLlamaFile(name) { + continue + } + f, err := llamaCacheOpen(dir, name, false) + if err != nil { + continue + } + info, err := f.Stat() + valid := err == nil && info.Size() > 0 && info.Size() <= llamaMaxModelSize && rejectSplitGGUF(f) == nil + _ = f.Close() + if valid { + models = append(models, llamaOwnedModel{ID: name, File: name, Size: info.Size()}) + } + } + if err := ctx.Err(); err != nil { + return nil, err + } + return json.Marshal(map[string]any{"models": models}) +} +func (e *Executor) deleteLlamaModel(ctx context.Context, id string) (json.RawMessage, error) { + if _, err := llamaRepoID(id); err != nil { + return nil, err + } + st, release, err := e.reserveLlamaCache(true) + if err != nil { + return nil, err + } + defer release() + dir, err := openLlamaCache(st.installDir+"-models", false) + if err != nil { + return nil, err + } + defer dir.Close() + manifest, err := readLlamaManifest(dir) + if err != nil { + return nil, err + } + for i, m := range manifest.Models { + if m.ID != id { + continue + } + if ctx.Err() != nil { + return nil, ctx.Err() + } + if err := llamaCacheRemoveOwned(dir, m.File, m.Identity, m.Size); err != nil { + return nil, err + } + manifest.Models = append(manifest.Models[:i], manifest.Models[i+1:]...) + if err := writeLlamaManifest(dir, manifest); err != nil { + return nil, err + } + return json.RawMessage(`{"ok":true}`), nil + } + return nil, fmt.Errorf("model was not downloaded by PAIR") +} + +// Idle transport reads are bounded independently of the 30-minute operation. +// Errors returned to callers intentionally omit URLs and remote body contents. +type llamaIdleConn struct{ net.Conn } + +func (c *llamaIdleConn) Read(p []byte) (int, error) { + if err := c.SetReadDeadline(time.Now().Add(30 * time.Second)); err != nil { + return 0, err + } + return c.Conn.Read(p) +} +func llamaDial(ctx context.Context, network, address string) (net.Conn, error) { + c, err := (&net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}).DialContext(ctx, network, address) + if err != nil { + return nil, err + } + return &llamaIdleConn{c}, nil +} + +// Resolve a managed selection by ID or its installed filename. Never promote +// a manually copied file to a PAIR-owned model through launch settings. +func recordedLlamaModelPath(cache, value string) (string, error) { + dir, err := openLlamaCache(cache, false) + if err != nil { + return "", err + } + defer dir.Close() + manifest, err := readLlamaManifest(dir) + if err != nil { + return "", err + } + for _, m := range manifest.Models { + if value == m.ID || value == m.File || value == strings.TrimSuffix(dir.Name(), "/")+"/"+m.File { + if err := llamaCacheMatch(dir, m.File, m.Identity, m.Size); err != nil { + return "", fmt.Errorf("recorded model changed; download it again: %w", err) + } + return m.File, nil + } + } + return "", fmt.Errorf("model must be a PAIR-recorded completed llama.cpp download") +} diff --git a/services/nvpair-engine-manager/llamacpp_download_test.go b/services/nvpair-engine-manager/llamacpp_download_test.go new file mode 100644 index 000000000..e12b876ee --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_download_test.go @@ -0,0 +1,301 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "context" + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + "time" +) + +const llamaTestID = "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M" + +func llamaTestBytes() []byte { + b := make([]byte, 24) + copy(b, "GGUF") + binary.LittleEndian.PutUint32(b[4:], 3) + return b +} +func llamaDownloadExecutor(t *testing.T, serve func(http.ResponseWriter, *http.Request), change func(*llamaHubRepo)) (*Executor, string) { + t.Helper() + m := llamaManifest(t) + p := m.Platforms["darwin/arm64"] + m.Platforms[runtime.GOOS+"/"+runtime.GOARCH] = p + reg := NewRegistry() + reg.engines["llamacpp"] = m + ex := NewExecutor(reg, NewReporter(nil), nil, t.TempDir()) + bytes := llamaTestBytes() + sum := sha256.Sum256(bytes) + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasPrefix(r.URL.Path, "/api/models/") { + repo := fixtureLlamaRepo(t) + repo.Siblings[0].LFS.SHA256 = hex.EncodeToString(sum[:]) + if change != nil { + change(&repo) + } + _ = json.NewEncoder(w).Encode(repo) + return + } + if serve != nil { + serve(w, r) + return + } + _, _ = w.Write(bytes) + })) + t.Cleanup(server.Close) + ex.llamaHub.endpoint = server.URL + ex.llamaHub.client.Transport = server.Client().Transport + ex.llamaFreeSpace = func(*os.File) (uint64, error) { return 1 << 40, nil } + return ex, filepath.Join(ex.baseDir, "llamacpp-models") +} +func pullLlamaTest(ex *Executor, id string) error { + _, err := ex.PullModelStream(context.Background(), "llamacpp", llamaTestID, nil, id) + return err +} +func TestLlamaInstalledInventoryHonorsCancellation(t *testing.T) { + ex, _ := llamaDownloadExecutor(t, nil, nil) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := ex.Action(ctx, "llamacpp", "list_downloaded", nil); err != context.Canceled { + t.Fatalf("ignored inventory cancellation: %v", err) + } +} + +func TestLlamaDownloadCommitSwapCannotClaimOrDeleteReplacement(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, nil) + ex.llamaCommit = func(dir *os.File, from, to string) error { + if err := llamaCacheCommit(dir, from, to); err != nil { + return err + } + path := filepath.Join(cache, to) + if err := os.Rename(path, path+".saved"); err != nil { + return err + } + return os.WriteFile(path, llamaTestBytes(), 0600) + } + if err := pullLlamaTest(ex, "commit-swap"); err == nil || !strings.Contains(err.Error(), "changed") { + t.Fatalf("commit replacement accepted: %v", err) + } + if _, err := ex.ModelDelete(context.Background(), "llamacpp", llamaTestID); err == nil { + t.Fatal("claimed and deleted a post-commit replacement") + } + if _, err := os.Stat(filepath.Join(cache, llamaOwnedFilename(llamaTestID))); err != nil { + t.Fatal("replacement disappeared") + } +} + +func TestLlamaDownloadChecksumAndDelete(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, nil) + if err := pullLlamaTest(ex, "success"); err != nil { + t.Fatal(err) + } + data, err := ex.Action(context.Background(), "llamacpp", "list_downloaded", nil) + if err != nil || !strings.Contains(string(data), llamaTestID) { + t.Fatalf("inventory: %s %v", data, err) + } + dir, err := openLlamaCache(cache, false) + if err != nil { + t.Fatal(err) + } + m, err := readLlamaManifest(dir) + dir.Close() + if err != nil || len(m.Models) != 1 || m.Models[0].FinalHost != "127.0.0.1" { + t.Fatalf("manifest: %+v %v", m, err) + } + st, _ := ex.state("llamacpp") + st.mu.Lock() + st.running = true + st.mu.Unlock() + if _, err := ex.ModelDelete(context.Background(), "llamacpp", llamaTestID); err == nil { + t.Fatal("deleted in-use model") + } + st.mu.Lock() + st.running = false + st.mu.Unlock() + if err := os.WriteFile(filepath.Join(cache, "unrecorded.gguf"), []byte("untouched"), 0600); err != nil { + t.Fatal(err) + } + if _, err := ex.ModelDelete(context.Background(), "llamacpp", "a/unrecorded:Q4_K_M"); err == nil { + t.Fatal("deleted unrecorded model") + } + if _, err := ex.ModelDelete(context.Background(), "llamacpp", llamaTestID); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(filepath.Join(cache, llamaOwnedFilename(llamaTestID))); !os.IsNotExist(err) { + t.Fatal("model remains") + } + b, err := os.ReadFile(filepath.Join(cache, "unrecorded.gguf")) + if err != nil || string(b) != "untouched" { + t.Fatal("unrecorded data changed") + } +} +func TestLlamaDownloadMismatchPreservesOtherFiles(t *testing.T) { + for _, which := range []string{"checksum", "size"} { + t.Run(which, func(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, func(r *llamaHubRepo) { + if which == "checksum" { + r.Siblings[0].LFS.SHA256 = strings.Repeat("b", 64) + } else { + r.Siblings[0].Size = 25 + r.Siblings[0].LFS.Size = 25 + } + }) + if err := os.MkdirAll(cache, 0700); err != nil { + t.Fatal(err) + } + other := filepath.Join(cache, "other.gguf") + _ = os.WriteFile(other, []byte("untouched"), 0600) + if err := pullLlamaTest(ex, "mismatch"); err == nil { + t.Fatal("mismatch accepted") + } + files, err := os.ReadDir(cache) + if err != nil || len(files) != 1 { + t.Fatalf("temp/final leaked: %v %v", files, err) + } + b, _ := os.ReadFile(other) + if string(b) != "untouched" { + t.Fatal("other file changed") + } + }) + } +} +func TestLlamaDownloadCancelDuplicateAndInterrupted(t *testing.T) { + for _, shutdown := range []bool{false, true} { + t.Run(map[bool]string{false: "cancel", true: "restart"}[shutdown], func(t *testing.T) { + started := make(chan struct{}) + ex, cache := llamaDownloadExecutor(t, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + _, _ = w.Write(llamaTestBytes()[:8]) + w.(http.Flusher).Flush() + close(started) + <-r.Context().Done() + }, nil) + events, unsub := ex.progress.subscribe("llamacpp") + defer unsub() + done := make(chan error, 1) + go func() { done <- pullLlamaTest(ex, "first") }() + select { + case <-started: + case <-time.After(3 * time.Second): + t.Fatal("not started") + } + _ = os.WriteFile(filepath.Join(cache, "other.gguf"), []byte("untouched"), 0600) + if err := pullLlamaTest(ex, "second"); err == nil { + t.Fatal("duplicate accepted") + } + before := time.Now() + if shutdown { + ex.shuttingDown.Store(true) + } + if !ex.CancelModelPull("llamacpp", llamaTestID, "first") { + t.Fatal("cancel refused") + } + if time.Since(before) > time.Second { + t.Fatal("acceptance blocked") + } + select { + case err := <-done: + if err != nil { + t.Fatal(err) + } + case <-time.After(3 * time.Second): + t.Fatal("cancel not confirmed") + } + expected := "cancelled" + if shutdown { + expected = "interrupted" + } + found := false + for len(events) > 0 { + ev := <-events + if ev.Stage == expected && ev.OpID == "first" { + found = true + } + } + if !found { + t.Fatalf("no %s state", expected) + } + files, err := os.ReadDir(cache) + if err != nil || len(files) != 1 || files[0].Name() != "other.gguf" { + t.Fatalf("cleanup wrong: %v %v", files, err) + } + }) + } +} +func TestLlamaDownloadCancelCleanupFailureIsNotConfirmed(t *testing.T) { + if os.Geteuid() == 0 { + t.Skip("requires an unprivileged user") + } + started := make(chan struct{}) + ex, cache := llamaDownloadExecutor(t, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + _, _ = w.Write(llamaTestBytes()[:8]) + w.(http.Flusher).Flush() + close(started) + <-r.Context().Done() + }, nil) + events, unsub := ex.progress.subscribe("llamacpp") + defer unsub() + done := make(chan error, 1) + go func() { done <- pullLlamaTest(ex, "cleanup") }() + select { + case <-started: + case <-time.After(3 * time.Second): + t.Fatal("not started") + } + if err := os.Chmod(cache, 0500); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = os.Chmod(cache, 0700) }) + if !ex.CancelModelPull("llamacpp", llamaTestID, "cleanup") { + t.Fatal("cancel refused") + } + select { + case err := <-done: + if err == nil || !strings.Contains(err.Error(), "could not be removed") { + t.Fatalf("cleanup failure hidden: %v", err) + } + case <-time.After(3 * time.Second): + t.Fatal("stop did not settle") + } + for len(events) > 0 { + if ev := <-events; ev.Stage == "cancelled" { + t.Fatal("claimed temp cleanup succeeded") + } + } +} + +func TestLlamaDownloadFreeSpaceAndHTTPRedirect(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, nil) + ex.llamaFreeSpace = func(*os.File) (uint64, error) { return 0, nil } + if err := pullLlamaTest(ex, "space"); err == nil || !strings.Contains(err.Error(), "free space") { + t.Fatalf("space: %v", err) + } + files, _ := os.ReadDir(cache) + if len(files) != 0 { + t.Fatal("space refusal created files") + } + reached := false + httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { reached = true })) + defer httpServer.Close() + ex, _ = llamaDownloadExecutor(t, func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, httpServer.URL, 302) }, nil) + if err := pullLlamaTest(ex, "redirect"); err == nil { + t.Fatal("unsafe redirect accepted") + } + if reached { + t.Fatal("followed HTTP redirect") + } +} diff --git a/services/nvpair-engine-manager/llamacpp_identity.go b/services/nvpair-engine-manager/llamacpp_identity.go new file mode 100644 index 000000000..7e932b349 --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_identity.go @@ -0,0 +1,55 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "fmt" + "os" + "strings" +) + +// Reserve the generated download namespace: an unrecorded or rejected commit +// must not reappear as a manually supplied model in inventory or launch. +func llamaManagedFilename(name string) bool { + name = strings.ToLower(name) + return strings.HasSuffix(name, ".gguf") && llamaSHA.MatchString(strings.TrimSuffix(name, ".gguf")) +} + +// Same-user processes can swap directory entries. macOS has no unlink-by-fd: +// keep this no-follow identity check immediately adjacent to unlinkat, leaving +// only the OS-minimum check-to-unlink window. Never unlink a known mismatch. +func llamaCacheRemoveOwned(dir *os.File, name, identity string, size int64) error { + if err := llamaCacheMatch(dir, name, identity, size); err != nil { + if os.IsNotExist(err) { + return err + } + return fmt.Errorf("cleanup skipped because the cache entry changed or is unavailable: %w", err) + } + return llamaCacheRemove(dir, name) +} + +// Compare the actual launch descriptor, not a pathname checked earlier. No +// re-hash of a potentially 64 GiB file is needed for post-commit replacements. +func validateLlamaRecordedFile(dir *os.File, name string, file *os.File) error { + manifest, err := readLlamaManifest(dir) + if err != nil { + return err + } + for _, m := range manifest.Models { + // Default macOS filesystems alias case variants of the same entry. + // Reserve and compare generated names case-insensitively on every OS. + if !strings.EqualFold(m.File, name) { + continue + } + info, err := file.Stat() + if err != nil || info.Size() != m.Size || llamaCacheIdentity(file) != m.Identity { + return fmt.Errorf("recorded model changed; refusing launch; delete the changed entry manually and download again") + } + return nil + } + if llamaManagedFilename(name) { + return fmt.Errorf("model is not a verified PAIR-recorded download; refusing launch") + } + return nil +} diff --git a/services/nvpair-engine-manager/llamacpp_identity_case_test.go b/services/nvpair-engine-manager/llamacpp_identity_case_test.go new file mode 100644 index 000000000..447ed41a7 --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_identity_case_test.go @@ -0,0 +1,44 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestLlamaRecordedIdentityCannotBeBypassedByFilenameCase(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, nil) + if err := pullLlamaTest(ex, "case-swap"); err != nil { + t.Fatal(err) + } + name := llamaOwnedFilename(llamaTestID) + path := filepath.Join(cache, name) + if err := os.Rename(path, path+".saved"); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, llamaTestBytes(), 0600); err != nil { + t.Fatal(err) + } + upper := filepath.Join(cache, strings.ToUpper(name)) + // Default macOS filesystems already alias this spelling. On Linux make the + // equivalent alias explicitly so the identity guard is exercised everywhere. + if err := os.Link(path, upper); err != nil && !os.IsExist(err) { + t.Fatal(err) + } + if _, err := llamaModelPath(cache, upper); err == nil { + t.Error("case alias bypassed selection identity") + } + if f, err := openLlamaModel(cache, upper); err == nil { + f.Close() + t.Error("case alias bypassed launch identity") + } + if !llamaManagedFilename(strings.ToUpper(name)) { + t.Error("case alias escaped reserved download namespace") + } +} diff --git a/services/nvpair-engine-manager/llamacpp_identity_test.go b/services/nvpair-engine-manager/llamacpp_identity_test.go new file mode 100644 index 000000000..7b5e21bd1 --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_identity_test.go @@ -0,0 +1,81 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestLlamaDownloadTempSwapAtCommit(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, nil) + ex.llamaCommit = func(dir *os.File, from, to string) error { + path := filepath.Join(cache, from) + if err := os.Rename(path, path+".saved"); err != nil { + return err + } + if err := os.WriteFile(path, llamaTestBytes(), 0600); err != nil { + return err + } + return llamaCacheCommit(dir, from, to) + } + if err := pullLlamaTest(ex, "temp-swap"); err == nil || !strings.Contains(err.Error(), "changed") { + t.Fatalf("swap accepted: %v", err) + } + dir, err := openLlamaCache(cache, false) + if err != nil { + t.Fatal(err) + } + defer dir.Close() + manifest, err := readLlamaManifest(dir) + if err != nil || len(manifest.Models) != 0 { + t.Fatalf("claimed replacement: %+v %v", manifest, err) + } + final := filepath.Join(cache, llamaOwnedFilename(llamaTestID)) + if _, err := os.Stat(final); err != nil { + t.Fatalf("replacement removed: %v", err) + } + data, err := ex.listLlamaModels(context.Background()) + if err != nil || strings.Contains(string(data), ".gguf") { + t.Fatalf("replacement selectable: %s %v", data, err) + } + if _, err := llamaModelPath(cache, final); err == nil { + t.Fatal("unrecorded managed replacement accepted for selection") + } + if f, err := openLlamaModel(cache, final); err == nil { + f.Close() + t.Fatal("unrecorded managed replacement accepted for launch") + } +} + +func TestLlamaRecordedReplacementRefusedAtLaunch(t *testing.T) { + ex, cache := llamaDownloadExecutor(t, nil, nil) + if err := pullLlamaTest(ex, "launch-swap"); err != nil { + t.Fatal(err) + } + name := llamaOwnedFilename(llamaTestID) + path := filepath.Join(cache, name) + if err := os.Rename(path, path+".saved"); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, llamaTestBytes(), 0600); err != nil { + t.Fatal(err) + } + for _, value := range []string{llamaTestID, name, path} { + t.Run(value, func(t *testing.T) { + if _, err := llamaModelPath(cache, value); err == nil { + t.Fatal("replacement accepted for selection") + } + if f, err := openLlamaModel(cache, value); err == nil { + f.Close() + t.Fatal("replacement accepted for launch") + } + }) + } +} diff --git a/services/nvpair-engine-manager/llamacpp_model.go b/services/nvpair-engine-manager/llamacpp_model.go index 64bd57853..0c7c4698a 100644 --- a/services/nvpair-engine-manager/llamacpp_model.go +++ b/services/nvpair-engine-manager/llamacpp_model.go @@ -23,6 +23,13 @@ func llamaModelPath(cache, value string) (string, error) { fail := func() (string, error) { return "", fmt.Errorf("model must be an existing regular .gguf file inside the PAIR llama.cpp model cache") } + if strings.Contains(value, ":Q4_K_M") { + file, err := recordedLlamaModelPath(cache, value) + if err != nil { + return "", err + } + value = file + } if cache == "" || value == "" || strings.Contains(value, "://") { return fail() } @@ -53,6 +60,21 @@ func llamaModelPath(cache, value string) (string, error) { if err != nil || !info.Mode().IsRegular() || !strings.EqualFold(filepath.Ext(path), ".gguf") { return fail() } + if filepath.Dir(path) == root { + dir, err := openLlamaCache(root, false) + if err != nil { + return "", err + } + defer dir.Close() + file, err := llamaCacheOpen(dir, filepath.Base(path), false) + if err != nil { + return "", err + } + defer file.Close() + if err := validateLlamaRecordedFile(dir, filepath.Base(path), file); err != nil { + return "", err + } + } return path, nil } diff --git a/services/nvpair-engine-manager/llamacpp_model_unix.go b/services/nvpair-engine-manager/llamacpp_model_unix.go index 85c8f407b..d8a865559 100644 --- a/services/nvpair-engine-manager/llamacpp_model_unix.go +++ b/services/nvpair-engine-manager/llamacpp_model_unix.go @@ -15,6 +15,13 @@ import ( ) func openLlamaModel(cache, value string) (*os.File, error) { + if strings.Contains(value, ":Q4_K_M") { + name, err := recordedLlamaModelPath(cache, value) + if err != nil { + return nil, err + } + value = name + } root, err := filepath.Abs(cache) if cache == "" || err != nil { return nil, fmt.Errorf("llama.cpp model cache must be an existing directory") @@ -82,6 +89,21 @@ func openLlamaModelAt(root, rel string) (*os.File, error) { _ = file.Close() return nil, fmt.Errorf("llama.cpp model must be a regular .gguf file inside the PAIR model cache") } + if len(parts) == 1 { + // Pin the same cache directory used for the launch open and manifest. + dup, err := unix.FcntlInt(uintptr(dirfd), unix.F_DUPFD_CLOEXEC, 0) + if err != nil { + _ = file.Close() + return nil, err + } + dir := os.NewFile(uintptr(dup), root) + err = validateLlamaRecordedFile(dir, rel, file) + _ = dir.Close() + if err != nil { + _ = file.Close() + return nil, err + } + } if err := rejectSplitGGUF(file); err != nil { _ = file.Close() return nil, err diff --git a/services/nvpair-engine-manager/llamacpp_remote_test.go b/services/nvpair-engine-manager/llamacpp_remote_test.go new file mode 100644 index 000000000..f7b53cafe --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_remote_test.go @@ -0,0 +1,26 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "encoding/json" + "strings" + "testing" +) + +func TestLlamaRemoteDownloadAndDeleteVerifiedPeer(t *testing.T) { + ex, _ := llamaDownloadExecutor(t, nil, nil) + srv, mesh := newToolControlPeer(t, ex, true) + params, _ := json.Marshal(map[string]string{"node": "host-server", "engine": "llamacpp", "model": llamaTestID, "opId": "llama-remote"}) + response := callToolRemote(t, srv, mesh, "engine:remote-pull-model", string(params)) + if response.Error != nil || !strings.Contains(string(response.Result), "sha256") { + t.Fatalf("remote pull: %+v", response) + } + response = callToolRemote(t, srv, mesh, "engine:remote-delete-model", string(params)) + if response.Error != nil || !strings.Contains(string(response.Result), `"ok":true`) { + t.Fatalf("remote delete: %+v", response) + } +} diff --git a/services/nvpair-engine-manager/llamacpp_space_test.go b/services/nvpair-engine-manager/llamacpp_space_test.go new file mode 100644 index 000000000..731b019cd --- /dev/null +++ b/services/nvpair-engine-manager/llamacpp_space_test.go @@ -0,0 +1,98 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "net/http" + "os" + "strings" + "testing" + "time" +) + +func TestLlamaDownloadFreeSpaceDropsDuringStreaming(t *testing.T) { + for _, interval := range []string{"bytes", "time"} { + t.Run(interval, func(t *testing.T) { + size := int64(257 << 20) + if interval == "time" { + size = 24 + } + ex, cache := llamaDownloadExecutor(t, func(w http.ResponseWriter, r *http.Request) { + if interval == "time" { + _, _ = w.Write(llamaTestBytes()[:8]) + w.(http.Flusher).Flush() + select { + case <-time.After(5100 * time.Millisecond): + case <-r.Context().Done(): + return + } + _, _ = w.Write(llamaTestBytes()[8:16]) + w.(http.Flusher).Flush() + } else { + chunk := make([]byte, 128<<10) + for i := 0; i < (256<<20)/len(chunk); i++ { + if _, err := w.Write(chunk); err != nil { + return + } + } + w.(http.Flusher).Flush() + } + // Never send EOF: refusal must occur during streaming, not only at commit. + <-r.Context().Done() + }, func(repo *llamaHubRepo) { + repo.Siblings[0].Size = size + repo.Siblings[0].LFS.Size = size + }) + checks := 0 + ex.llamaFreeSpace = func(f *os.File) (uint64, error) { + checks++ + if checks == 1 { + return 1 << 40, nil + } + info, err := f.Stat() + if err != nil || !info.Mode().IsRegular() { + t.Errorf("space recheck did not use transfer fd: %v", err) + } + return 0, nil + } + events, unsub := ex.progress.subscribe("llamacpp") + defer unsub() + observed := make(chan bool, 1) + go func() { + failed := false + for ev := range events { + if ev.Stage == "error" && strings.Contains(ev.Message, "not enough disk space") { + failed = true + } + } + observed <- failed + }() + done := make(chan error, 1) + go func() { done <- pullLlamaTest(ex, "space-drop") }() + select { + case err := <-done: + if err == nil || !strings.Contains(err.Error(), "not enough disk space") { + t.Fatalf("disk-space refusal missing: %v", err) + } + case <-time.After(12 * time.Second): + ex.CancelModelPull("llamacpp", llamaTestID, "space-drop") + <-done + t.Fatal("free space never rechecked during streaming") + } + if checks < 2 { + t.Fatal("no streaming statfs") + } + files, err := os.ReadDir(cache) + if err != nil || len(files) != 0 { + t.Fatalf("own temporary file remains: %v %v", files, err) + } + unsub() + if !<-observed { + t.Fatal("disk-space error absent from progress/state") + } + }) + } +} diff --git a/services/nvpair-engine-manager/manager.go b/services/nvpair-engine-manager/manager.go index ae586d059..64d2eec7c 100644 --- a/services/nvpair-engine-manager/manager.go +++ b/services/nvpair-engine-manager/manager.go @@ -274,6 +274,18 @@ func (m *Manager) handleMessage(ctx context.Context, msg *Message) { case "engine:errors": m.codec.Respond(msg.ID, map[string]any{"errors": m.exec.Errors()}) + case "engine:model-catalog": + var p struct { + Query string `json:"query,omitempty"` + } + if !m.parse(msg, &p) { + return + } + go func() { + rows, err := m.exec.llamaHub.catalog(ctx, p.Query) + m.respondOrErr(msg, map[string]any{"models": rows}, err) + }() + case "engine:models": go m.runModels(ctx, msg) diff --git a/services/nvpair-engine-manager/pull.go b/services/nvpair-engine-manager/pull.go index 235b6bf6d..be9e5c22a 100644 --- a/services/nvpair-engine-manager/pull.go +++ b/services/nvpair-engine-manager/pull.go @@ -65,6 +65,16 @@ func (e *Executor) PullModelStream(ctx context.Context, engine, model string, pa if err != nil { return nil, err } + if engine == "llamacpp" { + if _, err := llamaRepoID(model); err != nil { + return nil, err + } + id := "" + if len(operationID) > 0 { + id = operationID[0] + } + return e.trackedPull(ctx, engine, model, id, true, func(ctx context.Context) (json.RawMessage, error) { return e.pullLlamaModel(ctx, model) }) + } act, ok := st.manifest.Actions[pullModelAction] if !ok { return nil, fmt.Errorf("engine %q has no action %q", engine, pullModelAction) diff --git a/services/nvpair-engine-manager/pullcancel.go b/services/nvpair-engine-manager/pullcancel.go index e2077880a..af50d272c 100644 --- a/services/nvpair-engine-manager/pullcancel.go +++ b/services/nvpair-engine-manager/pullcancel.go @@ -14,8 +14,9 @@ import ( ) // Cancel only asks the adapter to stop. Confirmation is emitted by trackedPull -// after the adapter has released its transfer; no cache files are inspected or -// removed. A cancellation always addresses one operation, never its replacement. +// after the adapter has released its transfer. Shared engine caches are never +// inspected or removed; llama.cpp removes only its operation-owned temp file. +// A cancellation always addresses one operation, never its replacement. type activePull struct { engine, model, id string cancel context.CancelFunc @@ -48,6 +49,19 @@ func (e *Executor) CancelModelPull(engine, model, opID string) bool { return true } +// finishPull fences the adapter's terminal decision against operation-ID cancel +// acceptance. It never holds pullMu across filesystem or network I/O. +func (e *Executor) finishPull(ctx context.Context) bool { + p, ok := ctx.Value(activePullKey{}).(*activePull) + if !ok { + return false + } + e.pullMu.Lock() + defer e.pullMu.Unlock() + p.canCancel = false + return p.requested.Load() +} + func (e *Executor) trackedPull(ctx context.Context, engine, model, opID string, canCancel bool, run func(context.Context) (json.RawMessage, error)) (json.RawMessage, error) { if model == "" { return nil, fmt.Errorf("model is required") @@ -85,6 +99,9 @@ func (e *Executor) trackedPull(ctx context.Context, engine, model, opID string, } else if p.requested.Load() && errors.Is(err, context.Canceled) { stage = "cancelled" message = "Partial data is kept so the download can resume." + if engine == "llamacpp" { + message = "Download stopped; this operation's temporary file was removed." + } result, err = json.RawMessage(`{"status":"cancelled"}`), nil } else if err != nil { stage = "error" diff --git a/services/nvpair-engine-manager/remote_test.go b/services/nvpair-engine-manager/remote_test.go index 2c7143d77..5b4947e43 100644 --- a/services/nvpair-engine-manager/remote_test.go +++ b/services/nvpair-engine-manager/remote_test.go @@ -44,8 +44,13 @@ func callToolRemote(t *testing.T, srv *httptest.Server, mesh *clustertrust.Mesh, }() dec := json.NewDecoder(clientConn) var response Message - if err := dec.Decode(&response); err != nil { - t.Fatal(err) + for { + if err := dec.Decode(&response); err != nil { + t.Fatal(err) + } + if response.ID != nil { + break + } } select { case <-returned: diff --git a/services/nvpair-ui-broker/README.md b/services/nvpair-ui-broker/README.md index e203e4f67..fdc72c854 100644 --- a/services/nvpair-ui-broker/README.md +++ b/services/nvpair-ui-broker/README.md @@ -763,6 +763,11 @@ cache files nor synthesizes stopped state. Cancel acceptance is bounded by 15 s. `enginePullRelayTimeout` (31 min) bounds pull replies, outside the worker's 30 min action plus bounded stop and inside the clients' 32 min envelope. +`engine:model-catalog` relays the Go-owned llama.cpp catalog with a 25-second +budget outside the worker's 20-second metadata deadline. Existing action/pull, +cancel and remote methods also support PAIR-owned llama.cpp cache operations; +see the engine-manager [cache contract](../nvpair-engine-manager/LLAMACPP.md#go-model-hub-and-owned-cache). + Any other `engine:*` request is forwarded to `nvpair-engine-manager` verbatim and its response relayed straight back. This covers the whole engine control plane: `engine:get-installed`, `engine:describe`, `engine:status`, `engine:install`, `engine:uninstall`, `engine:start`, `engine:stop`, `engine:restart`, `engine:set-port`, `engine:action`, `engine:logs`, `engine:errors`, `engine:models`. (`engine:set-port` persists the engine's server port as a manifest override that survives a restart — see the `nvpair-engine-manager` README.) Lifecycle ops run for minutes (reporting progress via the `engine:install-progress` / `engine:state-changed` push events), so the relay imposes **no broker-side timeout** — fire the request and watch the event stream for the outcome. Error `-32000 "engine-manager not available"` when no engine-manager is supervised. The broker's prepare-shutdown request and subsequent stdin close join the same diff --git a/services/nvpair-ui-broker/broker.go b/services/nvpair-ui-broker/broker.go index e933a75c7..7635369af 100644 --- a/services/nvpair-ui-broker/broker.go +++ b/services/nvpair-ui-broker/broker.go @@ -3733,16 +3733,24 @@ func (b *Broker) relayToEngineNow(msg *Message) { ctx, cancel = context.WithTimeout(ctx, settings.LMStudioModelRelayBudget) } switch method { + case "engine:model-catalog": + ctx, cancel = context.WithTimeout(ctx, 25*time.Second) case "engine:cancel-pull", "engine:remote-cancel-pull": ctx, cancel = context.WithTimeout(ctx, engineCancelRelayTimeout) case "engine:remote-pull-model": ctx, cancel = context.WithTimeout(ctx, enginePullRelayTimeout) case "engine:action": var p struct { + Engine string `json:"engine"` Action string `json:"action"` } - if json.Unmarshal(msg.Params, &p) == nil && p.Action == "pull_model" { - ctx, cancel = context.WithTimeout(ctx, enginePullRelayTimeout) + if json.Unmarshal(msg.Params, &p) == nil { + if p.Action == "pull_model" { + ctx, cancel = context.WithTimeout(ctx, enginePullRelayTimeout) + } + if p.Engine == "llamacpp" && p.Action == "list_downloaded" { + ctx, cancel = context.WithTimeout(ctx, 10*time.Second) + } } } if cancel != nil { diff --git a/services/nvpair-ui-broker/llamacpp_downloads_test.go b/services/nvpair-ui-broker/llamacpp_downloads_test.go new file mode 100644 index 000000000..a7ae61ffd --- /dev/null +++ b/services/nvpair-ui-broker/llamacpp_downloads_test.go @@ -0,0 +1,37 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "encoding/json" + "strings" + "testing" +) + +func TestLlamaDownloadWorkerRestartInterruptsAndFencesLateEvents(t *testing.T) { + var out bytes.Buffer + b := &Broker{codec: NewCodec(writeOnlyReadWriter{&out}), engineSubscribed: true} + old := b.beginDownloadInstance() + worker := &rpcWorker{} + b.setEngineMgr(worker) + b.trackDownload(worker, &Message{Method: "engine:action", Params: json.RawMessage(`{"engine":"llamacpp","action":"pull_model","opId":"old","params":{"model":"a/b:Q4_K_M"}}`)}) + b.setEngineMgr(nil) + if !strings.Contains(out.String(), `"engine":"llamacpp"`) || !strings.Contains(out.String(), `"stage":"interrupted"`) { + t.Fatalf("no attributed interruption: %s", out.String()) + } + next := b.beginDownloadInstance() + replacement := &rpcWorker{} + b.setEngineMgr(replacement) + b.trackDownload(replacement, &Message{Method: "engine:action", Params: json.RawMessage(`{"engine":"llamacpp","action":"pull_model","opId":"new","params":{"model":"a/b:Q4_K_M"}}`)}) + before := out.Len() + b.forwardDownloadInstance(old, "engine:pull-progress", json.RawMessage(`{"engine":"llamacpp","model":"a/b:Q4_K_M","opId":"old","op":"pull","stage":"done"}`)) + if out.Len() != before || b.downloads["new"].owner != next { + t.Fatal("late old-worker event changed the replacement operation") + } + b.forwardDownloadInstance(next, "engine:pull-progress", json.RawMessage(`{"engine":"llamacpp","model":"a/b:Q4_K_M","opId":"new","op":"pull","stage":"done"}`)) + if len(b.downloads) != 0 { + t.Fatal("replacement did not finish") + } +}