From 8b57a999b650af2009f947bb3a3549a568174dc3 Mon Sep 17 00:00:00 2001 From: Erik LaBianca Date: Thu, 23 Jul 2026 16:15:10 -0400 Subject: [PATCH 01/25] =?UTF-8?q?feat(lucebox):=20hub=20CLI=20=E2=80=94=20?= =?UTF-8?q?install,=20serve,=20config,=20models?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit User-facing host wrapper + in-container Python CLI for launching and configuring the lucebox-hub image. Completes the docker-stack (#334) follow-up that intentionally shipped images without the CLI package. Host surface: - install.sh bootstrap + lucebox.sh wrapper (probe host, systemd unit, serve/pull/update/check/logs) - config.toml as system of record (env > file > defaults) In-container package (lucebox/): - check, pull, print-run, print-serve-argv - config {get,set,unset} - models {list,download} with VRAM-tier recommend_preset - VRAM-tier DFLASH_* heuristic (runtime_from_host); seeds config.toml on first models download --activate so 24 GB hosts get 98k/tq3_0 instead of the 16k class default - prefix_cache stays off by default (agent tool-prompt regression) Image/workspace: - COPY lucebox into CUDA and ROCm images; uv workspace member - entrypoint host_info missing-dir guard + pure-bash trim Tests/CI: lucebox pytest, wrapper sandbox scripts, lefthook. Deferred to follow-ups: autotune sweep/profiles, profile/smoke, agent-client harness adapters. --- .github/workflows/ci.yml | 29 + Dockerfile | 13 +- Dockerfile.rocm | 4 + install.sh | 138 +++ lefthook.yml | 59 + lucebox.sh | 1285 ++++++++++++++++++++++ lucebox/.gitignore | 3 + lucebox/README.md | 18 + lucebox/pyproject.toml | 54 + lucebox/src/lucebox/__init__.py | 16 + lucebox/src/lucebox/__main__.py | 6 + lucebox/src/lucebox/autotune.py | 80 ++ lucebox/src/lucebox/cli.py | 356 ++++++ lucebox/src/lucebox/config.py | 482 ++++++++ lucebox/src/lucebox/docker_run.py | 232 ++++ lucebox/src/lucebox/download.py | 515 +++++++++ lucebox/src/lucebox/host_check.py | 232 ++++ lucebox/src/lucebox/host_facts.py | 58 + lucebox/src/lucebox/py.typed | 1 + lucebox/src/lucebox/types.py | 140 +++ lucebox/tests/test_autotune.py | 48 + lucebox/tests/test_check.py | 118 ++ lucebox/tests/test_cli.py | 102 ++ lucebox/tests/test_config.py | 210 ++++ lucebox/tests/test_config_cli.py | 127 +++ lucebox/tests/test_docker_run.py | 254 +++++ lucebox/tests/test_download.py | 323 ++++++ lucebox/tests/test_models_cli.py | 142 +++ pyproject.toml | 10 +- scripts/check_lucebox_wrapper_sandbox.sh | 242 ++++ scripts/test_lucebox_sh.sh | 1130 +++++++++++++++++++ server/scripts/entrypoint.sh | 32 +- uv.lock | 32 + 33 files changed, 6474 insertions(+), 17 deletions(-) create mode 100755 install.sh create mode 100644 lefthook.yml create mode 100755 lucebox.sh create mode 100644 lucebox/.gitignore create mode 100644 lucebox/README.md create mode 100644 lucebox/pyproject.toml create mode 100644 lucebox/src/lucebox/__init__.py create mode 100644 lucebox/src/lucebox/__main__.py create mode 100644 lucebox/src/lucebox/autotune.py create mode 100644 lucebox/src/lucebox/cli.py create mode 100644 lucebox/src/lucebox/config.py create mode 100644 lucebox/src/lucebox/docker_run.py create mode 100644 lucebox/src/lucebox/download.py create mode 100644 lucebox/src/lucebox/host_check.py create mode 100644 lucebox/src/lucebox/host_facts.py create mode 100644 lucebox/src/lucebox/py.typed create mode 100644 lucebox/src/lucebox/types.py create mode 100644 lucebox/tests/test_autotune.py create mode 100644 lucebox/tests/test_check.py create mode 100644 lucebox/tests/test_cli.py create mode 100644 lucebox/tests/test_config.py create mode 100644 lucebox/tests/test_config_cli.py create mode 100644 lucebox/tests/test_docker_run.py create mode 100644 lucebox/tests/test_download.py create mode 100644 lucebox/tests/test_models_cli.py create mode 100755 scripts/check_lucebox_wrapper_sandbox.sh create mode 100755 scripts/test_lucebox_sh.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3cfe3e446..a6bdb964d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,6 +34,35 @@ jobs: - name: Lint Python surfaces touched by lucebox tooling run: uv run --frozen --extra dev ruff check . + - name: Install shellcheck (for bash test runner) + # ubuntu-latest typically ships shellcheck pre-installed, but pin + # the dependency explicitly so the bash test runner can always rely + # on `command -v shellcheck` succeeding. + run: | + if ! command -v shellcheck >/dev/null 2>&1; then + sudo apt-get update + sudo apt-get install -y shellcheck + fi + shellcheck --version | head -3 + + - name: Typecheck lucebox CLI + run: uv run --frozen --extra dev python -m mypy --package lucebox + + - name: Unit-test lucebox CLI + # The fast workspace sync above is enough: the suite mocks the + # docker / HTTP surfaces, so no torch wheel or GPU is needed. + # Keeps the lucebox Python honest on every push. + run: uv run --frozen --extra dev pytest lucebox -q + + - name: Smoke-test lucebox.sh wrapper + # Catches `set -u` regressions, syntax errors, and stale dispatch + # handlers in the host-side wrapper + the in-container entrypoint. + # Runs shellcheck --severity=error across every shipped .sh file, + # exercises every subcommand dispatch under `set -u`, and drives the + # entrypoint's draft-resolution block through every family-glob + # branch — all on the bare runner without docker/nvidia/systemd. + run: bash scripts/test_lucebox_sh.sh + build: name: Build (cmake + uv sync --extra megakernel) runs-on: ubuntu-latest diff --git a/Dockerfile b/Dockerfile index 954a4dc5e..6148f7ce5 100644 --- a/Dockerfile +++ b/Dockerfile @@ -118,14 +118,15 @@ RUN cd /src/server/build \ # of these reuses the cached CUDA layers above and only re-runs the # runtime stage's uv sync (~70s) instead of the full ~25-minute build. # -# Host-side Python tooling (lucebox/, harness/) is intentionally not copied -# here: this image is the server. Such tooling can layer on top later via a -# follow-up COPY directive or a runtime bind-mount during dev. +# The lucebox CLI is a workspace member (root pyproject + uv.lock), so its +# source must be present for `uv sync --frozen` in the runtime stage. It runs +# inside the container — the host wrapper `docker exec`s into it. COPY pyproject.toml uv.lock README.md /src/ COPY server/pyproject.toml server/README.md /src/server/ COPY server/scripts /src/server/scripts COPY optimizations/pflash /src/optimizations/pflash COPY optimizations/megakernel /src/optimizations/megakernel +COPY lucebox /src/lucebox # ─── Stage 2: runtime ─────────────────────────────────────────────────────── # Runtime image: ships nvidia driver libs but no nvcc / dev headers. Matches @@ -180,9 +181,9 @@ COPY --from=builder /src/optimizations/megakernel/pyproject.toml \ /src/optimizations/megakernel/README.md \ /opt/lucebox-hub/optimizations/megakernel/ -# Host-side Python tooling (lucebox/, harness/) is intentionally absent -# here: this image is the server base layer. Such tooling can layer on top -# later via a follow-up COPY directive or a runtime bind-mount during dev. +# The lucebox CLI package (a workspace member) — installed by the uv sync +# below and invoked in-container by the host wrapper via `docker exec`. +COPY --from=builder /src/lucebox /opt/lucebox-hub/lucebox # server: ship the entrypoint/benchmark scripts, the pyproject + README that uv # resolves against, and the pruned build tree (binaries + .so files from the diff --git a/Dockerfile.rocm b/Dockerfile.rocm index c5afbdd2f..ba8b8bb7e 100644 --- a/Dockerfile.rocm +++ b/Dockerfile.rocm @@ -123,6 +123,7 @@ COPY server/pyproject.toml server/README.md /src/server/ COPY server/scripts /src/server/scripts COPY optimizations/pflash /src/optimizations/pflash COPY optimizations/megakernel /src/optimizations/megakernel +COPY lucebox /src/lucebox # ─── Stage 2: runtime ─────────────────────────────────────────────────────── # Runtime reuses the ROCm base so the HIP runtime libs (libamdhip64, @@ -172,6 +173,9 @@ COPY --from=builder /src/optimizations/megakernel/pyproject.toml \ /src/optimizations/megakernel/README.md \ /opt/lucebox-hub/optimizations/megakernel/ +# The lucebox CLI is a workspace member and must be present before uv sync. +COPY --from=builder /src/lucebox /opt/lucebox-hub/lucebox + COPY --from=builder /src/server/scripts /opt/lucebox-hub/server/scripts COPY --from=builder /src/server/pyproject.toml /src/server/README.md \ /opt/lucebox-hub/server/ diff --git a/install.sh b/install.sh new file mode 100755 index 000000000..cff54a02e --- /dev/null +++ b/install.sh @@ -0,0 +1,138 @@ +#!/usr/bin/env bash +# install.sh — Bootstrap installer for the lucebox host wrapper. +# +# Canonical install (Luce-Org main, stable channel): +# +# curl -fsSL https://raw.githubusercontent.com/Luce-Org/lucebox-hub/main/install.sh | bash +# +# Install from a different fork / branch (dev channel). Note the env var +# is on the `bash` side of the pipe — `VAR=val curl … | bash` would attach +# it to the `curl` process, leaving `bash` with the canonical default: +# +# curl -fsSL https://raw.githubusercontent.com/easel/lucebox-hub/feat/lucebox-docker/install.sh | \ +# LUCEBOX_INSTALL_URL=https://raw.githubusercontent.com/easel/lucebox-hub/feat/lucebox-docker/lucebox.sh bash +# +# The installer bakes the source URL into the installed `lucebox.sh` as +# `LUCEBOX_INSTALLED_FROM=...`, so `lucebox update` later re-pulls from the +# same channel without the user having to remember which fork they used. +# +# Override the install destination via $LUCEBOX_INSTALL_DEST (default +# $HOME/.local/bin/lucebox). This is what `lucebox update` uses to replace +# the file in place. + +set -euo pipefail + +LUCEBOX_INSTALL_URL="${LUCEBOX_INSTALL_URL:-https://raw.githubusercontent.com/Luce-Org/lucebox-hub/main/lucebox.sh}" +DEST="${LUCEBOX_INSTALL_DEST:-$HOME/.local/bin/lucebox}" + +# ── helpers ─────────────────────────────────────────────────────────────── +C_OK=$'\033[1;32m' ; C_ERR=$'\033[1;31m' ; C_DIM=$'\033[2m' ; C_RST=$'\033[0m' +if [ ! -t 1 ] || [ "${NO_COLOR:-}" ]; then + C_OK="" ; C_ERR="" ; C_DIM="" ; C_RST="" +fi +info() { printf '%s[install]%s %s\n' "$C_DIM" "$C_RST" "$*"; } +ok() { printf '%s[install] ✓%s %s\n' "$C_OK" "$C_RST" "$*"; } +die() { printf '%s[install] ✗%s %s\n' "$C_ERR" "$C_RST" "$*" >&2; exit 1; } + +command -v curl >/dev/null 2>&1 || die "curl is required (apt-get install curl)" + +# ── fetch ───────────────────────────────────────────────────────────────── +tmp=$(mktemp -t lucebox.XXXXXX) || die "couldn't create temp file" +# shellcheck disable=SC2064 # we want $tmp expanded now, not at trap time +trap "rm -f '$tmp' '$tmp.bak'" EXIT +info "fetching $LUCEBOX_INSTALL_URL" +curl -fsSL "$LUCEBOX_INSTALL_URL" -o "$tmp" \ + || die "download failed from $LUCEBOX_INSTALL_URL" + +# ── sanity check ────────────────────────────────────────────────────────── +# Refuse to install something that isn't recognizably lucebox.sh. Catches +# 404 pages, redirects to HTML, and accidental URL typos. +head -1 "$tmp" | grep -q '^#!/usr/bin/env bash$' \ + || die "downloaded file does not look like a bash script (got: $(head -1 "$tmp"))" +grep -q '^VERSION=' "$tmp" \ + || die "downloaded file is missing VERSION marker — not lucebox.sh?" + +# ── decide what gets baked in as the persisted channel ─────────────────── +# `lucebox update` reads LUCEBOX_INSTALLED_FROM from the installed copy and +# re-fetches from it. Persisting a SHA-pinned URL is a footgun — every +# future update would re-install the same frozen SHA forever, defeating +# the point of `update`. So: +# +# 1. If $LUCEBOX_INSTALL_CHANNEL is set, that's the persisted URL +# (caller takes responsibility for picking a real branch URL). +# 2. Else if LUCEBOX_INSTALL_URL has a 40-char hex SHA segment, refuse +# to persist it — tell the user to set LUCEBOX_INSTALL_CHANNEL. +# Common case: someone curl'd from /raw// to bypass a stale CDN +# cache during dev; they meant for updates to track the branch. +# 3. Else persist LUCEBOX_INSTALL_URL as-is (branch or canonical main). +channel_url="${LUCEBOX_INSTALL_CHANNEL:-}" +if [ -z "$channel_url" ]; then + # Match a full 40-char hex SHA in the URL path, not the broader + # {7,40} range — a 7-39 char hex segment is more likely a branch + # name shaped like a short SHA (e.g. `feat/abc1234-hotfix`) than an + # actual SHA-pin. Keeping the gate at exactly 40 chars matches what + # `git rev-parse HEAD` emits and what `/raw//` URLs from + # GitHub's CDN actually carry. + if [[ "$LUCEBOX_INSTALL_URL" =~ /[0-9a-fA-F]{40}/[^/]+\.sh$ ]]; then + die "$(cat </install.sh | \\ + LUCEBOX_INSTALL_URL=/lucebox.sh \\ + LUCEBOX_INSTALL_CHANNEL=https://raw.githubusercontent.com////lucebox.sh \\ + bash +EOM +)" + fi + channel_url="$LUCEBOX_INSTALL_URL" +fi + +# Bake the channel URL into the file. Use a `|` delimiter since URLs +# contain `/`. The line is expected to exist in lucebox.sh with a `:-` +# default; we rewrite the whole assignment. +# +# The URL ends up inside a bash double-quoted literal in the installed +# script, so any of $ ` " \ in `channel_url` would break the installed +# file (or worse, allow command substitution to run at next sourcing). +# Validate that the URL is plain http(s)+ASCII-URL-safe characters; we +# don't expect arbitrary content here, only an upstream raw.github URL +# (or a forked equivalent). Escape the sed metachars (\&|) separately so +# the substitution itself round-trips. +case "$channel_url" in + *['"$`\']*) die "channel URL contains unsafe characters: $channel_url" ;; +esac +escaped_url=$(printf '%s' "$channel_url" | sed 's/[\\&|]/\\&/g') +sed "s|^LUCEBOX_INSTALLED_FROM=.*|LUCEBOX_INSTALLED_FROM=\"$escaped_url\"|" "$tmp" > "$tmp.baked" +mv "$tmp.baked" "$tmp" +grep -q "^LUCEBOX_INSTALLED_FROM=\"$escaped_url\"$" "$tmp" \ + || die "failed to bake install source into the downloaded script" + +# ── install ─────────────────────────────────────────────────────────────── +mkdir -p "$(dirname "$DEST")" +chmod +x "$tmp" +mv "$tmp" "$DEST" +trap - EXIT +ok "installed lucebox → $DEST" +info " fetched from: $LUCEBOX_INSTALL_URL" +info " update channel: $channel_url" +if [ "$LUCEBOX_INSTALL_URL" != "$channel_url" ]; then + info " (lucebox update will track the channel URL, not the fetch URL)" +fi + +# ── PATH hint ───────────────────────────────────────────────────────────── +case ":${PATH:-}:" in + *":$(dirname "$DEST"):"*) ;; + *) info " hint: add $(dirname "$DEST") to PATH so 'lucebox' is on the path" ;; +esac + +cat </dev/null || realpath "$0" 2>/dev/null || echo "$0")" +SCRIPT_NAME="$(basename "$SCRIPT_PATH")" + +# ── tunables / env overrides ─────────────────────────────────────────────── +# Host-side scalars (image registry+variant, port, container name, models +# dir). Resolution order, applied uniformly via _lucebox_resolve below: +# 1. $LUCEBOX_ per-invocation env override +# 2. config.toml
. persisted user choice (system of record) +# 3. derived / canonical default +# This keeps the wrapper and the in-container Python CLI agreeing on +# effective values — config.toml is the single source of truth, both +# sides read it. +UNIT_NAME="lucebox.service" +UNIT_PATH="${XDG_CONFIG_HOME:-$HOME/.config}/systemd/user/$UNIT_NAME" + +# CUDA driver floor for the prebuilt CUDA 12 image. +# shellcheck disable=SC2034 +MIN_DRIVER_CUDA12=525 + +# Canonical source of `lucebox.sh`. The bootstrap installer (`install.sh`) +# rewrites this line at install time to record which URL the user actually +# installed from — `lucebox update` then re-pulls from the same channel +# without losing track of forks. Falls back to the Luce-Org main branch +# when nothing was baked in (e.g. someone curl'd the script directly). +LUCEBOX_INSTALLED_FROM="${LUCEBOX_INSTALLED_FROM:-https://raw.githubusercontent.com/Luce-Org/lucebox-hub/main/lucebox.sh}" + +# Path to the persisted config.toml. Mirrors +# lucebox.config.default_config_path: $LUCEBOX_HOME/config.toml if set, +# else $HOME/.lucebox/config.toml. Read-only from this wrapper — the +# Python CLI is the writer. +_lucebox_config_path() { + if [ -n "${LUCEBOX_HOME:-}" ]; then + printf '%s/config.toml' "$LUCEBOX_HOME" + return + fi + printf '%s/.lucebox/config.toml' "$HOME" +} + +# Read a `
.` value from config.toml. Returns empty if the +# file is missing, the section/key is absent, or the value is empty. +# Handles the subset of TOML that lucebox writes: +# [section] +# key = "string" # surrounding double-quotes are stripped +# key = 8080 # bare scalars passed through verbatim +# key = true # same +# Inline `# comment` is honored. Arrays / inline tables / multi-line +# strings aren't written by the Python persister, so we don't parse them. +_lucebox_config_get() { + local dotted="$1" cfg + cfg="$(_lucebox_config_path)" + [ -f "$cfg" ] || return 0 + local section="${dotted%.*}" + local key="${dotted##*.}" + [ "$section" = "$dotted" ] && section="" + awk -v want_section="$section" -v want_key="$key" ' + BEGIN { current = "" } + /^[[:space:]]*\[/ { + t = $0 + sub(/^[[:space:]]*\[[[:space:]]*/, "", t) + sub(/[[:space:]]*\][[:space:]]*$/, "", t) + current = t + next + } + /^[[:space:]]*#/ { next } + /=/ { + if (current != want_section) next + line = $0 + sub(/#.*$/, "", line) + eq = index(line, "=") + if (eq == 0) next + k = substr(line, 1, eq - 1) + v = substr(line, eq + 1) + gsub(/^[[:space:]]+|[[:space:]]+$/, "", k) + gsub(/^[[:space:]]+|[[:space:]]+$/, "", v) + if (k != want_key) next + if (length(v) >= 2 && substr(v, 1, 1) == "\"" && substr(v, length(v), 1) == "\"") + v = substr(v, 2, length(v) - 2) + print v + exit + } + ' "$cfg" +} + +# Resolve a scalar through the precedence ladder. env_value comes from +# the caller (typically `"${LUCEBOX_FOO:-}"` — the `:-` matters under +# `set -u`). +_lucebox_resolve() { + local env_value="$1" toml_key="$2" default="$3" v + if [ -n "$env_value" ]; then + printf '%s' "$env_value" + return + fi + v="$(_lucebox_config_get "$toml_key")" + if [ -n "$v" ]; then + printf '%s' "$v" + return + fi + printf '%s' "$default" +} + +# Derive the default image URL from the install source so a fork install +# (e.g. easel/lucebox-hub) gets the fork's GHCR image automatically when +# config.toml hasn't pinned one yet. Pattern: +# https://raw.githubusercontent.com////lucebox.sh +# → ghcr.io// +# GHCR rejects mixed-case org paths so the org segment is lowercased; the +# repo name is preserved as-is. Falls back to the canonical Luce-Org image +# when the URL doesn't match the raw.githubusercontent.com pattern. +_lucebox_derive_image() { + # The ref segment can contain slashes (e.g. `feat/lucebox-docker`), so + # the middle `.+` greedily eats everything up to the trailing + # `/lucebox.sh`. The first two `[^/]+` capture org + repo, which are + # never slash-containing on GitHub. + local url="$1" org repo + if [[ "$url" =~ ^https?://raw\.githubusercontent\.com/([^/]+)/([^/]+)/.+/lucebox\.sh$ ]]; then + org=$(printf '%s' "${BASH_REMATCH[1]}" | tr '[:upper:]' '[:lower:]') + repo="${BASH_REMATCH[2]}" + printf 'ghcr.io/%s/%s' "$org" "$repo" + return + fi + printf 'ghcr.io/luce-org/lucebox-hub' +} + +# Effective scalars, env > config.toml > default. +CONTAINER_NAME=$(_lucebox_resolve "${LUCEBOX_CONTAINER:-}" runtime.container_name "lucebox") +DEFAULT_PORT=$(_lucebox_resolve "${LUCEBOX_PORT:-}" runtime.port "8080") +DEFAULT_MODELS_DIR=$(_lucebox_resolve "${LUCEBOX_MODELS:-}" paths.models "${XDG_DATA_HOME:-$HOME/.local/share}/lucebox/models") +IMAGE_BASE=$(_lucebox_resolve "${LUCEBOX_IMAGE:-}" image.registry "$(_lucebox_derive_image "$LUCEBOX_INSTALLED_FROM")") + +# ── LUCEBOX_HOST_* safe defaults (belt-and-suspenders) ──────────────────── +# `set -u` makes any unbound LUCEBOX_HOST_* read fatal. Historically this has +# been the #1 source of regressions in this wrapper: someone adds a code path +# that touches a LUCEBOX_HOST_* var before probe_host has run, the call sites +# that DO pre-probe still work, and the bug ships. To make the bug literally +# unrepresentable we seed every LUCEBOX_HOST_* with an explicit safe default +# at script-load time (these mirror probe_host's "nothing detected" state). +# probe_host then overwrites them with real values. Any future read — pre- or +# post-probe — is now well-defined. +: "${LUCEBOX_HOST_NPROC:=1}" +: "${LUCEBOX_HOST_RAM_GB:=0}" +: "${LUCEBOX_HOST_GPU_VENDOR:=none}" +: "${LUCEBOX_HOST_GPU_NAME:=}" +: "${LUCEBOX_HOST_GPU_COUNT:=0}" +: "${LUCEBOX_HOST_VRAM_GB:=0}" +: "${LUCEBOX_HOST_GPU_SM:=}" +: "${LUCEBOX_HOST_DRIVER_VERSION:=}" +: "${LUCEBOX_HOST_DRIVER_MAJOR:=0}" +: "${LUCEBOX_HOST_HAS_SYSTEMD:=0}" +: "${LUCEBOX_HOST_IS_WSL:=0}" +: "${LUCEBOX_HOST_HAS_DOCKER:=0}" +: "${LUCEBOX_HOST_DOCKER_VERSION:=}" +: "${LUCEBOX_HOST_HAS_CTK:=none}" +# Host-identity facts (item 1 — host-identity capture). These ride along +# the existing LUCEBOX_HOST_* convoy into the container so /opt/lucebox-hub/ +# HOST_INFO can be written without re-probing inside the container (where +# /proc and nvidia-smi see the container's view, not the rig's). +: "${LUCEBOX_HOST_OS_PRETTY:=}" +: "${LUCEBOX_HOST_KERNEL:=}" +: "${LUCEBOX_HOST_WSL_VERSION:=}" +: "${LUCEBOX_HOST_NVIDIA_CTK_VERSION:=}" +: "${LUCEBOX_HOST_CPU_MODEL:=}" +: "${LUCEBOX_HOST_GPU_LIST_CSV:=}" +: "${LUCEBOX_HOST_CUDA_VISIBLE_DEVICES:=}" +# Tracks whether probe_host has actually run; pieces of the code that need +# fresh host facts (e.g. cmd_check, cmd_serve) gate on this. Default 0. +: "${_LUCEBOX_HOST_PROBED:=0}" + +# ── output helpers ──────────────────────────────────────────────────────── +if [ -t 1 ] && [ -z "${NO_COLOR:-}" ]; then + C_INFO='\033[1;34m'; C_OK='\033[1;32m'; C_WARN='\033[1;33m' + C_ERR='\033[1;31m'; C_DIM='\033[2m'; C_RST='\033[0m' +else + C_INFO=''; C_OK=''; C_WARN=''; C_ERR=''; C_DIM=''; C_RST='' +fi + +info() { printf '%b[INFO]%b %s\n' "$C_INFO" "$C_RST" "$*"; } +ok() { printf '%b[OK]%b %s\n' "$C_OK" "$C_RST" "$*"; } +warn() { printf '%b[WARN]%b %s\n' "$C_WARN" "$C_RST" "$*"; } +err() { printf '%b[ERROR]%b %s\n' "$C_ERR" "$C_RST" "$*" >&2; } +hint() { printf ' %b%s%b\n' "$C_DIM" "$*" "$C_RST"; } +die() { err "$*"; exit 1; } + +# ── host probing ────────────────────────────────────────────────────────── +# Sets the LUCEBOX_HOST_* variables consumed by the in-container Python CLI +# (passed through with -e). The Python side trusts these and doesn't reprobe +# — it can't see the host's /proc anyway, only the container's. + +probe_host() { + LUCEBOX_HOST_NPROC=$(nproc 2>/dev/null || echo 1) + # RAM: try Linux /proc/meminfo first, then macOS/BSD sysctl, else 0. + LUCEBOX_HOST_RAM_GB=0 + if [ -r /proc/meminfo ]; then + LUCEBOX_HOST_RAM_GB=$(awk '/MemTotal/{printf "%.0f", $2/1024/1024}' /proc/meminfo 2>/dev/null || echo 0) + elif command -v sysctl &>/dev/null; then + mem_bytes=$(sysctl -n hw.memsize 2>/dev/null || echo 0) + LUCEBOX_HOST_RAM_GB=$(( mem_bytes / 1024 / 1024 / 1024 )) + fi + LUCEBOX_HOST_GPU_VENDOR="none" + LUCEBOX_HOST_GPU_NAME="" + LUCEBOX_HOST_GPU_COUNT=0 + LUCEBOX_HOST_VRAM_GB=0 + LUCEBOX_HOST_GPU_SM="" + LUCEBOX_HOST_DRIVER_VERSION="" + LUCEBOX_HOST_DRIVER_MAJOR=0 + + if command -v nvidia-smi &>/dev/null; then + local q + if q=$(nvidia-smi --query-gpu=name,memory.total,driver_version,compute_cap \ + --format=csv,noheader,nounits 2>/dev/null) && [ -n "$q" ]; then + LUCEBOX_HOST_GPU_VENDOR="nvidia" + LUCEBOX_HOST_GPU_NAME=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $1}') + local mem_mib + mem_mib=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $2}') + LUCEBOX_HOST_VRAM_GB=$((mem_mib / 1024)) + LUCEBOX_HOST_DRIVER_VERSION=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $3}') + LUCEBOX_HOST_DRIVER_MAJOR=${LUCEBOX_HOST_DRIVER_VERSION%%.*} + local cc + cc=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $4}') + LUCEBOX_HOST_GPU_SM="${cc//./}" + LUCEBOX_HOST_GPU_COUNT=$(printf '%s\n' "$q" | wc -l) + fi + # Multi-GPU enumeration for /props.host. The single-GPU vars + # above (GPU_NAME / GPU_SM / VRAM_GB / DRIVER_VERSION) keep + # describing GPU 0 for back-compat with cmd_check + autotune; + # the full per-GPU CSV rides along separately so HOST_INFO can + # emit the whole array. + LUCEBOX_HOST_GPU_LIST_CSV=$(nvidia-smi \ + --query-gpu=index,uuid,pci.bus_id,name,compute_cap,memory.total,power.limit \ + --format=csv,noheader 2>/dev/null || echo "") + fi + # CUDA_VISIBLE_DEVICES from the caller's env (empty default = "all GPUs"). + LUCEBOX_HOST_CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-}" + + # OS / kernel identity. /etc/os-release is the freedesktop spec for + # "what distro is this?" and we keep PRETTY_NAME verbatim (it already + # includes the version, e.g. "Ubuntu 22.04.3 LTS"). + LUCEBOX_HOST_OS_PRETTY="" + if [ -r /etc/os-release ]; then + # shellcheck source=/dev/null + LUCEBOX_HOST_OS_PRETTY=$(. /etc/os-release 2>/dev/null && printf '%s' "${PRETTY_NAME:-}") + fi + LUCEBOX_HOST_KERNEL=$(uname -r 2>/dev/null || echo "") + + # WSL version detection. "wsl2" matches the kernel-side string the + # MS-shipped WSL2 kernel embeds; "wsl1" is what the legacy translation + # layer writes. Anything else stays empty (= not WSL). + LUCEBOX_HOST_WSL_VERSION="" + if [ -r /proc/version ]; then + if grep -q "microsoft-standard-WSL2" /proc/version 2>/dev/null; then + LUCEBOX_HOST_WSL_VERSION="wsl2" + elif grep -qi "Microsoft" /proc/version 2>/dev/null; then + LUCEBOX_HOST_WSL_VERSION="wsl1" + fi + fi + + # CPU model — first "model name" hit in /proc/cpuinfo. Cheaper than + # lscpu and keeps the bash side dep-free. + LUCEBOX_HOST_CPU_MODEL="" + if [ -r /proc/cpuinfo ]; then + LUCEBOX_HOST_CPU_MODEL=$(awk -F': ' '/^model name/{print $2; exit}' /proc/cpuinfo 2>/dev/null || echo "") + fi + + LUCEBOX_HOST_HAS_SYSTEMD=0 + if command -v systemctl &>/dev/null && systemctl --user show-environment &>/dev/null; then + LUCEBOX_HOST_HAS_SYSTEMD=1 + fi + + LUCEBOX_HOST_IS_WSL=0 + if grep -qi microsoft /proc/version 2>/dev/null \ + || [ -e /proc/sys/fs/binfmt_misc/WSLInterop ]; then + LUCEBOX_HOST_IS_WSL=1 + fi + + LUCEBOX_HOST_HAS_DOCKER=0 + LUCEBOX_HOST_DOCKER_VERSION="" + if command -v docker &>/dev/null && docker ps &>/dev/null; then + LUCEBOX_HOST_HAS_DOCKER=1 + LUCEBOX_HOST_DOCKER_VERSION=$(timeout 5 docker version --format '{{.Server.Version}}' 2>/dev/null || echo "") + fi + + LUCEBOX_HOST_HAS_CTK="none" + if [ "$LUCEBOX_HOST_HAS_DOCKER" = "1" ]; then + if command -v nvidia-container-runtime &>/dev/null; then + LUCEBOX_HOST_HAS_CTK="runtime" + elif command -v nvidia-ctk &>/dev/null \ + && nvidia-ctk cdi list 2>/dev/null | grep -q 'nvidia.com/gpu'; then + LUCEBOX_HOST_HAS_CTK="cdi" + elif command -v nvidia-ctk &>/dev/null; then + LUCEBOX_HOST_HAS_CTK="installed-unwired" + fi + fi + + # NVIDIA Container Toolkit version (best-effort; empty when nvidia-ctk + # is not installed). nvidia-ctk --version prints "NVIDIA Container + # Toolkit CLI version 1.16.2" on a single line — extract the trailing + # token so the host-info JSON carries just the version, not the banner. + LUCEBOX_HOST_NVIDIA_CTK_VERSION="" + if command -v nvidia-ctk &>/dev/null; then + LUCEBOX_HOST_NVIDIA_CTK_VERSION=$(nvidia-ctk --version 2>/dev/null \ + | awk '/version/{print $NF; exit}' \ + || echo "") + fi + + export LUCEBOX_HOST_NPROC LUCEBOX_HOST_RAM_GB LUCEBOX_HOST_GPU_VENDOR + export LUCEBOX_HOST_GPU_NAME LUCEBOX_HOST_GPU_COUNT LUCEBOX_HOST_VRAM_GB + export LUCEBOX_HOST_GPU_SM LUCEBOX_HOST_DRIVER_VERSION LUCEBOX_HOST_DRIVER_MAJOR + export LUCEBOX_HOST_HAS_SYSTEMD LUCEBOX_HOST_IS_WSL + export LUCEBOX_HOST_HAS_DOCKER LUCEBOX_HOST_DOCKER_VERSION + export LUCEBOX_HOST_HAS_CTK + export LUCEBOX_HOST_OS_PRETTY LUCEBOX_HOST_KERNEL LUCEBOX_HOST_WSL_VERSION + export LUCEBOX_HOST_NVIDIA_CTK_VERSION LUCEBOX_HOST_CPU_MODEL + export LUCEBOX_HOST_GPU_LIST_CSV LUCEBOX_HOST_CUDA_VISIBLE_DEVICES + _LUCEBOX_HOST_PROBED=1 +} + +# Cheap idempotency wrapper. Anything that needs real host facts (vs the safe +# defaults seeded at script-load) calls this. Subcommands that go straight to +# `systemctl`/`journalctl` no longer need to remember to call probe_host. +ensure_probed() { + [ "$_LUCEBOX_HOST_PROBED" = "1" ] || probe_host +} + +pick_variant() { + # CUDA 12.8 is the supported image variant for this branch. Effective + # value goes through the same env > config.toml > default ladder as + # everything else so `config set image.variant=...` propagates. + _lucebox_resolve "${LUCEBOX_VARIANT:-}" image.variant "cuda12" +} + +# ── prereq checks (host-only) ───────────────────────────────────────────── +# Print-and-exit on anything that needs root to install. The Python CLI does +# the richer reporting; this is the bare minimum to make `docker run` viable. + +require_host_prereqs() { + local missing=0 + if ! command -v docker &>/dev/null; then + err "docker is not installed" + hint "Install: https://docs.docker.com/engine/install/" + missing=1 + elif ! docker ps &>/dev/null; then + err "docker daemon not reachable" + hint "sudo systemctl start docker (or: add your user to the 'docker' group, then re-login)" + missing=1 + fi + + if ! command -v nvidia-smi &>/dev/null; then + err "nvidia-smi not found — no NVIDIA driver detected" + hint "Install the NVIDIA driver: https://www.nvidia.com/Download/index.aspx" + missing=1 + elif ! nvidia-smi --query-gpu=name --format=csv,noheader &>/dev/null; then + err "nvidia-smi present but NVML calls fail — likely a driver/library mismatch" + hint "Reboot, or reinstall the matching NVIDIA driver package" + missing=1 + fi + + [ "$missing" = "0" ] || exit 1 +} + +require_ctk() { + case "$LUCEBOX_HOST_HAS_CTK" in + runtime|cdi) return 0 ;; + installed-unwired) + err "NVIDIA Container Toolkit installed but not wired into docker" + hint "sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker" + hint " or generate a CDI spec: sudo nvidia-ctk cdi generate --output=/etc/cdi/nvidia.yaml" + exit 1 ;; + none|*) + err "NVIDIA Container Toolkit not installed" + hint "Install: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html" + hint "Then register with docker:" + hint " sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker" + exit 1 ;; + esac +} + +require_systemd() { + # Earlier versions of this wrapper had `start`/`stop`/`logs`/etc. drop + # straight into cmd_systemctl_passthrough without probing first, which + # tripped `set -u` on the reference below. Two layers of defence now: + # 1) top-of-script seeds LUCEBOX_HOST_HAS_SYSTEMD=0 unconditionally, so + # no read can be unbound even if probe_host is bypassed entirely. + # 2) ensure_probed runs probe_host on first call so we still get the + # real answer for the require_systemd error path. + ensure_probed + if [ "$LUCEBOX_HOST_HAS_SYSTEMD" != "1" ]; then + err "user systemd is not available — required for $1" + hint "On WSL: set 'systemd=true' under [boot] in /etc/wsl.conf, then 'wsl --shutdown'." + hint "Otherwise: install systemd, or run '$SCRIPT_NAME serve' to run in the foreground without systemd." + exit 1 + fi +} + +# ── docker run construction ─────────────────────────────────────────────── +# All the Python-CLI subcommands share the same docker run incantation: +# mount the host docker socket (so the in-container CLI can spawn server / +# bench containers on the host daemon), mount $HOME at the same path (so +# paths look identical in and out), and pass host facts via env. When an +# NVIDIA GPU is detected we also pass --gpus all so the orchestrator can +# call nvidia-smi during profile snapshot export; without it nvidia_smi_csv (and +# any downstream power/utilization fields) come back empty. + +DOCKER_SOCK_PATH="${DOCKER_HOST:-/var/run/docker.sock}" +DOCKER_SOCK_PATH="${DOCKER_SOCK_PATH#unix://}" + +# Append `-e LUCEBOX_HOST_=` for every exported host fact onto the +# named docker-argv array (bash 4.3+ nameref). The Python side reads these +# instead of reprobing — see build_orchestrator_argv / cmd_exec_in_container. +_append_host_env() { + # shellcheck disable=SC2178 # nameref to a caller's array, not a string + local -n _arr="$1" + local var + for var in $(compgen -e | grep '^LUCEBOX_HOST_' || true); do + _arr+=(-e "$var=${!var}") + done +} + +# Append the LUCEBOX_* scalar overrides (image/variant/port/container/models) +# plus the optional HF_TOKEN guard onto the named docker-argv array. Shared +# by the docker-run (build_orchestrator_argv) and docker-exec +# (cmd_exec_in_container) paths so both forward an identical env subset. +_append_scalar_env() { + # shellcheck disable=SC2178 # nameref to a caller's array, not a string + local -n _arr="$1" + local variant="$2" + _arr+=(-e "LUCEBOX_IMAGE=$IMAGE_BASE") + _arr+=(-e "LUCEBOX_VARIANT=$variant") + _arr+=(-e "LUCEBOX_PORT=$DEFAULT_PORT") + _arr+=(-e "LUCEBOX_CONTAINER=$CONTAINER_NAME") + _arr+=(-e "LUCEBOX_MODELS=$DEFAULT_MODELS_DIR") + [ -n "${HF_TOKEN:-}" ] && _arr+=(-e "HF_TOKEN=$HF_TOKEN") + return 0 +} + +# Pick docker's interactive flags: -it on a real tty, -i otherwise. +# Writes into a caller-supplied array via nameref. This MUST run in the +# caller's scope (not a subshell or `< <(...)` process substitution): the +# `[ -t 1 ]` test inspects fd 1, and inside a process substitution fd 1 is +# the pipe to the consumer, not the terminal — which would force -i even on +# a real tty and break the interactive client TUIs (lucebox claude, etc.). +_set_tty_flags() { # usage: _set_tty_flags arrayname + # shellcheck disable=SC2178 + local -n _a="$1" + if [ -t 0 ] && [ -t 1 ]; then + _a=(-it) + else + _a=(-i) + fi +} + +build_orchestrator_argv() { + local variant="$1"; shift + local tty=() + _set_tty_flags tty + local argv=(docker run --rm "${tty[@]}") + if [ "${LUCEBOX_HOST_GPU_VENDOR:-none}" = "nvidia" ]; then + argv+=(--gpus all) + fi + argv+=(--name "${CONTAINER_NAME}-cli-$$") + argv+=(--user "$(id -u):$(id -g)") + # Only bind-mount the docker socket when DOCKER_HOST actually points + # at a unix socket on this host. With DOCKER_HOST=tcp://… or ssh://… + # the path we'd construct is `tcp` or empty, and `docker run -v` would + # bark with an "invalid mount" error before the orchestrator even + # starts. The orchestrator-in-container relies on docker access only + # when actually needed; pulling that mount when the host talks to + # docker over TCP/SSH is fine. + if [ -S "$DOCKER_SOCK_PATH" ]; then + argv+=(--group-add "$(stat -c '%g' "$DOCKER_SOCK_PATH")") + argv+=(-v "$DOCKER_SOCK_PATH:/var/run/docker.sock") + fi + argv+=(-v "$HOME:$HOME") + # Bind-mount the XDG models dir explicitly (host = container path) so + # paths line up in/out. The $HOME mount above already covers it when + # XDG_DATA_HOME is unset, but an explicit -v is required when the user + # points XDG_DATA_HOME outside $HOME. + mkdir -p "$DEFAULT_MODELS_DIR" + argv+=(-v "$DEFAULT_MODELS_DIR:$DEFAULT_MODELS_DIR") + argv+=(-w "$PWD") + argv+=(-e "HOME=$HOME") + # Host facts — Python side reads these instead of reprobing. + _append_host_env argv + # User overrides for image/port/container/models scalars + HF_TOKEN. + # Always exports the resolved models dir so the in-container CLI sees + # the same path the wrapper mounts (the XDG default flows through too). + _append_scalar_env argv "$variant" + + argv+=("${IMAGE_BASE}:${variant}") + # `lucebox` is the entrypoint subcommand handled by server/scripts/entrypoint.sh + # — it execs `python -m lucebox` with whatever args we pass on. + argv+=(lucebox "$@") + printf '%s\n' "${argv[@]}" +} + +# ── subcommand implementations ──────────────────────────────────────────── + +cmd_serve() { + # Long-running foreground server. Also what systemd's ExecStart= calls. + # + # Two-stage so config.toml takes effect: + # 1. Run an ephemeral orchestrator container that emits the canonical + # server docker-run argv from .lucebox/config.toml (one arg per + # line on stdout). + # 2. Exec that argv. + # + # If stage 1 fails (image not pulled yet, no config), fall back to a + # conservative docker run — the container's own VRAM-tiered autotune + # picks reasonable defaults from there. + require_host_prereqs + ensure_probed + require_ctk + local variant + variant=$(pick_variant) + + # Pre-flight: refuse to stomp on something that's already serving this + # slot. Three states to distinguish, because silently `docker rm -f`-ing + # whatever is there hides real bugs (e.g. the user forgot they had a + # systemd unit up, and we'd happily race two servers on the same port): + # + # 1. systemd unit active → refuse, redirect to `logs`/`stop` + # 2. container running (no systemd)→ refuse, redirect to `docker logs` + # 3. container present but stopped → orphan from a SIGKILLed previous + # run (docker run --rm only cleans up on clean exit). Remove it, + # but TELL the user — they need to know their last run died dirty. + # CRITICAL: when systemd invokes US as the unit's ExecStart, is-active + # returns true *because of us* — refusing here would deadlock the unit + # in a restart loop (and historically did — commit a30dbe5 shipped this + # bug). systemd sets $INVOCATION_ID in every service exec, so its + # presence is the unambiguous "I am running as the systemd ExecStart" + # signal. Skip the unit-active check in that case; the container-state + # check below still catches a stale container holding the slot. + if [ -z "${INVOCATION_ID:-}" ] \ + && systemctl --user is-active --quiet "$UNIT_NAME" 2>/dev/null; then + err "${UNIT_NAME} is already running under systemd." + hint " $SCRIPT_NAME logs # follow the journal" + hint " $SCRIPT_NAME restart # bounce the service" + hint " $SCRIPT_NAME stop # stop the service" + exit 1 + fi + local container_state + container_state=$(docker inspect --format '{{.State.Status}}' "$CONTAINER_NAME" 2>/dev/null || echo absent) + case "$container_state" in + absent) + ;; + running|restarting) + err "Container '$CONTAINER_NAME' is already running (outside systemd)." + hint " docker logs -f $CONTAINER_NAME # follow output" + hint " $SCRIPT_NAME stop # stop it" + exit 1 + ;; + exited|created|paused|dead) + info "Removing stale '$CONTAINER_NAME' container (state=$container_state, likely from a previous unclean exit)" + docker rm -f "$CONTAINER_NAME" >/dev/null + ;; + *) + warn "Container '$CONTAINER_NAME' is in unexpected state '$container_state' — removing" + docker rm -f "$CONTAINER_NAME" >/dev/null + ;; + esac + + local orch_argv server_argv + mapfile -t orch_argv < <(build_orchestrator_argv "$variant" print-serve-argv) + + if mapfile -t server_argv < <("${orch_argv[@]}" 2>/dev/null) \ + && [ "${#server_argv[@]}" -gt 0 ] \ + && [ "${server_argv[0]}" = "docker" ]; then + info "Starting lucebox server (variant=$variant, from config.toml)" + _serve_and_track "${server_argv[@]}" + return $? + fi + + warn "Couldn't fetch server argv from container (image not pulled?) — using fallback" + info "Starting lucebox server (variant=$variant, port=$DEFAULT_PORT, defaults only)" + local fallback_models="$DEFAULT_MODELS_DIR" + mkdir -p "$fallback_models" + # Forward host facts even on the fallback path so the in-container + # entrypoint can still write /opt/lucebox-hub/HOST_INFO from the host's + # view of the rig. Matches the orchestrator path (see + # build_orchestrator_argv) — without it, HOST_INFO would be written + # with "source: unknown" any time print-serve-argv fails. + local fallback_argv=(docker run --rm + --name "$CONTAINER_NAME" + --gpus all + -p "$DEFAULT_PORT:8080" + -v "$HOME:$HOME" + -v "$fallback_models:/opt/lucebox-hub/server/models") + _append_host_env fallback_argv + fallback_argv+=("${IMAGE_BASE}:${variant}") + _serve_and_track "${fallback_argv[@]}" +} + +# Foreground server runner with controlling-process lifetime semantics: +# the docker daemon owns containers independently of the CLI, so a bare +# `exec docker run` leaves the container alive after the wrapper's parent +# (a terminal, a systemd unit, anything) goes away. `docker run --rm` only +# cleans up on the container's own clean exit, not on our death. +# +# Fix: run docker as a child, install signal traps that issue `docker stop` +# before exiting. Now `lucebox serve` behaves like a normal foreground +# program — close the terminal, kill the wrapper, send SIGTERM from +# systemd, the container goes down with it. +# +# Stops also from EXIT so even a `set -e` propagation cleans up. +_serve_and_track() { + "$@" & + local docker_pid=$! + # shellcheck disable=SC2317 # called via trap, not "unreachable" + _serve_stop() { + trap - HUP INT TERM EXIT + # Best-effort: container may already be exiting / never started. + # `docker stop` blocks up to -t seconds for graceful shutdown + # (server handles SIGTERM), then SIGKILLs. 10s is enough for the + # in-flight request to finish on a typical decode. + docker stop -t 10 "$CONTAINER_NAME" >/dev/null 2>&1 || true + # docker_pid may be out of scope when this fires as the EXIT trap + # after _serve_and_track has unwound (e.g. the server fast-failed); + # guard so `set -u` doesn't turn cleanup into its own error. + [ -n "${docker_pid:-}" ] && wait "$docker_pid" 2>/dev/null || true + } + trap _serve_stop HUP INT TERM EXIT + wait "$docker_pid" + local rc=$? + trap - HUP INT TERM EXIT + return $rc +} + +cmd_systemd_install() { + require_host_prereqs + ensure_probed + require_systemd "service install" + local docker_bin + docker_bin=$(command -v docker) + + mkdir -p "$(dirname "$UNIT_PATH")" + # Capture the user's resolved env at install time so the unit launches + # with the same image/variant/port/models the user expected when they + # ran `lucebox install`. Systemd's user-session env is sparse — without + # this block, the wrapper inside the unit would fall back to the + # in-script defaults and silently pick a different image or models + # directory than the user's interactive session uses. + # + # ExecStartPre cleans up any orphaned container with the target name + # left behind by a previous crash (docker's `--rm` only fires on clean + # exit — a SIGKILL or daemon restart leaves the name claimed, and the + # next ExecStart would die with "name already in use" while systemd + # reports a useless "exit code 125"). + cat > "$UNIT_PATH" </dev/null | awk -F= '/^Linger=/{print $2}') + if [ "$linger" != "yes" ]; then + warn "Linger is off for $USER — the service will stop when you log out" + hint "To enable (requires sudo): sudo loginctl enable-linger \"$USER\"" + fi + + printf '\nNext:\n' + hint " $SCRIPT_NAME start # start now" + hint " $SCRIPT_NAME enable # start at every login" + hint " $SCRIPT_NAME logs # follow the journal" +} + +cmd_systemd_uninstall() { + require_systemd "service uninstall" + if systemctl --user is-active --quiet "$UNIT_NAME" 2>/dev/null; then + info "Stopping $UNIT_NAME" + systemctl --user stop "$UNIT_NAME" || true + fi + if systemctl --user is-enabled --quiet "$UNIT_NAME" 2>/dev/null; then + info "Disabling $UNIT_NAME" + systemctl --user disable "$UNIT_NAME" || true + fi + if [ -f "$UNIT_PATH" ]; then + rm -f "$UNIT_PATH" + ok "Removed $UNIT_PATH" + else + info "No unit at $UNIT_PATH — nothing to remove" + fi + systemctl --user daemon-reload + hint "Config and models are left in place. Remove them by hand if you want." +} + +cmd_systemctl_passthrough() { + local action="$1" + require_systemd "$action" + if [ ! -f "$UNIT_PATH" ]; then + err "$UNIT_NAME is not installed — run '$SCRIPT_NAME install' first" + exit 1 + fi + case "$action" in + start|restart) + # `systemctl start` is fire-and-forget for Type=exec: it returns + # success as soon as execve() completes, even if the wrapper + # exits 1 a millisecond later. That gave us the worst possible + # UX — `lucebox start` reports no error but no container ever + # binds port 8080. Poll is-active for a few seconds and dump + # status + recent journal lines so the user sees the real cause. + local current + current=$(systemctl --user is-active "$UNIT_NAME" 2>/dev/null || true) + # `start` against an already-active unit: systemctl returns 0 + # silently. That's polite for scripts but confusing for humans + # — say so explicitly. For `restart` always run through. + if [ "$action" = "start" ] && [ "$current" = "active" ]; then + ok "$UNIT_NAME is already active" + hint "logs: $SCRIPT_NAME logs" + hint "smoke: curl -s http://localhost:$DEFAULT_PORT/v1/models" + hint "(use \`$SCRIPT_NAME restart\` to bounce, \`$SCRIPT_NAME stop\` to halt)" + return 0 + fi + # `start` against a unit stuck in restart-loop ("activating") is + # the symptom of a broken ExecStart — calling start would just + # block waiting for active that never comes. Surface this + # specifically so the user goes to `lucebox logs` to find the + # exit reason rather than waiting for the poll to give up. + if [ "$action" = "start" ] && [ "$current" = "activating" ]; then + err "$UNIT_NAME is in restart-loop (state=activating)" + hint "the unit is failing and being auto-restarted by systemd" + hint " $SCRIPT_NAME stop # halt the loop first" + hint " $SCRIPT_NAME logs # find the exit reason" + exit 1 + fi + info "$action $UNIT_NAME" + if ! systemctl --user "$action" "$UNIT_NAME"; then + err "systemctl --user $action $UNIT_NAME failed" + systemctl --user status "$UNIT_NAME" --no-pager -n 30 || true + exit 1 + fi + local i state + for i in 1 2 3 4 5 6 7 8 9 10; do + state=$(systemctl --user is-active "$UNIT_NAME" 2>/dev/null || true) + case "$state" in + active) break ;; # already up — no need to keep polling + activating) ;; # still booting; keep waiting + *) break ;; # failed / inactive — fall through to error path + esac + sleep 1 + done + state=$(systemctl --user is-active "$UNIT_NAME" 2>/dev/null || true) + if [ "$state" != "active" ]; then + err "$UNIT_NAME did not reach active state (current: ${state:-unknown})" + if [ "$state" = "activating" ]; then + hint "the unit is in a restart loop — \`$SCRIPT_NAME stop\` to halt it" + fi + hint "status:" + systemctl --user status "$UNIT_NAME" --no-pager -n 30 || true + hint "recent journal:" + journalctl --user -u "$UNIT_NAME" -n 30 --no-pager || true + exit 1 + fi + ok "$UNIT_NAME is active" + hint "logs: $SCRIPT_NAME logs" + hint "smoke: curl -s http://localhost:$DEFAULT_PORT/v1/models" + ;; + stop|enable|disable) + exec systemctl --user "$action" "$UNIT_NAME" ;; + status) + exec systemctl --user status "$UNIT_NAME" --no-pager ;; + *) + die "unknown systemctl passthrough: $action" ;; + esac +} + +cmd_logs() { + require_systemd "logs" + # Pure passthrough: any flags the user wants (-f, -n, --since, ...) go + # straight to journalctl. Default is follow. + if [ $# -eq 0 ]; then + exec journalctl --user -u "$UNIT_NAME" -f + fi + exec journalctl --user -u "$UNIT_NAME" "$@" +} + +cmd_pull() { + # Pull has to run on the host. Delegating this into the container creates a + # stale-image trap: docker may start an old local tag before the fresh tag + # has been pulled. + require_host_prereqs + local variant + variant=$(pick_variant) + info "Pulling ${IMAGE_BASE}:${variant}" + exec docker pull "${IMAGE_BASE}:${variant}" +} + +cmd_update() { + # Re-run the bootstrap installer against the channel we were installed + # from. The installer is the source of truth for "how do you install + # lucebox correctly" — chmod, atomic mv, validation, baking the source + # URL back into the new copy so the channel is preserved across + # upgrades. Keeping the logic in install.sh means it can evolve + # independently (sha verify, signature check, etc.) and the installed + # `lucebox update` picks those changes up on the next run. + # + # The installer URL is derived from LUCEBOX_INSTALLED_FROM by swapping + # `lucebox.sh` → `install.sh` in the same directory, so forks don't + # need a separate registration. Override the source channel via + # $LUCEBOX_INSTALL_URL (e.g. to switch from canonical to a dev fork). + local source_url installer_url target + source_url="${LUCEBOX_INSTALL_URL:-$LUCEBOX_INSTALLED_FROM}" + if [[ "$source_url" != */lucebox.sh ]]; then + die "LUCEBOX_INSTALLED_FROM doesn't end in /lucebox.sh: $source_url" + fi + installer_url="${source_url%/lucebox.sh}/install.sh" + target=$(realpath "$SCRIPT_PATH") + + info "Updating lucebox via $installer_url" + info " source: $source_url" + info " target: $target" + + # Pass the URLs through to install.sh via env. The installer reads + # $LUCEBOX_INSTALL_URL (which we set to source_url) and + # $LUCEBOX_INSTALL_DEST (the realpath of *this* file, so a symlinked + # install replaces the actual file behind the link). + LUCEBOX_INSTALL_URL="$source_url" \ + LUCEBOX_INSTALL_DEST="$target" \ + bash -c "$(curl -fsSL "$installer_url")" \ + || die "update failed (installer exited non-zero)" +} + +cmd_completion() { + # Print shell completion script for bash / zsh / fish. Usage: + # + # # bash (in ~/.bashrc): + # source <(lucebox completion bash) + # + # # zsh (in ~/.zshrc, before `compinit`): + # source <(lucebox completion zsh) + # + # # fish: + # lucebox completion fish | source + # + # Keep this in sync with the dispatch table in main() and the sub-app + # verbs (config get/set/unset, models list/download). Adding a new + # top-level command means adding it here too. + local shell="${1:-}" + case "$shell" in + bash) + cat <<'BASH' +# lucebox bash completion. Source from ~/.bashrc: +# source <(lucebox completion bash) +_lucebox_complete() { + local cur prev cmds config_verbs models_verbs completion_shells + COMPREPLY=() + cur="${COMP_WORDS[COMP_CWORD]}" + prev="${COMP_WORDS[COMP_CWORD-1]}" + cmds="install uninstall start stop restart enable disable status logs \ + serve pull update check completion config models \ + print-run help version" + config_verbs="get set unset" + models_verbs="list download" + completion_shells="bash zsh fish" + + # Sub-app verbs / shell args. + case "$prev" in + config) COMPREPLY=( $(compgen -W "$config_verbs" -- "$cur") ); return ;; + models) COMPREPLY=( $(compgen -W "$models_verbs" -- "$cur") ); return ;; + completion) COMPREPLY=( $(compgen -W "$completion_shells" -- "$cur") ); return ;; + esac + + # Top-level command. + if [ "$COMP_CWORD" = 1 ]; then + COMPREPLY=( $(compgen -W "$cmds" -- "$cur") ) + return + fi +} +complete -F _lucebox_complete lucebox lucebox.sh +BASH + ;; + zsh) + # Bash-compat shim: zsh sources our bash completion through + # bashcompinit. Users who prefer native zsh _arguments-style + # completion can write their own; this gets `` working + # in two lines for free. + cat <<'ZSH' +# lucebox zsh completion. Source from ~/.zshrc (after compinit): +# source <(lucebox completion zsh) +autoload -Uz compinit bashcompinit +compinit +bashcompinit +ZSH + cmd_completion bash + ;; + fish) + cat <<'FISH' +# lucebox fish completion. Source from ~/.config/fish/config.fish: +# lucebox completion fish | source +complete -c lucebox -f +set -l __lucebox_cmds install uninstall start stop restart enable disable \ + status logs serve pull update check completion config models \ + print-run help version +for cmd in $__lucebox_cmds + complete -c lucebox -n "not __fish_seen_subcommand_from $__lucebox_cmds" -a $cmd +end +complete -c lucebox -n "__fish_seen_subcommand_from config" -a "get set unset" +complete -c lucebox -n "__fish_seen_subcommand_from models" -a "list download" +complete -c lucebox -n "__fish_seen_subcommand_from completion" -a "bash zsh fish" +FISH + ;; + ""|--help|-h) + cat </dev/null; then + _row 0 "docker daemon" "installed but unreachable — start the daemon or add user to 'docker' group" + else + _row 0 "docker daemon" "not installed — https://docs.docker.com/engine/install/" + fi + + # nvidia container toolkit + case "$LUCEBOX_HOST_HAS_CTK" in + runtime) _row 1 "nvidia ctk" "wired into docker (runtime)" ;; + cdi) _row 1 "nvidia ctk" "wired via CDI (nvidia.com/gpu)" ;; + installed-unwired) _row warn "nvidia ctk" "installed but not registered with docker — sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker" ;; + none|*) _row 0 "nvidia ctk" "not installed — https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html" ;; + esac + + # nvidia-smi + driver + if [ "$LUCEBOX_HOST_GPU_VENDOR" = "nvidia" ]; then + if [ "$LUCEBOX_HOST_DRIVER_MAJOR" -ge "$MIN_DRIVER_CUDA12" ]; then + _row 1 "nvidia driver" "$LUCEBOX_HOST_DRIVER_VERSION (≥ $MIN_DRIVER_CUDA12 required for cuda12)" + else + _row 0 "nvidia driver" "$LUCEBOX_HOST_DRIVER_VERSION (< $MIN_DRIVER_CUDA12 — cuda12 image will fail)" + fi + elif command -v nvidia-smi &>/dev/null; then + _row 0 "nvidia driver" "nvidia-smi present but NVML calls fail — driver/library mismatch, try reboot" + else + _row 0 "nvidia driver" "nvidia-smi not found — install the NVIDIA driver" + fi + + # GPU detail + if [ "$LUCEBOX_HOST_GPU_VENDOR" = "nvidia" ]; then + _row 1 "gpu" "$LUCEBOX_HOST_GPU_NAME × $LUCEBOX_HOST_GPU_COUNT (sm_$LUCEBOX_HOST_GPU_SM, ${LUCEBOX_HOST_VRAM_GB} GB VRAM)" + # cuda12 image arch coverage: sm_75;80;86;89;90;120 (see docker-bake.hcl) + case "$LUCEBOX_HOST_GPU_SM" in + 75|80|86|89|90|120) _row 1 "cuda12 arch" "sm_$LUCEBOX_HOST_GPU_SM covered by image" ;; + "") _row warn "cuda12 arch" "compute_cap not detected" ;; + *) _row warn "cuda12 arch" "sm_$LUCEBOX_HOST_GPU_SM not in image arch list (75;80;86;89;90;120)" ;; + esac + fi + + # systemd + if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ]; then + _row 1 "user systemd" "available (needed for '$SCRIPT_NAME install')" + elif [ "$LUCEBOX_HOST_IS_WSL" = "1" ]; then + _row warn "user systemd" "WSL detected — set 'systemd=true' under [boot] in /etc/wsl.conf, then 'wsl --shutdown'" + else + _row warn "user systemd" "not available — '$SCRIPT_NAME install' (service unit) won't work; '$SCRIPT_NAME serve' (foreground) will" + fi + + # image we'd pull — marked ✗ when the host clearly can't run cuda12 + # (no nvidia driver, or no CTK wired into docker). It's still useful + # to print the line so the user knows what would be pulled, but a + # green ✓ would be misleading. + if [ "$LUCEBOX_HOST_GPU_VENDOR" != "nvidia" ]; then + _row 0 "image" "${IMAGE_BASE}:${variant} — requires NVIDIA driver" + elif [ "$LUCEBOX_HOST_HAS_CTK" = "none" ] || [ "$LUCEBOX_HOST_HAS_CTK" = "installed-unwired" ]; then + _row 0 "image" "${IMAGE_BASE}:${variant} — needs NVIDIA Container Toolkit wired into docker" + else + _row 1 "image" "${IMAGE_BASE}:${variant}" + fi + # RAM / cores (informational) + _row 1 "host" "${LUCEBOX_HOST_NPROC} cpus, ${LUCEBOX_HOST_RAM_GB} GB RAM" +} + +cmd_in_container() { + # Generic dispatcher: anything that isn't a systemd action goes here. + # Runs the in-container Python CLI with the supplied argv. + require_host_prereqs + ensure_probed + # CTK isn't strictly required for every subcommand (e.g. `config get` + # or `autotune` only touch local files), but the server-spawning + # subcommands need it. + # Letting docker error its own way is fine for the no-CTK case. + local variant + variant=$(pick_variant) + local argv + mapfile -t argv < <(build_orchestrator_argv "$variant" "$@") + exec "${argv[@]}" +} + +# Is the long-running lucebox container currently up? Used by the dispatcher +# to decide between `docker exec` into it (cheap, shares the running server's +# network namespace so localhost:8080 reaches the server) vs. `docker run` +# (cold start, isolated network — can't reach the live server). +# +# `docker ps -q -f name=^$` prints the container id when running, +# empty otherwise. The anchored regex avoids matching `lucebox-cli-12345` +# style ephemeral siblings. +_lucebox_container_running() { + # No docker on PATH → definitely not running. Don't even probe. + command -v docker >/dev/null 2>&1 || return 1 + local id + id=$(docker ps -q -f "name=^${CONTAINER_NAME}\$" 2>/dev/null || true) + [ -n "$id" ] +} + +# `docker exec` variant of cmd_in_container. Same calling convention, but: +# - shares the running container's network namespace (localhost:8080 → the +# server), filesystem, and mounts — no bind mounts needed. +# - skips the ~1-3s cold-start cost of a fresh `docker run --rm`. +# - only safe for steady-state / read-only / config-only subcommands. Any +# command that restarts the lucebox service (autotune --sweep, serve) +# would kill the very container the exec is in — caller must route those +# to cmd_in_container instead. +# +# Pass through the same env-var subset the run path uses so the in-container +# CLI sees consistent overrides whichever route it took: HOME, every +# LUCEBOX_HOST_*, the image/port/container/models scalars, and HF_TOKEN. +cmd_exec_in_container() { + require_host_prereqs + ensure_probed + local variant + variant=$(pick_variant) + local tty=() + _set_tty_flags tty + local argv=(docker exec "${tty[@]}") + argv+=(--user "$(id -u):$(id -g)") + argv+=(-w "$PWD") + argv+=(-e "HOME=$HOME") + _append_host_env argv + _append_scalar_env argv "$variant" + # The image has no top-level `lucebox` binary on PATH — that name only + # works as the first arg to /opt/lucebox-hub/server/scripts/entrypoint.sh, + # which then `exec uv run ... python -m lucebox`s. docker exec bypasses + # the image's ENTRYPOINT, so we invoke the entrypoint shim explicitly + # with `lucebox` as its SUBCMD and the user's argv tail. Keeps the + # exec path bit-for-bit equivalent to what docker run does on the + # SUBCMD=lucebox branch. + argv+=("$CONTAINER_NAME" /opt/lucebox-hub/server/scripts/entrypoint.sh lucebox "$@") + exec "${argv[@]}" +} + +# Decide whether a given (subcommand, argv) pair is safe to run via +# `docker exec` into the live container. Returns 0 (yes, prefer exec) or 1 +# (no, must use docker run / host-side). +# +# The safe-to-exec set is exactly the steady-state / read-only / hits-the- +# running-server subcommands. Anything that restarts the service, mutates +# images, or is itself the long-running service must stay on cmd_in_container. +# +_lucebox_prefer_exec() { + local cmd="$1"; shift + case "$cmd" in + config|models|check|print-run|print-serve-argv) + return 0 + ;; + *) + return 1 + ;; + esac +} + +# Top-level routing for the in-container Python CLI. Picks between exec +# (cheap, shares the live server's namespace) and run (cold start, isolated). +# +# Decision tree: +# 1. LUCEBOX_NO_EXEC=1 / --no-exec was set → always run, never exec. +# Useful for debugging the wrapper or when the in-container Python is +# stale relative to the image. +# 2. cmd is not in the prefer-exec list → run (sweep, service mutators). +# 3. container is running → exec (the fast path, hits the live server). +# 4. container is not running → run (fall back so first-run / pre-install +# flows still work without a live service). +cmd_route_to_container() { + local cmd="$1"; shift + if [ "${LUCEBOX_NO_EXEC:-0}" = "1" ]; then + cmd_in_container "$cmd" "$@" + return + fi + if _lucebox_prefer_exec "$cmd" "$@" && _lucebox_container_running; then + cmd_exec_in_container "$cmd" "$@" + return + fi + cmd_in_container "$cmd" "$@" +} + +usage() { + cat < print shell completion script (bash / zsh / fish) + models list / download / activate model presets + config read / write keys in .lucebox/config.toml + print-run print the docker-run command for the server + +Misc: + help, --help, -h this message + version, --version print version + +Environment overrides: + LUCEBOX_IMAGE image name without tag (default: ghcr.io/luce-org/lucebox-hub) + LUCEBOX_VARIANT image tag to pull/run (default: cuda12) + LUCEBOX_PORT host port for the server (default: 8080) + LUCEBOX_CONTAINER server container name (default: lucebox) + LUCEBOX_MODELS host model directory (default: \$XDG_DATA_HOME/lucebox/models + LUCEBOX_NO_EXEC=1 force docker-run for in-container subcommands even + when the container is up (equivalent to --no-exec) + HF_TOKEN propagated to \`models download\` for gated HF repos + +Container routing: + When the long-running '$CONTAINER_NAME' container is up, steady-state + subcommands (config, models, check, print-run, print-serve-argv) + 'docker exec' into it instead of starting a fresh container. This avoids + the ~1-3s docker-run cold-start AND shares the live server's network + namespace so localhost:\$LUCEBOX_PORT reaches the server. Service-restarting + commands (serve, pull, update, install, etc.) stay on the host-side / + docker-run path. Pass --no-exec (or LUCEBOX_NO_EXEC=1) to force docker-run. +EOF +} + +# ── dispatch ────────────────────────────────────────────────────────────── + +main() { + # Global flag pass: `--no-exec` anywhere before the subcommand forces the + # docker-run path even if the container is up. Equivalent to + # `LUCEBOX_NO_EXEC=1 lucebox ...`. We pop it out of argv up-front so the + # rest of dispatch doesn't have to know about it. + local args=() + while [ $# -gt 0 ]; do + case "$1" in + --no-exec) export LUCEBOX_NO_EXEC=1; shift ;; + *) args+=("$1"); shift ;; + esac + done + set -- "${args[@]}" + + local cmd="${1:-help}" + [ $# -gt 0 ] && shift + case "$cmd" in + # Systemd surface + install) cmd_systemd_install "$@" ;; + uninstall) cmd_systemd_uninstall "$@" ;; + start|stop|restart|enable|disable|status) + cmd_systemctl_passthrough "$cmd" "$@" ;; + logs) cmd_logs "$@" ;; + + # Direct server + serve) cmd_serve "$@" ;; + pull) cmd_pull "$@" ;; + + # Self-update — re-runs the bootstrap installer against the channel + # this script was installed from (LUCEBOX_INSTALLED_FROM). + update) cmd_update "$@" ;; + + # Host-only readiness check — pure shell, never enters the container. + check) cmd_check "$@" ;; + + # Shell completion — print a script the user sources into their rc + # file. Bash and zsh share the bash-style emitter (zsh users add a + # `bashcompinit; complete` shim); fish is native. + completion) cmd_completion "$@" ;; + + # Help / version + help|--help|-h) usage ;; + version|--version) printf '%s\n' "$VERSION" ;; + + # Everything else → in-container Python CLI. cmd_route_to_container + # picks between `docker exec` into the live container (cheap, shares + # the running server's network namespace) and `docker run` (cold, + # isolated) based on container state + the safe-to-exec command set. + *) cmd_route_to_container "$cmd" "$@" ;; + esac +} + +main "$@" diff --git a/lucebox/.gitignore b/lucebox/.gitignore new file mode 100644 index 000000000..15f95f0e7 --- /dev/null +++ b/lucebox/.gitignore @@ -0,0 +1,3 @@ + +# Generated by hatch-vcs at build time from git tags. +src/lucebox/_version.py diff --git a/lucebox/README.md b/lucebox/README.md new file mode 100644 index 000000000..747a49eef --- /dev/null +++ b/lucebox/README.md @@ -0,0 +1,18 @@ +# lucebox — host CLI for the lucebox-hub container + +This package ships *inside* the `ghcr.io/luce-org/lucebox-hub` Docker image +and is invoked from the host via the [`lucebox.sh`](../lucebox.sh) wrapper: + + lucebox.sh check # `docker run … lucebox check` + lucebox.sh config get + lucebox.sh print-run + +The wrapper is the only thing that runs on the host; everything else (host +checks, TOML config, docker daemon calls, model download) is Python in the +container. Host facts (driver, GPU, RAM, VRAM, systemd availability) are +passed in via `LUCEBOX_HOST_*` environment variables so the Python side +doesn't reprobe. The autotune sweep, profiling, and agent-client launchers +land in follow-up PRs. + +Subcommands are defined in [`lucebox/cli.py`](src/lucebox/cli.py). See the +top-level [README.md](../README.md) for the user-facing flow. diff --git a/lucebox/pyproject.toml b/lucebox/pyproject.toml new file mode 100644 index 000000000..5277b268d --- /dev/null +++ b/lucebox/pyproject.toml @@ -0,0 +1,54 @@ +[project] +name = "lucebox" +# Version is derived from git tags via hatch-vcs (see [tool.hatch.version] +# below). Tag `lucebox-v0.2.1` → release version `0.2.1`. Commits past a +# tag get a `.devN+g` suffix so dev installs are visibly distinct +# from releases. Single source of truth: the git tag. +dynamic = ["version"] +description = "Host-side CLI for the lucebox-hub container: launch, config, model download" +readme = "README.md" +requires-python = ">=3.11" +authors = [{ name = "Lucebox" }] +license = { text = "Apache-2.0" } + +# Kept intentionally narrow. typer pulls click+rich; tomli-w gives us TOML +# writes (stdlib tomllib only reads). httpx for the smoke + readiness probes. +# huggingface_hub for download-models — used directly (not via subprocess) +# so we can drive a Rich progress bar + verify sha256 against the repo +# metadata before re-fetching multi-GB GGUFs. +dependencies = [ + "typer>=0.12", + "rich>=13", + "httpx>=0.27", + "tomli-w>=1.0", + "huggingface_hub>=0.27", + # luce-bench is consumed lazily by the autotune sweep scorer + # (agent_replay_pass_rate in sweep.py does a function-local + # `from lucebench.areas.agent_recorded import ...` wrapped in try/except). + # It's deliberately NOT a hard dep here because the workspace can't lock + # against it until #337 (luce-bench in-tree) lands. Install with + # `uv pip install luce-bench` on the host running the scorer. +] + +[project.scripts] +lucebox = "lucebox.cli:app" + +[build-system] +requires = ["hatchling", "hatch-vcs"] +build-backend = "hatchling.build" + +[tool.hatch.version] +source = "vcs" +# Untagged checkouts (e.g. fresh clone before tagging lucebox-v0.2.1) +# resolve to this rather than 0.0.0.dev0. +fallback-version = "0.2.1.dev0" +raw-options.tag_regex = '''^lucebox-v(?P\d+\.\d+\.\d+)$''' + +[tool.hatch.build.hooks.vcs] +# Build hook writes the resolved version into src/lucebox/_version.py +# so `__init__.py` can `from lucebox._version import __version__`. +# Generated file — see lucebox/.gitignore. +version-file = "src/lucebox/_version.py" + +[tool.hatch.build.targets.wheel] +packages = ["src/lucebox"] diff --git a/lucebox/src/lucebox/__init__.py b/lucebox/src/lucebox/__init__.py new file mode 100644 index 000000000..8ca821024 --- /dev/null +++ b/lucebox/src/lucebox/__init__.py @@ -0,0 +1,16 @@ +"""lucebox — host-side CLI for the lucebox-hub container. + +Runs inside the container; the host wrapper at ../lucebox.sh handles `docker +run` plumbing and systemd integration. This package owns: TOML config, the +host-derived DFLASH_* serve heuristic, docker daemon calls (via the mounted +socket), and model download. The empirical autotune sweep, profiling, and +agent-client launchers land in follow-up PRs. +""" + +# Version is generated by hatch-vcs at build time into _version.py. +# Fresh source-tree checkouts before any build will not yet have the +# file — fall back to a dev marker so imports don't break. +try: + from lucebox._version import __version__ +except ImportError: + __version__ = "0.0.0.dev0+unbuilt" diff --git a/lucebox/src/lucebox/__main__.py b/lucebox/src/lucebox/__main__.py new file mode 100644 index 000000000..128e2ca87 --- /dev/null +++ b/lucebox/src/lucebox/__main__.py @@ -0,0 +1,6 @@ +"""Entry point for `python -m lucebox`.""" + +from lucebox.cli import app + +if __name__ == "__main__": + app() diff --git a/lucebox/src/lucebox/autotune.py b/lucebox/src/lucebox/autotune.py new file mode 100644 index 000000000..51b4f8424 --- /dev/null +++ b/lucebox/src/lucebox/autotune.py @@ -0,0 +1,80 @@ +"""Heuristic autotune: VRAM tier → DflashRuntime defaults. + +The recommended runtime is computed from HostFacts (VRAM, is_wsl) — stateless: +it takes HostFacts in and returns a fresh DflashRuntime. ``config.live_config`` +applies it so ``lucebox print-serve-argv`` / ``docker_run`` bake conservative +DFLASH_* defaults into the serve command for the detected VRAM tier. + +The empirical sweep + per-workload profiles (``lucebox autotune --sweep``) +live in a follow-up PR; this module keeps only the host-derived heuristic +that the serve path depends on. +""" + +from __future__ import annotations + +from lucebox.types import DflashRuntime, HostFacts + + +def runtime_from_host(host: HostFacts) -> DflashRuntime: + """Pick a conservative DflashRuntime that 'should work' on this VRAM tier. + + Tiers (NVIDIA, baseline = Qwen3.6-27B Q4_K_M ~18 GB total): + <12 GB — too small for 27B; pick min ctx as a floor so a fallback + start at least gets an error from the daemon rather than + a silent OOM. + 12-21 — fits but tight; cap ctx. + 22-31 — 24 GB-class consumer flagships (3090/4090/5090/5090-Laptop). + 98 K with tq3_0 KV (~2 GB KV + ~18 GB model ≈ 20 GB). + Confirmed on bragi (RTX 5090 Laptop, 23 GB VRAM) 2026-05-31. + 32-47 — RTX 6000 Ada / A100 40 GB. Full 128 K. + ≥48 — A100 80 GB / H100 / RTX 6000 Pro. Full 128 K. + + Prefix cache remains an explicit sweep tunable, but the automatic + baseline keeps it off because tool prompts currently exercise a daemon + snapshot path that is not reliable with prefix slots enabled. + Empirically confirmed on bragi 2026-05-31: prefix_cache_slots=32 + caused -19pp regression on agent_recorded (23.1% vs 42.3% baseline). + 5 previously-passing cases regressed; 0 new cases unlocked. See + docs/experiments/qwen3.6-27b-prefix-cache-regression-bragi-2026-05-31.md. + + On `lazy`: the C++ server requires `--prefill-drafter` (and `--draft`) + to be set for `--lazy-draft` to take effect, and silently ignores it + otherwise (`--lazy-draft ignored: requires both --prefill-drafter and + --draft`). Since the heuristic path does NOT set `prefill_drafter`, + we default `lazy=False` here — "what we say" matches "what runs". + Users who explicitly opt in via config.toml will be warned at server + startup that the flag is being dropped (see entrypoint.sh). + """ + if host.vram_gb <= 0: + return DflashRuntime() # no VRAM signal — stick with class defaults + + if host.vram_gb < 12: + return DflashRuntime(max_ctx=4096) + if host.vram_gb < 22: + return DflashRuntime(max_ctx=32768) + if host.vram_gb < 32: + # 22-31 GB cards. tq3_0 KV is required at 98K: model (~18-19 GB) + + # q8_0 KV at 98K (~5-6 GB) = 24-25 GB → OOM, while tq3_0 KV (~2 GB) + # leaves ~3 GB headroom. Confirmed on bragi (RTX 5090 Laptop, 23 GB + # VRAM) 2026-05-30 — q8_0 timed out on every 98K cell; all tq3_0 cells + # passed. Preset-size-aware capping (large models → 32K) lives with the + # autotune sweep in a follow-up PR. + if host.is_wsl: + # Bumped from max_ctx=65536 → 98304 on 2026-05-30 after the + # coding-agent-loop sweep on sindri proved 98K serves real + # 90K-token agentic prompts with ~3 GB VRAM headroom and no + # CUDA VMM failures. See + # docs/experiments/gemma4-26b-coding-agent-loop-sweep-2026-05-30.md. + # The original 65K cap cited unverified VMM failures — + # bisect history showed no commit reproducing them. + return DflashRuntime( + budget=16, max_ctx=98304, + cache_type_k="tq3_0", cache_type_v="tq3_0", + ) + return DflashRuntime( + max_ctx=98304, + cache_type_k="tq3_0", cache_type_v="tq3_0", + ) + if host.vram_gb < 48: + return DflashRuntime(max_ctx=131072) + return DflashRuntime(max_ctx=131072) diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py new file mode 100644 index 000000000..87cc29366 --- /dev/null +++ b/lucebox/src/lucebox/cli.py @@ -0,0 +1,356 @@ +"""Typer app — the user-facing subcommands. + +Layout follows the host wrapper's dispatch table. Anything `lucebox` +doesn't intercept (everything outside the systemd surface) ends up here. + +Subcommand inventory: + check — readiness report + config get/set/unset — read / write a single key in config.toml + pull — docker pull the cuda12 image + print-run — emit the docker-run command for the server + print-serve-argv — same, raw argv lines (consumed by `lucebox serve`) + models — list / download presets, activate one +""" + +from __future__ import annotations + +import os +import sys +from dataclasses import replace +from pathlib import Path +from typing import Annotated + +import typer +from rich.console import Console +from rich.table import Table + +import lucebox.config as config_mod +import lucebox.docker_run as docker_run +import lucebox.download as download_mod +import lucebox.host_check as host_check +from lucebox import __version__ +from lucebox.config import config_get, config_set, config_unset, live_config +from lucebox.host_facts import from_env + +app = typer.Typer( + name="lucebox", + help="Host CLI for the lucebox-hub container. Invoked by lucebox.sh.", + no_args_is_help=True, + add_completion=False, +) +console = Console() + + +# ── helpers ──────────────────────────────────────────────────────────────── + + +def _load_or_build() -> config_mod.Config: # type: ignore[name-defined] + """env > config.toml > dataclass defaults — the canonical precedence. + + Without the env-overlay step below, `config_mod.load()` returned the + persisted config verbatim and `LUCEBOX_IMAGE` / `LUCEBOX_VARIANT` / + `LUCEBOX_PORT` / `LUCEBOX_CONTAINER` / `LUCEBOX_MODELS` from the + systemd unit's `Environment=` (or any one-shot shell export) were + silently dropped. That contradicted the precedence lucebox.sh + documents and applies — and bit sindri when its config.toml had + `[image]` without `registry`, so the dataclass default + `ghcr.io/luce-org/lucebox-hub` won over the unit's + `LUCEBOX_IMAGE=ghcr.io/easel/lucebox-hub`. + + Fix: overlay env on top of the loaded config (or the live_config + fallback when config.toml is absent). Only the five top-level + scalars have env hooks — dflash/host/model don't, by design. + """ + cfg = config_mod.load() + if cfg is None: + cfg = live_config() + # Overlay live host facts. When ``config.toml`` exists without a + # ``[host]`` block (the common case — operators don't hand-edit + # host facts), ``cfg.host`` defaults to a zero-filled ``HostFacts`` + # and the DFLASH_* serve heuristic silently falls through to the + # "no VRAM signal" path. Re-probe from env so the wrapper-exported + # LUCEBOX_HOST_* facts always win over the persisted (possibly + # absent) snapshot. + live_host = from_env() + host = live_host if live_host.vram_gb > 0 or live_host.nproc > 0 else cfg.host + return replace( + cfg, + variant=os.environ.get("LUCEBOX_VARIANT", cfg.variant), + image=os.environ.get("LUCEBOX_IMAGE", cfg.image), + container_name=os.environ.get("LUCEBOX_CONTAINER", cfg.container_name), + port=int(os.environ.get("LUCEBOX_PORT", str(cfg.port))), + models_dir=Path(os.environ.get("LUCEBOX_MODELS", str(cfg.models_dir))), + host=host, + ) + + +# ── subcommands ──────────────────────────────────────────────────────────── + + +@app.command() +def check() -> None: + """Print a readiness report (driver, docker, CTK, RAM, VRAM, systemd).""" + host = from_env() + results = host_check.run_checks(host) + worst = host_check.render(console, host, results) + if worst == "fail": + raise typer.Exit(code=1) + + +@app.command() +def pull() -> None: + """`docker pull` the image variant from config.toml.""" + cfg = _load_or_build() + tag = f"{cfg.image}:{cfg.variant}" + console.print(f"[bold]Pulling {tag}[/bold] (~14 GB; takes a while)…") + rc = docker_run.docker_pull(tag) + if rc != 0: + raise typer.Exit(code=rc) + + +@app.command("print-run") +def print_run() -> None: + """Print the docker-run command for the server (copy-pasteable).""" + cfg = _load_or_build() + spec = docker_run.server_run_spec(cfg) + print(spec.printable()) + + +@app.command("print-serve-argv") +def print_serve_argv() -> None: + """Emit the server docker-run argv, one token per line. + + Consumed by lucebox.sh's `serve` subcommand and the systemd unit. Kept as + a separate command from `print-run` so the bash side has a guaranteed + machine-readable contract that's independent of the pretty formatter. + """ + cfg = _load_or_build() + spec = docker_run.server_run_spec(cfg) + for tok in spec.argv(): + print(tok) + + +# ── config sub-app ───────────────────────────────────────────────────────── + + +config_app = typer.Typer(no_args_is_help=True, help="Read/write keys in config.toml.") +app.add_typer(config_app, name="config") + + +@config_app.command("get") +def config_get_cmd( + key: Annotated[str, typer.Argument(help="Dotted key (omit to list every key).")] = "", +) -> None: + """Print a single key (or every reachable key) with its origin annotation.""" + try: + entries = config_get(key or None) + except KeyError as exc: + console.print(f"[red]{exc}[/red]") + raise typer.Exit(code=2) from exc + for k, (value, origin) in entries.items(): + console.print(f"{k} = {value!r} ([dim]from {origin}[/dim])") + + +@config_app.command("set") +def config_set_cmd( + kv: Annotated[str, typer.Argument(help='"key=value" pair (e.g. "model.preset=qwen3.6-27b")')], +) -> None: + """Set one dotted key. Auto-creates config.toml when missing. + + Only the named key is written — other on-disk keys are preserved + untouched, unset keys stay implicit. Use `lucebox config unset` to + remove a key (next read falls back to the live default). + """ + if "=" not in kv: + console.print("[red]argument must be key=value[/red]") + raise typer.Exit(code=2) + key, _, value = kv.partition("=") + key = key.strip() + value = value.strip() + try: + config_set(key, value) + except (KeyError, ValueError) as exc: + console.print(f"[red]{exc}[/red]") + raise typer.Exit(code=2) from exc + console.print(f"[green]Set[/green] {key} = {value}") + + +@config_app.command("unset") +def config_unset_cmd( + key: Annotated[str, typer.Argument(help="Dotted key to remove from config.toml.")], +) -> None: + """Remove a key from config.toml. Next read uses the live default.""" + try: + changed = config_unset(key) + except KeyError as exc: + console.print(f"[red]{exc}[/red]") + raise typer.Exit(code=2) from exc + if changed: + console.print(f"[green]Unset[/green] {key}") + else: + console.print(f"[dim]{key} was not in config.toml; nothing to do[/dim]") + + +# ── models sub-app ───────────────────────────────────────────────────────── + + +models_app = typer.Typer( + no_args_is_help=False, help="Manage local model presets (list, download, activate)." +) +app.add_typer(models_app, name="models") + + +def _print_installed_presets() -> None: + cfg = _load_or_build() + installed = download_mod.installed_presets(cfg) + active = cfg.model.preset + console.print(f"Models dir: [bold]{cfg.models_dir}[/bold]") + if not installed: + console.print("[dim]No presets installed yet — try `lucebox models download`.[/dim]") + return + table = Table() + table.add_column("preset") + table.add_column("status") + table.add_column("size (GB)") + for pres in installed: + marker = "* " if pres.name == active else " " + size_gb = download_mod.installed_size_gb(cfg, pres) + table.add_row(f"{marker}{pres.name}", "installed", f"{size_gb:.1f}") + console.print(table) + total = sum(download_mod.installed_size_gb(cfg, p) for p in installed) + console.print(f"[dim]Total disk usage: {total:.1f} GB[/dim]") + + +@models_app.callback(invoke_without_command=True) +def models_default(ctx: typer.Context) -> None: + """Default action: list installed presets, mark active with `*`.""" + if ctx.invoked_subcommand is None: + _print_installed_presets() + + +@models_app.command("list") +def models_list() -> None: + """Show every registered preset (installed or not) with status + size.""" + cfg = _load_or_build() + active = cfg.model.preset + table = Table() + table.add_column("preset") + table.add_column("status") + table.add_column("size (GB)") + table.add_column("description") + for name in sorted(download_mod.PRESETS): + pres = download_mod.PRESETS[name] + marker = "* " if name == active else " " + status = download_mod.installed_status(cfg, pres) + size = download_mod.installed_size_gb(cfg, pres) + size_text = f"{size:.1f}" if size > 0 else f"~{pres.approx_total_gb}*" + table.add_row(f"{marker}{name}", status, size_text, pres.description or "") + console.print(table) + + +@models_app.command("download") +def models_download( + preset: Annotated[str, typer.Argument(help="Preset name (empty = recommend)")] = "", + activate: Annotated[ + bool, typer.Option("--activate", help="Also set as active preset (model.preset).") + ] = False, +) -> None: + """Fetch a preset's GGUFs into the models dir. + + With no argument and no preset configured, recommends one for this + host's VRAM tier and auto-activates it (the first-install path). + Otherwise the named preset is downloaded; pass ``--activate`` to + also flip `model.preset` to it. + """ + cfg = _load_or_build() + if not preset: + if cfg.model.preset: + console.print( + "[yellow]No preset specified and one is already active. " + "Pass an explicit preset name (or use --activate to switch).[/yellow]" + ) + raise typer.Exit(code=2) + recommended = download_mod.recommend_preset(cfg.host) + if recommended is None: + console.print( + "[red]Cannot recommend a preset for this host. " + "Run `lucebox models list` and pick one explicitly.[/red]" + ) + raise typer.Exit(code=2) + preset = recommended + activate = True + console.print( + f"[bold]Recommended preset: {preset}[/bold] " + "(no preset configured; auto-activating after download)" + ) + + try: + pres = download_mod.resolve_preset(preset) + except KeyError as exc: + console.print(f"[red]{exc}[/red]") + raise typer.Exit(code=2) from exc + + current = download_mod.status(cfg, pres) + console.print(f"Models dir: [bold]{cfg.models_dir}[/bold]") + console.print(f"Preset: [bold]{pres.name}[/bold]") + console.print( + f" target ({pres.target_repo}/{pres.target_file}):" + f" {'present' if current['target_present'] else 'will download'}" + ) + if pres.has_draft: + console.print( + f" draft ({pres.draft_repo}/{pres.draft_file}):" + f" {'present' if current['draft_present'] else 'will download'}" + ) + else: + console.print(" draft [dim](none — target-only preset)[/dim]") + + if current["target_present"] and current["draft_present"]: + console.print("[green]Already present.[/green]") + else: + console.print(f"[bold]Downloading[/bold] (~{pres.approx_total_gb} GB total)…") + rc = download_mod.download_preset(cfg, pres) + if rc != 0: + raise typer.Exit(code=rc) + console.print("[green]Done.[/green]") + + if activate: + config_set("model.preset", preset) + if pres.target_file: + config_set("model.target_file", pres.target_file) + if pres.has_draft and pres.draft_file: + config_set("model.draft_file", pres.draft_file) + else: + # Drop any stale draft_file from a previous activation; the + # active preset has no draft. + config_unset("model.draft_file") + console.print(f"[green]Activated:[/green] model.preset = {preset}") + # First-time setup: bake the VRAM-tier DFLASH_* heuristic into + # config.toml so `lucebox serve` is auto-tuned to this host instead + # of falling back to the conservative class defaults. Never clobbers + # an existing [dflash] section. + if config_mod.seed_dflash_from_host(cfg.host): + max_ctx = config_get("dflash.max_ctx")["dflash.max_ctx"][0] + console.print( + f"[green]Auto-tuned:[/green] VRAM-tier DFLASH_* defaults " + f"(max_ctx={max_ctx}) written to config.toml" + ) + + +@app.command() +def version() -> None: + """Print lucebox version.""" + print(__version__) + + +def main() -> None: + """Module entrypoint — `python -m lucebox`.""" + try: + app() + except KeyboardInterrupt: + console.print("\n[dim]interrupted[/dim]") + sys.exit(130) + + +if __name__ == "__main__": + main() diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py new file mode 100644 index 000000000..b3f63eb0d --- /dev/null +++ b/lucebox/src/lucebox/config.py @@ -0,0 +1,482 @@ +"""Sparse TOML persistence for .lucebox/config.toml. + +Single source of truth for user-overridden configuration. We track which +dotted keys were explicitly set by the user (or by commands acting on +their behalf) and serialize ONLY those keys back to disk — defaults +stay implicit, so `config.toml` reads like a diff against live defaults +and upgrades that add new fields don't gratuitously rewrite every file. + +The dotted-key surface area is small and flat: + model.preset, model.target_file, model.draft_file + port, models_dir, variant, image, container_name + dflash. for each of the 11 DflashRuntime knobs + think_max + +Load resolves the TOML file → ``Config`` object, with anything absent +filled from ``Config()`` defaults. Save writes back only the keys that +appear in the TOML doc (tracked on ``Config._user_set``). The TOML doc +itself is a plain ``dict[str, Any]`` carrying only the set keys. +""" + +from __future__ import annotations + +import os +import re +import tomllib +from collections.abc import Callable +from dataclasses import asdict, replace +from datetime import UTC +from pathlib import Path +from typing import Any + +import tomli_w + +from lucebox.types import ( + Config, + DflashRuntime, + HostFacts, + ModelMeta, + Variant, + default_models_dir, +) + + +def default_config_path() -> Path: + """Where .lucebox/config.toml lives. + + Convention: under $LUCEBOX_HOME if set, otherwise $HOME/.lucebox. Lives in + the bind-mounted host home dir so the config survives container teardown + and is editable from the host. + """ + base = os.environ.get("LUCEBOX_HOME") + if base: + return Path(base) / "config.toml" + return Path.home() / ".lucebox" / "config.toml" + + +# ── dotted-key registry ──────────────────────────────────────────────────── + +def _cast_prefill_mode(v: Any) -> str: + s = str(v) + if s not in {"off", "auto", "always"}: + raise ValueError(f"prefill_mode must be off/auto/always, got {s!r}") + return s + + +def _cast_bool(v: Any) -> bool: + """Strict-ish boolean coercion for config values. + + - Native booleans pass through. + - Strings: 1/true/yes/on → True; 0/false/no/off/"" → False (case-insensitive). + - Anything else raises ``ValueError`` rather than silently coercing, + because that's what bit ``dflash.debug_thinking_logits`` — the + built-in ``bool`` caster turned ``"false"`` into ``True``. + """ + if isinstance(v, bool): + return v + if isinstance(v, str): + s = v.strip().lower() + if s in ("1", "true", "yes", "on"): + return True + if s in ("0", "false", "no", "off", ""): + return False + raise ValueError(f"cannot parse boolean: {v!r}") + if isinstance(v, int): + return bool(v) + raise ValueError(f"cannot parse boolean: {v!r}") + + +# Each entry: dotted-key → (toml_path, type_caster, default_getter). +# ``toml_path`` is the (section, field) pair on disk; ``"_root"`` means the +# key lives at the top level (no [section]). ``default_getter`` returns the +# in-memory default so ``config get`` can annotate origin. +KEY_REGISTRY: dict[str, tuple[tuple[str, str], Callable[[Any], Any]]] = { + "variant": (("image", "variant"), str), + "image": (("image", "registry"), str), + "container_name": (("runtime", "container_name"), str), + "port": (("runtime", "port"), int), + "models_dir": (("paths", "models"), str), + "model.preset": (("model", "preset"), str), + "model.target_file": (("model", "target_file"), str), + "model.draft_file": (("model", "draft_file"), str), + "dflash.budget": (("dflash", "budget"), int), + "dflash.max_ctx": (("dflash", "max_ctx"), int), + "dflash.lazy": (("dflash", "lazy"), _cast_bool), + "dflash.prefix_cache_slots": (("dflash", "prefix_cache_slots"), int), + "dflash.prefill_cache_slots": (("dflash", "prefill_cache_slots"), int), + "dflash.cache_type_k": (("dflash", "cache_type_k"), str), + "dflash.cache_type_v": (("dflash", "cache_type_v"), str), + "dflash.prefill_mode": (("dflash", "prefill_mode"), _cast_prefill_mode), + "dflash.prefill_keep_ratio": (("dflash", "prefill_keep_ratio"), float), + "dflash.prefill_threshold": (("dflash", "prefill_threshold"), int), + "dflash.prefill_drafter": (("dflash", "prefill_drafter"), str), + "dflash.think_max": (("dflash", "think_max"), int), + "dflash.fa_window": (("dflash", "fa_window"), int), + "dflash.think_soft_close_min_ratio": ( + ("dflash", "think_soft_close_min_ratio"), float), + "dflash.debug_thinking_logits": ( + ("dflash", "debug_thinking_logits"), _cast_bool), +} + + +def _doc_get(doc: dict[str, Any], section: str, field: str) -> Any: + if section == "_root": + return doc.get(field) + sub = doc.get(section) + if isinstance(sub, dict): + return sub.get(field) + return None + + +def _doc_set(doc: dict[str, Any], section: str, field: str, value: Any) -> None: + if section == "_root": + doc[field] = value + return + doc.setdefault(section, {})[field] = value + + +def _doc_unset(doc: dict[str, Any], section: str, field: str) -> bool: + """Remove a dotted key from the doc. Returns True iff something was removed.""" + if section == "_root": + if field in doc: + del doc[field] + return True + return False + sub = doc.get(section) + if isinstance(sub, dict) and field in sub: + del sub[field] + if not sub: + del doc[section] + return True + return False + + +# ── load ─────────────────────────────────────────────────────────────────── + + +def load(path: Path | None = None) -> Config | None: + """Load config.toml, or return None if missing. + + If a legacy `.env` sits next to it (or in place of it), migrate that + first and write back as TOML. + """ + path = path or default_config_path() + if path.exists(): + return _load_toml(path) + + legacy = path.with_suffix(".env") + if legacy.exists(): + cfg, doc = _load_legacy_env(legacy) + save(cfg, path, doc=doc) + return cfg + + return None + + +def _load_toml(path: Path) -> Config: + raw = tomllib.loads(path.read_text()) + return _from_dict(raw) + + +def load_doc(path: Path | None = None) -> dict[str, Any]: + """Return the raw TOML doc (a dict). Empty when no file or empty file.""" + path = path or default_config_path() + if not path.exists(): + return {} + return tomllib.loads(path.read_text()) + + +_LEGACY_KEY_MAP: dict[str, tuple[str, str, Callable[[str], Any]]] = { + "DFLASH_BUDGET": ("dflash", "budget", int), + "DFLASH_MAX_CTX": ("dflash", "max_ctx", int), + "DFLASH_LAZY": ("dflash", "lazy", + lambda v: str(v).strip().lower() in ("1", "true", "yes", "on")), + "DFLASH_PREFIX_CACHE_SLOTS": ("dflash", "prefix_cache_slots", int), + "DFLASH_PORT": ("runtime", "port", int), + "LUCEBOX_VARIANT": ("image", "variant", str), + "LUCEBOX_IMAGE": ("image", "registry", str), + "LUCEBOX_MODELS": ("paths", "models", str), +} + + +def _load_legacy_env(path: Path) -> tuple[Config, dict[str, Any]]: + """Best-effort migration from the bash-era .lucebox/config.env.""" + raw: dict[str, Any] = {} + line_re = re.compile(r"^([A-Z_][A-Z0-9_]*)=(.*)$") + for line in path.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + m = line_re.match(line) + if not m: + continue + key, val = m.group(1), m.group(2).strip().strip('"').strip("'") + if key not in _LEGACY_KEY_MAP: + continue + section, field, cast_fn = _LEGACY_KEY_MAP[key] + try: + raw.setdefault(section, {})[field] = cast_fn(val) + except (TypeError, ValueError): + continue + return _from_dict(raw), raw + + +def _from_dict(raw: dict[str, Any]) -> Config: + img = raw.get("image", {}) + variant: Variant = str(img.get("variant", "cuda12")) + registry = img.get("registry", "ghcr.io/luce-org/lucebox-hub") + + runtime = raw.get("runtime", {}) + port = int(runtime.get("port", 8080)) + container_name = str(runtime.get("container_name", "lucebox")) + + paths = raw.get("paths", {}) + models_dir = Path(paths.get("models", str(default_models_dir()))) + + df = raw.get("dflash", {}) + dflash = DflashRuntime( + budget=int(df.get("budget", 22)), + max_ctx=int(df.get("max_ctx", 16384)), + lazy=bool(df.get("lazy", False)), + prefix_cache_slots=int(df.get("prefix_cache_slots", 0)), + prefill_cache_slots=int(df.get("prefill_cache_slots", 0)), + cache_type_k=str(df.get("cache_type_k", "")), + cache_type_v=str(df.get("cache_type_v", "")), + prefill_mode=df.get("prefill_mode", "off"), + prefill_keep_ratio=float(df.get("prefill_keep_ratio", 0.05)), + prefill_threshold=int(df.get("prefill_threshold", 32000)), + prefill_drafter=str(df.get("prefill_drafter", "")), + think_max=int(df.get("think_max", 15488)), + fa_window=int(df.get("fa_window", 0)), + think_soft_close_min_ratio=float( + df.get("think_soft_close_min_ratio", 0.0)), + debug_thinking_logits=bool(df.get("debug_thinking_logits", False)), + ) + + host_raw = raw.get("host", {}) + host = HostFacts( + nproc=int(host_raw.get("nproc", 0)), + ram_gb=int(host_raw.get("ram_gb", 0)), + gpu_vendor=host_raw.get("gpu_vendor", "none"), + gpu_name=str(host_raw.get("gpu_name", "")), + gpu_count=int(host_raw.get("gpu_count", 0)), + vram_gb=int(host_raw.get("vram_gb", 0)), + gpu_sm=str(host_raw.get("gpu_sm", "")), + driver_version=str(host_raw.get("driver_version", "")), + driver_major=int(host_raw.get("driver_major", 0)), + has_systemd=bool(host_raw.get("has_systemd", False)), + is_wsl=bool(host_raw.get("is_wsl", False)), + has_docker=bool(host_raw.get("has_docker", False)), + docker_version=str(host_raw.get("docker_version", "")), + ctk=host_raw.get("ctk", "none"), + ) + + # `[model]` is optional — legacy configs (pre-multi-model) carry no + # such section and we want them to keep working unchanged. If + # `preset` is set but `target_file` / `draft_file` isn't, derive + # them from the registry so users only have to write one key. + mdl = raw.get("model", {}) + preset_name = str(mdl.get("preset", "")) + target_file = str(mdl.get("target_file", "")) + draft_file = str(mdl.get("draft_file", "")) + if preset_name and (not target_file or not draft_file): + from lucebox.download import PRESETS + + if preset_name in PRESETS: + pres = PRESETS[preset_name] + if not target_file: + target_file = pres.target_file + if not draft_file and pres.has_draft and pres.draft_file: + draft_file = pres.draft_file + model = ModelMeta(preset=preset_name, target_file=target_file, draft_file=draft_file) + + return Config( + variant=variant, + image=registry, + container_name=container_name, + port=port, + models_dir=models_dir, + dflash=dflash, + host=host, + model=model, + ) + + +# ── save ─────────────────────────────────────────────────────────────────── + + +def _atomic_write_doc(path: Path, doc: dict[str, Any]) -> None: + """Serialize ``doc`` to TOML and write it to ``path`` atomically. + + Write to a sibling ``.toml.tmp`` then ``replace`` so a crash mid-write + never leaves a truncated config.toml. Caller ensures ``path.parent`` exists. + """ + tmp = path.with_suffix(".toml.tmp") + tmp.write_bytes(tomli_w.dumps(doc).encode("utf-8")) + tmp.replace(path) + + +def save(cfg: Config, path: Path | None = None, *, doc: dict[str, Any] | None = None) -> Path: + """Persist a Config to ``path``. Only keys present in ``doc`` are written. + + ``doc`` is the raw TOML mapping returned by ``load_doc`` — it carries + exactly the keys the user (or a command on their behalf) has set. When + ``doc=None`` and the file exists we re-use the on-disk doc; when both + are absent we write an empty file. + """ + path = path or default_config_path() + path.parent.mkdir(parents=True, exist_ok=True) + if doc is None: + doc = load_doc(path) + _atomic_write_doc(path, doc) + # Silence unused-arg: cfg is the on-disk representation's source of + # truth for callers that want to round-trip through a Config object, + # but the sparse write never re-derives keys from it. + del cfg + return path + + +def seed_dflash_from_host(host: HostFacts, *, path: Path | None = None) -> bool: + """Persist the VRAM-tier DFLASH_* heuristic to config.toml on first setup. + + Returns True when it wrote. No-op (returns False) when a ``[dflash]`` + section already exists, so it never clobbers values a prior tune or the + user set. Called when a preset is first activated: without it a fresh + install serves at the conservative ``DflashRuntime`` class defaults + (``load()`` returns those for a config.toml that has no ``[dflash]``), + ignoring the host's VRAM tier. The ``live_config`` heuristic only fires + when there is no config.toml at all, which stops being true the moment a + model is activated — so the heuristic is persisted here instead. + """ + from datetime import datetime + + import lucebox.autotune as autotune_mod + + path = path or default_config_path() + doc = load_doc(path) + if "dflash" in doc: + return False + runtime = autotune_mod.runtime_from_host(host) + for field, value in asdict(runtime).items(): + _doc_set(doc, "dflash", field, _value_to_toml(value)) + _doc_set(doc, "autotune", "source", "heuristic") + _doc_set( + doc, "autotune", "timestamp", + datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"), + ) + path.parent.mkdir(parents=True, exist_ok=True) + _atomic_write_doc(path, doc) + return True + + +# ── dotted-key API ───────────────────────────────────────────────────────── + + +def _value_to_toml(value: Any) -> Any: + """Make a Python value safe for tomli_w (no None, Path→str).""" + if isinstance(value, Path): + return str(value) + return value + + +def _live_default(key: str) -> Any: + """Return the in-memory default for ``key`` (from a fresh Config()).""" + cfg = Config() + section_field = KEY_REGISTRY[key][0] + section, field = section_field + if section == "image": + return {"variant": cfg.variant, "registry": cfg.image}[field] + if section == "runtime": + return {"port": cfg.port, "container_name": cfg.container_name}[field] + if section == "paths": + return str(cfg.models_dir) if field == "models" else None + if section == "dflash": + return getattr(cfg.dflash, field) + if section == "model": + return getattr(cfg.model, field) + return None + + +def config_set(key: str, value: Any, *, path: Path | None = None) -> None: + """Set one dotted key and write the file. Auto-creates a missing file.""" + if key not in KEY_REGISTRY: + raise KeyError(f"unknown config key {key!r}; known: {sorted(KEY_REGISTRY)}") + section_field, caster = KEY_REGISTRY[key] + section, field = section_field + try: + cast_value = caster(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"cannot coerce {value!r} for {key}: {exc}") from exc + path = path or default_config_path() + doc = load_doc(path) if path.exists() else {} + _doc_set(doc, section, field, _value_to_toml(cast_value)) + path.parent.mkdir(parents=True, exist_ok=True) + _atomic_write_doc(path, doc) + + +def config_unset(key: str, *, path: Path | None = None) -> bool: + """Remove a dotted key from the file. Returns True if something changed.""" + if key not in KEY_REGISTRY: + raise KeyError(f"unknown config key {key!r}; known: {sorted(KEY_REGISTRY)}") + section_field, _ = KEY_REGISTRY[key] + section, field = section_field + path = path or default_config_path() + if not path.exists(): + return False + doc = load_doc(path) + changed = _doc_unset(doc, section, field) + if changed: + # Leave the file in place even when empty — `config set` will + # repopulate; deleting would surprise users who expect their + # config dir to exist. + _atomic_write_doc(path, doc) + return changed + + +def config_get(key: str | None = None, *, path: Path | None = None) -> dict[str, tuple[Any, str]]: + """Return ``{key: (value, origin)}``. ``origin`` is ``"file"`` or ``"default"``. + + When ``key`` is None or empty, every registered key is returned. + Otherwise just that one key (still as a single-item dict, for caller + uniformity). + """ + path = path or default_config_path() + doc = load_doc(path) if path.exists() else {} + keys = [key] if key else list(KEY_REGISTRY) + out: dict[str, tuple[Any, str]] = {} + for k in keys: + if k not in KEY_REGISTRY: + raise KeyError(f"unknown config key {k!r}; known: {sorted(KEY_REGISTRY)}") + section_field, _ = KEY_REGISTRY[k] + section, field = section_field + in_file = _doc_get(doc, section, field) + if in_file is not None: + out[k] = (in_file, "file") + else: + out[k] = (_live_default(k), "default") + return out + + +def live_config() -> Config: + """Build a fresh Config from current host facts + the DFLASH_* heuristic. + + Used as the no-config fallback in ``cli._load_or_build`` and reused by + the ``models`` sub-app, so the host probe + heuristic + env-override + logic lives in one place rather than being duplicated per caller. + """ + # Lazy import to avoid the autotune ↔ config import cycle the importer + # would hit if this moved to module scope. + import lucebox.autotune as autotune_mod + from lucebox.host_facts import from_env + + host = from_env() + default = Config() + return replace( + default, + variant=os.environ.get("LUCEBOX_VARIANT", "cuda12"), + image=os.environ.get("LUCEBOX_IMAGE", default.image), + container_name=os.environ.get("LUCEBOX_CONTAINER", default.container_name), + port=int(os.environ.get("LUCEBOX_PORT", str(default.port))), + models_dir=Path(os.environ.get("LUCEBOX_MODELS", str(default.models_dir))), + dflash=autotune_mod.runtime_from_host(host), + host=host, + ) diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py new file mode 100644 index 000000000..ff3615b26 --- /dev/null +++ b/lucebox/src/lucebox/docker_run.py @@ -0,0 +1,232 @@ +"""Build and execute `docker run` argv for the server and download containers. + +We shell out to the `docker` CLI rather than using the docker SDK because +(a) the CLI is the user-visible contract — errors look the same whether +issued by lucebox or the user; (b) zero import cost; (c) trivially mockable +via subprocess in tests. Wrap everything in one module so swapping to the +SDK later is a single-file change. +""" + +from __future__ import annotations + +import os +import shlex +import subprocess +from dataclasses import dataclass +from pathlib import Path + +from lucebox.types import Config + + +def _host_facts_env() -> list[tuple[str, str]]: + """Forward LUCEBOX_HOST_* from the orchestrator's env into the server. + + lucebox.sh's probe_host() exports every host-identity fact (OS, + kernel, GPU list CSV, CTK version, …) before invoking ``docker run`` + on the orchestrator. The orchestrator inherits them and we pass + them through verbatim so the server entrypoint can write + /opt/lucebox-hub/HOST_INFO without re-probing inside the container + (where /proc and nvidia-smi see the container's view, not the + rig's). See entrypoint.sh::write_host_info and http_server.cpp's + /props.host block. + """ + out: list[tuple[str, str]] = [] + for key, value in sorted(os.environ.items()): + if key.startswith("LUCEBOX_HOST_"): + out.append((key, value)) + return out + + +def _resolve_model_files(cfg: Config) -> tuple[str, str, str]: + """Return (target_file, draft_file, draft_dir) for DFLASH_TARGET / DFLASH_DRAFT. + + Resolution order — first non-empty wins per field: + 1. cfg.model.target_file / draft_file (explicit override in config.toml) + 2. PRESETS[cfg.model.preset].target_file / draft_file / speculator_dir (registry) + 3. "" (entrypoint autodetect path runs unchanged). + + ``draft_dir`` is a directory name under ``models/draft/`` holding a + safetensors speculator (e.g. ``laguna-xs2-speculator``). It is only set + when the preset declares one AND the directory exists on disk; otherwise + it is empty. When non-empty, docker_run_spec uses it as DFLASH_DRAFT + (a directory path) instead of the GGUF-file path, allowing the entrypoint + to discover the safetensors file inside it. + + Imported lazily to avoid the lucebox.types ↔ lucebox.download circular + import that surfaces when this module is imported from ``__init__``. + """ + target = cfg.model.target_file + draft = cfg.model.draft_file + draft_dir = "" + if (not target or not draft) and cfg.model.preset: + from lucebox.download import PRESETS + + pres = PRESETS.get(cfg.model.preset) + if pres is not None: + if not target: + target = pres.target_file + if not draft and pres.has_draft and pres.draft_file: + draft = pres.draft_file + if not draft and pres.speculator_dir: + spec_path = cfg.models_dir / "draft" / pres.speculator_dir + if spec_path.is_dir(): + draft_dir = pres.speculator_dir + return target, draft, draft_dir + + +def _runtime_volumes(cfg: Config) -> tuple[tuple[str, str], ...]: + """Mount models plus $HOME so absolute symlink targets remain valid.""" + home = str(Path.home()) + models = str(cfg.models_dir) + volumes = [(models, "/opt/lucebox-hub/server/models")] + if home != models: + volumes.append((home, home)) + return tuple(volumes) + + +@dataclass(frozen=True, slots=True) +class DockerRunSpec: + """Pre-render of a docker-run command. Render via `argv()` or `printable()`.""" + + image: str + name: str + gpus: bool = True + detach: bool = False + remove: bool = True + port_publish: tuple[int, int] | None = None # (host, container) + volumes: tuple[tuple[str, str], ...] = () + env: tuple[tuple[str, str], ...] = () + entrypoint_args: tuple[str, ...] = () + extra: tuple[str, ...] = () + + def argv(self) -> list[str]: + out = ["docker", "run"] + if self.remove: + out.append("--rm") + if self.detach: + out.append("-d") + out += ["--name", self.name] + if self.gpus: + out += ["--gpus", "all"] + if self.port_publish is not None: + host, container = self.port_publish + out += ["-p", f"{host}:{container}"] + for host_path, container_path in self.volumes: + out += ["-v", f"{host_path}:{container_path}"] + for k, v in self.env: + out += ["-e", f"{k}={v}"] + out += list(self.extra) + out.append(self.image) + out += list(self.entrypoint_args) + return out + + def printable(self) -> str: + """Human-readable, one-flag-per-line docker run. Copy-pasteable.""" + argv = self.argv() + if not argv: + return "" + out = argv[0] + i = 1 + while i < len(argv): + tok = argv[i] + out += " \\\n " + tok + # Glue value-taking flags onto the same line. + if tok in { + "-p", + "-v", + "-e", + "--name", + "--gpus", + "--env", + "--volume", + "--publish", + "--entrypoint", + } and i + 1 < len(argv): + i += 1 + out += " " + shlex.quote(argv[i]) + i += 1 + return out + + +# ── server argv from Config ──────────────────────────────────────────────── + + +def server_run_spec(cfg: Config) -> DockerRunSpec: + """Long-running OpenAI-compatible server. Foreground (systemd manages + lifecycle), --gpus all, models bind-mounted, DFLASH_* propagated. + """ + # LUCEBOX_HOST_* first so they ride out front in the rendered argv, + # making it obvious in `print-run` output what host facts get forwarded. + env: list[tuple[str, str]] = list(_host_facts_env()) + env += [ + ("DFLASH_BUDGET", str(cfg.dflash.budget)), + ("DFLASH_MAX_CTX", str(cfg.dflash.max_ctx)), + ("DFLASH_PREFIX_CACHE_SLOTS", str(cfg.dflash.prefix_cache_slots)), + ("DFLASH_PREFILL_CACHE_SLOTS", str(cfg.dflash.prefill_cache_slots)), + ("DFLASH_THINK_MAX", str(cfg.dflash.think_max)), + ("DFLASH_PORT", "8080"), + ] + # Resolve target/draft GGUFs in priority order: + # 1. cfg.model.target_file / draft_file (explicit override in config.toml) + # 2. PRESETS[cfg.model.preset].target_file / draft_file / speculator_dir (registry) + # 3. unset — entrypoint's autodetect path runs unchanged. + # Container view of the models dir is /opt/lucebox-hub/server/models + # (see _runtime_volumes); the entrypoint reads DFLASH_TARGET / DFLASH_DRAFT. + # draft_dir is a subdirectory of models/draft/ holding a safetensors speculator; + # it takes effect only when draft_file is empty and the directory exists on disk. + target_file, draft_file, draft_dir = _resolve_model_files(cfg) + if target_file: + env.append(("DFLASH_TARGET", f"/opt/lucebox-hub/server/models/{target_file}")) + if draft_file: + env.append(("DFLASH_DRAFT", f"/opt/lucebox-hub/server/models/draft/{draft_file}")) + elif draft_dir: + env.append(("DFLASH_DRAFT", f"/opt/lucebox-hub/server/models/draft/{draft_dir}")) + if cfg.dflash.lazy: + env.append(("DFLASH_LAZY", "1")) + if cfg.dflash.cache_type_k: + env.append(("DFLASH_CACHE_TYPE_K", cfg.dflash.cache_type_k)) + if cfg.dflash.cache_type_v: + env.append(("DFLASH_CACHE_TYPE_V", cfg.dflash.cache_type_v)) + if cfg.dflash.prefill_mode != "off": + env += [ + ("DFLASH_PREFILL_MODE", cfg.dflash.prefill_mode), + ("DFLASH_PREFILL_KEEP", str(cfg.dflash.prefill_keep_ratio)), + ("DFLASH_PREFILL_THRESHOLD", str(cfg.dflash.prefill_threshold)), + ] + if cfg.dflash.prefill_drafter: + env.append(("DFLASH_PREFILL_DRAFTER", cfg.dflash.prefill_drafter)) + # fa_window=0 is the server's own default (full attention); only emit + # the env when the operator has selected a sparse decode window. The + # entrypoint mirrors this guard so an unset env reproduces the + # server's stock behavior. + if cfg.dflash.fa_window > 0: + env.append(("DFLASH_FA_WINDOW", str(cfg.dflash.fa_window))) + # Soft-close ratio: 0.0 is server-side disabled (byte-identical + # to pre-PR-#326 behavior). Emit only when nonzero to keep the + # docker env minimal and mirror the entrypoint's `case` guard. + if cfg.dflash.think_soft_close_min_ratio > 0.0: + env.append(( + "DFLASH_THINK_SOFT_CLOSE_MIN_RATIO", + f"{cfg.dflash.think_soft_close_min_ratio:g}", + )) + if cfg.dflash.debug_thinking_logits: + env.append(("DFLASH_DEBUG_THINKING_LOGITS", "1")) + + return DockerRunSpec( + image=f"{cfg.image}:{cfg.variant}", + name=cfg.container_name, + gpus=True, + remove=True, + detach=False, + port_publish=(cfg.port, 8080), + volumes=_runtime_volumes(cfg), + env=tuple(env), + ) + + +# ── subprocess helpers ───────────────────────────────────────────────────── + + +def docker_pull(image_tag: str) -> int: + """Pull an image, streaming progress. Returns docker's exit code.""" + return subprocess.call(["docker", "pull", image_tag]) diff --git a/lucebox/src/lucebox/download.py b/lucebox/src/lucebox/download.py new file mode 100644 index 000000000..df1d5c96f --- /dev/null +++ b/lucebox/src/lucebox/download.py @@ -0,0 +1,515 @@ +"""Model download orchestration. + +Runs *inside* the orchestrator container. Uses `huggingface_hub` directly +(no subprocess) so we can: + + * drive a Rich progress bar based on real byte counts (the previous + `uvx hf download` subprocess produced no visible progress inside the + container — hf-xet's TTY detection misfires there), + * verify each candidate file's size and sha256 against the repo + metadata BEFORE downloading, so a re-run on a host that already has + the target GGUF (e.g. previous download into the same models_dir) + skips the multi-GB fetch entirely. + +The :data:`PRESETS` registry encodes the canonical (target_repo, +target_file, draft_repo, draft_file) tuple per model — selectable via +``lucebox models download ``. ``DEFAULT_PRESET`` stays pinned to +Qwen3.6-27B for back-compat with callers that pre-date the registry. +Drafts are optional: presets that have no published DFlash draft +(e.g. Laguna's speculator is safetensors, not GGUF) carry +``draft_repo=None`` and run target-only. +""" + +from __future__ import annotations + +import hashlib +import os +import threading +import time +from dataclasses import dataclass +from pathlib import Path + +# hf-xet (huggingface_hub ≥ 1.16) streams the entire file in one final +# burst — the polling-based progress bar sits at 0% for ~14 minutes +# then snaps to 100% on a 17 GB GGUF. Force the chunked Python +# downloader instead so bytes grow continuously and the Rich bar tracks +# reality. Set before importing hf_hub_download so the import picks +# the env up. `setdefault` lets a user override on the command line. +os.environ.setdefault("HF_HUB_DISABLE_XET", "1") + +from huggingface_hub import HfApi, hf_hub_download # noqa: E402 +from huggingface_hub._local_folder import get_local_download_paths # noqa: E402 +from rich.console import Console +from rich.progress import ( + BarColumn, + DownloadColumn, + Progress, + TextColumn, + TimeRemainingColumn, + TransferSpeedColumn, +) + +from lucebox.types import Config, HostFacts + + +@dataclass(frozen=True, slots=True) +class ModelPreset: + """Canonical (target, draft) repo+filename pair for a supported model. + + ``draft_repo`` and ``draft_file`` may both be ``None`` for models + where no GGUF DFlash draft is published (e.g. Laguna's safetensors + speculator). In that case the entrypoint runs target-only — DFlash + speculative decoding is disabled but the server still works. + + ``speculator_dir`` names a directory under ``models/draft/`` that holds + a safetensors-format speculator (e.g. ``model.safetensors``). When + present on disk the server launch sets ``DFLASH_DRAFT`` to that + directory; absent, the server runs target-only. Unlike ``draft_file`` + (which marks the preset as incomplete when missing), ``speculator_dir`` + is optional supplementary hardware and doesn't affect installed_status. + """ + + name: str + target_repo: str + target_file: str + draft_repo: str | None + draft_file: str | None + approx_total_gb: int + description: str = "" + speculator_dir: str | None = None + + @property + def has_draft(self) -> bool: + return bool(self.draft_repo and self.draft_file) + + +# Registry of supported models. Keyed by preset name; the CLI surface +# exposes these via `lucebox models download ` and the +# `lucebox models list` table. The values come straight from the model +# cards under share/model_cards/ — keep them in sync. +PRESETS: dict[str, ModelPreset] = { + "qwen3.6-27b": ModelPreset( + name="qwen3.6-27b", + target_repo="unsloth/Qwen3.6-27B-GGUF", + target_file="Qwen3.6-27B-Q4_K_M.gguf", + draft_repo="spiritbuun/Qwen3.6-27B-DFlash-GGUF", + draft_file="dflash-draft-3.6-q4_k_m.gguf", + approx_total_gb=17, + description="Qwen3.6 27B dense (Q4_K_M) + Qwen3.6 DFlash draft. Lucebox default.", + ), + "gemma-4-26b": ModelPreset( + name="gemma-4-26b", + target_repo="bartowski/google_gemma-4-26B-A4B-it-GGUF", + target_file="google_gemma-4-26B-A4B-it-Q4_K_M.gguf", + draft_repo="Lucebox/gemma-4-26B-A4B-it-DFlash-GGUF", + draft_file="gemma-4-26B-A4B-it-DFlash-q8_0.gguf", + approx_total_gb=18, + description="Gemma 4 26B-A4B IT MoE (Q4_K_M) + Lucebox DFlash q8_0 draft.", + ), + "gemma-4-31b": ModelPreset( + name="gemma-4-31b", + target_repo="bartowski/google_gemma-4-31B-it-GGUF", + target_file="google_gemma-4-31B-it-Q4_K_M.gguf", + draft_repo="Lucebox/gemma-4-31B-it-DFlash-GGUF", + draft_file="gemma-4-31B-it-DFlash-q8_0.gguf", + approx_total_gb=21, + description="Gemma 4 31B IT dense (Q4_K_M) + Lucebox DFlash q8_0 draft.", + ), + "laguna-xs.2": ModelPreset( + name="laguna-xs.2", + target_repo="Lucebox/Laguna-XS.2-GGUF", + target_file="laguna-xs2-Q4_K_M.gguf", + # Laguna's DFlash speculator is safetensors-format + # (poolside/Laguna-XS.2-speculator.dflash), downloaded manually + # into models/draft/laguna-xs2-speculator/. The download command + # doesn't fetch it automatically — it's opt-in. When present, + # speculator_dir wires it into DFLASH_DRAFT at server launch. + draft_repo=None, + draft_file=None, + speculator_dir="laguna-xs2-speculator", + approx_total_gb=20, + description=( + "Laguna-XS.2 MoE code model (Q4_K_M). " + "DFlash safetensors speculator in draft/laguna-xs2-speculator/ " + "is used automatically when present." + ), + ), + "qwen3.6-moe": ModelPreset( + name="qwen3.6-moe", + target_repo="unsloth/Qwen3.6-35B-A3B-GGUF", + # Unsloth's MoE repo publishes both a "UD" (dynamic) and a plain + # Q4_K_M family. Verified 2026-05-28 via HfApi.repo_info: the + # `-UD-Q4_K_M.gguf` variant (22.1 GB) is the canonical Q4_K_M + # release — there is no plain `Q4_K_M.gguf` on the MoE repo. + target_file="Qwen3.6-35B-A3B-UD-Q4_K_M.gguf", + # No DFlash draft GGUF has been published for the MoE variant + # (probed Lucebox/* and spiritbuun/* repos 2026-05-28 — none + # exist). Target-only, mirroring laguna-xs.2's wiring. The + # lucebox C++ server speaks the `qwen35moe` arch natively + # (server/src/qwen35moe/) so this runs without a draft. + draft_repo=None, + draft_file=None, + approx_total_gb=22, + description=( + "Qwen3.6 35B-A3B MoE (3B active per token), Q4_K_M unsloth " + "dynamic quant. Target-only — no DFlash MoE draft published " + "yet. Uses lucebox's qwen35moe arch backend." + ), + ), +} + +DEFAULT_PRESET = PRESETS["qwen3.6-27b"] + + +def resolve_preset(name: str | None) -> ModelPreset: + """Look up a preset by name, with a friendly error on typos. + + ``None`` (or empty string) resolves to :data:`DEFAULT_PRESET` so + callers and the CLI default both flow through one code path. + """ + if not name: + return DEFAULT_PRESET + if name in PRESETS: + return PRESETS[name] + # Build a suggestion list — show every known preset; the user's + # search space is small (4 entries today) so listing them all is + # cheaper and clearer than a fuzzy-match heuristic. + known = ", ".join(sorted(PRESETS.keys())) + raise KeyError(f"unknown preset {name!r}. Known presets: {known}") + + +def _file_meta(api: HfApi, repo_id: str, filename: str) -> tuple[int, str | None]: + """Return (expected_size, lfs_sha256_or_None) for filename in repo_id.""" + info = api.model_info(repo_id, files_metadata=True) + for sib in info.siblings or []: + if sib.rfilename == filename: + sha = getattr(sib.lfs, "sha256", None) if sib.lfs else None + return int(sib.size or 0), sha + raise FileNotFoundError(f"{filename} not present in repo {repo_id}") + + +def _sha256(path: Path, chunk_mb: int = 16) -> str: + h = hashlib.sha256() + chunk = chunk_mb * 1024 * 1024 + with path.open("rb") as f: + while buf := f.read(chunk): + h.update(buf) + return h.hexdigest() + + +def _local_matches(path: Path, size: int, sha256: str | None, console: Console) -> bool: + """True iff a local file at `path` matches the expected size + sha256. + + Size mismatch shortcircuits (cheap). Sha256 is verified for LFS files + (multi-GB GGUFs always carry one) and skipped when the repo doesn't + expose a hash. Hashing 17 GB takes ~30s on a fast SSD — worth it to + avoid a multi-GB re-download on rate-limited / metered links. + """ + if not path.exists(): + return False + actual_size = path.stat().st_size + if actual_size != size: + console.print( + f" [yellow]✗[/yellow] {path.name} present but size {actual_size:,} != " + f"expected {size:,} — will re-download" + ) + return False + if sha256: + console.print(f" [dim]verifying sha256 of {path.name} ({actual_size / 1e9:.1f} GB)…[/dim]") + actual_sha = _sha256(path) + if actual_sha != sha256: + console.print( + f" [yellow]✗[/yellow] {path.name} sha256 {actual_sha[:12]}… != " + f"expected {sha256[:12]}… — will re-download" + ) + return False + return True + + +def _incomplete_path_candidates(local_dir: Path, filename: str, etag: str | None) -> list[Path]: + """Return likely paths of the partial file currently being written. + + huggingface_hub 1.x (with hf-xet) stages downloads under + ``{local_dir}/.cache/huggingface/download/`` using a *hashed* name — + ``{short_hash(metadata_filename)}.{etag}.incomplete`` — so a naive + ``{filename}.incomplete`` poll never sees any growth and the + progress bar sits at 0 % for the whole multi-GB transfer. + + We get the *exact* expected staging path from + ``get_local_download_paths().incomplete_path(etag)`` when we already + know the LFS sha256 (which acts as the etag for Xet downloads), and + fall back to globbing every ``*.incomplete`` in the staging dir + otherwise. The legacy non-Xet downloader writes a ``.incomplete`` + next to the destination blob in ``~/.cache/huggingface/hub`` — but + when ``local_dir`` is set hf-hub always uses the local staging dir, + so the two candidates above cover every code path we hit. + """ + paths = get_local_download_paths(local_dir, filename) + candidates: list[Path] = [] + if etag: + candidates.append(paths.incomplete_path(etag)) + # Fallback: every .incomplete file in the staging dir. This is what + # rescues us when sha256 is unknown (non-LFS file) or when hf-hub + # changes the etag derivation again in some future release. + candidates.append(paths.metadata_path.parent) # sentinel: glob this dir + return candidates + + +def _current_bytes(target: Path, candidates: list[Path]) -> int: + """Best-effort byte count of the file currently being written.""" + if target.exists(): + try: + return target.stat().st_size + except OSError: + pass + for c in candidates: + if c.is_dir(): + # Glob every .incomplete in the staging dir; return the + # largest (there's typically only one in-flight transfer). + largest = 0 + try: + for p in c.glob("*.incomplete"): + try: + largest = max(largest, p.stat().st_size) + except OSError: + continue + except OSError: + continue + if largest: + return largest + else: + try: + if c.exists(): + return c.stat().st_size + except OSError: + continue + return 0 + + +def _download_with_progress( + repo_id: str, + filename: str, + local_dir: Path, + expected_size: int, + console: Console, + etag: str | None = None, +) -> Path: + """Download a single HF file with a Rich progress bar. + + Runs hf_hub_download in a worker thread; the main thread polls the + growing file size and updates the Rich progress bar. The polled + target is computed via ``get_local_download_paths`` so we hit the + actual hf-xet staging path (a hashed filename under + ``.cache/huggingface/download/``), not a guess. + """ + local_dir.mkdir(parents=True, exist_ok=True) + target = local_dir / filename + candidates = _incomplete_path_candidates(local_dir, filename, etag) + + result: list[str | None] = [None] + error: list[BaseException | None] = [None] + + def _worker() -> None: + try: + result[0] = hf_hub_download( + repo_id=repo_id, + filename=filename, + local_dir=str(local_dir), + ) + except BaseException as exc: # propagate to main thread + error[0] = exc + + t = threading.Thread(target=_worker, daemon=True) + t.start() + + with Progress( + TextColumn("[cyan]{task.description}"), + BarColumn(bar_width=40), + DownloadColumn(), + TransferSpeedColumn(), + TimeRemainingColumn(), + console=console, + transient=False, + ) as progress: + task = progress.add_task(filename, total=expected_size or 1) + while t.is_alive(): + current = _current_bytes(target, candidates) + # Always tick the bar — even at 0 bytes — so Rich repaints + # the spinner/ETA and the user sees the UI is alive within + # the first poll tick rather than a blank "Downloading…" line. + progress.update(task, completed=min(current, expected_size or current or 1)) + time.sleep(0.5) + # Final tick after the worker finishes so the bar paints 100%. + if target.exists(): + progress.update(task, completed=target.stat().st_size) + + t.join(timeout=5) + if error[0] is not None: + raise error[0] + if result[0] is None: + raise RuntimeError(f"hf_hub_download returned no path for {filename}") + return Path(result[0]) + + +def _fetch( + api: HfApi, + repo_id: str, + filename: str, + local_dir: Path, + console: Console, +) -> Path: + """Verify-or-download a single file. Skips when the local copy matches.""" + size, sha = _file_meta(api, repo_id, filename) + target = local_dir / filename + if _local_matches(target, size, sha, console): + console.print(f" [green]✓[/green] {filename} already present (size + sha256 match)") + return target + # `sha` doubles as the etag for hf-xet's staging path + # ({local_dir}/.cache/huggingface/download/{hash}.{etag}.incomplete); + # passing it through is what makes the Rich progress bar see real + # byte counts during the multi-GB transfer. + return _download_with_progress(repo_id, filename, local_dir, size, console, etag=sha) + + +def download_preset(cfg: Config, preset: ModelPreset | None = None) -> int: + """Fetch the target GGUF + (optional) DFlash draft into cfg.models_dir. + + Returns 0 on success, non-zero on failure. Verifies each file's size + and (LFS) sha256 against the repo metadata before downloading, so a + repeat run with the files already on disk is a no-op + sha256 walk. + + ``preset=None`` resolves to :data:`DEFAULT_PRESET` for back-compat; + presets with ``has_draft=False`` (e.g. Laguna) skip the draft fetch + entirely and let the server run target-only. + """ + preset = preset or DEFAULT_PRESET + console = Console() + api = HfApi() + models = cfg.models_dir + models.mkdir(parents=True, exist_ok=True) + draft = models / "draft" + draft.mkdir(exist_ok=True) + + try: + _fetch(api, preset.target_repo, preset.target_file, models, console) + if preset.has_draft: + # Narrow the optionals for the type-checker — has_draft is + # exactly the predicate that proves these aren't None. + assert preset.draft_repo is not None and preset.draft_file is not None + _fetch(api, preset.draft_repo, preset.draft_file, draft, console) + else: + console.print( + f" [dim]no DFlash draft published for {preset.name} — running target-only[/dim]" + ) + except Exception as exc: + console.print(f"[red]download failed:[/red] {exc}") + return 1 + return 0 + + +def _local_target_path(cfg: Config, preset: ModelPreset) -> Path: + return cfg.models_dir / preset.target_file + + +def _local_draft_path(cfg: Config, preset: ModelPreset) -> Path | None: + if not (preset.has_draft and preset.draft_file): + return None + return cfg.models_dir / "draft" / preset.draft_file + + +def installed_status(cfg: Config, preset: ModelPreset) -> str: + """Return ``"installed"`` / ``"partial"`` / ``"absent"`` for a preset. + + Size-only — doesn't hash. ``"installed"`` requires the target (and + draft when one is published) to exist on disk; ``"partial"`` means + at least one of the two is present but the set is incomplete. + """ + target_exists = _local_target_path(cfg, preset).exists() + draft_path = _local_draft_path(cfg, preset) + if draft_path is None: + return "installed" if target_exists else "absent" + draft_exists = draft_path.exists() + if target_exists and draft_exists: + return "installed" + if target_exists or draft_exists: + return "partial" + return "absent" + + +def installed_size_gb(cfg: Config, preset: ModelPreset) -> float: + """Sum of on-disk byte sizes for the preset's files, in GB (binary 1e9).""" + total = 0 + target = _local_target_path(cfg, preset) + if target.exists(): + try: + total += target.stat().st_size + except OSError: + pass + draft = _local_draft_path(cfg, preset) + if draft is not None and draft.exists(): + try: + total += draft.stat().st_size + except OSError: + pass + return total / 1e9 + + +def installed_presets(cfg: Config) -> list[ModelPreset]: + """Return every preset whose files are currently present in cfg.models_dir. + + "Present" follows ``installed_status`` — fully installed only. + Partial states (target without draft, etc.) are excluded so the + default ``lucebox models`` view stays uncluttered. + """ + out: list[ModelPreset] = [] + for name in sorted(PRESETS): + pres = PRESETS[name] + if installed_status(cfg, pres) == "installed": + out.append(pres) + return out + + +def status(cfg: Config, preset: ModelPreset | None = None) -> dict[str, bool]: + """Quick presence check — what's already on disk? Size-only, no sha256. + + For presets without a published DFlash draft, ``draft_present`` is + reported as ``True`` (nothing to fetch → nothing missing). That + keeps the "all present, nothing to do" UX path uniform whether or + not a draft exists. + """ + preset = preset or DEFAULT_PRESET + api = HfApi() + out: dict[str, bool] = {} + try: + size, _ = _file_meta(api, preset.target_repo, preset.target_file) + local = cfg.models_dir / preset.target_file + out["target_present"] = local.exists() and local.stat().st_size == size + except Exception: + out["target_present"] = False + + if preset.has_draft: + assert preset.draft_repo is not None and preset.draft_file is not None + try: + size, _ = _file_meta(api, preset.draft_repo, preset.draft_file) + local = cfg.models_dir / "draft" / preset.draft_file + out["draft_present"] = local.exists() and local.stat().st_size == size + except Exception: + out["draft_present"] = False + else: + out["draft_present"] = True + return out + + +def recommend_preset(host: HostFacts) -> str | None: + """Pick a default preset for first-run install. None = ask the user. + + Tiers follow the model size catalog: 22 GB+ → Qwen3.6-27B (the + Lucebox default), 16-21 GB → Laguna-XS.2 (small target-only). Below + 16 GB we punt and let the user pick explicitly — the registered + presets all need at least 16 GB to run usefully. + """ + if host.vram_gb >= 22: + return "qwen3.6-27b" + if host.vram_gb >= 16: + return "laguna-xs.2" + return None diff --git a/lucebox/src/lucebox/host_check.py b/lucebox/src/lucebox/host_check.py new file mode 100644 index 000000000..2ce8d3889 --- /dev/null +++ b/lucebox/src/lucebox/host_check.py @@ -0,0 +1,232 @@ +"""Readiness check: aggregate HostFacts (provided by lucebox.sh) with the +docker-daemon checks we can do from inside the container via the mounted +socket. Prints a status report and returns an aggregate severity. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Literal + +from rich.console import Console + +from lucebox.types import HostFacts + +Severity = Literal["ok", "warn", "fail"] +_SEVERITY_ORDER: dict[Severity, int] = {"ok": 0, "warn": 1, "fail": 2} + + +@dataclass(frozen=True, slots=True) +class CheckResult: + name: str + severity: Severity + message: str + hint: str | None = None + + +def run_checks(host: HostFacts) -> list[CheckResult]: + return [ + _check_docker(host), + _check_nvidia_driver(host), + _check_ctk(host), + _check_ram(host), + _check_vram(host), + _check_systemd(host), + ] + + +def _check_docker(host: HostFacts) -> CheckResult: + if not host.has_docker: + return CheckResult( + "docker", + "fail", + "docker daemon unreachable", + "sudo systemctl start docker, or add your user to the 'docker' group", + ) + return CheckResult("docker", "ok", f"daemon reachable ({host.docker_version})") + + +def _check_nvidia_driver(host: HostFacts) -> CheckResult: + if host.gpu_vendor != "nvidia": + if host.gpu_vendor == "amd": + return CheckResult( + "gpu", + "fail", + "AMD GPU detected — prebuilt images are NVIDIA-only", + "Build dflash from source with HIP; see dflash/README.md", + ) + return CheckResult("gpu", "fail", "no NVIDIA GPU detected") + if not host.driver_version: + return CheckResult( + "driver", + "warn", + "nvidia-smi present but NVML query failed (likely driver/library mismatch)", + "reboot, or reinstall the matching NVIDIA driver", + ) + if host.driver_major < 525: + return CheckResult( + "driver", + "fail", + f"driver r{host.driver_major} too old (need r525+ for cuda12)", + "upgrade the NVIDIA driver", + ) + return CheckResult("driver", "ok", f"nvidia r{host.driver_major} ({host.driver_version})") + + +def _check_ctk(host: HostFacts) -> CheckResult: + match host.ctk: + case "runtime": + return CheckResult("ctk", "ok", "NVIDIA Container Toolkit registered as docker runtime") + case "cdi": + return CheckResult("ctk", "ok", "NVIDIA Container Toolkit available via CDI") + case "installed-unwired": + return CheckResult( + "ctk", + "warn", + "NVIDIA Container Toolkit installed but not wired into docker", + "sudo nvidia-ctk runtime configure --runtime=docker && " + "sudo systemctl restart docker", + ) + case _: + return CheckResult( + "ctk", + "fail", + "NVIDIA Container Toolkit not installed", + "https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html", + ) + + +def _check_ram(host: HostFacts) -> CheckResult: + if host.ram_gb == 0: + return CheckResult("ram", "warn", "RAM unknown") + if host.ram_gb < 16: + return CheckResult("ram", "warn", f"{host.ram_gb} GB RAM — model load may swap") + return CheckResult("ram", "ok", f"{host.ram_gb} GB RAM") + + +def _check_vram(host: HostFacts) -> CheckResult: + if host.vram_gb == 0: + return CheckResult("vram", "warn", "VRAM unknown") + if host.vram_gb < 12: + return CheckResult( + "vram", + "fail", + f"VRAM {host.vram_gb} GB < 12 GB — 27B target won't fit", + "use a smaller model preset or larger GPU", + ) + if host.vram_gb < 22: + return CheckResult( + "vram", + "warn", + f"VRAM {host.vram_gb} GB — 27B fits but max_ctx will be capped near 32K", + ) + return CheckResult("vram", "ok", f"VRAM {host.vram_gb} GB ({host.gpu_name})") + + +def _check_systemd(host: HostFacts) -> CheckResult: + if not host.has_systemd: + return CheckResult( + "systemd", + "warn", + "user systemd not available", + "WSL: enable systemd in /etc/wsl.conf; otherwise 'lucebox serve' " + "still works in the foreground", + ) + return CheckResult("systemd", "ok", "user systemd available") + + +def aggregate(results: list[CheckResult]) -> Severity: + worst: Severity = "ok" + for r in results: + if _SEVERITY_ORDER[r.severity] > _SEVERITY_ORDER[worst]: + worst = r.severity + return worst + + +def render(console: Console, host: HostFacts, results: list[CheckResult]) -> Severity: + """Print a status block, return the worst severity.""" + summary = f"[bold]Host:[/bold] {host.nproc} CPUs · {host.ram_gb} GB RAM" + if host.gpu_vendor == "nvidia" and host.gpu_name: + summary += f" · {host.gpu_name} · {host.vram_gb} GB VRAM" + ( + f" (sm_{host.gpu_sm})" if host.gpu_sm else "" + ) + if host.is_wsl: + summary += " · WSL2" + console.print(summary) + console.print() + + sev_style = { + "ok": "[green]OK[/green]", + "warn": "[yellow]WARN[/yellow]", + "fail": "[red]FAIL[/red]", + } + for r in results: + console.print(f" {sev_style[r.severity]:<22} {r.name:<8} {r.message}") + if r.hint: + console.print(f" {'':<22} {'':<8} [dim]{r.hint}[/dim]") + + render_host_facts(console) + + worst = aggregate(results) + console.print() + if worst == "ok": + console.print("[green]All checks passed.[/green]") + elif worst == "warn": + console.print("[yellow]Checks passed with warnings.[/yellow]") + else: + console.print( + "[red]Critical checks failed — fix the issues above before 'lucebox start'.[/red]" + ) + return worst + + +def render_host_facts(console: Console) -> None: + """Print a pretty 'Host facts' section sourced from LUCEBOX_HOST_*. + + Same data that ends up in /opt/lucebox-hub/HOST_INFO inside the + container — printed here so the operator can sanity-check the + rig classification BEFORE starting a long bench run, and so the + CI exit-code gate (the pass/fail checks above) stays orthogonal + to the informational host facts. + + Reads from the same LUCEBOX_HOST_* env the host wrapper exports + (see lucebox.sh::probe_host). Quiet — emits the section header + even when most facts are unset, since "no host facts probed at + all" is itself a useful signal. + """ + console.print() + console.print("[bold]Host facts[/bold] (LUCEBOX_HOST_*, surfaced as /props.host)") + facts = [ + ("os", os.environ.get("LUCEBOX_HOST_OS_PRETTY", "")), + ("kernel", os.environ.get("LUCEBOX_HOST_KERNEL", "")), + ("wsl_version", os.environ.get("LUCEBOX_HOST_WSL_VERSION", "")), + ("docker", os.environ.get("LUCEBOX_HOST_DOCKER_VERSION", "")), + ("nvidia_driver", os.environ.get("LUCEBOX_HOST_DRIVER_VERSION", "")), + ("nvidia_ctk", os.environ.get("LUCEBOX_HOST_NVIDIA_CTK_VERSION", "")), + ("cpu", os.environ.get("LUCEBOX_HOST_CPU_MODEL", "")), + ("cuda_visible_devices", os.environ.get("LUCEBOX_HOST_CUDA_VISIBLE_DEVICES", "")), + ] + for key, value in facts: + display = value if value else "[dim](unset)[/dim]" + console.print(f" {key:<22} {display}") + + # Multi-GPU table — one line per device. LUCEBOX_HOST_GPU_LIST_CSV + # carries the verbatim nvidia-smi CSV the host wrapper probed. + csv = os.environ.get("LUCEBOX_HOST_GPU_LIST_CSV", "") + if csv: + console.print(" gpus:") + for line in csv.splitlines(): + line = line.strip() + if not line: + continue + parts = [c.strip() for c in line.split(",")] + if len(parts) >= 7: + idx, _uuid, _pci, name, sm, mem, plimit = parts[:7] + console.print( + f" [{idx}] {name} (sm_{sm}, {mem}, {plimit})" + ) + else: + console.print(f" {line}") + else: + console.print(" gpus [dim](none — nvidia-smi unavailable)[/dim]") diff --git a/lucebox/src/lucebox/host_facts.py b/lucebox/src/lucebox/host_facts.py new file mode 100644 index 000000000..5deb6721a --- /dev/null +++ b/lucebox/src/lucebox/host_facts.py @@ -0,0 +1,58 @@ +"""Read HostFacts from the LUCEBOX_HOST_* env vars that lucebox.sh exports. + +We deliberately don't try to detect anything ourselves on the Python side — +inside the container, /proc/meminfo reports the container's view, not the +host's, and nvidia-smi may or may not be available depending on how the +caller invoked us. The host wrapper is the only thing that can see the +truth, and it's already paid for the probe. +""" + +from __future__ import annotations + +import os +from typing import cast + +from lucebox.types import CtkStatus, GpuVendor, HostFacts + + +def _env_int(key: str, default: int = 0) -> int: + raw = os.environ.get(key, "").strip() + if not raw: + return default + try: + return int(raw) + except ValueError: + return default + + +def _env_bool(key: str) -> bool: + return os.environ.get(key, "").strip() in {"1", "true", "yes", "on"} + + +def from_env() -> HostFacts: + vendor: GpuVendor = "none" + raw_vendor = os.environ.get("LUCEBOX_HOST_GPU_VENDOR", "none") + if raw_vendor in {"nvidia", "amd", "none"}: + vendor = cast(GpuVendor, raw_vendor) + + ctk: CtkStatus = "none" + raw_ctk = os.environ.get("LUCEBOX_HOST_HAS_CTK", "none") + if raw_ctk in {"runtime", "cdi", "installed-unwired", "none"}: + ctk = cast(CtkStatus, raw_ctk) + + return HostFacts( + nproc=_env_int("LUCEBOX_HOST_NPROC"), + ram_gb=_env_int("LUCEBOX_HOST_RAM_GB"), + gpu_vendor=vendor, + gpu_name=os.environ.get("LUCEBOX_HOST_GPU_NAME", ""), + gpu_count=_env_int("LUCEBOX_HOST_GPU_COUNT"), + vram_gb=_env_int("LUCEBOX_HOST_VRAM_GB"), + gpu_sm=os.environ.get("LUCEBOX_HOST_GPU_SM", ""), + driver_version=os.environ.get("LUCEBOX_HOST_DRIVER_VERSION", ""), + driver_major=_env_int("LUCEBOX_HOST_DRIVER_MAJOR"), + has_systemd=_env_bool("LUCEBOX_HOST_HAS_SYSTEMD"), + is_wsl=_env_bool("LUCEBOX_HOST_IS_WSL"), + has_docker=_env_bool("LUCEBOX_HOST_HAS_DOCKER"), + docker_version=os.environ.get("LUCEBOX_HOST_DOCKER_VERSION", ""), + ctk=ctk, + ) diff --git a/lucebox/src/lucebox/py.typed b/lucebox/src/lucebox/py.typed new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/lucebox/src/lucebox/py.typed @@ -0,0 +1 @@ + diff --git a/lucebox/src/lucebox/types.py b/lucebox/src/lucebox/types.py new file mode 100644 index 000000000..e1d3620d7 --- /dev/null +++ b/lucebox/src/lucebox/types.py @@ -0,0 +1,140 @@ +"""Shared dataclasses passed between modules. + +HostFacts is populated from the LUCEBOX_HOST_* env vars set by lucebox.sh. +Config is what we serialize to/from .lucebox/config.toml. Both are frozen so +mistakes (e.g. mutating a config after autotune wrote it) fail loudly. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Literal + +Variant = str +CtkStatus = Literal["runtime", "cdi", "installed-unwired", "none"] + + +def default_models_dir() -> Path: + """Resolve the default models directory under the XDG Base Directory spec. + + $XDG_DATA_HOME (default ~/.local/share) is the conventional location for + user-specific data files on Linux + macOS. Lucebox nests its model store + under that so downloads live alongside other per-user app data instead + of cluttering $HOME directly. The host wrapper bind-mounts this path + into the container so paths line up in and out of the image. + """ + base = os.environ.get("XDG_DATA_HOME") or str(Path.home() / ".local" / "share") + return Path(base) / "lucebox" / "models" + + +GpuVendor = Literal["nvidia", "amd", "none"] + + +@dataclass(frozen=True, slots=True) +class HostFacts: + """Probed once by lucebox.sh, passed in via env vars. Single source of + truth on the Python side — we never reprobe (we can't see host /proc).""" + + nproc: int = 0 + ram_gb: int = 0 + gpu_vendor: GpuVendor = "none" + gpu_name: str = "" + gpu_count: int = 0 + vram_gb: int = 0 + gpu_sm: str = "" # e.g. "120" — matches docker-bake arch lists + driver_version: str = "" # e.g. "595.71.05" + driver_major: int = 0 + has_systemd: bool = False + is_wsl: bool = False + has_docker: bool = False + docker_version: str = "" + ctk: CtkStatus = "none" + + +@dataclass(frozen=True, slots=True) +class DflashRuntime: + """The DFLASH_* knobs as typed values. Serialized under [dflash] in TOML + and emitted as -e DFLASH_FOO=bar args to docker run. + + The 11 fields below (budget through prefill_drafter) form the strict + allowlist mirrored by lucebench's snapshot config.json — keep both + in lockstep. ``think_max`` is a separate phase-1 thinking cap that + isn't part of the runtime snapshot allowlist (it's per-request, not + per-server). + """ + + budget: int = 22 + max_ctx: int = 16384 + lazy: bool = False + prefix_cache_slots: int = 0 + prefill_cache_slots: int = 0 + cache_type_k: str = "" + cache_type_v: str = "" + prefill_mode: Literal["off", "auto", "always"] = "off" + prefill_keep_ratio: float = 0.05 + prefill_threshold: int = 32000 + prefill_drafter: str = "" + # Phase-1 (thinking) cap when a request opts into thinking. Default mirrors + # antirez/ds4 ds4_eval.c: think_max_tokens = max_tokens - hard_limit_reply + # budget = 16000 - 512 = 15488. The server's own hardcoded default is 10000. + think_max: int = 15488 + # Flash-attention sliding-window on full-attention layers. 0 = full + # attention (server default). On gemma4's hybrid iSWA the full-attn + # layers grow KV linearly with max_ctx; a sparse fa_window keeps + # decode compute bounded on long prompts without changing the KV + # footprint. Q: passed through to the server's `--fa-window ` + # flag (see server/src/server/server_main.cpp). + fa_window: int = 0 + # Soft-close thinking termination dial (PR #326 in lucebox-hub). + # Lets the AR loop force early when the close-token logit + # comes within this probability ratio of the chosen-token logit. + # Range [0.0, 1.0]; 0.0 = disabled (byte-identical to pre-change + # behaviour). 0.5 = close when close-token prob >= 0.5 * chosen-token + # prob; 0.9 = aggressive. Qwen3.5/3.6 AR path only in v1. Surfaced + # to the server via DFLASH_THINK_SOFT_CLOSE_MIN_RATIO → + # --think-soft-close-min-ratio. + think_soft_close_min_ratio: float = 0.0 + # Diagnostic: when True, surface --debug-thinking-logits to the + # server CLI via DFLASH_DEBUG_THINKING_LOGITS=1, producing one + # stderr line per thinking AR step recording the close-vs-chosen + # logit gap. Used to fit a sliding-ratio curve from real trajectory + # data. Heavy stderr (one line per thinking token across all + # in-flight requests); leave off in production. + debug_thinking_logits: bool = False + + +@dataclass(frozen=True, slots=True) +class ModelMeta: + """Which preset the operator picked at configure/download time. + + Persisted under ``[model]`` in config.toml so `lucebox serve` can + pass ``DFLASH_TARGET=/opt/lucebox-hub/server/models/`` and + ``DFLASH_DRAFT`` for the draft GGUF (when one is published for the + preset). The entrypoint's "multiple candidate GGUFs" branch never + has to guess which one to load. + + ``target_file`` and ``draft_file`` are advanced overrides — when set + they win over the preset's registry default. Empty strings mean + "fall back to the registry value for [model] preset, then to the + entrypoint's autodetect". + """ + + preset: str = "" + target_file: str = "" + draft_file: str = "" + + +@dataclass(frozen=True, slots=True) +class Config: + """The whole config.toml, materialized.""" + + variant: Variant = "cuda12" + image: str = "ghcr.io/luce-org/lucebox-hub" + container_name: str = "lucebox" + port: int = 8080 + models_dir: Path = field(default_factory=default_models_dir) + dflash: DflashRuntime = field(default_factory=DflashRuntime) + host: HostFacts = field(default_factory=HostFacts) + model: ModelMeta = field(default_factory=ModelMeta) diff --git a/lucebox/tests/test_autotune.py b/lucebox/tests/test_autotune.py new file mode 100644 index 000000000..c47c17c9e --- /dev/null +++ b/lucebox/tests/test_autotune.py @@ -0,0 +1,48 @@ +from lucebox.autotune import runtime_from_host +from lucebox.types import HostFacts + + +def test_wsl_24gb_defaults_leave_cuda_headroom() -> None: + runtime = runtime_from_host(HostFacts(vram_gb=24, is_wsl=True)) + + assert runtime.budget == 16 + # Bumped 65536 → 98304 on 2026-05-30 after the gemma4-26b coding- + # agent-loop sweep proved 98K serves 90K-token agentic prompts + # with ~3 GB VRAM headroom and no CUDA VMM failures on the 3090 Ti + # WSL configuration (see + # docs/experiments/gemma4-26b-coding-agent-loop-sweep-2026-05-30.md). + assert runtime.max_ctx == 98304 + # lazy is False because the heuristic path does NOT set prefill_drafter, + # and the C++ server silently ignores --lazy-draft without it. Flipping + # to False makes the host config match runtime behaviour. See the + # `entrypoint.sh` warning emitted when the two are out-of-sync. + assert runtime.lazy is False + assert runtime.prefix_cache_slots == 0 + + +def test_native_24gb_caps_context_below_vmm_failure_boundary() -> None: + runtime = runtime_from_host(HostFacts(vram_gb=24, is_wsl=False)) + + assert runtime.budget == 22 + assert runtime.max_ctx == 98304 + assert runtime.lazy is False # see WSL test above + assert runtime.prefix_cache_slots == 0 + + +def test_no_heuristic_tier_sets_lazy_without_prefill_drafter() -> None: + """Regression for the `--lazy-draft ignored` silent no-op. + + The C++ dflash_server drops `--lazy-draft` unless `--prefill-drafter` + is also passed. The heuristic doesn't set `prefill_drafter`, so any + tier that sets `lazy=True` would produce a host config that doesn't + match what actually ran — exactly the mismatch the sindri decode + sweep tripped over (every docker.stderr contained the warning). + """ + for vram in (0, 8, 16, 24, 40, 80): + for is_wsl in (False, True): + rt = runtime_from_host(HostFacts(vram_gb=vram, is_wsl=is_wsl)) + if rt.lazy: + assert rt.prefill_drafter, ( + f"vram={vram} is_wsl={is_wsl}: lazy=True without " + f"prefill_drafter → silent no-op on the C++ server" + ) diff --git a/lucebox/tests/test_check.py b/lucebox/tests/test_check.py new file mode 100644 index 000000000..3fdd469d9 --- /dev/null +++ b/lucebox/tests/test_check.py @@ -0,0 +1,118 @@ +"""Tests for ``lucebox check`` — readiness report. + +The check command has two surfaces that must stay independent: + + * pass/fail checks → drive the exit code, so the command is usable + as a CI exit-code gate; + * Host facts section → informational, prints the LUCEBOX_HOST_* + convoy that gets baked into /opt/lucebox-hub/HOST_INFO inside + the container. +""" + +from __future__ import annotations + +import pytest +from lucebox.cli import app +from lucebox.types import HostFacts +from rich.console import Console +from typer.testing import CliRunner + +from lucebox import host_check + + +def test_check_prints_host_facts_section(monkeypatch: pytest.MonkeyPatch) -> None: + """`lucebox check` includes a Host facts block sourced from LUCEBOX_HOST_*.""" + monkeypatch.setenv("LUCEBOX_HOST_OS_PRETTY", "Ubuntu 22.04.3 LTS") + monkeypatch.setenv("LUCEBOX_HOST_KERNEL", "6.6.87.2-microsoft-standard-WSL2") + monkeypatch.setenv("LUCEBOX_HOST_WSL_VERSION", "wsl2") + monkeypatch.setenv("LUCEBOX_HOST_DOCKER_VERSION", "29.1.3") + monkeypatch.setenv("LUCEBOX_HOST_DRIVER_VERSION", "596.36") + monkeypatch.setenv("LUCEBOX_HOST_NVIDIA_CTK_VERSION", "1.16.2") + monkeypatch.setenv("LUCEBOX_HOST_CPU_MODEL", "Intel Test CPU") + monkeypatch.setenv( + "LUCEBOX_HOST_GPU_LIST_CSV", + "0, GPU-abc, 00000000:01:00.0, NVIDIA RTX 5090, 12.0, 24576 MiB, 175.00 W", + ) + # Stub HostFacts so the pass/fail checks succeed at least minimally. + # `cli.check` imports `from_env` into its module namespace, so patch + # both names. + def stub() -> HostFacts: + return HostFacts( + nproc=24, + ram_gb=64, + gpu_vendor="nvidia", + gpu_name="NVIDIA RTX 5090", + gpu_count=1, + vram_gb=24, + gpu_sm="120", + driver_version="596.36", + driver_major=596, + has_systemd=True, + is_wsl=True, + has_docker=True, + docker_version="29.1.3", + ctk="runtime", + ) + monkeypatch.setattr("lucebox.host_facts.from_env", stub) + monkeypatch.setattr("lucebox.cli.from_env", stub) + result = CliRunner().invoke(app, ["check"]) + # The pass/fail half of `check` should still exit 0 on this stubbed host. + assert result.exit_code == 0, result.stdout + assert "Host facts" in result.stdout + assert "Ubuntu 22.04.3 LTS" in result.stdout + assert "wsl2" in result.stdout + assert "1.16.2" in result.stdout + assert "Intel Test CPU" in result.stdout + # Multi-GPU table line. + assert "NVIDIA RTX 5090" in result.stdout + + +def test_render_host_facts_unset_env_shows_placeholders( + monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + """All LUCEBOX_HOST_* unset → section still renders with explicit (unset) markers.""" + for k in list(__import__("os").environ): + if k.startswith("LUCEBOX_HOST_"): + monkeypatch.delenv(k, raising=False) + console = Console(force_terminal=False, no_color=True, record=True) + host_check.render_host_facts(console) + text = console.export_text() + assert "Host facts" in text + # Multi-line section renders even when no env was passed in. + assert "(unset)" in text + assert "gpus" in text + + +def test_check_exit_code_independent_of_host_facts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Host facts section must not change the exit-code semantics of check. + + Drives the pass/fail logic through a known-fail HostFacts (no docker) + and asserts the exit code is still 1, regardless of what the Host + facts block prints. + """ + monkeypatch.setenv("LUCEBOX_HOST_OS_PRETTY", "Bare Linux") + def stub() -> HostFacts: + return HostFacts( + nproc=8, + ram_gb=16, + gpu_vendor="nvidia", + gpu_name="X", + gpu_count=1, + vram_gb=24, + gpu_sm="86", + driver_version="555.00", + driver_major=555, + has_systemd=False, + is_wsl=False, + has_docker=False, # → fail + docker_version="", + ctk="none", # also fail + ) + monkeypatch.setattr("lucebox.host_facts.from_env", stub) + monkeypatch.setattr("lucebox.cli.from_env", stub) + result = CliRunner().invoke(app, ["check"]) + assert result.exit_code == 1 + # Host facts block still printed despite the failure. + assert "Host facts" in result.stdout diff --git a/lucebox/tests/test_cli.py b/lucebox/tests/test_cli.py new file mode 100644 index 000000000..f7628e8b7 --- /dev/null +++ b/lucebox/tests/test_cli.py @@ -0,0 +1,102 @@ +"""Tests for the top-level Typer surface.""" + +from __future__ import annotations + +import os + +import pytest +from lucebox.cli import app +from typer.testing import CliRunner + + +def test_config_subcommand_is_registered() -> None: + result = CliRunner().invoke(app, ["config", "--help"]) + assert result.exit_code == 0 + assert "get" in result.output + assert "set" in result.output + assert "unset" in result.output + + +def test_models_subcommand_is_registered() -> None: + result = CliRunner().invoke(app, ["models", "--help"]) + assert result.exit_code == 0 + assert "list" in result.output + assert "download" in result.output + + +@pytest.mark.parametrize( + "verb", + [ + "autotune", + "sweep", + "profile", + "smoke", + "claude", + "codex", + "opencode", + "hermes", + "pi", + "openclaw", + ], +) +def test_deferred_verbs_are_not_registered(verb: str) -> None: + """autotune/sweep, profile/smoke and the client launchers are deferred to + follow-up PRs — this core CLI (launch / serve / install / download) must + not expose them.""" + result = CliRunner().invoke(app, [verb, "--help"]) + assert result.exit_code != 0 + + +def test_core_verbs_present_in_app() -> None: + """The core launch/serve surface stays wired into the Typer command table.""" + registered = { + c.name or (c.callback.__name__ if c.callback else "") + for c in app.registered_commands + } + for verb in ("check", "pull", "print-run", "print-serve-argv", "version"): + assert verb in registered + + +def test_legacy_subcommands_are_removed() -> None: + """`configure` and `download-models` were folded into config/models.""" + cfg = CliRunner().invoke(app, ["configure", "--help"]) + assert cfg.exit_code != 0 + dl = CliRunner().invoke(app, ["download-models", "--help"]) + assert dl.exit_code != 0 + + +def test_server_run_spec_forwards_lucebox_host_env(monkeypatch) -> None: + """server_run_spec carries LUCEBOX_HOST_* from the orchestrator into the server. + + lucebox.sh exports the LUCEBOX_HOST_* convoy before `docker run` on the + orchestrator; the orchestrator inherits them and we forward each one + as ``-e KEY=VALUE`` to the server container so entrypoint.sh's + write_host_info() can populate /opt/lucebox-hub/HOST_INFO. + """ + import lucebox.docker_run as docker_run + from lucebox.config import live_config + + # Scrub any pre-existing LUCEBOX_HOST_* env so the test sees only what we set. + for k in list(os.environ): + if k.startswith("LUCEBOX_HOST_"): + monkeypatch.delenv(k, raising=False) + monkeypatch.setenv("LUCEBOX_HOST_OS_PRETTY", "Ubuntu 22.04.3 LTS") + monkeypatch.setenv("LUCEBOX_HOST_KERNEL", "6.6.87.2-microsoft-standard-WSL2") + monkeypatch.setenv("LUCEBOX_HOST_WSL_VERSION", "wsl2") + monkeypatch.setenv( + "LUCEBOX_HOST_GPU_LIST_CSV", + "0, GPU-x, 00000000:01:00.0, NVIDIA RTX 5090, 12.0, 24576 MiB, 175.00 W", + ) + + cfg = live_config() + spec = docker_run.server_run_spec(cfg) + env_keys = {k for k, _ in spec.env} + assert "LUCEBOX_HOST_OS_PRETTY" in env_keys + assert "LUCEBOX_HOST_KERNEL" in env_keys + assert "LUCEBOX_HOST_WSL_VERSION" in env_keys + assert "LUCEBOX_HOST_GPU_LIST_CSV" in env_keys + # DFLASH_* still present. + assert "DFLASH_BUDGET" in env_keys + # Values surface verbatim. + env_map = dict(spec.env) + assert env_map["LUCEBOX_HOST_OS_PRETTY"] == "Ubuntu 22.04.3 LTS" diff --git a/lucebox/tests/test_config.py b/lucebox/tests/test_config.py new file mode 100644 index 000000000..2ec795a12 --- /dev/null +++ b/lucebox/tests/test_config.py @@ -0,0 +1,210 @@ +"""Tests for the sparse TOML config persistence layer.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from lucebox.config import config_get, config_set, config_unset + +from lucebox import config + + +def test_legacy_env_migration_skips_invalid_values(tmp_path: Path) -> None: + legacy = tmp_path / "config.env" + legacy.write_text("DFLASH_BUDGET=not-an-int\nDFLASH_MAX_CTX=65536\nDFLASH_LAZY=true\n") + + cfg, _doc = config._load_legacy_env(legacy) + + assert cfg.dflash.budget == 22 + assert cfg.dflash.max_ctx == 65536 + assert cfg.dflash.lazy is True + + +def test_image_variant_round_trips_from_toml(tmp_path: Path) -> None: + path = tmp_path / "config.toml" + path.write_text( + "[image]\n" + 'registry = "ghcr.io/luce-org/lucebox-hub"\n' + 'variant = "integration-props-uv-squared-clean-cuda12"\n' + ) + + cfg = config._load_toml(path) + + assert cfg.image == "ghcr.io/luce-org/lucebox-hub" + assert cfg.variant == "integration-props-uv-squared-clean-cuda12" + + +def test_model_preset_round_trips_through_set_and_load(tmp_path: Path) -> None: + """Setting model.preset writes a sparse TOML doc that loads back correctly.""" + path = tmp_path / "config.toml" + config_set("model.preset", "gemma-4-26b", path=path) + config_set("model.target_file", "google_gemma-4-26B-A4B-it-Q4_K_M.gguf", path=path) + + cfg = config._load_toml(path) + assert cfg.model.preset == "gemma-4-26b" + assert cfg.model.target_file == "google_gemma-4-26B-A4B-it-Q4_K_M.gguf" + + +def test_legacy_config_without_model_section_stays_unpinned(tmp_path: Path) -> None: + """Legacy configs (no [model] section) must NOT silently pin to qwen.""" + path = tmp_path / "config.toml" + path.write_text('[image]\nvariant = "cuda12"\n') + + cfg = config._load_toml(path) + + assert cfg.model.preset == "" + assert cfg.model.target_file == "" + assert cfg.model.draft_file == "" + + +def test_model_section_picks_target_file_from_registry(tmp_path: Path) -> None: + """A bare [model] preset="..." entry pulls target_file from the registry.""" + path = tmp_path / "config.toml" + path.write_text('[model]\npreset = "gemma-4-31b"\n') + + cfg = config._load_toml(path) + + assert cfg.model.preset == "gemma-4-31b" + assert cfg.model.target_file == "google_gemma-4-31B-it-Q4_K_M.gguf" + + +def test_model_section_picks_draft_file_from_registry(tmp_path: Path) -> None: + """When preset has a published draft GGUF, [model] preset="..." picks draft_file too.""" + path = tmp_path / "config.toml" + path.write_text('[model]\npreset = "qwen3.6-27b"\n') + + cfg = config._load_toml(path) + assert cfg.model.preset == "qwen3.6-27b" + assert cfg.model.draft_file == "dflash-draft-3.6-q4_k_m.gguf" + + +def test_config_set_writes_only_named_key(tmp_path: Path) -> None: + """Sparse persistence: setting one key does NOT serialize every default.""" + path = tmp_path / "config.toml" + config_set("dflash.budget", 16, path=path) + body = path.read_text() + # The only [dflash] field that should appear is budget — none of the others. + assert "[dflash]" in body + assert "budget = 16" in body + assert "max_ctx" not in body # not user-set, must not appear + assert "lazy" not in body + assert "[host]" not in body # whole section absent + assert "[image]" not in body # not touched either + + +def test_config_set_preserves_existing_keys(tmp_path: Path) -> None: + """Setting a new key leaves previously-set keys intact.""" + path = tmp_path / "config.toml" + config_set("dflash.budget", 16, path=path) + config_set("model.preset", "qwen3.6-27b", path=path) + body = path.read_text() + assert "budget = 16" in body + assert 'preset = "qwen3.6-27b"' in body + + +def test_config_unset_removes_one_key(tmp_path: Path) -> None: + """Unset removes the named key and leaves siblings alone.""" + path = tmp_path / "config.toml" + config_set("dflash.budget", 16, path=path) + config_set("dflash.max_ctx", 65536, path=path) + changed = config_unset("dflash.budget", path=path) + assert changed is True + body = path.read_text() + assert "budget" not in body + assert "max_ctx = 65536" in body + + +def test_config_unset_drops_empty_section(tmp_path: Path) -> None: + """Unsetting the last key in a section drops the empty section.""" + path = tmp_path / "config.toml" + config_set("dflash.budget", 16, path=path) + config_unset("dflash.budget", path=path) + body = path.read_text() + # The section may still exist as an empty table but `[dflash]` shouldn't. + assert "[dflash]" not in body + + +def test_config_get_reports_origin(tmp_path: Path) -> None: + """Each key carries an origin label — `file` when overridden, `default` otherwise.""" + path = tmp_path / "config.toml" + config_set("dflash.budget", 9, path=path) + entries = config_get(path=path) + assert entries["dflash.budget"] == (9, "file") + # max_ctx wasn't set so should report the live default. + value, origin = entries["dflash.max_ctx"] + assert origin == "default" + assert value == 16384 # DflashRuntime.max_ctx default + + +def test_config_get_rejects_unknown_key(tmp_path: Path) -> None: + path = tmp_path / "config.toml" + with pytest.raises(KeyError): + config_get("not.a.key", path=path) + + +def test_config_set_rejects_unknown_key(tmp_path: Path) -> None: + path = tmp_path / "config.toml" + with pytest.raises(KeyError): + config_set("not.a.key", 1, path=path) + + +def test_config_set_auto_creates_file(tmp_path: Path) -> None: + """`config set` creates a missing config.toml on first write.""" + path = tmp_path / "config.toml" + assert not path.exists() + config_set("port", 9090, path=path) + assert path.exists() + assert "port = 9090" in path.read_text() + + +def test_save_writes_sparse_doc(tmp_path: Path) -> None: + """`save` writes whatever doc is handed in — no defaults serialized.""" + path = tmp_path / "config.toml" + cfg = config._from_dict({}) + config.save(cfg, path, doc={"dflash": {"budget": 9}}) + body = path.read_text() + assert "budget = 9" in body + assert "max_ctx" not in body + + +def test_live_config_uses_recommend_preset_indirectly(tmp_path: Path) -> None: + """``live_config()`` returns a Config — no implicit preset when none given.""" + # The function probes the env-provided HostFacts; with no preset arg + # we must NOT silently pin one (that would surprise legacy installs). + cfg = config.live_config() + assert cfg.model.preset == "" + + +def test_seed_dflash_writes_heuristic_when_absent(tmp_path: Path) -> None: + """First-time activate seeds the VRAM-tier heuristic into config.toml. + + A config.toml with only [model] (what `models download --activate` + writes) would otherwise load with DflashRuntime class defaults + (max_ctx=16384), ignoring the host tier. Seeding writes the heuristic. + """ + from lucebox.types import HostFacts + + path = tmp_path / "config.toml" + path.write_text('[model]\npreset = "qwen3.6-27b"\n') + wrote = config.seed_dflash_from_host(HostFacts(vram_gb=24, is_wsl=True), path=path) + assert wrote is True + loaded = config.load(path) + assert loaded is not None + # 24 GB tier heuristic → 98304, not the 16384 class default. + assert loaded.dflash.max_ctx == 98304 + # Provenance recorded; [model] preserved. + doc = config.load_doc(path) + assert doc["autotune"]["source"] == "heuristic" + assert doc["model"]["preset"] == "qwen3.6-27b" + + +def test_seed_dflash_is_noop_when_dflash_present(tmp_path: Path) -> None: + """Never clobber a [dflash] the user or a prior tune already wrote.""" + from lucebox.types import HostFacts + + path = tmp_path / "config.toml" + path.write_text("[dflash]\nmax_ctx = 4096\n") + wrote = config.seed_dflash_from_host(HostFacts(vram_gb=80), path=path) + assert wrote is False + assert config.load(path).dflash.max_ctx == 4096 diff --git a/lucebox/tests/test_config_cli.py b/lucebox/tests/test_config_cli.py new file mode 100644 index 000000000..446ab41b6 --- /dev/null +++ b/lucebox/tests/test_config_cli.py @@ -0,0 +1,127 @@ +"""Tests for the ``lucebox config`` sub-app CLI.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from lucebox.cli import app +from typer.testing import CliRunner + + +def _set_config_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + monkeypatch.setenv("LUCEBOX_HOME", str(tmp_path)) + return tmp_path / "config.toml" + + +def test_config_set_then_get_round_trip( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cfg_path = _set_config_path(tmp_path, monkeypatch) + set_result = CliRunner().invoke(app, ["config", "set", "dflash.budget=12"]) + assert set_result.exit_code == 0 + assert cfg_path.exists() + get_result = CliRunner().invoke(app, ["config", "get", "dflash.budget"]) + assert get_result.exit_code == 0 + assert "12" in get_result.stdout + assert "from file" in get_result.stdout + + +def test_config_get_with_no_key_lists_every_registered_key( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _set_config_path(tmp_path, monkeypatch) + result = CliRunner().invoke(app, ["config", "get"]) + assert result.exit_code == 0 + # Every registered dotted key shows up at least once. + for key in ("model.preset", "dflash.budget", "port"): + assert key in result.stdout + + +def test_config_unset_drops_key( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cfg_path = _set_config_path(tmp_path, monkeypatch) + CliRunner().invoke(app, ["config", "set", "dflash.budget=9"]) + assert "budget = 9" in cfg_path.read_text() + unset_result = CliRunner().invoke(app, ["config", "unset", "dflash.budget"]) + assert unset_result.exit_code == 0 + body = cfg_path.read_text() + assert "budget" not in body + + +def test_config_set_unknown_key_errors( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _set_config_path(tmp_path, monkeypatch) + result = CliRunner().invoke(app, ["config", "set", "totally.unknown=1"]) + assert result.exit_code == 2 + + +def test_config_set_rejects_missing_equals( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _set_config_path(tmp_path, monkeypatch) + result = CliRunner().invoke(app, ["config", "set", "dflash.budget"]) + assert result.exit_code == 2 + + +def test_config_set_creates_file_when_missing( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cfg_path = _set_config_path(tmp_path, monkeypatch) + assert not cfg_path.exists() + CliRunner().invoke(app, ["config", "set", "port=9090"]) + assert cfg_path.exists() + assert "port = 9090" in cfg_path.read_text() + + +def test_load_or_build_env_overrides_persisted_config( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """LUCEBOX_* env vars must win over config.toml. + + Regression test for the precedence bug fixed in this commit: prior + to the fix, `_load_or_build()` returned `config_mod.load()`'s result + verbatim when config.toml existed, so the systemd unit's + `Environment=LUCEBOX_IMAGE=...` was silently ignored. Sindri's + config.toml had `[image]` without `registry`, which made the + dataclass default `ghcr.io/luce-org/lucebox-hub` win over the + intended easel image. + """ + from lucebox.cli import _load_or_build + + cfg_path = _set_config_path(tmp_path, monkeypatch) + # Write a config.toml WITHOUT an image.registry line — the + # bug-trigger shape on sindri. + cfg_path.write_text( + '[image]\nvariant = "cuda12"\n[runtime]\nport = 9090\n' + '[dflash]\nbudget = 22\n' + ) + # Env should override what config.toml says (and what dataclass + # defaults fill in for missing keys). + monkeypatch.setenv("LUCEBOX_IMAGE", "ghcr.io/myfork/lucebox-hub") + monkeypatch.setenv("LUCEBOX_PORT", "7777") + monkeypatch.setenv("LUCEBOX_CONTAINER", "lucebox-test") + cfg = _load_or_build() + assert cfg.image == "ghcr.io/myfork/lucebox-hub" # env beats dataclass default + assert cfg.port == 7777 # env beats config.toml + assert cfg.container_name == "lucebox-test" # env applied + # variant is in config.toml — config.toml value (no env override). + assert cfg.variant == "cuda12" + # dflash IS persisted in config.toml — env doesn't touch it (no DFLASH_* + # env hooks at this layer). + assert cfg.dflash.budget == 22 + + +def test_load_or_build_no_toml_env_overrides_defaults( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """When config.toml is absent, env must still override defaults.""" + from lucebox.cli import _load_or_build + + _set_config_path(tmp_path, monkeypatch) + # Don't write a config.toml — exercise the live_config() fallback. + monkeypatch.setenv("LUCEBOX_IMAGE", "ghcr.io/myfork/lucebox-hub") + cfg = _load_or_build() + assert cfg.image == "ghcr.io/myfork/lucebox-hub" diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py new file mode 100644 index 000000000..bb888514f --- /dev/null +++ b/lucebox/tests/test_docker_run.py @@ -0,0 +1,254 @@ +"""Tests for the docker-run serve-argv builder. + +This is the core's whole job: turn a Config into the exact `docker run` +command (and DFLASH_* env) that launches the server. The argv contract is +what `lucebox serve` / the systemd unit / `print-run` all consume, so it is +pinned field-by-field here rather than only smoke-tested. +""" + +from __future__ import annotations + +from pathlib import Path + +from lucebox.download import PRESETS +from lucebox.types import Config, DflashRuntime, ModelMeta + +from lucebox import docker_run + + +def _env(spec) -> dict[str, str]: + return dict(spec.env) + + +# ── DockerRunSpec.argv ─────────────────────────────────────────────────────── + + +def test_argv_minimal_defaults() -> None: + spec = docker_run.DockerRunSpec(image="img:tag", name="box") + argv = spec.argv() + assert argv[:2] == ["docker", "run"] + assert "--rm" in argv # remove defaults True + assert ["--name", "box"] == argv[argv.index("--name") : argv.index("--name") + 2] + assert ["--gpus", "all"] == argv[argv.index("--gpus") : argv.index("--gpus") + 2] + # image is the last positional (no entrypoint_args here) + assert argv[-1] == "img:tag" + assert "-d" not in argv # detach defaults False + + +def test_argv_flags_and_ordering() -> None: + spec = docker_run.DockerRunSpec( + image="img:tag", + name="box", + gpus=False, + detach=True, + remove=False, + port_publish=(8080, 8080), + volumes=(("/host/models", "/opt/lucebox-hub/server/models"),), + env=(("DFLASH_BUDGET", "22"),), + entrypoint_args=("serve",), + extra=("--shm-size", "1g"), + ) + argv = spec.argv() + assert "--rm" not in argv # remove=False + assert "-d" in argv # detach + assert "--gpus" not in argv # gpus=False + assert ["-p", "8080:8080"] == argv[argv.index("-p") : argv.index("-p") + 2] + assert ["-v", "/host/models:/opt/lucebox-hub/server/models"] == argv[ + argv.index("-v") : argv.index("-v") + 2 + ] + assert ["-e", "DFLASH_BUDGET=22"] == argv[argv.index("-e") : argv.index("-e") + 2] + # extra flags precede the image; entrypoint_args follow it. + assert argv[-1] == "serve" + assert argv[-2] == "img:tag" + assert argv.index("--shm-size") < argv.index("img:tag") + + +def test_printable_glues_value_taking_flags() -> None: + spec = docker_run.DockerRunSpec( + image="img:tag", + name="box", + port_publish=(8080, 8080), + env=(("K", "v"),), + ) + out = spec.printable() + # one flag per line, continued with backslash-newline + assert out.startswith("docker \\\n run") + # value-taking flags keep their value on the same line + assert "--name box" in out + assert "--gpus all" in out + assert "-p 8080:8080" in out + assert "-e K=v" in out + + +# ── _runtime_volumes ───────────────────────────────────────────────────────── + + +def test_runtime_volumes_mounts_models_and_home(tmp_path: Path) -> None: + cfg = Config(models_dir=tmp_path / "models") + vols = docker_run._runtime_volumes(cfg) + assert (str(tmp_path / "models"), "/opt/lucebox-hub/server/models") in vols + # $HOME is also mounted so absolute symlink targets resolve in-container. + assert any(host == str(Path.home()) for host, _ in vols) + + +def test_runtime_volumes_dedupes_when_models_is_home(monkeypatch, tmp_path: Path) -> None: + monkeypatch.setattr(Path, "home", staticmethod(lambda: tmp_path)) + cfg = Config(models_dir=tmp_path) + vols = docker_run._runtime_volumes(cfg) + # models_dir == home → only the models mount, no duplicate home mount. + assert len(vols) == 1 + + +# ── _resolve_model_files ───────────────────────────────────────────────────── + + +def test_resolve_model_files_explicit_override_wins(tmp_path: Path) -> None: + cfg = Config( + models_dir=tmp_path, + model=ModelMeta(preset="qwen3.6-27b", target_file="custom.gguf", draft_file="d.gguf"), + ) + target, draft, draft_dir = docker_run._resolve_model_files(cfg) + assert target == "custom.gguf" + assert draft == "d.gguf" + assert draft_dir == "" + + +def test_resolve_model_files_falls_back_to_preset_registry(tmp_path: Path) -> None: + pres = PRESETS["qwen3.6-27b"] + cfg = Config(models_dir=tmp_path, model=ModelMeta(preset="qwen3.6-27b")) + target, draft, draft_dir = docker_run._resolve_model_files(cfg) + assert target == pres.target_file + assert draft == (pres.draft_file or "") + assert draft_dir == "" # no speculator dir on disk + + +def test_resolve_model_files_no_preset_no_override(tmp_path: Path) -> None: + cfg = Config(models_dir=tmp_path) # ModelMeta() defaults: all empty + assert docker_run._resolve_model_files(cfg) == ("", "", "") + + +# ── server_run_spec ────────────────────────────────────────────────────────── + + +def test_server_run_spec_top_level_shape(tmp_path: Path) -> None: + cfg = Config( + image="ghcr.io/x/lucebox-hub", + variant="cuda12", + container_name="lucebox", + port=9000, + models_dir=tmp_path, + ) + spec = docker_run.server_run_spec(cfg) + assert spec.image == "ghcr.io/x/lucebox-hub:cuda12" + assert spec.name == "lucebox" + assert spec.gpus is True + assert spec.remove is True + assert spec.detach is False + assert spec.port_publish == (9000, 8080) + assert (str(tmp_path), "/opt/lucebox-hub/server/models") in spec.volumes + + +def test_server_run_spec_always_emits_core_dflash_env(tmp_path: Path) -> None: + cfg = Config(models_dir=tmp_path, dflash=DflashRuntime(budget=22, max_ctx=32768)) + env = _env(docker_run.server_run_spec(cfg)) + assert env["DFLASH_BUDGET"] == "22" + assert env["DFLASH_MAX_CTX"] == "32768" + assert env["DFLASH_PREFIX_CACHE_SLOTS"] == "0" + assert env["DFLASH_PREFILL_CACHE_SLOTS"] == "0" + assert env["DFLASH_THINK_MAX"] == "15488" + assert env["DFLASH_PORT"] == "8080" + + +def test_server_run_spec_optional_env_off_by_default(tmp_path: Path) -> None: + env = _env(docker_run.server_run_spec(Config(models_dir=tmp_path))) + for absent in ( + "DFLASH_LAZY", + "DFLASH_CACHE_TYPE_K", + "DFLASH_CACHE_TYPE_V", + "DFLASH_PREFILL_MODE", + "DFLASH_FA_WINDOW", + "DFLASH_THINK_SOFT_CLOSE_MIN_RATIO", + "DFLASH_DEBUG_THINKING_LOGITS", + "DFLASH_TARGET", + "DFLASH_DRAFT", + ): + assert absent not in env + + +def test_server_run_spec_optional_env_emitted_when_set(tmp_path: Path) -> None: + cfg = Config( + models_dir=tmp_path, + dflash=DflashRuntime( + lazy=True, + cache_type_k="tq3_0", + cache_type_v="tq3_0", + prefill_mode="auto", + prefill_keep_ratio=0.1, + prefill_threshold=20000, + prefill_drafter="drafter.gguf", + fa_window=512, + think_soft_close_min_ratio=0.5, + debug_thinking_logits=True, + ), + ) + env = _env(docker_run.server_run_spec(cfg)) + assert env["DFLASH_LAZY"] == "1" + assert env["DFLASH_CACHE_TYPE_K"] == "tq3_0" + assert env["DFLASH_CACHE_TYPE_V"] == "tq3_0" + assert env["DFLASH_PREFILL_MODE"] == "auto" + assert env["DFLASH_PREFILL_KEEP"] == "0.1" + assert env["DFLASH_PREFILL_THRESHOLD"] == "20000" + assert env["DFLASH_PREFILL_DRAFTER"] == "drafter.gguf" + assert env["DFLASH_FA_WINDOW"] == "512" + assert env["DFLASH_THINK_SOFT_CLOSE_MIN_RATIO"] == "0.5" + assert env["DFLASH_DEBUG_THINKING_LOGITS"] == "1" + + +def test_server_run_spec_resolves_target_and_draft_paths(tmp_path: Path) -> None: + pres = PRESETS["qwen3.6-27b"] + cfg = Config(models_dir=tmp_path, model=ModelMeta(preset="qwen3.6-27b")) + env = _env(docker_run.server_run_spec(cfg)) + assert env["DFLASH_TARGET"] == f"/opt/lucebox-hub/server/models/{pres.target_file}" + if pres.draft_file: + assert env["DFLASH_DRAFT"] == ( + f"/opt/lucebox-hub/server/models/draft/{pres.draft_file}" + ) + + +def test_server_run_spec_forwards_host_env(monkeypatch, tmp_path: Path) -> None: + monkeypatch.setenv("LUCEBOX_HOST_OS_PRETTY", "Ubuntu 22.04") + monkeypatch.setenv("LUCEBOX_HOST_GPU_NAME", "RTX 5090") + env = _env(docker_run.server_run_spec(Config(models_dir=tmp_path))) + assert env["LUCEBOX_HOST_OS_PRETTY"] == "Ubuntu 22.04" + assert env["LUCEBOX_HOST_GPU_NAME"] == "RTX 5090" + + +def test_large_preset_serves_at_safe_default_ctx(tmp_path: Path) -> None: + """Regression guard for the preset-cap analysis (#5). + + Activating a preset writes only [model], never [dflash], so a loaded + Config keeps the conservative DflashRuntime() floor (max_ctx=16384). + The VRAM-tier heuristic's higher caps only apply via `autotune --apply` + (which threads cfg.model.preset and is a separate PR). This test pins + that a large preset does NOT silently serve at a high, OOM-prone ctx + through the default serve path. + """ + cfg = Config(models_dir=tmp_path, model=ModelMeta(preset="qwen3.6-27b")) + env = _env(docker_run.server_run_spec(cfg)) + assert env["DFLASH_MAX_CTX"] == "16384" + + +# ── docker_pull ────────────────────────────────────────────────────────────── + + +def test_docker_pull_shells_out_and_returns_code(monkeypatch) -> None: + seen: dict[str, list[str]] = {} + + def fake_call(argv: list[str]) -> int: + seen["argv"] = argv + return 7 + + monkeypatch.setattr(docker_run.subprocess, "call", fake_call) + rc = docker_run.docker_pull("img:tag") + assert rc == 7 + assert seen["argv"] == ["docker", "pull", "img:tag"] diff --git a/lucebox/tests/test_download.py b/lucebox/tests/test_download.py new file mode 100644 index 000000000..8b69e96b8 --- /dev/null +++ b/lucebox/tests/test_download.py @@ -0,0 +1,323 @@ +"""Tests for the model-download orchestration. + +The downloader now drives `huggingface_hub.hf_hub_download` directly +(no subprocess) and verifies size + sha256 against the repo metadata +before re-fetching. The tests stub out the network calls so the +behavior contract — what gets requested, when downloads are skipped — +stays pinned without actually talking to the Hub. +""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest +from lucebox.download import ( + DEFAULT_PRESET, + PRESETS, + recommend_preset, + resolve_preset, + status, +) +from lucebox.types import HostFacts + +from lucebox import download + + +def test_default_preset_uses_quantized_gguf_draft(): + assert DEFAULT_PRESET.draft_repo == "spiritbuun/Qwen3.6-27B-DFlash-GGUF" + assert DEFAULT_PRESET.draft_file == "dflash-draft-3.6-q4_k_m.gguf" + + +def test_default_preset_is_registered_under_qwen_name(): + assert DEFAULT_PRESET is PRESETS["qwen3.6-27b"] + assert DEFAULT_PRESET.name == "qwen3.6-27b" + + +def test_resolve_preset_returns_default_on_none(): + assert resolve_preset(None) is DEFAULT_PRESET + assert resolve_preset("") is DEFAULT_PRESET + + +def test_resolve_preset_picks_gemma_target_and_draft(): + pres = resolve_preset("gemma-4-26b") + assert pres.name == "gemma-4-26b" + assert pres.target_repo == "bartowski/google_gemma-4-26B-A4B-it-GGUF" + assert pres.target_file == "google_gemma-4-26B-A4B-it-Q4_K_M.gguf" + assert pres.draft_repo == "Lucebox/gemma-4-26B-A4B-it-DFlash-GGUF" + assert pres.draft_file == "gemma-4-26B-A4B-it-DFlash-q8_0.gguf" + assert pres.has_draft + + +def test_resolve_preset_supports_target_only_laguna(): + pres = resolve_preset("laguna-xs.2") + assert pres.target_repo == "Lucebox/Laguna-XS.2-GGUF" + assert pres.draft_repo is None + assert not pres.has_draft + + +def test_resolve_preset_picks_qwen36_moe_target_only(): + """Qwen3.6 MoE preset routes to unsloth's UD-Q4_K_M file, no draft. + + The MoE variant has no published DFlash draft GGUF (verified against + HfApi.repo_info 2026-05-28), so it runs target-only like Laguna. The + file stem is `Qwen3.6-35B-A3B-UD-Q4_K_M.gguf` — the unsloth repo only + publishes the UD ("unsloth dynamic") family at Q4_K_M, not a plain + `Q4_K_M.gguf`. + """ + pres = resolve_preset("qwen3.6-moe") + assert pres.name == "qwen3.6-moe" + assert pres.target_repo == "unsloth/Qwen3.6-35B-A3B-GGUF" + assert pres.target_file == "Qwen3.6-35B-A3B-UD-Q4_K_M.gguf" + assert pres.draft_repo is None + assert pres.draft_file is None + assert not pres.has_draft + + +def test_download_preset_target_only_qwen36_moe_skips_draft(tmp_path, monkeypatch): + """qwen3.6-moe behaves identically to laguna-xs.2: target only, no draft fetch.""" + cfg = SimpleNamespace(models_dir=tmp_path) + pres = resolve_preset("qwen3.6-moe") + assert not pres.has_draft + fetches: list[tuple[str, str]] = [] + + def _meta(_api, repo_id: str, filename: str) -> tuple[int, None]: + return 10, None + + def _stub_fetch(api, repo_id, filename, local_dir, console): # noqa: ARG001 + fetches.append((repo_id, filename)) + out = local_dir / filename + out.parent.mkdir(parents=True, exist_ok=True) + with out.open("wb") as f: + f.truncate(10) + return out + + monkeypatch.setattr(download, "_file_meta", _meta) + monkeypatch.setattr(download, "_fetch", _stub_fetch) + + assert download.download_preset(cfg, pres) == 0 + # Only the target — no draft attempt at all. + assert fetches == [(pres.target_repo, pres.target_file)] + + +def test_status_qwen36_moe_reports_draft_present_when_target_only(tmp_path, monkeypatch): + """No published draft → status reports draft_present=True (nothing to fetch).""" + cfg = SimpleNamespace(models_dir=tmp_path) + pres = resolve_preset("qwen3.6-moe") + + def _meta(_api, repo_id: str, filename: str) -> tuple[int, None]: + return 22 * 10**9, None + + monkeypatch.setattr(download, "_file_meta", _meta) + # Target absent → target_present False, draft_present True (no draft). + assert status(cfg, pres) == {"target_present": False, "draft_present": True} + + +def test_resolve_preset_unknown_name_lists_known_options(): + with pytest.raises(KeyError) as exc_info: + resolve_preset("qwen-99b") + msg = str(exc_info.value) + # Every registered preset must appear in the suggestion list so the + # user can copy-paste the right name. + for name in PRESETS: + assert name in msg + + +def _stub_file_meta(target_size: int, draft_size: int): + """Build a `_file_meta` replacement that returns (size, None) per repo+file. + + sha256 is left None so tests don't need to compute real hashes; the + real metadata path is exercised by the live `models download` + invocation, not the unit tests. + """ + + def _meta(_api, repo_id: str, filename: str) -> tuple[int, None]: + if repo_id == DEFAULT_PRESET.target_repo and filename == DEFAULT_PRESET.target_file: + return target_size, None + if repo_id == DEFAULT_PRESET.draft_repo and filename == DEFAULT_PRESET.draft_file: + return draft_size, None + raise FileNotFoundError(f"unexpected ({repo_id}, {filename})") + + return _meta + + +def test_status_checks_default_draft_gguf(tmp_path, monkeypatch): + cfg = SimpleNamespace(models_dir=tmp_path) + draft_dir = tmp_path / "draft" + draft_dir.mkdir() + target = tmp_path / DEFAULT_PRESET.target_file + draft = draft_dir / DEFAULT_PRESET.draft_file + + monkeypatch.setattr(download, "_file_meta", _stub_file_meta(target_size=1024, draft_size=512)) + + # Neither file exists yet. + assert status(cfg) == {"target_present": False, "draft_present": False} + + # Write files at the expected sizes. + with target.open("wb") as f: + f.truncate(1024) + with draft.open("wb") as f: + f.truncate(512) + assert status(cfg) == {"target_present": True, "draft_present": True} + + +def test_status_rejects_partial_model_files(tmp_path, monkeypatch): + cfg = SimpleNamespace(models_dir=tmp_path) + draft_dir = tmp_path / "draft" + draft_dir.mkdir() + target = tmp_path / DEFAULT_PRESET.target_file + draft = draft_dir / DEFAULT_PRESET.draft_file + target.write_bytes(b"partial") + draft.write_bytes(b"partial") + + # Repo says the target is 1 GB; a 7-byte file is partial, not present. + monkeypatch.setattr( + download, "_file_meta", _stub_file_meta(target_size=10**9, draft_size=10**6) + ) + assert status(cfg) == {"target_present": False, "draft_present": False} + + +def test_current_bytes_reads_xet_staging_path(tmp_path): + """Regression: progress polling must see hf-xet's hashed staging file. + + huggingface_hub 1.x writes partial Xet downloads to + ``{local_dir}/.cache/huggingface/download/{short_hash}.{etag}.incomplete`` + — NOT to ``{local_dir}/{filename}.incomplete``. Before the fix the + polling code only checked the latter (which never appears) so the + Rich progress bar sat at 0 bytes for the entire transfer. + """ + filename = "model.gguf" + etag = "abc123" + candidates = download._incomplete_path_candidates(tmp_path, filename, etag) + # The first candidate must point at the actual hf-xet staging path. + xet_path: Path = candidates[0] + assert xet_path.parent == tmp_path / ".cache" / "huggingface" / "download" + assert xet_path.name.endswith(f".{etag}.incomplete") + + # Now: writing to that path must be observed by _current_bytes. + xet_path.parent.mkdir(parents=True, exist_ok=True) + xet_path.write_bytes(b"x" * 4096) + target = tmp_path / filename + assert download._current_bytes(target, candidates) == 4096 + + +def test_current_bytes_falls_back_to_glob_without_etag(tmp_path): + """When sha256 is unknown we still find growing .incomplete files.""" + filename = "model.gguf" + candidates = download._incomplete_path_candidates(tmp_path, filename, etag=None) + target = tmp_path / filename + + staging = tmp_path / ".cache" / "huggingface" / "download" + staging.mkdir(parents=True, exist_ok=True) + (staging / "deadbeef.deadbeef.incomplete").write_bytes(b"x" * 8192) + assert download._current_bytes(target, candidates) == 8192 + + +def test_current_bytes_prefers_final_target_when_complete(tmp_path): + filename = "model.gguf" + candidates = download._incomplete_path_candidates(tmp_path, filename, etag="abc") + target = tmp_path / filename + target.write_bytes(b"x" * 1234) + assert download._current_bytes(target, candidates) == 1234 + + +def test_download_preset_fetches_exact_draft_file(tmp_path, monkeypatch): + cfg = SimpleNamespace(models_dir=tmp_path) + fetches: list[tuple[str, str, str]] = [] + + monkeypatch.setattr(download, "_file_meta", _stub_file_meta(target_size=10, draft_size=10)) + + # Stub the actual download to record what was requested + create a stub + # file of the expected size so `_local_matches` would pass on a re-run. + def _stub_fetch(api, repo_id, filename, local_dir, console): # noqa: ARG001 + fetches.append((repo_id, filename, str(local_dir))) + target = local_dir / filename + target.parent.mkdir(parents=True, exist_ok=True) + with target.open("wb") as f: + f.truncate(10) + return target + + monkeypatch.setattr(download, "_fetch", _stub_fetch) + + assert download.download_preset(cfg) == 0 + assert (DEFAULT_PRESET.target_repo, DEFAULT_PRESET.target_file, str(tmp_path)) in fetches + assert ( + DEFAULT_PRESET.draft_repo, + DEFAULT_PRESET.draft_file, + str(tmp_path / "draft"), + ) in fetches + + +def test_download_preset_routes_gemma_preset_to_correct_repos(tmp_path, monkeypatch): + cfg = SimpleNamespace(models_dir=tmp_path) + pres = resolve_preset("gemma-4-26b") + fetches: list[tuple[str, str, str]] = [] + + def _meta(_api, repo_id: str, filename: str) -> tuple[int, None]: + return 10, None + + def _stub_fetch(api, repo_id, filename, local_dir, console): # noqa: ARG001 + fetches.append((repo_id, filename, str(local_dir))) + out = local_dir / filename + out.parent.mkdir(parents=True, exist_ok=True) + with out.open("wb") as f: + f.truncate(10) + return out + + monkeypatch.setattr(download, "_file_meta", _meta) + monkeypatch.setattr(download, "_fetch", _stub_fetch) + + assert download.download_preset(cfg, pres) == 0 + assert (pres.target_repo, pres.target_file, str(tmp_path)) in fetches + assert (pres.draft_repo, pres.draft_file, str(tmp_path / "draft")) in fetches + + +def test_download_preset_target_only_skips_draft_fetch(tmp_path, monkeypatch): + cfg = SimpleNamespace(models_dir=tmp_path) + pres = resolve_preset("laguna-xs.2") + assert not pres.has_draft + fetches: list[tuple[str, str]] = [] + + def _meta(_api, repo_id: str, filename: str) -> tuple[int, None]: + return 10, None + + def _stub_fetch(api, repo_id, filename, local_dir, console): # noqa: ARG001 + fetches.append((repo_id, filename)) + out = local_dir / filename + out.parent.mkdir(parents=True, exist_ok=True) + with out.open("wb") as f: + f.truncate(10) + return out + + monkeypatch.setattr(download, "_file_meta", _meta) + monkeypatch.setattr(download, "_fetch", _stub_fetch) + + assert download.download_preset(cfg, pres) == 0 + # Target fetched, no draft fetch attempted at all. + assert fetches == [(pres.target_repo, pres.target_file)] + + +def test_status_target_only_preset_reports_draft_as_present(tmp_path, monkeypatch): + cfg = SimpleNamespace(models_dir=tmp_path) + pres = resolve_preset("laguna-xs.2") + + def _meta(_api, repo_id: str, filename: str) -> tuple[int, None]: + return 1024, None + + monkeypatch.setattr(download, "_file_meta", _meta) + # Target absent → target_present False, draft_present True (nothing to download). + assert status(cfg, pres) == {"target_present": False, "draft_present": True} + + +def test_recommend_preset_tiers() -> None: + """First-run preset recommendation is a pure VRAM-tier function. + + 22 GB+ → the Lucebox default (qwen3.6-27b); 16-21 GB → laguna-xs.2; + below 16 GB → None (the registered presets need ≥16 GB, so we punt to + an explicit choice rather than recommend something that can't run). + """ + assert recommend_preset(HostFacts(vram_gb=24)) == "qwen3.6-27b" + assert recommend_preset(HostFacts(vram_gb=22)) == "qwen3.6-27b" + assert recommend_preset(HostFacts(vram_gb=20)) == "laguna-xs.2" + assert recommend_preset(HostFacts(vram_gb=16)) == "laguna-xs.2" + assert recommend_preset(HostFacts(vram_gb=12)) is None + assert recommend_preset(HostFacts(vram_gb=0)) is None diff --git a/lucebox/tests/test_models_cli.py b/lucebox/tests/test_models_cli.py new file mode 100644 index 000000000..f44583044 --- /dev/null +++ b/lucebox/tests/test_models_cli.py @@ -0,0 +1,142 @@ +"""Tests for the ``lucebox models`` sub-app.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from lucebox.cli import app +from lucebox.download import PRESETS +from lucebox.types import HostFacts +from typer.testing import CliRunner + +from lucebox import config as config_mod +from lucebox import download as download_mod + + +def _set_config_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + monkeypatch.setenv("LUCEBOX_HOME", str(tmp_path)) + monkeypatch.setenv("LUCEBOX_MODELS", str(tmp_path / "models")) + return tmp_path / "config.toml" + + +def _stub_host(monkeypatch: pytest.MonkeyPatch, vram_gb: int) -> None: + monkeypatch.setattr("lucebox.host_facts.from_env", lambda: HostFacts(vram_gb=vram_gb)) + monkeypatch.setattr("lucebox.cli.from_env", lambda: HostFacts(vram_gb=vram_gb)) + + +def test_models_list_shows_every_registered_preset( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + result = CliRunner().invoke(app, ["models", "list"]) + assert result.exit_code == 0 + for name in PRESETS: + assert name in result.stdout + + +def test_models_default_view_lists_only_installed( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + # No models on disk → default view says "no presets installed". + result = CliRunner().invoke(app, ["models"]) + assert result.exit_code == 0 + assert "No presets installed" in result.stdout or "Models dir" in result.stdout + + +def test_models_download_recommends_when_empty( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """No preset configured + nothing on argv → auto-recommend + auto-activate.""" + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + + # Stub the network calls so the test doesn't try to talk to HF. + monkeypatch.setattr(download_mod, "download_preset", lambda cfg, pres: 0) + monkeypatch.setattr( + download_mod, + "status", + lambda cfg, pres: {"target_present": True, "draft_present": True}, + ) + + result = CliRunner().invoke(app, ["models", "download"]) + assert result.exit_code == 0 + assert "Recommended preset" in result.stdout + assert cfg_path.exists() + # The active preset should now be model.preset = qwen3.6-27b. + entries = config_mod.config_get(path=cfg_path) + assert entries["model.preset"] == ("qwen3.6-27b", "file") + + +def test_models_download_refuses_silent_switch( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """When a preset is already active, `download` with no arg refuses.""" + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + config_mod.config_set("model.preset", "qwen3.6-27b", path=cfg_path) + + result = CliRunner().invoke(app, ["models", "download"]) + assert result.exit_code == 2 + assert "already active" in result.stdout.lower() + + +def test_models_download_explicit_preset_no_activate( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Passing a preset without --activate downloads but doesn't flip model.preset.""" + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + monkeypatch.setattr(download_mod, "download_preset", lambda cfg, pres: 0) + monkeypatch.setattr( + download_mod, + "status", + lambda cfg, pres: {"target_present": False, "draft_present": False}, + ) + + result = CliRunner().invoke(app, ["models", "download", "gemma-4-26b"]) + assert result.exit_code == 0 + if cfg_path.exists(): + entries = config_mod.config_get(path=cfg_path) + assert entries["model.preset"] == ("", "default") + + +def test_models_download_explicit_preset_with_activate( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + monkeypatch.setattr(download_mod, "download_preset", lambda cfg, pres: 0) + monkeypatch.setattr( + download_mod, + "status", + lambda cfg, pres: {"target_present": False, "draft_present": False}, + ) + + result = CliRunner().invoke(app, ["models", "download", "gemma-4-26b", "--activate"]) + assert result.exit_code == 0 + entries = config_mod.config_get(path=cfg_path) + assert entries["model.preset"] == ("gemma-4-26b", "file") + + +def test_installed_helpers_track_presence( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """``installed_status`` / ``installed_size_gb`` reflect on-disk byte counts.""" + _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + from lucebox.config import live_config + + cfg = live_config() + cfg.models_dir.mkdir(parents=True, exist_ok=True) + laguna = PRESETS["laguna-xs.2"] + assert download_mod.installed_status(cfg, laguna) == "absent" + + target = cfg.models_dir / laguna.target_file + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(b"x" * (5 * 10**9)) + assert download_mod.installed_status(cfg, laguna) == "installed" + assert download_mod.installed_size_gb(cfg, laguna) == pytest.approx(5.0, rel=0.01) diff --git a/pyproject.toml b/pyproject.toml index 56ae2bf4f..520838041 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,7 @@ license = { text = "Apache-2.0" } authors = [{ name = "Lucebox" }] dependencies = [ + "lucebox", "lucebox-dflash", "pflash", ] @@ -23,7 +24,7 @@ line-length = 100 # server-internal and optimization Python (server/scripts, optimizations/*) # carries pre-existing style debt and is added to `include` as it is cleaned # up. Vendored deps stay excluded permanently (extend-exclude below). -include = ["harness/**/*.py", "scripts/**/*.py"] +include = ["harness/**/*.py", "scripts/**/*.py", "lucebox/**/*.py"] extend-exclude = [ "dflash/deps", "megakernel", @@ -51,11 +52,12 @@ package = false no-build-isolation-package = ["qwen35-megakernel-bf16"] [tool.uv.workspace] -# Workspace members. Keeping the list to the packages that live in this -# repo lets `uv lock --check` / `uv sync --frozen` pass. -members = ["server", "optimizations/megakernel", "optimizations/pflash"] +# Workspace members. PR adds the lucebox/ package alongside the existing +# server / megakernel / pflash members. +members = ["lucebox", "server", "optimizations/megakernel", "optimizations/pflash"] [tool.uv.sources] +lucebox = { workspace = true } lucebox-dflash = { workspace = true } pflash = { workspace = true } qwen35-megakernel-bf16 = { workspace = true } diff --git a/scripts/check_lucebox_wrapper_sandbox.sh b/scripts/check_lucebox_wrapper_sandbox.sh new file mode 100755 index 000000000..df2b2b9bc --- /dev/null +++ b/scripts/check_lucebox_wrapper_sandbox.sh @@ -0,0 +1,242 @@ +#!/usr/bin/env bash +# Exercise the host-side lucebox.sh installer/wrapper from an isolated prefix. +# +# The script intentionally runs from a throwaway HOME, XDG_CONFIG_HOME, +# LUCEBOX_HOME, model directory, and working directory. That catches accidental +# dependencies on the checkout or the user's real ~/.lucebox while keeping the +# test reproducible enough to paste into a bug report. + +set -euo pipefail + +IMAGE="${LUCEBOX_TEST_IMAGE:-ghcr.io/easel/lucebox-hub}" +VARIANT="${LUCEBOX_TEST_VARIANT:-integration-props-uv-squared-clean-cuda12}" +WRAPPER_SOURCE="${LUCEBOX_TEST_WRAPPER_SOURCE:-local}" +RUN_PULL="${LUCEBOX_TEST_RUN_PULL:-1}" +RUN_CONTAINER_CLI="${LUCEBOX_TEST_RUN_CONTAINER_CLI:-1}" +KEEP_SANDBOX="${LUCEBOX_TEST_KEEP_SANDBOX:-0}" + +ROOT="" +LOG="" + +usage() { + cat <&2; usage >&2; exit 2 ;; + esac +done + +die() { + echo "[FAIL] $*" >&2 + if [ -n "$LOG" ] && [ -f "$LOG" ]; then + echo "[FAIL] transcript: $LOG" >&2 + fi + exit 1 +} + +note() { + printf '[INFO] %s\n' "$*" +} + +pass() { + printf '[PASS] %s\n' "$*" +} + +assert_file() { + [ -f "$1" ] || die "missing file: $1" + pass "file exists: $1" +} + +assert_contains() { + local file="$1" + local pattern="$2" + if ! grep -Fq "$pattern" "$file"; then + echo "----- $file -----" >&2 + sed -n '1,220p' "$file" >&2 || true + echo "-----------------" >&2 + die "expected '$pattern' in $file" + fi + pass "$file contains: $pattern" +} + +run_logged() { + note "run: $*" + { + printf '\n===== %s =====\n' "$*" + "$@" + printf '===== exit=0 =====\n' + } 2>&1 | tee -a "$LOG" +} + +run_logged_capture() { + local out="$1" + shift + note "run: $* > $out" + { + printf '\n===== %s > %s =====\n' "$*" "$out" + "$@" + local rc=$? + printf '===== exit=%s =====\n' "$rc" + return "$rc" + } 2>&1 | tee "$out" | tee -a "$LOG" >/dev/null +} + +cleanup() { + if [ -n "$ROOT" ] && [ "$KEEP_SANDBOX" != "1" ]; then + rm -rf "$ROOT" + elif [ -n "$ROOT" ]; then + note "kept sandbox: $ROOT" + note "transcript: $LOG" + fi +} +trap cleanup EXIT + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +ROOT="$(mktemp -d "${TMPDIR:-/tmp}/lucebox-wrapper-sandbox.XXXXXX")" +LOG="$ROOT/transcript.log" + +HOME_DIR="$ROOT/home" +BIN_DIR="$ROOT/bin" +XDG_DIR="$ROOT/xdg" +MODELS_DIR="$ROOT/models" +WORK_DIR="$ROOT/work" +mkdir -p "$HOME_DIR" "$BIN_DIR" "$XDG_DIR" "$MODELS_DIR" "$WORK_DIR" + +note "sandbox: $ROOT" +note "transcript: $LOG" + +case "$WRAPPER_SOURCE" in + local) + cp "$REPO_ROOT/lucebox.sh" "$BIN_DIR/lucebox" + ;; + http://*|https://*) + curl -fsSL "$WRAPPER_SOURCE" -o "$BIN_DIR/lucebox" + ;; + *) + cp "$WRAPPER_SOURCE" "$BIN_DIR/lucebox" + ;; +esac +chmod +x "$BIN_DIR/lucebox" + +FIRST_LINE="$(head -n 1 "$BIN_DIR/lucebox")" +[ "$FIRST_LINE" = "#!/usr/bin/env bash" ] || die "unexpected shebang: $FIRST_LINE" +pass "wrapper has expected shebang" + +export HOME="$HOME_DIR" +export XDG_CONFIG_HOME="$XDG_DIR" +export LUCEBOX_HOME="$HOME_DIR/.lucebox" +export LUCEBOX_MODELS="$MODELS_DIR" +export LUCEBOX_IMAGE="$IMAGE" +export LUCEBOX_VARIANT="$VARIANT" +export LUCEBOX_CONTAINER="lucebox-sandbox" +export LUCEBOX_PORT="18080" +export PATH="$BIN_DIR:$PATH" + +cd "$WORK_DIR" +[ "$PWD" = "$WORK_DIR" ] || die "failed to enter sandbox workdir" +pass "working directory isolated: $PWD" + +run_logged_capture "$ROOT/version.out" lucebox version +assert_contains "$ROOT/version.out" "0.2.0" + +run_logged_capture "$ROOT/help.out" lucebox help +assert_contains "$ROOT/help.out" "LUCEBOX_VARIANT" +assert_contains "$ROOT/help.out" "LUCEBOX_IMAGE" + +docker manifest inspect "${IMAGE}:${VARIANT}" >/dev/null +pass "image manifest exists: ${IMAGE}:${VARIANT}" + +if [ "$RUN_PULL" = "1" ]; then + run_logged_capture "$ROOT/pull.out" lucebox pull + assert_contains "$ROOT/pull.out" "${IMAGE}:${VARIANT}" +fi + +if [ "$RUN_CONTAINER_CLI" = "1" ]; then + run_logged_capture "$ROOT/check.out" lucebox check + # Sparse persistence: `config set` creates config.toml with only the + # named key. Replaces the old `configure --overwrite` path. + run_logged_capture "$ROOT/config-image.out" lucebox config set "image=$IMAGE" + run_logged_capture "$ROOT/config-variant.out" lucebox config set "variant=$VARIANT" + assert_file "$LUCEBOX_HOME/config.toml" + [ "$(stat -c '%u' "$LUCEBOX_HOME/config.toml")" = "$(id -u)" ] \ + || die "config.toml is not owned by the invoking user" + pass "config.toml ownership matches invoking user" + assert_contains "$LUCEBOX_HOME/config.toml" "registry = \"$IMAGE\"" + assert_contains "$LUCEBOX_HOME/config.toml" "variant = \"$VARIANT\"" + + run_logged_capture "$ROOT/print-run.out" lucebox print-run + assert_contains "$ROOT/print-run.out" "${IMAGE}:${VARIANT}" + assert_contains "$ROOT/print-run.out" "$MODELS_DIR:/opt/lucebox-hub/dflash/models" + if grep -Fq "$REPO_ROOT" "$ROOT/print-run.out"; then + die "print-run leaked repository path: $REPO_ROOT" + fi + pass "print-run did not reference repository checkout" +fi + +# Exercise `lucebox install` without allowing it to call real systemctl, +# loginctl, docker, or nvidia-smi. The generated user unit must land under the +# sandbox XDG_CONFIG_HOME and point ExecStart at the sandbox-installed wrapper. +SHIM_DIR="$ROOT/shims" +mkdir -p "$SHIM_DIR" +cat > "$SHIM_DIR/docker" <<'EOF' +#!/usr/bin/env bash +case "${1:-}" in + info) exit 0 ;; + version) echo "25.0.0"; exit 0 ;; + stop) exit 0 ;; + *) echo "docker shim: $*" >&2; exit 0 ;; +esac +EOF +cat > "$SHIM_DIR/nvidia-smi" <<'EOF' +#!/usr/bin/env bash +case "$*" in + *"--query-gpu=name,memory.total,driver_version,compute_cap"*) + echo "Fake GPU, 24576, 555.42.01, 8.6"; exit 0 ;; + *"--query-gpu=name"*) + echo "Fake GPU"; exit 0 ;; + *) echo "Fake GPU"; exit 0 ;; +esac +EOF +cat > "$SHIM_DIR/systemctl" <<'EOF' +#!/usr/bin/env bash +if [ "$1" = "--user" ] && [ "$2" = "show-environment" ]; then exit 0; fi +if [ "$1" = "--user" ] && [ "$2" = "daemon-reload" ]; then exit 0; fi +echo "systemctl shim: $*" >&2 +exit 0 +EOF +cat > "$SHIM_DIR/loginctl" <<'EOF' +#!/usr/bin/env bash +echo "Linger=no" +EOF +chmod +x "$SHIM_DIR/docker" "$SHIM_DIR/nvidia-smi" "$SHIM_DIR/systemctl" "$SHIM_DIR/loginctl" + +PATH="$SHIM_DIR:$BIN_DIR:$PATH" run_logged_capture "$ROOT/install.out" lucebox install +UNIT="$XDG_CONFIG_HOME/systemd/user/lucebox.service" +assert_file "$UNIT" +assert_contains "$UNIT" "ExecStart=$BIN_DIR/lucebox serve" +assert_contains "$UNIT" "ExecStop=$SHIM_DIR/docker stop -t 30 lucebox-sandbox" +assert_contains "$ROOT/install.out" "Installed $UNIT" + +pass "sandbox wrapper check completed" +note "summary: image=${IMAGE}:${VARIANT} wrapper_source=${WRAPPER_SOURCE}" diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh new file mode 100755 index 000000000..3960760f6 --- /dev/null +++ b/scripts/test_lucebox_sh.sh @@ -0,0 +1,1130 @@ +#!/usr/bin/env bash +# scripts/test_lucebox_sh.sh — smoke tests for the host-side wrapper + +# every other bash script we ship. +# +# Catches regressions like: +# * syntax errors (bash -n) +# * shellcheck error-level findings across every shipped bash script +# * `set -u` violations in command paths that don't need docker/nvidia — +# each subcommand dispatch is exercised in isolation to verify no +# LUCEBOX_HOST_* or DFLASH_* read fires before the helper that should +# populate it has run. +# * missing dispatch handlers (help, version, check, usage) +# * stale references to subcommands removed from main's case +# +# The wrapper is shell + has zero non-coreutils deps for the host-only +# commands, so this script doesn't need docker/nvidia/systemd present — +# probe_host degrades cleanly when those aren't found, and the +# formatter must render fine for the "everything is missing" case too. +# +# Run from anywhere: scripts/test_lucebox_sh.sh + +set -euo pipefail + +# Resolve repo root + script under test. +ROOT="$(git rev-parse --show-toplevel 2>/dev/null || (cd "$(dirname "$0")/.." && pwd))" +SCRIPT="$ROOT/lucebox.sh" +ENTRYPOINT="$ROOT/server/scripts/entrypoint.sh" +INSTALLER="$ROOT/install.sh" + +if [ ! -f "$SCRIPT" ]; then + echo "FAIL: lucebox.sh not found at $SCRIPT" >&2 + exit 1 +fi + +# entrypoint.sh ships with the docker-stack PR (#334). When it's absent +# (e.g. on the lucebox-cli branch in isolation), skip the entire suite — +# every section below either references $ENTRYPOINT in shellcheck targets, +# parses it with `bash -n`, or sources/dispatches into it directly. The +# host-only lucebox.sh wrapper itself is covered by lucebox.sh's own unit +# tests; this script's value is the wrapper↔entrypoint contract. +if [ ! -f "$ENTRYPOINT" ]; then + echo "Skipping entrypoint tests: server/scripts/entrypoint.sh not present (provided by #334 docker-stack)" + exit 0 +fi + +fail=0 +pass=0 +report() { + if [ "$1" = "ok" ]; then + printf ' \033[1;32m✓\033[0m %s\n' "$2" + pass=$((pass + 1)) + else + printf ' \033[1;31m✗\033[0m %s\n' "$2" + if [ -n "${3:-}" ]; then + printf ' %s\n' "$3" + fi + fail=$((fail + 1)) + fi +} + +# Helper: run the wrapper with strict bash, capture stdout+stderr, check for +# (a) zero exit code, (b) substring match. NO_COLOR is set so colour codes +# don't pollute substring matches. +assert_runs() { + local label="$1" cmd="$2" expect="${3:-}" + local out rc + out=$(NO_COLOR=1 bash -c "$cmd" 2>&1) + rc=$? + if [ "$rc" -ne 0 ]; then + report fail "$label" "exit $rc; output: $(printf '%s' "$out" | head -3)" + return + fi + if [ -n "$expect" ] && ! grep -qF "$expect" <<<"$out"; then + report fail "$label" "missing expected substring '$expect'; got: $(printf '%s' "$out" | head -3)" + return + fi + report ok "$label" +} + +# Helper: run a subcommand whose successful completion would normally need +# docker / nvidia / systemd. We only care that the bash dispatch up to the +# point of the missing dependency does NOT trip `set -u`. Exit code is +# allowed to be non-zero; what we forbid is a raw "unbound variable" / +# "syntax error" / "line N:" leak in the captured output. +# +# Wrapped in `timeout` so subcommands that exec into a follow-style binary +# (logs → journalctl -f, status when systemd is healthy, etc.) don't hang +# the test runner on a dev box where the underlying tools succeed. +assert_no_set_u_leak() { + local label="$1" + shift + local out + out=$(NO_COLOR=1 timeout 5 bash "$@" 2>&1 || true) + # The "line N:" pattern is anchored to a script-path prefix to avoid + # false positives from journalctl output ("systemd[1385106]:") which + # contains a similar shape but isn't a bash error. Bash always emits + # the source filename before the line number, e.g. + # /tmp/lbh-flat/lucebox.sh: line 200: VAR: unbound variable + if grep -qE 'unbound variable|syntax error|\.sh: line [0-9]+:' <<<"$out"; then + report fail "$label" "raw bash error leaked: $(head -3 <<<"$out")" + else + report ok "$label" + fi +} + +echo "[test_lucebox_sh] running against $SCRIPT" + +# ── 1. shellcheck ───────────────────────────────────────────────────────── +# Run shellcheck across every bash script we ship (the wrapper, the +# in-container entrypoint, and every helper under scripts/). Error-level +# findings fail the build; warnings are informational only — those have +# been triaged and the SC2034/SC2155/SC2164 hits in sweep_ds4_2case.sh +# aren't user-visible bugs. +SHELLCHECK_TARGETS=( + "$SCRIPT" + "$ENTRYPOINT" + "$INSTALLER" +) +# Add every scripts/*.sh except this one (don't recurse into our own tests). +while IFS= read -r -d '' f; do + [ "$f" = "${BASH_SOURCE[0]}" ] && continue + SHELLCHECK_TARGETS+=("$f") +done < <(find "$ROOT/scripts" -maxdepth 1 -name '*.sh' -type f -print0 2>/dev/null) +SHELLCHECK_TARGETS+=("${BASH_SOURCE[0]}") + +if command -v shellcheck >/dev/null 2>&1; then + sc_out=$(shellcheck --severity=error "${SHELLCHECK_TARGETS[@]}" 2>&1) || sc_rc=$? + sc_rc="${sc_rc:-0}" + if [ "$sc_rc" -eq 0 ]; then + report ok "shellcheck --severity=error (${#SHELLCHECK_TARGETS[@]} files)" + else + report fail "shellcheck --severity=error" "$(printf '%s' "$sc_out" | head -10)" + fi +else + report fail "shellcheck not installed" "install via 'apt-get install -y shellcheck' (Ubuntu) or 'brew install shellcheck'" +fi + +# ── 2. Syntax / parse ───────────────────────────────────────────────────── +if bash -n "$SCRIPT"; then report ok "bash -n lucebox.sh parses cleanly" +else report fail "bash -n lucebox.sh"; fi +if bash -n "$ENTRYPOINT"; then report ok "bash -n entrypoint.sh parses cleanly" +else report fail "bash -n entrypoint.sh"; fi + +# ── 3. Trivial subcommands (zero-exit expected) ─────────────────────────── +assert_runs "help" "bash '$SCRIPT' help" "host-side wrapper" +assert_runs "--help" "bash '$SCRIPT' --help" "host-side wrapper" +assert_runs "-h" "bash '$SCRIPT' -h" "host-side wrapper" +assert_runs "version" "bash '$SCRIPT' version" "" +assert_runs "--version" "bash '$SCRIPT' --version" "" + +# ── 4. check — host-only, must run to completion even without docker/nvidia. +# This is the path that broke last time (multi-byte glyph + set -u). +assert_runs "check" "bash '$SCRIPT' check" "host readiness report" + +# ── 5. systemd-surface subcommands — every one of these used to crash with +# `LUCEBOX_HOST_HAS_SYSTEMD: unbound variable` because cmd_systemctl_passthrough +# / cmd_logs / cmd_systemd_uninstall reached require_systemd without first +# calling probe_host. The fix routes through require_systemd → probe_host +# when the var is unset; these tests pin that invariant. +# +# On the bare runner there is no user systemd, no installed unit, and no +# docker — so every command is expected to exit non-zero with a CLEAN error +# message. What we forbid is a raw bash "unbound variable" leak. +for sub in start stop restart enable disable status install uninstall; do + assert_no_set_u_leak "$sub dispatch (no set -u leak)" "$SCRIPT" "$sub" +done +# `logs` is special: it execs `journalctl -f` which streams every historical +# journal record for the unit. On a dev box where the lucebox service has +# actually run, that stream contains every past error — including the very +# bugs this test exists to prevent — and we'd false-positive on them. Pass +# `-n 0 --no-pager` so we only see new entries (none, in the test window). +assert_no_set_u_leak "logs dispatch (no set -u leak)" "$SCRIPT" logs -n 0 --no-pager + +# ── 6. server-spawning subcommands — exercise the dispatch up to where +# the missing docker daemon stops them. `serve` is intentionally skipped +# because on a host with a working docker + the cuda12 image already +# pulled, it would actually exec into the container — at which point +# we'd be testing the image's entrypoint, not the wrapper. `pull` just +# execs `docker pull`, so we still smoke its host-side dispatch. +assert_no_set_u_leak "pull dispatch (no set -u leak)" "$SCRIPT" pull + +# ── 7. Unknown subcommand → cmd_in_container fallback path. Same rule: +# clean error, no raw bash leak. +assert_no_set_u_leak "unknown subcommand dispatch" "$SCRIPT" no-such-subcommand + +# ── 8. Pre-populated LUCEBOX_HOST_* env (simulates an already-probed host +# whose vars are passed in from a parent process). Useful in CI matrices +# where we want to mock a "good host" without nvidia-smi/docker on PATH. +out=$( + NO_COLOR=1 \ + LUCEBOX_HOST_HAS_SYSTEMD=0 \ + LUCEBOX_HOST_HAS_DOCKER=0 \ + LUCEBOX_HOST_HAS_CTK=none \ + LUCEBOX_HOST_GPU_VENDOR=none \ + LUCEBOX_HOST_GPU_NAME="" \ + LUCEBOX_HOST_GPU_COUNT=0 \ + LUCEBOX_HOST_VRAM_GB=0 \ + LUCEBOX_HOST_GPU_SM="" \ + LUCEBOX_HOST_DRIVER_VERSION="" \ + LUCEBOX_HOST_DRIVER_MAJOR=0 \ + LUCEBOX_HOST_NPROC=1 \ + LUCEBOX_HOST_RAM_GB=0 \ + LUCEBOX_HOST_IS_WSL=0 \ + LUCEBOX_HOST_DOCKER_VERSION="" \ + timeout 5 bash "$SCRIPT" start 2>&1 || true +) +if grep -qE 'unbound variable|syntax error' <<<"$out"; then + report fail "start with pre-populated LUCEBOX_HOST_* env" "leak: $(head -3 <<<"$out")" +else + report ok "start with pre-populated LUCEBOX_HOST_* env" +fi + +# ── 8b. PIN the top-of-script LUCEBOX_HOST_* safe-default seeds. Even with +# probe_host short-circuited to a no-op (the worst-case bug recurrence: a +# future refactor accidentally deletes the call from a dispatch path) the +# wrapper must not leak `unbound variable` on `start`. We achieve "probe_host +# is a no-op" by exporting `_LUCEBOX_HOST_PROBED=1` so ensure_probed skips +# the real probe — equivalent to a future refactor that calls ensure_probed +# but mis-implements the gate. +out=$( + NO_COLOR=1 \ + _LUCEBOX_HOST_PROBED=1 \ + timeout 5 bash "$SCRIPT" start 2>&1 || true +) +if grep -qE 'unbound variable|syntax error' <<<"$out"; then + report fail "start with probe_host bypassed (seed defaults must catch this)" "leak: $(head -3 <<<"$out")" +else + report ok "start with probe_host bypassed (seed defaults intact)" +fi + +# Same for every other systemd-surface subcommand, since the seed defaults +# are the only thing keeping these safe under `set -u` if probe_host is ever +# bypassed. +for sub in stop restart enable disable status install uninstall logs; do + out=$( + NO_COLOR=1 \ + _LUCEBOX_HOST_PROBED=1 \ + timeout 5 bash "$SCRIPT" "$sub" -n 0 --no-pager 2>&1 || true + ) + if grep -qE 'unbound variable|syntax error' <<<"$out"; then + report fail "$sub with probe_host bypassed" "leak: $(head -3 <<<"$out")" + else + report ok "$sub with probe_host bypassed" + fi +done + +# ── 8c. Install path writes a robust unit file. Use a sandbox HOME so we +# don't clobber the developer's real ~/.config/systemd/user/lucebox.service, +# and verify the generated unit contains the Environment= / ExecStartPre= +# hardening that Bug 2 ("systemctl start succeeds but no container") added. +# The install runs in a host with no real systemd (the sandbox doesn't have +# `systemctl --user`), so we pre-seed LUCEBOX_HOST_HAS_SYSTEMD=1 to slip past +# the require_systemd gate, then stub out the `systemctl` binary itself so +# daemon-reload is a no-op. +test_install_writes_robust_unit() { + local label="install writes hardened unit file" + local sandbox shim_dir + sandbox=$(mktemp -d) + shim_dir="$sandbox/bin" + mkdir -p "$shim_dir" + # Stub systemctl + docker + nvidia-smi + loginctl so the install's + # require_host_prereqs and daemon-reload calls all succeed. + for binname in systemctl docker nvidia-smi loginctl; do + cat > "$shim_dir/$binname" <<'STUB' +#!/usr/bin/env bash +case "$1" in + ps|version) exit 0 ;; + show-user) echo "Linger=no" ;; + --query-gpu=*) echo "Fake, 24576, 550.00, 8.9" ;; +esac +exit 0 +STUB + chmod +x "$shim_dir/$binname" + done + local out rc unit_path + unit_path="$sandbox/.config/systemd/user/lucebox.service" + out=$( + set +e + HOME="$sandbox" \ + XDG_CONFIG_HOME="$sandbox/.config" \ + XDG_DATA_HOME="$sandbox/.local/share" \ + PATH="$shim_dir:$PATH" \ + LUCEBOX_HOST_HAS_SYSTEMD=1 \ + LUCEBOX_HOST_HAS_DOCKER=1 \ + LUCEBOX_HOST_HAS_CTK=runtime \ + LUCEBOX_HOST_GPU_VENDOR=nvidia \ + _LUCEBOX_HOST_PROBED=1 \ + NO_COLOR=1 \ + timeout 10 bash "$SCRIPT" install 2>&1 + echo "RC=$?" + ) + rc=$(grep -oE 'RC=[0-9]+$' <<<"$out" | tail -1 | sed 's/^RC=//') + rc="${rc:-99}" + if [ "$rc" != "0" ]; then + report fail "$label" "exit $rc; output: $(head -10 <<<"$out")" + rm -rf "$sandbox" + return + fi + if [ ! -f "$unit_path" ]; then + report fail "$label" "unit file not written at $unit_path" + rm -rf "$sandbox" + return + fi + # Required hardening — each line is a Bug-2 root-cause defence: + # ExecStartPre=…docker rm -f … → clear orphaned container name + # Environment=PATH=… → systemd user-session PATH is sparse + # Environment=LUCEBOX_IMAGE=… → pin the image the user installed against + # SuccessExitStatus=143 → `serve` exits 143 on SIGTERM; a normal + # `systemctl stop` must land "inactive", not "failed" + local missing="" + for needle in \ + "ExecStartPre=" \ + "Environment=PATH=" \ + "Environment=LUCEBOX_IMAGE=" \ + "Environment=LUCEBOX_VARIANT=" \ + "Environment=LUCEBOX_PORT=" \ + "Environment=LUCEBOX_MODELS=" \ + "SuccessExitStatus=143" \ + ; do + grep -qF "$needle" "$unit_path" || missing="$missing $needle" + done + if [ -n "$missing" ]; then + report fail "$label" "unit missing required directives:$missing" + rm -rf "$sandbox" + return + fi + report ok "$label" + rm -rf "$sandbox" +} +test_install_writes_robust_unit + +# ── 9. entrypoint.sh dispatch — confirm the in-container dispatch routes +# trivial subcommands (shell, an unknown passthrough) without firing +# `set -u` on DFLASH_* / DRAFT_* vars that only get assigned on the +# serve path. We can't fully exec the serve path here (it needs nvidia +# and the compiled binary) but we can confirm the early dispatch is clean. +# +# Each `exec` would actually try to run the underlying binary, which we +# don't have — so we shim it by overriding `exec` via a wrapper script. +# Easier: just confirm `bash -n` parses and run a tiny subset. +out=$(NO_COLOR=1 SUBCMD=help bash -c " + cd '$ROOT' + # Simulate 'docker run ... lucebox-hub:cuda12 shell echo ok' — entrypoint + # gets SUBCMD=shell and execs /bin/bash with the rest of argv. We replace + # exec via PATH so we don't actually exec. + tmpdir=\$(mktemp -d) + trap 'rm -rf \$tmpdir' EXIT + cat > \$tmpdir/uv <<'STUB' +#!/usr/bin/env bash +echo \"uv stub: \$*\" +exit 0 +STUB + chmod +x \$tmpdir/uv + PATH=\$tmpdir:\$PATH bash $ENTRYPOINT shell -c 'echo entrypoint-shell-dispatched' +" 2>&1 || true) +if grep -qE 'unbound variable|syntax error' <<<"$out"; then + report fail "entrypoint shell dispatch (no set -u leak)" "leak: $(head -5 <<<"$out")" +else + report ok "entrypoint shell dispatch (no set -u leak)" +fi + +# ── 10. entrypoint.sh serve-path under `set -u` — drive the REAL +# server/scripts/entrypoint.sh through its full draft-resolution block by +# sandboxing it with a synthetic DFLASH_DIR layout and a `dflash_server` +# shim that captures argv instead of execing the native binary. The +# `DRAFT_FAMILY_GLOB: unbound variable` bug fired precisely here — the +# previous version of this test inlined the block instead of sourcing +# the real file, and silently passed even when the shipped script was +# broken. So this test invokes server/scripts/entrypoint.sh directly. +# Build the shared entrypoint-serve sandbox: a synthetic DFLASH_DIR layout +# plus the `dflash_server` + `nvidia-smi` shims used by the three serve-path +# tests below. Assigns sandbox/models_dir/draft_dir/bin_dir/shim_dir into the +# CALLER'S scope (bash dynamic scoping) — the caller must `local`-declare +# them first. Mirrors the _make_docker_shim factoring above. +_make_entrypoint_sandbox() { + sandbox=$(mktemp -d) + models_dir="$sandbox/models" + draft_dir="$models_dir/draft" + bin_dir="$sandbox/build" + shim_dir="$sandbox/bin" + mkdir -p "$draft_dir" "$bin_dir" "$shim_dir" + # `dflash_server` shim — print argv and exit 0 instead of running. + cat > "$bin_dir/dflash_server" <<'STUB' +#!/usr/bin/env bash +printf '[shim] dflash_server' +for a in "$@"; do printf ' %q' "$a"; done +printf '\n' +exit 0 +STUB + chmod +x "$bin_dir/dflash_server" + # `nvidia-smi` shim — pretend we have a 24 GB GPU so the autotune + # block runs but doesn't pick the under-12-GB warn tier. + cat > "$shim_dir/nvidia-smi" <<'STUB' +#!/usr/bin/env bash +case "$*" in + *"--query-gpu=memory.total"*) echo 24576 ;; + -L|*-L*) echo "GPU 0: Fake (UUID: 0)" ;; + *) echo "ok" ;; +esac +exit 0 +STUB + chmod +x "$shim_dir/nvidia-smi" +} + +test_entrypoint_serve_path() { + local label="$1" target_name="$2" draft_file="$3" + local sandbox draft_dir models_dir bin_dir shim_dir + _make_entrypoint_sandbox + # Synthetic target (must be a real file at least 5 GB to pass the + # auto-detect block, OR we set DFLASH_TARGET explicitly to skip it). + touch "$models_dir/$target_name" + touch "$draft_dir/$draft_file" + + local out rc + out=$( + set +e + PATH="$shim_dir:$PATH" \ + DFLASH_DIR="$sandbox" \ + DFLASH_SERVER_BIN="$bin_dir/dflash_server" \ + DFLASH_TARGET="$models_dir/$target_name" \ + DFLASH_DRAFT="$draft_dir" \ + timeout 10 bash "$ENTRYPOINT" serve 2>&1 + echo "RC=$?" + ) + rc=$(grep -oE 'RC=[0-9]+$' <<<"$out" | tail -1 | sed 's/^RC=//') + rc="${rc:-99}" + rm -rf "$sandbox" + if grep -qE 'unbound variable|syntax error' <<<"$out"; then + report fail "$label" "leak: $(head -5 <<<"$out")" + elif [ "$rc" != "0" ]; then + report fail "$label" "exit $rc; output: $(head -5 <<<"$out")" + elif ! grep -qF "[shim] dflash_server" <<<"$out"; then + report fail "$label" "shim never executed; output: $(head -5 <<<"$out")" + else + report ok "$label" + fi +} + +# Exercise three branches of the family-glob logic: qwen3.6 + gemma-4 (the +# two families with family-specific globs) and an unknown target that +# triggers the empty-FAMILY_GLOBS fallback to the generic glob list. +test_entrypoint_serve_path "entrypoint serve: qwen3.6 family match" \ + "Qwen3.6-27B-Q4_K_M.gguf" "dflash-draft-3.6-test.gguf" +test_entrypoint_serve_path "entrypoint serve: gemma-4-31b family match" \ + "gemma-4-31B-it-Q8_0.gguf" "gemma-4-31b-dflash-q8.gguf" +test_entrypoint_serve_path "entrypoint serve: generic fallback" \ + "Mystery-Model-7B.gguf" "model.gguf" + +# ── 11. entrypoint.sh serve-path with MULTIPLE target-sized GGUFs in +# models/. The single-candidate fixture in test 10 doesn't exercise the +# auto-detect path that picks "first alphabetically" when more than one +# target ≥5 GB lives in the models dir — that path is what the sindri +# decode sweep tripped over after the user added the qwen3.6-moe preset +# (commit 4b6bced) alongside the existing Qwen3.6-27B target. The crash +# manifested as `DRAFT_FAMILY_GLOB: unbound variable`, and the partial +# fix in a87bb93 didn't survive a recurrence. +# +# Uses sparse files (`truncate -s 6G`) so the test stays cheap on disk — +# the 6 GB virtual size is enough to clear the find ... -size +5G filter +# without consuming actual blocks. Skip if truncate is missing (e.g. +# minimal busybox CI image). +test_entrypoint_multi_target() { + local label="$1" + shift + if ! command -v truncate &>/dev/null; then + report ok "$label (skipped: truncate not available)" + return + fi + local sandbox draft_dir models_dir bin_dir shim_dir + _make_entrypoint_sandbox + # Two qwen3.6-shaped targets ≥5 GB each — exactly the layout that + # broke on sindri (Qwen3.6-27B + Qwen3.6-35B-A3B-UD-Q4_K_M). + truncate -s 6G "$models_dir/Qwen3.6-27B-Q4_K_M.gguf" + truncate -s 6G "$models_dir/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf" + touch "$draft_dir/dflash-draft-3.6-test.gguf" + + local out rc + out=$( + set +e + # NOTE: deliberately NOT setting DFLASH_TARGET — the test must + # exercise the auto-detect block (line ~151). The explicit-config + # workaround from the bug report would skip the bug entirely. + PATH="$shim_dir:$PATH" \ + DFLASH_DIR="$sandbox" \ + DFLASH_SERVER_BIN="$bin_dir/dflash_server" \ + DFLASH_DRAFT="$draft_dir" \ + timeout 10 bash "$ENTRYPOINT" serve 2>&1 + echo "RC=$?" + ) + rc=$(grep -oE 'RC=[0-9]+$' <<<"$out" | tail -1 | sed 's/^RC=//') + rc="${rc:-99}" + rm -rf "$sandbox" + # The auto-detect block is entered (so any `set -u` regression on + # DRAFT_FAMILY_GLOB will trip) and then the entrypoint refuses to + # auto-pick — the deliberate safety added in PR #334's cubic round. + # We require: no set-u leak, the refuse warn fired, a non-zero exit + # (so a future regression that logs the warning but still returns 0 + # cannot slip past — the container MUST fail to start, not silently + # auto-pick a stale GGUF), and the shim was NOT exec'd. + if grep -qE 'unbound variable|syntax error' <<<"$out"; then + report fail "$label" "leak: $(grep -E 'unbound variable|syntax error' <<<"$out" | head -3)" + elif ! grep -qF "Refusing to auto-select" <<<"$out"; then + report fail "$label" "refuse-to-auto-pick warn missing — did the auto-detect block fire? rc=$rc output: $(head -5 <<<"$out")" + elif [ "$rc" = "0" ]; then + report fail "$label" "refuse warn fired but rc=0 — entrypoint must exit non-zero on multi-target refuse" + elif grep -qF "[shim] dflash_server" <<<"$out"; then + report fail "$label" "shim was exec'd despite multi-target refuse" + else + report ok "$label" + fi +} + +# Drive the regression: the sindri layout that broke (post-moe-preset). +test_entrypoint_multi_target "entrypoint serve: multi-target auto-detect (no DRAFT_FAMILY_GLOB leak)" + +# Also drive the DFLASH_DRAFT-is-a-file path. The init at entrypoint.sh:257 +# sits inside `if [ -d "$DFLASH_DRAFT" ]; then` — when DRAFT is a file the +# block is skipped, and any future read of DRAFT_FAMILY_GLOB outside the +# block would trip set -u. The defensive `:-` guard at the read site is +# meant to survive that refactor; this test guarantees it. +test_entrypoint_draft_is_file() { + local label="$1" + local sandbox draft_dir models_dir bin_dir shim_dir + _make_entrypoint_sandbox + touch "$models_dir/Qwen3.6-27B-Q4_K_M.gguf" + # DFLASH_DRAFT points at a FILE (not a directory). + touch "$draft_dir/dflash-draft-3.6-test.gguf" + + local out rc + out=$( + set +e + PATH="$shim_dir:$PATH" \ + DFLASH_DIR="$sandbox" \ + DFLASH_SERVER_BIN="$bin_dir/dflash_server" \ + DFLASH_TARGET="$models_dir/Qwen3.6-27B-Q4_K_M.gguf" \ + DFLASH_DRAFT="$draft_dir/dflash-draft-3.6-test.gguf" \ + timeout 10 bash "$ENTRYPOINT" serve 2>&1 + echo "RC=$?" + ) + rc=$(grep -oE 'RC=[0-9]+$' <<<"$out" | tail -1 | sed 's/^RC=//') + rc="${rc:-99}" + rm -rf "$sandbox" + if grep -qE 'unbound variable|syntax error' <<<"$out"; then + report fail "$label" "leak: $(grep -E 'unbound variable|syntax error' <<<"$out" | head -3)" + elif [ "$rc" != "0" ]; then + report fail "$label" "exit $rc; output: $(head -5 <<<"$out")" + else + report ok "$label" + fi +} +test_entrypoint_draft_is_file "entrypoint serve: DFLASH_DRAFT is a file (no DRAFT_FAMILY_GLOB leak)" + +# ── 12. entrypoint.sh writes HOST_INFO atomically on the serve path. The +# C++ server reads /opt/lucebox-hub/HOST_INFO into ServerConfig.host_info +# and surfaces it under /props.host. We can't write to /opt/lucebox-hub +# from the test runner, so override the path by sourcing the helpers and +# calling _build_host_info_json directly. The full entrypoint runs in +# test 10/11 already; this test pins the JSON shape independently. +test_entrypoint_host_info_json() { + local label="$1" + # Source the helper functions from the real entrypoint.sh. + # shellcheck disable=SC1090 + source <(awk '/^_json_escape\(\) \{/,/^\}/' "$ENTRYPOINT") + # shellcheck disable=SC1090 + source <(awk '/^_json_str_or_null\(\) \{/,/^\}/' "$ENTRYPOINT") + # shellcheck disable=SC1090 + source <(awk '/^_json_int_or_null\(\) \{/,/^\}/' "$ENTRYPOINT") + # shellcheck disable=SC1090 + source <(awk '/^_trim\(\) \{/,/^\}/' "$ENTRYPOINT") + # shellcheck disable=SC1090 + source <(awk '/^_emit_gpu_array\(\) \{/,/^\}/' "$ENTRYPOINT") + # shellcheck disable=SC1090 + source <(awk '/^_build_host_info_json\(\) \{/,/^\}/' "$ENTRYPOINT") + + local out + LUCEBOX_HOST_OS_PRETTY="Ubuntu 22.04.3 LTS" \ + LUCEBOX_HOST_KERNEL="6.6.87.2-microsoft-standard-WSL2" \ + LUCEBOX_HOST_WSL_VERSION="wsl2" \ + LUCEBOX_HOST_DOCKER_VERSION="29.1.3" \ + LUCEBOX_HOST_DRIVER_VERSION="596.36" \ + LUCEBOX_HOST_NVIDIA_CTK_VERSION="1.16.2" \ + LUCEBOX_HOST_CPU_MODEL='Intel(R) Core(TM) Ultra 9 275HX' \ + LUCEBOX_HOST_NPROC=24 \ + LUCEBOX_HOST_RAM_GB=64 \ + LUCEBOX_HOST_GPU_LIST_CSV="0, GPU-abc, 00000000:01:00.0, NVIDIA RTX 5090, 12.0, 24576 MiB, 175.00 W" \ + LUCEBOX_HOST_CUDA_VISIBLE_DEVICES="0" \ + out=$(_build_host_info_json "lucebox.sh" "lucebox.sh" "2026-05-28T20:31:42Z") + if ! python3 -c "import json,sys; d=json.loads(sys.argv[1]); assert d['os_pretty']=='Ubuntu 22.04.3 LTS'; assert d['wsl_version']=='wsl2'; assert d['nvidia_ctk_version']=='1.16.2'; assert d['source']=='lucebox.sh'; assert d['gpus'][0]['vram_gb']==24; assert d['gpus'][0]['name']=='NVIDIA RTX 5090'" "$out" >/dev/null 2>&1; then + report fail "$label (populated)" "JSON shape mismatch: $out" + return + fi + # Now drive the unknown path: every LUCEBOX_HOST_* unset → nulls and source=unknown. + out=$(env -i bash -c " + set -u + $(declare -f _json_escape _json_str_or_null _json_int_or_null _emit_gpu_array _build_host_info_json) + _build_host_info_json 'unknown' 'entrypoint.sh' '2026-05-28T20:31:42Z' + ") + if ! python3 -c "import json,sys; d=json.loads(sys.argv[1]); assert d['source']=='unknown'; assert d['gpus']==[]; assert d['os_pretty'] is None" "$out" >/dev/null 2>&1; then + report fail "$label (unknown)" "JSON shape mismatch: $out" + return + fi + report ok "$label" +} +test_entrypoint_host_info_json "entrypoint HOST_INFO JSON shape (populated + unknown)" + +# ── install.sh end-to-end ───────────────────────────────────────────────── +# Drive install.sh against a file:// URL pointing at a fixture lucebox.sh, +# verify the installed copy has LUCEBOX_INSTALLED_FROM rewritten to the +# fetched URL — that's the contract that `lucebox update` depends on to +# preserve the user's channel across upgrades. +test_install_sh_bakes_source_url() { + local label="$1" + local tmp dest_dir dest_path src_url out rc + tmp=$(mktemp -d -t lucebox-install.XXXXXX) + # Use the real lucebox.sh as the "remote" file — `file://` works with + # curl out of the box and exercises the same install.sh code path as + # an https fetch would. + src_url="file://$SCRIPT" + dest_dir="$tmp/bin" + dest_path="$dest_dir/lucebox" + out=$(LUCEBOX_INSTALL_URL="$src_url" LUCEBOX_INSTALL_DEST="$dest_path" \ + NO_COLOR=1 bash "$INSTALLER" 2>&1) || rc=$? + rc="${rc:-0}" + if [ "$rc" -ne 0 ]; then + rm -rf "$tmp" + report fail "$label" "installer exited $rc; output: $(printf '%s' "$out" | head -3)" + return + fi + if [ ! -x "$dest_path" ]; then + rm -rf "$tmp" + report fail "$label" "installed file missing or not executable at $dest_path" + return + fi + if ! grep -q "^LUCEBOX_INSTALLED_FROM=\"$src_url\"$" "$dest_path"; then + rm -rf "$tmp" + report fail "$label" "LUCEBOX_INSTALLED_FROM not rewritten in installed copy" + return + fi + rm -rf "$tmp" + report ok "$label" +} +test_install_sh_bakes_source_url "install.sh bakes LUCEBOX_INSTALLED_FROM into installed copy" + +# ── update dispatch ─────────────────────────────────────────────────────── +# `lucebox update` must dispatch to cmd_update — verify it's wired in the +# main case statement and appears in --help. We can't actually run the +# update (it'd curl + replace this very script) so the test is parse-level. +test_update_subcommand_wired() { + local label="$1" + local out + out=$(LUCEBOX_HOST_HAS_SYSTEMD=0 "$SCRIPT" --help 2>&1) + if ! grep -q '^ update ' <<<"$out"; then + report fail "$label" "update command missing from --help output" + return + fi + if ! grep -q '^[[:space:]]*update)[[:space:]]*cmd_update' "$SCRIPT"; then + report fail "$label" "update) → cmd_update dispatch not wired" + return + fi + report ok "$label" +} +test_update_subcommand_wired "lucebox update subcommand is wired" + +# ── IMAGE_BASE derived from install source ──────────────────────────────── +# Source lucebox.sh in a subshell with LUCEBOX_INSTALLED_FROM pointing at +# various URLs, then check that IMAGE_BASE comes out right. Uses +# `set -e; return` early so we don't actually run the wrapper's main(). +test_image_base_derives_from_install_url() { + local label="$1" url expected got + for case in \ + "https://raw.githubusercontent.com/easel/lucebox-hub/feat/lucebox-docker/lucebox.sh|ghcr.io/easel/lucebox-hub" \ + "https://raw.githubusercontent.com/Luce-Org/lucebox-hub/main/lucebox.sh|ghcr.io/luce-org/lucebox-hub" \ + "https://raw.githubusercontent.com/easel/lucebox-hub/601ab52/lucebox.sh|ghcr.io/easel/lucebox-hub" \ + "https://example.com/bogus|ghcr.io/luce-org/lucebox-hub" + do + url="${case%%|*}" + expected="${case##*|}" + # Extract the derivation function from the script and run it in + # isolation — sourcing the whole script triggers main() and side + # effects we don't want under a test harness. + got=$(bash -c ' + '"$(sed -n "/^_lucebox_derive_image()/,/^}/p" "$SCRIPT")"' + _lucebox_derive_image "$1" + ' bash "$url") + if [ "$got" != "$expected" ]; then + report fail "$label" "url=$url expected=$expected got=$got" + return + fi + done + report ok "$label" +} +test_image_base_derives_from_install_url "IMAGE_BASE derived from LUCEBOX_INSTALLED_FROM (4 URL shapes)" + +# ── config.toml reader + resolver ───────────────────────────────────────── +# Drive _lucebox_config_get + _lucebox_resolve against a fixture +# config.toml in a tmp $LUCEBOX_HOME. Verifies the wrapper agrees with +# the Python CLI on every scalar that lives in [image]/[runtime]/[paths]. +test_config_toml_reader_and_resolve() { + local label="$1" tmp got + tmp=$(mktemp -d -t lucebox-cfg.XXXXXX) + cat > "$tmp/config.toml" <<'TOML' +[image] +variant = "cuda13" +registry = "ghcr.io/myorg/forkedhub" + +[runtime] +port = 9090 +container_name = "luce-test" + +[paths] +models = "/srv/models" + +[dflash] +budget = 22 +lazy = false +TOML + + # Exercise both helpers + the resolver via a subshell that sources + # the relevant snippets out of lucebox.sh. Each case is a triple: + # env_value | toml_key | default | expected + local cases=( + "|image.registry|ghcr.io/luce-org/lucebox-hub|ghcr.io/myorg/forkedhub" + "|image.variant|cuda12|cuda13" + "|runtime.port|8080|9090" + "|runtime.container_name|lucebox|luce-test" + "|paths.models|/var/lib/lucebox|/srv/models" + "OVERRIDE|image.registry|ghcr.io/luce-org/lucebox-hub|OVERRIDE" + "|missing.key|fallback-default|fallback-default" + ) + local case env_value toml_key default expected + for case in "${cases[@]}"; do + IFS='|' read -r env_value toml_key default expected <<<"$case" + got=$(LUCEBOX_HOME="$tmp" bash -c ' + '"$(sed -n "/^_lucebox_config_path()/,/^}/p" "$SCRIPT")"' + '"$(sed -n "/^_lucebox_config_get()/,/^}/p" "$SCRIPT")"' + '"$(sed -n "/^_lucebox_resolve()/,/^}/p" "$SCRIPT")"' + _lucebox_resolve "$1" "$2" "$3" + ' bash "$env_value" "$toml_key" "$default") + if [ "$got" != "$expected" ]; then + rm -rf "$tmp" + report fail "$label" "env=$env_value key=$toml_key default=$default expected=$expected got=$got" + return + fi + done + rm -rf "$tmp" + report ok "$label" +} +test_config_toml_reader_and_resolve "config.toml reader + env > toml > default resolution (7 cases)" + +# ── cmd_serve under systemd: INVOCATION_ID short-circuits is-active ────── +# When systemd invokes the wrapper as a unit's ExecStart, it sets +# $INVOCATION_ID. The wrapper must NOT then refuse "already running under +# systemd" — that's a self-defeating check that turns into a restart loop. +# Verify the guard is present in the source (the actual behavior requires +# a running systemd unit to test end-to-end, which the harness can't do). +test_cmd_serve_invocation_id_guard() { + local label="$1" + if ! grep -q 'INVOCATION_ID' "$SCRIPT"; then + report fail "$label" "INVOCATION_ID guard missing from cmd_serve preflight" + return + fi + # The guard must be the AND-condition gating the is-active check. + # If grep finds the is-active line WITHOUT INVOCATION_ID nearby, + # the guard isn't wired correctly. + if ! awk ' + /INVOCATION_ID/ { saw_guard = NR } + /is-active --quiet "\$UNIT_NAME"/ { + if (saw_guard && NR - saw_guard <= 3) found = 1 + } + END { exit (found ? 0 : 1) } + ' "$SCRIPT"; then + report fail "$label" "INVOCATION_ID not adjacent to is-active check (guard not wired)" + return + fi + report ok "$label" +} +test_cmd_serve_invocation_id_guard "cmd_serve has INVOCATION_ID guard on systemd is-active check" + +# ── cmd_systemctl_passthrough: smart start ─────────────────────────────── +# Verify the source has the "already active" + "restart loop" short +# circuits for the start action. Behavior-level testing requires a real +# unit; this is a source-level guarantee that the branches exist. +test_cmd_start_already_active_shortcircuit() { + local label="$1" + if ! grep -q 'is already active' "$SCRIPT"; then + report fail "$label" "already-active short-circuit missing" + return + fi + if ! grep -q 'is in restart-loop' "$SCRIPT"; then + report fail "$label" "restart-loop short-circuit missing" + return + fi + report ok "$label" +} +test_cmd_start_already_active_shortcircuit "lucebox start has already-active + restart-loop short-circuits" + +# ── install.sh SHA-pin refusal + CHANNEL override ──────────────────────── +# A SHA-pinned LUCEBOX_INSTALL_URL with no LUCEBOX_INSTALL_CHANNEL must +# refuse — otherwise `lucebox update` would re-fetch that frozen SHA +# forever. With CHANNEL set, the bake-in uses the channel URL, not the +# fetch URL. +test_install_sha_pin_refusal_and_channel_override() { + local label="$1" tmp got rc + tmp=$(mktemp -d -t lucebox-sha.XXXXXX) + + # Case 1: SHA-pinned URL without CHANNEL → must refuse + LUCEBOX_INSTALL_URL="https://raw.githubusercontent.com/easel/lucebox-hub/abc1234567/lucebox.sh" \ + LUCEBOX_INSTALL_DEST="$tmp/lucebox1" \ + NO_COLOR=1 \ + bash "$INSTALLER" >/dev/null 2>&1 && rc=0 || rc=$? + if [ "$rc" -eq 0 ]; then + rm -rf "$tmp" + report fail "$label" "SHA-pinned URL without CHANNEL should have refused (rc=$rc, got success)" + return + fi + if [ -f "$tmp/lucebox1" ]; then + rm -rf "$tmp" + report fail "$label" "SHA-pinned URL refusal still wrote $tmp/lucebox1" + return + fi + + # Case 2: SHA-pinned URL WITH CHANNEL → installs, bakes CHANNEL + LUCEBOX_INSTALL_URL="file://$SCRIPT" \ + LUCEBOX_INSTALL_CHANNEL="https://raw.githubusercontent.com/easel/lucebox-hub/feat/lucebox-docker/lucebox.sh" \ + LUCEBOX_INSTALL_DEST="$tmp/lucebox2" \ + NO_COLOR=1 \ + bash "$INSTALLER" >/dev/null 2>&1 || rc=$? + got=$(grep '^LUCEBOX_INSTALLED_FROM=' "$tmp/lucebox2" 2>/dev/null || echo missing) + if [ "$got" != 'LUCEBOX_INSTALLED_FROM="https://raw.githubusercontent.com/easel/lucebox-hub/feat/lucebox-docker/lucebox.sh"' ]; then + rm -rf "$tmp" + report fail "$label" "CHANNEL not baked; got: $got" + return + fi + + rm -rf "$tmp" + report ok "$label" +} +test_install_sha_pin_refusal_and_channel_override "install.sh refuses SHA-pin without CHANNEL + honors CHANNEL override" + +# ── lucebox completion ─────────────────────────────────────────────────── +# The completion script must source cleanly and complete a known prefix. +test_completion_bash() { + local label="$1" out + out=$(LUCEBOX_HOST_HAS_SYSTEMD=0 bash -c ' + source <("$1" completion bash 2>/dev/null) + COMP_WORDS=(lucebox conf) + COMP_CWORD=1 + _lucebox_complete + printf "%s\n" "${COMPREPLY[@]}" + ' bash "$SCRIPT") + if ! grep -qx 'config' <<<"$out"; then + report fail "$label" "completion didn't suggest 'config' for prefix 'conf'; got: $(printf '%s' "$out" | tr '\n' ' ')" + return + fi + report ok "$label" +} +test_completion_bash "lucebox completion bash completes a known prefix" + +# ── docker exec routing ─────────────────────────────────────────────────── +# When the lucebox container is running, steady-state subcommands must +# `docker exec` into it (cheap + shares the live server's net namespace) and +# service-restarting subcommands (serve, pull, ...) must stay on +# `docker run`. We mock docker via a PATH shim that: +# - on `docker ps -q -f name=^lucebox$` prints a fake container id +# (signals "container is running") iff DOCKER_FAKE_RUNNING=1. +# - on any other call (run, exec, pull, ...) echoes its argv on stdout and +# exits 0. The test then asserts on the captured first-token (run vs exec) +# and trailing argv. +# +# nvidia-smi is stubbed too so probe_host doesn't barf, but the captured argv +# we care about is the docker invocation downstream of dispatch. +_make_docker_shim() { + local sandbox="$1" running="$2" + local shim_dir="$sandbox/bin" + mkdir -p "$shim_dir" + # docker shim: dispatch on first arg. Important: ps -q -f name=^lucebox$ + # must print a fake id when DOCKER_FAKE_RUNNING=1 and nothing otherwise. + # All other invocations (run, exec, pull) print "DOCKER_INVOKED " + # on stdout so the caller can grep it. + cat > "$shim_dir/docker" < "$shim_dir/nvidia-smi" <<'STUB' +#!/usr/bin/env bash +case "$*" in + *"--query-gpu="*) echo "Fake GPU, 24576, 550.00, 8.9" ;; + *) echo "ok" ;; +esac +exit 0 +STUB + chmod +x "$shim_dir/nvidia-smi" +} + +# Drive the wrapper through the dispatch case under test and capture the +# docker invocation it would have exec'd. Because `cmd_in_container` / +# `cmd_exec_in_container` call `exec docker ...` we replace `exec` semantics +# by running the wrapper in a subshell — the docker shim prints what it was +# called with and the captured stdout is the proof. +_run_wrapper_capture_docker() { + local sandbox="$1"; shift + local shim_dir="$sandbox/bin" + set +e + HOME="$sandbox" \ + XDG_CONFIG_HOME="$sandbox/.config" \ + XDG_DATA_HOME="$sandbox/.local/share" \ + LUCEBOX_HOME="$sandbox/.lucebox" \ + PATH="$shim_dir:$PATH" \ + LUCEBOX_HOST_HAS_DOCKER=1 \ + LUCEBOX_HOST_HAS_CTK=runtime \ + LUCEBOX_HOST_GPU_VENDOR=nvidia \ + LUCEBOX_HOST_DRIVER_MAJOR=550 \ + LUCEBOX_HOST_DRIVER_VERSION="550.00" \ + LUCEBOX_HOST_GPU_NAME="Fake GPU" \ + LUCEBOX_HOST_GPU_COUNT=1 \ + LUCEBOX_HOST_VRAM_GB=24 \ + LUCEBOX_HOST_GPU_SM="89" \ + LUCEBOX_HOST_NPROC=8 \ + LUCEBOX_HOST_RAM_GB=64 \ + LUCEBOX_HOST_HAS_SYSTEMD=0 \ + LUCEBOX_HOST_IS_WSL=0 \ + LUCEBOX_HOST_DOCKER_VERSION="29.1.3" \ + _LUCEBOX_HOST_PROBED=1 \ + NO_COLOR=1 \ + timeout 10 bash "$SCRIPT" "$@" 2>&1 + set -e +} + +test_routes_to_exec_when_running() { + local label="$1" sandbox out + sandbox=$(mktemp -d -t lucebox-route.XXXXXX) + _make_docker_shim "$sandbox" 1 + out=$(_run_wrapper_capture_docker "$sandbox" config get model.preset || true) + rm -rf "$sandbox" + if ! grep -q '^DOCKER_INVOKED exec' <<<"$out"; then + report fail "$label" "expected 'docker exec' invocation; got: $(head -3 <<<"$out")" + return + fi + if grep -q '^DOCKER_INVOKED run' <<<"$out"; then + report fail "$label" "got 'docker run' when container is up — should have exec'd" + return + fi + # Sanity: the exec line ends with `lucebox config get model.preset`. + if ! grep -qE 'lucebox config get model.preset' <<<"$out"; then + report fail "$label" "exec argv missing tail 'lucebox config get model.preset'; got: $(head -3 <<<"$out")" + return + fi + # The exec path must forward the LUCEBOX_* scalar env subset (shared + # with the docker-run path via _append_scalar_env). Pin LUCEBOX_IMAGE= + # so a regression in that helper is caught here. + if ! grep -q 'LUCEBOX_IMAGE=' <<<"$out"; then + report fail "$label" "exec argv missing 'LUCEBOX_IMAGE=' scalar env; got: $(head -3 <<<"$out")" + return + fi + report ok "$label" +} +test_routes_to_exec_when_running "config get routes to docker exec when container running" + +test_routes_to_run_when_not_running() { + local label="$1" sandbox out + sandbox=$(mktemp -d -t lucebox-route.XXXXXX) + _make_docker_shim "$sandbox" 0 + out=$(_run_wrapper_capture_docker "$sandbox" config get model.preset || true) + rm -rf "$sandbox" + if ! grep -q '^DOCKER_INVOKED run' <<<"$out"; then + report fail "$label" "expected 'docker run' invocation (container not running); got: $(head -3 <<<"$out")" + return + fi + if grep -q '^DOCKER_INVOKED exec' <<<"$out"; then + report fail "$label" "got 'docker exec' but container is not running — should fall back to run" + return + fi + report ok "$label" +} +test_routes_to_run_when_not_running "config get falls back to docker run when container not running" + +test_no_exec_flag_forces_run() { + local label="$1" sandbox out + sandbox=$(mktemp -d -t lucebox-route.XXXXXX) + _make_docker_shim "$sandbox" 1 + # --no-exec must override the prefer-exec path even when container is up. + out=$(_run_wrapper_capture_docker "$sandbox" --no-exec config get model.preset || true) + rm -rf "$sandbox" + if grep -q '^DOCKER_INVOKED exec' <<<"$out"; then + report fail "$label" "--no-exec failed to force run path; got exec" + return + fi + if ! grep -q '^DOCKER_INVOKED run' <<<"$out"; then + report fail "$label" "expected 'docker run' under --no-exec; got: $(head -3 <<<"$out")" + return + fi + report ok "$label" +} +test_no_exec_flag_forces_run "--no-exec flag forces docker run even when container is up" + +test_no_exec_env_forces_run() { + local label="$1" sandbox out + sandbox=$(mktemp -d -t lucebox-route.XXXXXX) + _make_docker_shim "$sandbox" 1 + out=$( + LUCEBOX_NO_EXEC=1 _run_wrapper_capture_docker "$sandbox" config get model.preset || true + ) + rm -rf "$sandbox" + if grep -q '^DOCKER_INVOKED exec' <<<"$out"; then + report fail "$label" "LUCEBOX_NO_EXEC=1 failed to force run path; got exec" + return + fi + if ! grep -q '^DOCKER_INVOKED run' <<<"$out"; then + report fail "$label" "expected 'docker run' under LUCEBOX_NO_EXEC=1; got: $(head -3 <<<"$out")" + return + fi + report ok "$label" +} +test_no_exec_env_forces_run "LUCEBOX_NO_EXEC=1 env override forces docker run" + +test_models_routes_to_exec() { + local label="$1" sandbox out + sandbox=$(mktemp -d -t lucebox-route.XXXXXX) + _make_docker_shim "$sandbox" 1 + out=$(_run_wrapper_capture_docker "$sandbox" models list || true) + rm -rf "$sandbox" + if ! grep -q '^DOCKER_INVOKED exec' <<<"$out"; then + report fail "$label" "expected 'docker exec' for models when running; got: $(head -3 <<<"$out")" + return + fi + # Confirm the exec'd command tail is `lucebox models list` — the + # in-container CLI's argv must NOT be polluted with dispatcher bookkeeping. + if ! grep -qE 'lucebox models list' <<<"$out"; then + report fail "$label" "exec'd argv missing 'lucebox models list' tail" + return + fi + report ok "$label" +} +test_models_routes_to_exec "models list routes to docker exec when container running" + +# ── usage mentions exec-when-running ────────────────────────────────────── +test_usage_mentions_exec_routing() { + local label="$1" out + out=$(NO_COLOR=1 bash "$SCRIPT" --help 2>&1) + if ! grep -qi 'docker exec\|--no-exec' <<<"$out"; then + report fail "$label" "usage doesn't mention the exec routing / --no-exec flag" + return + fi + report ok "$label" +} +test_usage_mentions_exec_routing "usage documents docker exec routing + --no-exec flag" + +# ── TTY flag selection. Regression guard for the process-substitution bug: +# _set_tty_flags must run in the CALLER's scope so `[ -t 1 ]` inspects the +# real terminal. If it is ever moved back behind `< <(...)` or `$(...)`, +# fd 1 becomes a pipe and it emits -i even on a real tty, silently dropping +# docker's -t and breaking the interactive client TUIs (lucebox claude …). +# The rest of this suite runs non-tty, so only this test exercises the -it +# branch — via a real PTY allocated by python's pty.fork. +test_tty_flags_selection() { + local label="$1" fn out + fn=$(awk '/^_set_tty_flags\(\) \{/,/^\}/' "$SCRIPT") + + # (a) non-tty (stdin /dev/null, stdout a pipe) → -i + out=$(bash -c "$fn"$'\n''f=(); _set_tty_flags f; printf "%s" "${f[*]}"' /dev/null) + if [ "$out" != "-i" ]; then + report fail "$label" "non-tty expected -i, got '$out'" + return + fi + + # (b) real tty on fd0+fd1 (python pty.fork) → -it + out=$(python3 - "$SCRIPT" <<'PY' 2>/dev/null +import os, pty, re, sys +src = open(sys.argv[1]).read() +fn = re.search(r'^_set_tty_flags\(\) \{.*?^\}', src, re.S | re.M).group(0) +script = fn + '\nf=(); _set_tty_flags f; printf "TTYFLAG=%s\\n" "${f[*]}"\n' +pid, fd = pty.fork() +if pid == 0: + os.execvp("bash", ["bash", "-c", script]) +buf = b"" +try: + while True: + chunk = os.read(fd, 1024) + if not chunk: + break + buf += chunk +except OSError: + pass +os.waitpid(pid, 0) +m = re.search(rb"TTYFLAG=(\S+)", buf) +sys.stdout.write(m.group(1).decode() if m else "NONE") +PY +) + if [ "$out" != "-it" ]; then + report fail "$label" "real tty expected -it, got '$out'" + return + fi + report ok "$label" +} +test_tty_flags_selection "_set_tty_flags: -it on a real tty, -i otherwise" + +echo +if [ "$fail" -eq 0 ]; then + echo "[test_lucebox_sh] $pass passed, 0 failed" + exit 0 +else + echo "[test_lucebox_sh] $pass passed, $fail failed" >&2 + exit 1 +fi diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index 2602517f5..cd0b5184e 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -73,6 +73,15 @@ esac # write-failure (read-only FS, etc.) gets a warning and we continue. write_host_info() { local target="/opt/lucebox-hub/HOST_INFO" + # If the target dir doesn't exist (e.g. running the entrypoint outside + # the canonical container layout: unit tests, plain `docker run` without + # a bind mount), don't try to write — bash's own "No such file or + # directory" complaint on the `> "$tmp"` redirect below would leak to + # stderr regardless of `2>/dev/null` (that suppresses the command's + # stderr, not the redirect itself). HOST_INFO is informational. + if [ ! -d "$(dirname "$target")" ]; then + return 0 + fi local tmp="${target}.tmp.$$" local collected_at collected_at=$(date -u +%FT%TZ 2>/dev/null || echo "") @@ -158,6 +167,15 @@ _json_int_or_null() { # `nvidia-smi --query-gpu=index,uuid,pci.bus_id,name,compute_cap,memory.total,power.limit # --format=csv,noheader` produced on the host) into a JSON # array. Empty CSV → "[]". Each row becomes one object. +# Strip leading/trailing whitespace from a string. Pure bash (no sed fork) +# via prefix/suffix removal of the longest run of spaces or tabs. +_trim() { + local s="$1" + s="${s#"${s%%[![:space:]]*}"}" # leading + s="${s%"${s##*[![:space:]]}"}" # trailing + printf '%s' "$s" +} + _emit_gpu_array() { local csv="${LUCEBOX_HOST_GPU_LIST_CSV:-}" if [ -z "$csv" ]; then @@ -173,13 +191,13 @@ _emit_gpu_array() { # split on `,` alone and trim whitespace per field so both forms parse. local idx uuid pci name cc mem plimit IFS=',' read -r idx uuid pci name cc mem plimit <<<"$line" - idx=$(printf '%s' "$idx" | sed 's/^[[:space:]]*//; s/[[:space:]]*$//') - uuid=$(printf '%s' "$uuid" | sed 's/^[[:space:]]*//; s/[[:space:]]*$//') - pci=$(printf '%s' "$pci" | sed 's/^[[:space:]]*//; s/[[:space:]]*$//') - name=$(printf '%s' "$name" | sed 's/^[[:space:]]*//; s/[[:space:]]*$//') - cc=$(printf '%s' "$cc" | sed 's/^[[:space:]]*//; s/[[:space:]]*$//') - mem=$(printf '%s' "$mem" | sed 's/^[[:space:]]*//; s/[[:space:]]*$//') - plimit=$(printf '%s' "$plimit" | sed 's/^[[:space:]]*//; s/[[:space:]]*$//') + idx=$(_trim "$idx") + uuid=$(_trim "$uuid") + pci=$(_trim "$pci") + name=$(_trim "$name") + cc=$(_trim "$cc") + mem=$(_trim "$mem") + plimit=$(_trim "$plimit") # Strip units. "24576 MiB" → 24576; "175.00 W" → 175 (truncate). local mem_mib vram_gb power_w mem_mib=$(printf '%s' "$mem" | awk '{print $1+0}') diff --git a/uv.lock b/uv.lock index fee8de0df..ba16922d8 100644 --- a/uv.lock +++ b/uv.lock @@ -9,6 +9,7 @@ resolution-markers = [ [manifest] members = [ + "lucebox", "lucebox-dflash", "lucebox-hub", "pflash", @@ -429,6 +430,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/59/67/a6739ac96e28b7855808bdb0370e250606104a859750d209e5a0716fe7ab/librt-0.11.0-cp312-cp312-win_arm64.whl", hash = "sha256:2f10cf143e4a9bb0f4f5af568a00df94a2d69ef41c2579584454bb0fe5cc642c", size = 103470, upload-time = "2026-05-10T18:16:10.369Z" }, ] +[[package]] +name = "lucebox" +source = { editable = "lucebox" } +dependencies = [ + { name = "httpx" }, + { name = "huggingface-hub" }, + { name = "rich" }, + { name = "tomli-w" }, + { name = "typer" }, +] + +[package.metadata] +requires-dist = [ + { name = "httpx", specifier = ">=0.27" }, + { name = "huggingface-hub", specifier = ">=0.27" }, + { name = "rich", specifier = ">=13" }, + { name = "tomli-w", specifier = ">=1.0" }, + { name = "typer", specifier = ">=0.12" }, +] + [[package]] name = "lucebox-dflash" version = "0.1.0" @@ -466,6 +487,7 @@ name = "lucebox-hub" version = "0.0.0" source = { virtual = "." } dependencies = [ + { name = "lucebox" }, { name = "lucebox-dflash" }, { name = "pflash" }, ] @@ -482,6 +504,7 @@ megakernel = [ [package.metadata] requires-dist = [ + { name = "lucebox", editable = "lucebox" }, { name = "lucebox-dflash", virtual = "server" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.10,<2" }, { name = "pflash", editable = "optimizations/pflash" }, @@ -1124,6 +1147,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/72/f4/0de46cfa12cdcbcd464cc59fde36912af405696f687e53a091fb432f694c/tokenizers-0.22.2-cp39-abi3-win_arm64.whl", hash = "sha256:9ce725d22864a1e965217204946f830c37876eee3b2ba6fc6255e8e903d5fcbc", size = 2612133, upload-time = "2026-01-05T10:45:17.232Z" }, ] +[[package]] +name = "tomli-w" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/19/75/241269d1da26b624c0d5e110e8149093c759b7a286138f4efd61a60e75fe/tomli_w-1.2.0.tar.gz", hash = "sha256:2dd14fac5a47c27be9cd4c976af5a12d87fb1f0b4512f81d69cce3b35ae25021", size = 7184, upload-time = "2025-01-15T12:07:24.262Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/18/c86eb8e0202e32dd3df50d43d7ff9854f8e0603945ff398974c1d91ac1ef/tomli_w-1.2.0-py3-none-any.whl", hash = "sha256:188306098d013b691fcadc011abd66727d3c414c571bb01b1a174ba8c983cf90", size = 6675, upload-time = "2025-01-15T12:07:22.074Z" }, +] + [[package]] name = "torch" version = "2.11.0+cu128" From a083605c10e9f144db65b424089147241e8a31f5 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 27 Jul 2026 11:55:05 +0200 Subject: [PATCH 02/25] feat(lucebox): support AMD and heterogeneous GPU builds --- install.sh | 2 +- lucebox.sh | 402 +++++++++++++++++++++++++----- lucebox/README.md | 7 +- lucebox/src/lucebox/autotune.py | 2 +- lucebox/src/lucebox/cli.py | 2 +- lucebox/src/lucebox/config.py | 8 +- lucebox/src/lucebox/docker_run.py | 36 ++- lucebox/src/lucebox/host_check.py | 75 ++++-- lucebox/src/lucebox/host_facts.py | 7 +- lucebox/src/lucebox/types.py | 12 +- lucebox/tests/test_check.py | 50 ++++ lucebox/tests/test_config.py | 18 ++ lucebox/tests/test_docker_run.py | 35 ++- lucebox/tests/test_models_cli.py | 5 +- scripts/test_lucebox_sh.sh | 142 +++++++++++ server/scripts/entrypoint.sh | 68 ++++- 16 files changed, 767 insertions(+), 104 deletions(-) diff --git a/install.sh b/install.sh index cff54a02e..c44f92712 100755 --- a/install.sh +++ b/install.sh @@ -131,7 +131,7 @@ esac cat </dev/null || echo 1) # RAM: try Linux /proc/meminfo first, then macOS/BSD sysctl, else 0. @@ -232,18 +300,30 @@ probe_host() { LUCEBOX_HOST_RAM_GB=$(( mem_bytes / 1024 / 1024 / 1024 )) fi LUCEBOX_HOST_GPU_VENDOR="none" + LUCEBOX_HOST_HAS_NVIDIA_GPU=0 + LUCEBOX_HOST_HAS_AMD_GPU=0 LUCEBOX_HOST_GPU_NAME="" LUCEBOX_HOST_GPU_COUNT=0 LUCEBOX_HOST_VRAM_GB=0 LUCEBOX_HOST_GPU_SM="" LUCEBOX_HOST_DRIVER_VERSION="" LUCEBOX_HOST_DRIVER_MAJOR=0 + LUCEBOX_HOST_GPU_LIST_CSV="" + LUCEBOX_HOST_ROCM_VERSION="" + LUCEBOX_HOST_HAS_KFD=0 + LUCEBOX_HOST_HAS_DRI=0 + LUCEBOX_HOST_AMD_GPU_NAME="" + LUCEBOX_HOST_AMD_GPU_COUNT=0 + LUCEBOX_HOST_AMD_VRAM_GB=0 + LUCEBOX_HOST_AMD_GPU_ARCH="" + LUCEBOX_HOST_AMD_GPU_LIST_CSV="" if command -v nvidia-smi &>/dev/null; then local q if q=$(nvidia-smi --query-gpu=name,memory.total,driver_version,compute_cap \ --format=csv,noheader,nounits 2>/dev/null) && [ -n "$q" ]; then LUCEBOX_HOST_GPU_VENDOR="nvidia" + LUCEBOX_HOST_HAS_NVIDIA_GPU=1 LUCEBOX_HOST_GPU_NAME=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $1}') local mem_mib mem_mib=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $2}') @@ -264,8 +344,96 @@ probe_host() { --query-gpu=index,uuid,pci.bus_id,name,compute_cap,memory.total,power.limit \ --format=csv,noheader 2>/dev/null || echo "") fi + + # Probe AMD independently even on mixed NVIDIA + Strix systems. A working + # NVIDIA GPU remains the default backend (RTX 3090 + Strix → cuda12), but + # recording the AMD companion prevents the APU from confusing readiness + # reporting and lets an explicit rocm variant remain possible. + local amd_csv="" amd_rows="" + if command -v amd-smi &>/dev/null; then + amd_csv=$(amd-smi static --asic --vram --csv 2>/dev/null || echo "") + if [ -n "$amd_csv" ]; then + amd_rows=$(printf '%s\n' "$amd_csv" | _parse_amd_smi_csv) + fi + fi + if [ -z "$amd_rows" ] && command -v rocm-smi &>/dev/null; then + amd_csv=$(rocm-smi --showproductname --showmeminfo vram --csv 2>/dev/null || echo "") + if [ -n "$amd_csv" ]; then + amd_rows=$(printf '%s\n' "$amd_csv" | _parse_rocm_smi_csv) + fi + fi + # Minimal fallback for ROCm installations without either SMI frontend. + if [ -z "$amd_rows" ] && [ -e /dev/kfd ] && command -v rocminfo &>/dev/null; then + local roc_arches + roc_arches=$(rocminfo 2>/dev/null \ + | awk '/^[[:space:]]*Name:[[:space:]]+gfx[0-9a-z]+/{print $2}' \ + | awk '!seen[$0]++' || echo "") + if [ -n "$roc_arches" ]; then + local amd_idx=0 arch + while IFS= read -r arch; do + amd_rows+="${amd_idx}|AMD GPU|${arch}|0"$'\n' + amd_idx=$((amd_idx + 1)) + done <<<"$roc_arches" + amd_rows=${amd_rows%$'\n'} + fi + fi + + if [ -n "$amd_rows" ]; then + LUCEBOX_HOST_HAS_AMD_GPU=1 + LUCEBOX_HOST_AMD_GPU_COUNT=$(printf '%s\n' "$amd_rows" | awk 'NF{n++} END{print n+0}') + local amd_primary + amd_primary=$(printf '%s\n' "$amd_rows" | sort -t'|' -k4,4nr | head -1) + local amd_idx amd_name amd_arch amd_mem_mib + IFS='|' read -r amd_idx amd_name amd_arch amd_mem_mib <<<"$amd_primary" + LUCEBOX_HOST_AMD_GPU_NAME="$amd_name" + LUCEBOX_HOST_AMD_GPU_ARCH="$amd_arch" + LUCEBOX_HOST_AMD_VRAM_GB=$((amd_mem_mib / 1024)) + LUCEBOX_HOST_AMD_GPU_LIST_CSV=$(printf '%s\n' "$amd_rows" \ + | awk -F'|' '{printf "%s, , , %s, %s, %s MiB,\n", $1, $2, $3, $4}') + # Strix Halo exposes most memory as unified system RAM, while SMI may + # report only a 512 MiB carve-out. Use host RAM as the effective model + # capacity on a Strix-only build; a discrete R9700 remains primary on + # the R9700 + Strix build because it has the largest physical VRAM. + if [ "$amd_arch" = "gfx1151" ] \ + && [ "$LUCEBOX_HOST_AMD_VRAM_GB" -lt 12 ] \ + && [ "$LUCEBOX_HOST_RAM_GB" -ge 32 ]; then + LUCEBOX_HOST_AMD_VRAM_GB=$LUCEBOX_HOST_RAM_GB + fi + + if command -v amd-smi &>/dev/null; then + LUCEBOX_HOST_ROCM_VERSION=$(amd-smi version 2>/dev/null \ + | sed -n 's/.*ROCm version: \([^ |]*\).*/\1/p' \ + | head -1 || echo "") + fi + if [ -z "$LUCEBOX_HOST_ROCM_VERSION" ] && command -v hipconfig &>/dev/null; then + LUCEBOX_HOST_ROCM_VERSION=$(hipconfig --version 2>/dev/null \ + | sed 's/-.*//' | head -1 || echo "") + fi + + if [ "$LUCEBOX_HOST_GPU_VENDOR" = "none" ]; then + LUCEBOX_HOST_GPU_VENDOR="amd" + LUCEBOX_HOST_GPU_NAME="$LUCEBOX_HOST_AMD_GPU_NAME" + LUCEBOX_HOST_GPU_COUNT=$LUCEBOX_HOST_AMD_GPU_COUNT + LUCEBOX_HOST_VRAM_GB=$LUCEBOX_HOST_AMD_VRAM_GB + LUCEBOX_HOST_GPU_SM="$LUCEBOX_HOST_AMD_GPU_ARCH" + LUCEBOX_HOST_GPU_LIST_CSV="$LUCEBOX_HOST_AMD_GPU_LIST_CSV" + fi + fi + + if [ -r /dev/kfd ] && [ -w /dev/kfd ]; then + LUCEBOX_HOST_HAS_KFD=1 + fi + local render_node + for render_node in /dev/dri/renderD*; do + [ -e "$render_node" ] || continue + if [ -r "$render_node" ] && [ -w "$render_node" ]; then + LUCEBOX_HOST_HAS_DRI=1 + break + fi + done # CUDA_VISIBLE_DEVICES from the caller's env (empty default = "all GPUs"). LUCEBOX_HOST_CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-}" + LUCEBOX_HOST_HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-}" # OS / kernel identity. /etc/os-release is the freedesktop spec for # "what distro is this?" and we keep PRETTY_NAME verbatim (it already @@ -338,14 +506,20 @@ probe_host() { fi export LUCEBOX_HOST_NPROC LUCEBOX_HOST_RAM_GB LUCEBOX_HOST_GPU_VENDOR + export LUCEBOX_HOST_HAS_NVIDIA_GPU LUCEBOX_HOST_HAS_AMD_GPU export LUCEBOX_HOST_GPU_NAME LUCEBOX_HOST_GPU_COUNT LUCEBOX_HOST_VRAM_GB export LUCEBOX_HOST_GPU_SM LUCEBOX_HOST_DRIVER_VERSION LUCEBOX_HOST_DRIVER_MAJOR export LUCEBOX_HOST_HAS_SYSTEMD LUCEBOX_HOST_IS_WSL export LUCEBOX_HOST_HAS_DOCKER LUCEBOX_HOST_DOCKER_VERSION export LUCEBOX_HOST_HAS_CTK + export LUCEBOX_HOST_ROCM_VERSION LUCEBOX_HOST_HAS_KFD LUCEBOX_HOST_HAS_DRI + export LUCEBOX_HOST_AMD_GPU_NAME LUCEBOX_HOST_AMD_GPU_COUNT + export LUCEBOX_HOST_AMD_VRAM_GB LUCEBOX_HOST_AMD_GPU_ARCH + export LUCEBOX_HOST_AMD_GPU_LIST_CSV export LUCEBOX_HOST_OS_PRETTY LUCEBOX_HOST_KERNEL LUCEBOX_HOST_WSL_VERSION export LUCEBOX_HOST_NVIDIA_CTK_VERSION LUCEBOX_HOST_CPU_MODEL export LUCEBOX_HOST_GPU_LIST_CSV LUCEBOX_HOST_CUDA_VISIBLE_DEVICES + export LUCEBOX_HOST_HIP_VISIBLE_DEVICES _LUCEBOX_HOST_PROBED=1 } @@ -357,10 +531,41 @@ ensure_probed() { } pick_variant() { - # CUDA 12.8 is the supported image variant for this branch. Effective - # value goes through the same env > config.toml > default ladder as - # everything else so `config set image.variant=...` propagates. - _lucebox_resolve "${LUCEBOX_VARIANT:-}" image.variant "cuda12" + # Explicit env/config always wins. On a fresh install choose the backend + # from hardware: a working NVIDIA GPU takes priority on RTX + Strix builds; + # otherwise an AMD GPU selects ROCm (R9700 + Strix and Strix-only builds). + local configured + if [ -n "${LUCEBOX_VARIANT:-}" ]; then + printf '%s' "$LUCEBOX_VARIANT" + return + fi + configured=$(_lucebox_config_get image.variant) + if [ -n "$configured" ]; then + printf '%s' "$configured" + return + fi + ensure_probed + if [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ]; then + printf 'cuda12' + elif [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ]; then + printf 'rocm' + else + # Keep the historical default so `lucebox check` can still tell a + # GPU-less host what image it would otherwise use. + printf 'cuda12' + fi +} + +_variant_is_rocm() { + # Variant names are not limited to the moving `rocm` tag. Releases and + # CI produce tags such as `0.3.0-rocm` and `pr-335-rocm`; treat any tag + # containing the backend marker as ROCm. Spell out case-insensitivity + # instead of using Bash 4's `${value,,}` so host-only commands such as + # `check` keep working under the older Bash bundled with macOS too. + case "$1" in + *[Rr][Oo][Cc][Mm]*) return 0 ;; + *) return 1 ;; + esac } # ── prereq checks (host-only) ───────────────────────────────────────────── @@ -368,6 +573,9 @@ pick_variant() { # the richer reporting; this is the bare minimum to make `docker run` viable. require_host_prereqs() { + ensure_probed + local variant="${1:-}" + [ -n "$variant" ] || variant=$(pick_variant) local missing=0 if ! command -v docker &>/dev/null; then err "docker is not installed" @@ -379,20 +587,36 @@ require_host_prereqs() { missing=1 fi - if ! command -v nvidia-smi &>/dev/null; then - err "nvidia-smi not found — no NVIDIA driver detected" - hint "Install the NVIDIA driver: https://www.nvidia.com/Download/index.aspx" - missing=1 - elif ! nvidia-smi --query-gpu=name --format=csv,noheader &>/dev/null; then - err "nvidia-smi present but NVML calls fail — likely a driver/library mismatch" - hint "Reboot, or reinstall the matching NVIDIA driver package" - missing=1 + if _variant_is_rocm "$variant"; then + if [ "$LUCEBOX_HOST_HAS_AMD_GPU" != "1" ]; then + err "ROCm image selected but no working AMD GPU was detected" + hint "Install ROCm/amd-smi, or choose LUCEBOX_VARIANT=cuda12 on an NVIDIA build." + missing=1 + fi + if [ "$LUCEBOX_HOST_HAS_KFD" != "1" ]; then + err "/dev/kfd is missing or not accessible" + hint "Add the user to the render group, then re-login: sudo usermod -aG render \"$USER\"" + missing=1 + fi + if [ "$LUCEBOX_HOST_HAS_DRI" != "1" ]; then + err "no accessible /dev/dri/renderD* device was found" + hint "Add the user to the render and video groups, then re-login." + missing=1 + fi + else + if [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" != "1" ]; then + err "CUDA image selected but no working NVIDIA GPU was detected" + hint "Install the NVIDIA driver, or choose LUCEBOX_VARIANT=rocm on an AMD build." + missing=1 + fi fi [ "$missing" = "0" ] || exit 1 } require_ctk() { + local variant="${1:-$(pick_variant)}" + _variant_is_rocm "$variant" && return 0 case "$LUCEBOX_HOST_HAS_CTK" in runtime|cdi) return 0 ;; installed-unwired) @@ -430,10 +654,9 @@ require_systemd() { # All the Python-CLI subcommands share the same docker run incantation: # mount the host docker socket (so the in-container CLI can spawn server / # bench containers on the host daemon), mount $HOME at the same path (so -# paths look identical in and out), and pass host facts via env. When an -# NVIDIA GPU is detected we also pass --gpus all so the orchestrator can -# call nvidia-smi during profile snapshot export; without it nvidia_smi_csv (and -# any downstream power/utilization fields) come back empty. +# paths look identical in and out), and pass host facts via env. The selected +# image gets its native accelerator contract: --gpus all for +# CUDA; /dev/kfd + /dev/dri and the render/video groups for ROCm. DOCKER_SOCK_PATH="${DOCKER_HOST:-/var/run/docker.sock}" DOCKER_SOCK_PATH="${DOCKER_SOCK_PATH#unix://}" @@ -467,6 +690,27 @@ _append_scalar_env() { return 0 } +# Append the Docker accelerator contract for the chosen image variant. +# CUDA and ROCm use fundamentally different runtime flags; keeping this in +# one helper prevents the orchestrator, canonical server argv, and fallback +# server path from drifting apart. +_append_gpu_args() { # usage: _append_gpu_args arrayname variant + # shellcheck disable=SC2178 + local -n _gpu_arr="$1" + local variant="$2" + if _variant_is_rocm "$variant"; then + _gpu_arr+=( + --device /dev/kfd + --device /dev/dri + --group-add video + --group-add render + --security-opt seccomp=unconfined + ) + else + _gpu_arr+=(--gpus all) + fi +} + # Pick docker's interactive flags: -it on a real tty, -i otherwise. # Writes into a caller-supplied array via nameref. This MUST run in the # caller's scope (not a subshell or `< <(...)` process substitution): the @@ -488,9 +732,7 @@ build_orchestrator_argv() { local tty=() _set_tty_flags tty local argv=(docker run --rm "${tty[@]}") - if [ "${LUCEBOX_HOST_GPU_VENDOR:-none}" = "nvidia" ]; then - argv+=(--gpus all) - fi + _append_gpu_args argv "$variant" argv+=(--name "${CONTAINER_NAME}-cli-$$") argv+=(--user "$(id -u):$(id -g)") # Only bind-mount the docker socket when DOCKER_HOST actually points @@ -541,11 +783,11 @@ cmd_serve() { # If stage 1 fails (image not pulled yet, no config), fall back to a # conservative docker run — the container's own VRAM-tiered autotune # picks reasonable defaults from there. - require_host_prereqs ensure_probed - require_ctk local variant variant=$(pick_variant) + require_host_prereqs "$variant" + require_ctk "$variant" # Pre-flight: refuse to stomp on something that's already serving this # slot. Three states to distinguish, because silently `docker rm -f`-ing @@ -615,10 +857,10 @@ cmd_serve() { # with "source: unknown" any time print-serve-argv fails. local fallback_argv=(docker run --rm --name "$CONTAINER_NAME" - --gpus all -p "$DEFAULT_PORT:8080" -v "$HOME:$HOME" -v "$fallback_models:/opt/lucebox-hub/server/models") + _append_gpu_args fallback_argv "$variant" _append_host_env fallback_argv fallback_argv+=("${IMAGE_BASE}:${variant}") _serve_and_track "${fallback_argv[@]}" @@ -660,8 +902,10 @@ _serve_and_track() { } cmd_systemd_install() { - require_host_prereqs ensure_probed + local variant + variant=$(pick_variant) + require_host_prereqs "$variant" require_systemd "service install" local docker_bin docker_bin=$(command -v docker) @@ -696,7 +940,7 @@ RestartSec=10 SuccessExitStatus=143 SIGTERM Environment=PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin Environment=LUCEBOX_IMAGE=$IMAGE_BASE -Environment=LUCEBOX_VARIANT=$(pick_variant) +Environment=LUCEBOX_VARIANT=$variant Environment=LUCEBOX_PORT=$DEFAULT_PORT Environment=LUCEBOX_CONTAINER=$CONTAINER_NAME Environment=LUCEBOX_MODELS=$DEFAULT_MODELS_DIR @@ -839,9 +1083,10 @@ cmd_pull() { # Pull has to run on the host. Delegating this into the container creates a # stale-image trap: docker may start an old local tag before the fresh tag # has been pulled. - require_host_prereqs + ensure_probed local variant variant=$(pick_variant) + require_host_prereqs "$variant" info "Pulling ${IMAGE_BASE}:${variant}" exec docker pull "${IMAGE_BASE}:${variant}" } @@ -1011,36 +1256,67 @@ cmd_check() { _row 0 "docker daemon" "not installed — https://docs.docker.com/engine/install/" fi - # nvidia container toolkit - case "$LUCEBOX_HOST_HAS_CTK" in - runtime) _row 1 "nvidia ctk" "wired into docker (runtime)" ;; - cdi) _row 1 "nvidia ctk" "wired via CDI (nvidia.com/gpu)" ;; - installed-unwired) _row warn "nvidia ctk" "installed but not registered with docker — sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker" ;; - none|*) _row 0 "nvidia ctk" "not installed — https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html" ;; - esac + if [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ]; then + # nvidia container toolkit + case "$LUCEBOX_HOST_HAS_CTK" in + runtime) _row 1 "nvidia ctk" "wired into docker (runtime)" ;; + cdi) _row 1 "nvidia ctk" "wired via CDI (nvidia.com/gpu)" ;; + installed-unwired) _row warn "nvidia ctk" "installed but not registered with docker — sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker" ;; + none|*) _row 0 "nvidia ctk" "not installed — https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html" ;; + esac - # nvidia-smi + driver - if [ "$LUCEBOX_HOST_GPU_VENDOR" = "nvidia" ]; then if [ "$LUCEBOX_HOST_DRIVER_MAJOR" -ge "$MIN_DRIVER_CUDA12" ]; then _row 1 "nvidia driver" "$LUCEBOX_HOST_DRIVER_VERSION (≥ $MIN_DRIVER_CUDA12 required for cuda12)" else _row 0 "nvidia driver" "$LUCEBOX_HOST_DRIVER_VERSION (< $MIN_DRIVER_CUDA12 — cuda12 image will fail)" fi - elif command -v nvidia-smi &>/dev/null; then - _row 0 "nvidia driver" "nvidia-smi present but NVML calls fail — driver/library mismatch, try reboot" - else - _row 0 "nvidia driver" "nvidia-smi not found — install the NVIDIA driver" - fi - - # GPU detail - if [ "$LUCEBOX_HOST_GPU_VENDOR" = "nvidia" ]; then - _row 1 "gpu" "$LUCEBOX_HOST_GPU_NAME × $LUCEBOX_HOST_GPU_COUNT (sm_$LUCEBOX_HOST_GPU_SM, ${LUCEBOX_HOST_VRAM_GB} GB VRAM)" + _row 1 "nvidia gpu" "$LUCEBOX_HOST_GPU_NAME × $LUCEBOX_HOST_GPU_COUNT (sm_$LUCEBOX_HOST_GPU_SM, ${LUCEBOX_HOST_VRAM_GB} GB VRAM)" # cuda12 image arch coverage: sm_75;80;86;89;90;120 (see docker-bake.hcl) case "$LUCEBOX_HOST_GPU_SM" in 75|80|86|89|90|120) _row 1 "cuda12 arch" "sm_$LUCEBOX_HOST_GPU_SM covered by image" ;; - "") _row warn "cuda12 arch" "compute_cap not detected" ;; + "") _row warn "cuda12 arch" "compute capability not detected" ;; *) _row warn "cuda12 arch" "sm_$LUCEBOX_HOST_GPU_SM not in image arch list (75;80;86;89;90;120)" ;; esac + elif command -v nvidia-smi &>/dev/null; then + _row 0 "nvidia driver" "nvidia-smi present but NVML calls fail — driver/library mismatch, try reboot" + fi + + if [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ]; then + if [ -n "$LUCEBOX_HOST_ROCM_VERSION" ]; then + _row 1 "rocm" "$LUCEBOX_HOST_ROCM_VERSION" + else + _row warn "rocm" "GPU detected, userspace version unavailable" + fi + _row 1 "amd gpu" "${LUCEBOX_HOST_AMD_GPU_COUNT} device(s); primary ${LUCEBOX_HOST_AMD_GPU_NAME} (${LUCEBOX_HOST_AMD_GPU_ARCH}, ${LUCEBOX_HOST_AMD_VRAM_GB} GB effective)" + local amd_line amd_device_idx amd_device_name amd_device_arch amd_device_mem + while IFS= read -r amd_line; do + [ -n "$amd_line" ] || continue + amd_device_idx=$(printf '%s' "$amd_line" | awk -F',' '{gsub(/^[[:space:]]+|[[:space:]]+$/, "", $1); print $1}') + amd_device_name=$(printf '%s' "$amd_line" | awk -F',' '{gsub(/^[[:space:]]+|[[:space:]]+$/, "", $4); print $4}') + amd_device_arch=$(printf '%s' "$amd_line" | awk -F',' '{gsub(/^[[:space:]]+|[[:space:]]+$/, "", $5); print $5}') + amd_device_mem=$(printf '%s' "$amd_line" | awk -F',' '{gsub(/^[[:space:]]+|[[:space:]]+$/, "", $6); print $6}') + _row 1 " amd[$amd_device_idx]" "$amd_device_name ($amd_device_arch, $amd_device_mem physical)" + done <<<"$LUCEBOX_HOST_AMD_GPU_LIST_CSV" + case "$LUCEBOX_HOST_AMD_GPU_ARCH" in + gfx1100|gfx1151|gfx1200|gfx1201) _row 1 "rocm arch" "$LUCEBOX_HOST_AMD_GPU_ARCH covered by published image" ;; + "") _row warn "rocm arch" "gfx architecture not detected" ;; + *) _row warn "rocm arch" "$LUCEBOX_HOST_AMD_GPU_ARCH not in published image arch list" ;; + esac + if [ "$LUCEBOX_HOST_HAS_KFD" = "1" ]; then + _row 1 "amd /dev/kfd" "accessible" + else + _row 0 "amd /dev/kfd" "missing or inaccessible — user needs the render group" + fi + if [ "$LUCEBOX_HOST_HAS_DRI" = "1" ]; then + _row 1 "amd /dev/dri" "render node accessible" + else + _row 0 "amd /dev/dri" "no accessible renderD* node — user needs render/video groups" + fi + fi + + if [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" != "1" ] \ + && [ "$LUCEBOX_HOST_HAS_AMD_GPU" != "1" ]; then + _row 0 "gpu" "no supported NVIDIA or AMD GPU detected" fi # systemd @@ -1052,16 +1328,24 @@ cmd_check() { _row warn "user systemd" "not available — '$SCRIPT_NAME install' (service unit) won't work; '$SCRIPT_NAME serve' (foreground) will" fi - # image we'd pull — marked ✗ when the host clearly can't run cuda12 - # (no nvidia driver, or no CTK wired into docker). It's still useful - # to print the line so the user knows what would be pulled, but a - # green ✓ would be misleading. - if [ "$LUCEBOX_HOST_GPU_VENDOR" != "nvidia" ]; then - _row 0 "image" "${IMAGE_BASE}:${variant} — requires NVIDIA driver" - elif [ "$LUCEBOX_HOST_HAS_CTK" = "none" ] || [ "$LUCEBOX_HOST_HAS_CTK" = "installed-unwired" ]; then - _row 0 "image" "${IMAGE_BASE}:${variant} — needs NVIDIA Container Toolkit wired into docker" + # Selected backend and image. On heterogeneous builds the unselected APU + # stays visible in the inventory above but does not change this decision. + if _variant_is_rocm "$variant"; then + if [ "$LUCEBOX_HOST_HAS_AMD_GPU" != "1" ]; then + _row 0 "image" "${IMAGE_BASE}:${variant} — requires an AMD GPU" + elif [ "$LUCEBOX_HOST_HAS_KFD" != "1" ] || [ "$LUCEBOX_HOST_HAS_DRI" != "1" ]; then + _row 0 "image" "${IMAGE_BASE}:${variant} — needs accessible /dev/kfd and /dev/dri" + else + _row 1 "image" "${IMAGE_BASE}:${variant} (AMD selected)" + fi else - _row 1 "image" "${IMAGE_BASE}:${variant}" + if [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" != "1" ]; then + _row 0 "image" "${IMAGE_BASE}:${variant} — requires an NVIDIA GPU" + elif [ "$LUCEBOX_HOST_HAS_CTK" = "none" ] || [ "$LUCEBOX_HOST_HAS_CTK" = "installed-unwired" ]; then + _row 0 "image" "${IMAGE_BASE}:${variant} — needs NVIDIA Container Toolkit wired into docker" + else + _row 1 "image" "${IMAGE_BASE}:${variant} (NVIDIA selected)" + fi fi # RAM / cores (informational) _row 1 "host" "${LUCEBOX_HOST_NPROC} cpus, ${LUCEBOX_HOST_RAM_GB} GB RAM" @@ -1070,7 +1354,6 @@ cmd_check() { cmd_in_container() { # Generic dispatcher: anything that isn't a systemd action goes here. # Runs the in-container Python CLI with the supplied argv. - require_host_prereqs ensure_probed # CTK isn't strictly required for every subcommand (e.g. `config get` # or `autotune` only touch local files), but the server-spawning @@ -1078,6 +1361,7 @@ cmd_in_container() { # Letting docker error its own way is fine for the no-CTK case. local variant variant=$(pick_variant) + require_host_prereqs "$variant" local argv mapfile -t argv < <(build_orchestrator_argv "$variant" "$@") exec "${argv[@]}" @@ -1112,10 +1396,10 @@ _lucebox_container_running() { # CLI sees consistent overrides whichever route it took: HOME, every # LUCEBOX_HOST_*, the image/port/container/models scalars, and HF_TOKEN. cmd_exec_in_container() { - require_host_prereqs ensure_probed local variant variant=$(pick_variant) + require_host_prereqs "$variant" local tty=() _set_tty_flags tty local argv=(docker exec "${tty[@]}") @@ -1196,7 +1480,7 @@ Direct server invocation (foreground, no systemd): Provisioning + workloads (delegated to the in-container Python CLI): check host + docker readiness report - pull docker pull the cuda12 image + pull docker pull the auto-selected CUDA or ROCm image update re-run the bootstrap installer to upgrade this script completion print shell completion script (bash / zsh / fish) models list / download / activate model presets @@ -1209,7 +1493,7 @@ Misc: Environment overrides: LUCEBOX_IMAGE image name without tag (default: ghcr.io/luce-org/lucebox-hub) - LUCEBOX_VARIANT image tag to pull/run (default: cuda12) + LUCEBOX_VARIANT image tag override (default: cuda12 on NVIDIA, rocm on AMD) LUCEBOX_PORT host port for the server (default: 8080) LUCEBOX_CONTAINER server container name (default: lucebox) LUCEBOX_MODELS host model directory (default: \$XDG_DATA_HOME/lucebox/models diff --git a/lucebox/README.md b/lucebox/README.md index 747a49eef..6af3d2a02 100644 --- a/lucebox/README.md +++ b/lucebox/README.md @@ -7,9 +7,12 @@ and is invoked from the host via the [`lucebox.sh`](../lucebox.sh) wrapper: lucebox.sh config get lucebox.sh print-run -The wrapper is the only thing that runs on the host; everything else (host +The wrapper is the only thing that runs on the host. It selects CUDA for a +working NVIDIA GPU (including RTX 3090 + Strix builds) and ROCm for AMD builds +(including R9700 + Strix), then applies the matching Docker device flags. +Everything else (host checks, TOML config, docker daemon calls, model download) is Python in the -container. Host facts (driver, GPU, RAM, VRAM, systemd availability) are +container. Host facts (driver/runtime, GPU, RAM, VRAM, systemd availability) are passed in via `LUCEBOX_HOST_*` environment variables so the Python side doesn't reprobe. The autotune sweep, profiling, and agent-client launchers land in follow-up PRs. diff --git a/lucebox/src/lucebox/autotune.py b/lucebox/src/lucebox/autotune.py index 51b4f8424..5541cf271 100644 --- a/lucebox/src/lucebox/autotune.py +++ b/lucebox/src/lucebox/autotune.py @@ -18,7 +18,7 @@ def runtime_from_host(host: HostFacts) -> DflashRuntime: """Pick a conservative DflashRuntime that 'should work' on this VRAM tier. - Tiers (NVIDIA, baseline = Qwen3.6-27B Q4_K_M ~18 GB total): + Tiers (selected CUDA/ROCm device, baseline = Qwen3.6-27B Q4_K_M ~18 GB total): <12 GB — too small for 27B; pick min ctx as a floor so a fallback start at least gets an error from the daemon rather than a silent OOM. diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index 87cc29366..e3c6101bd 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -6,7 +6,7 @@ Subcommand inventory: check — readiness report config get/set/unset — read / write a single key in config.toml - pull — docker pull the cuda12 image + pull — docker pull the selected CUDA or ROCm image print-run — emit the docker-run command for the server print-serve-argv — same, raw argv lines (consumed by `lucebox serve`) models — list / download presets, activate one diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py index b3f63eb0d..6e3d113c8 100644 --- a/lucebox/src/lucebox/config.py +++ b/lucebox/src/lucebox/config.py @@ -257,12 +257,17 @@ def _from_dict(raw: dict[str, Any]) -> Config: nproc=int(host_raw.get("nproc", 0)), ram_gb=int(host_raw.get("ram_gb", 0)), gpu_vendor=host_raw.get("gpu_vendor", "none"), + has_nvidia_gpu=bool(host_raw.get("has_nvidia_gpu", False)), + has_amd_gpu=bool(host_raw.get("has_amd_gpu", False)), gpu_name=str(host_raw.get("gpu_name", "")), gpu_count=int(host_raw.get("gpu_count", 0)), vram_gb=int(host_raw.get("vram_gb", 0)), gpu_sm=str(host_raw.get("gpu_sm", "")), driver_version=str(host_raw.get("driver_version", "")), driver_major=int(host_raw.get("driver_major", 0)), + rocm_version=str(host_raw.get("rocm_version", "")), + has_kfd=bool(host_raw.get("has_kfd", False)), + has_dri=bool(host_raw.get("has_dri", False)), has_systemd=bool(host_raw.get("has_systemd", False)), is_wsl=bool(host_raw.get("is_wsl", False)), has_docker=bool(host_raw.get("has_docker", False)), @@ -470,9 +475,10 @@ def live_config() -> Config: host = from_env() default = Config() + default_variant = "rocm" if host.gpu_vendor == "amd" else "cuda12" return replace( default, - variant=os.environ.get("LUCEBOX_VARIANT", "cuda12"), + variant=os.environ.get("LUCEBOX_VARIANT", default_variant), image=os.environ.get("LUCEBOX_IMAGE", default.image), container_name=os.environ.get("LUCEBOX_CONTAINER", default.container_name), port=int(os.environ.get("LUCEBOX_PORT", str(default.port))), diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index ff3615b26..8e61e355f 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -15,7 +15,7 @@ from dataclasses import dataclass from pathlib import Path -from lucebox.types import Config +from lucebox.types import Config, GpuVendor def _host_facts_env() -> list[tuple[str, str]]: @@ -91,6 +91,7 @@ class DockerRunSpec: image: str name: str gpus: bool = True + gpu_vendor: GpuVendor = "nvidia" detach: bool = False remove: bool = True port_publish: tuple[int, int] | None = None # (host, container) @@ -106,8 +107,21 @@ def argv(self) -> list[str]: if self.detach: out.append("-d") out += ["--name", self.name] - if self.gpus: + if self.gpus and self.gpu_vendor == "nvidia": out += ["--gpus", "all"] + elif self.gpus and self.gpu_vendor == "amd": + out += [ + "--device", + "/dev/kfd", + "--device", + "/dev/dri", + "--group-add", + "video", + "--group-add", + "render", + "--security-opt", + "seccomp=unconfined", + ] if self.port_publish is not None: host, container = self.port_publish out += ["-p", f"{host}:{container}"] @@ -141,6 +155,9 @@ def printable(self) -> str: "--volume", "--publish", "--entrypoint", + "--device", + "--group-add", + "--security-opt", } and i + 1 < len(argv): i += 1 out += " " + shlex.quote(argv[i]) @@ -153,7 +170,7 @@ def printable(self) -> str: def server_run_spec(cfg: Config) -> DockerRunSpec: """Long-running OpenAI-compatible server. Foreground (systemd manages - lifecycle), --gpus all, models bind-mounted, DFLASH_* propagated. + lifecycle), vendor-specific GPU devices, models bind-mounted, DFLASH_* propagated. """ # LUCEBOX_HOST_* first so they ride out front in the rendered argv, # making it obvious in `print-run` output what host facts get forwarded. @@ -212,10 +229,23 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: if cfg.dflash.debug_thinking_logits: env.append(("DFLASH_DEBUG_THINKING_LOGITS", "1")) + # The chosen image variant is the runtime contract. This matters on a + # heterogeneous RTX + Strix host: the probe records both vendors, while an + # explicit ``variant=rocm`` must still receive AMD device flags. Unknown + # custom variants fall back to the probe's selected/default vendor. + variant_lower = cfg.variant.lower() + if "rocm" in variant_lower: + gpu_vendor: GpuVendor = "amd" + elif "cuda" in variant_lower: + gpu_vendor = "nvidia" + else: + gpu_vendor = cfg.host.gpu_vendor if cfg.host.gpu_vendor != "none" else "nvidia" + return DockerRunSpec( image=f"{cfg.image}:{cfg.variant}", name=cfg.container_name, gpus=True, + gpu_vendor=gpu_vendor, remove=True, detach=False, port_publish=(cfg.port, 8080), diff --git a/lucebox/src/lucebox/host_check.py b/lucebox/src/lucebox/host_check.py index 2ce8d3889..8f2d8d8d8 100644 --- a/lucebox/src/lucebox/host_check.py +++ b/lucebox/src/lucebox/host_check.py @@ -26,14 +26,15 @@ class CheckResult: def run_checks(host: HostFacts) -> list[CheckResult]: - return [ - _check_docker(host), - _check_nvidia_driver(host), - _check_ctk(host), - _check_ram(host), - _check_vram(host), - _check_systemd(host), - ] + results = [_check_docker(host)] + if host.gpu_vendor == "nvidia": + results += [_check_nvidia_driver(host), _check_ctk(host)] + elif host.gpu_vendor == "amd": + results += [_check_amd_driver(host), _check_amd_devices(host)] + else: + results.append(CheckResult("gpu", "fail", "no supported NVIDIA or AMD GPU detected")) + results += [_check_ram(host), _check_vram(host), _check_systemd(host)] + return results def _check_docker(host: HostFacts) -> CheckResult: @@ -48,15 +49,6 @@ def _check_docker(host: HostFacts) -> CheckResult: def _check_nvidia_driver(host: HostFacts) -> CheckResult: - if host.gpu_vendor != "nvidia": - if host.gpu_vendor == "amd": - return CheckResult( - "gpu", - "fail", - "AMD GPU detected — prebuilt images are NVIDIA-only", - "Build dflash from source with HIP; see dflash/README.md", - ) - return CheckResult("gpu", "fail", "no NVIDIA GPU detected") if not host.driver_version: return CheckResult( "driver", @@ -74,6 +66,33 @@ def _check_nvidia_driver(host: HostFacts) -> CheckResult: return CheckResult("driver", "ok", f"nvidia r{host.driver_major} ({host.driver_version})") +def _check_amd_driver(host: HostFacts) -> CheckResult: + if not host.rocm_version: + return CheckResult( + "rocm", + "warn", + "AMD GPU detected but ROCm userspace version is unknown", + "install amd-smi or hipconfig so Lucebox can verify the ROCm runtime", + ) + return CheckResult("rocm", "ok", f"ROCm {host.rocm_version} ({host.gpu_sm or 'gfx?'})") + + +def _check_amd_devices(host: HostFacts) -> CheckResult: + missing: list[str] = [] + if not host.has_kfd: + missing.append("/dev/kfd") + if not host.has_dri: + missing.append("/dev/dri/renderD*") + if missing: + return CheckResult( + "devices", + "fail", + f"missing or inaccessible: {', '.join(missing)}", + "add the user to the render and video groups, then re-login", + ) + return CheckResult("devices", "ok", "/dev/kfd and /dev/dri are accessible") + + def _check_ctk(host: HostFacts) -> CheckResult: match host.ctk: case "runtime": @@ -147,10 +166,13 @@ def aggregate(results: list[CheckResult]) -> Severity: def render(console: Console, host: HostFacts, results: list[CheckResult]) -> Severity: """Print a status block, return the worst severity.""" summary = f"[bold]Host:[/bold] {host.nproc} CPUs · {host.ram_gb} GB RAM" - if host.gpu_vendor == "nvidia" and host.gpu_name: - summary += f" · {host.gpu_name} · {host.vram_gb} GB VRAM" + ( - f" (sm_{host.gpu_sm})" if host.gpu_sm else "" - ) + if host.gpu_name: + arch = host.gpu_sm + if arch and host.gpu_vendor == "nvidia": + arch = f"sm_{arch}" + summary += f" · {host.gpu_name} · {host.vram_gb} GB VRAM" + if arch: + summary += f" ({arch})" if host.is_wsl: summary += " · WSL2" console.print(summary) @@ -202,10 +224,13 @@ def render_host_facts(console: Console) -> None: ("kernel", os.environ.get("LUCEBOX_HOST_KERNEL", "")), ("wsl_version", os.environ.get("LUCEBOX_HOST_WSL_VERSION", "")), ("docker", os.environ.get("LUCEBOX_HOST_DOCKER_VERSION", "")), + ("gpu_vendor", os.environ.get("LUCEBOX_HOST_GPU_VENDOR", "")), ("nvidia_driver", os.environ.get("LUCEBOX_HOST_DRIVER_VERSION", "")), ("nvidia_ctk", os.environ.get("LUCEBOX_HOST_NVIDIA_CTK_VERSION", "")), + ("rocm", os.environ.get("LUCEBOX_HOST_ROCM_VERSION", "")), ("cpu", os.environ.get("LUCEBOX_HOST_CPU_MODEL", "")), ("cuda_visible_devices", os.environ.get("LUCEBOX_HOST_CUDA_VISIBLE_DEVICES", "")), + ("hip_visible_devices", os.environ.get("LUCEBOX_HOST_HIP_VISIBLE_DEVICES", "")), ] for key, value in facts: display = value if value else "[dim](unset)[/dim]" @@ -223,10 +248,10 @@ def render_host_facts(console: Console) -> None: parts = [c.strip() for c in line.split(",")] if len(parts) >= 7: idx, _uuid, _pci, name, sm, mem, plimit = parts[:7] - console.print( - f" [{idx}] {name} (sm_{sm}, {mem}, {plimit})" - ) + arch = sm if sm.startswith("gfx") else f"sm_{sm}" + detail = ", ".join(part for part in (arch, mem, plimit) if part) + console.print(f" [{idx}] {name} ({detail})") else: console.print(f" {line}") else: - console.print(" gpus [dim](none — nvidia-smi unavailable)[/dim]") + console.print(" gpus [dim](none — GPU probe unavailable)[/dim]") diff --git a/lucebox/src/lucebox/host_facts.py b/lucebox/src/lucebox/host_facts.py index 5deb6721a..c9556d246 100644 --- a/lucebox/src/lucebox/host_facts.py +++ b/lucebox/src/lucebox/host_facts.py @@ -2,7 +2,7 @@ We deliberately don't try to detect anything ourselves on the Python side — inside the container, /proc/meminfo reports the container's view, not the -host's, and nvidia-smi may or may not be available depending on how the +host's, and nvidia-smi/amd-smi may or may not be available depending on how the caller invoked us. The host wrapper is the only thing that can see the truth, and it's already paid for the probe. """ @@ -44,12 +44,17 @@ def from_env() -> HostFacts: nproc=_env_int("LUCEBOX_HOST_NPROC"), ram_gb=_env_int("LUCEBOX_HOST_RAM_GB"), gpu_vendor=vendor, + has_nvidia_gpu=_env_bool("LUCEBOX_HOST_HAS_NVIDIA_GPU"), + has_amd_gpu=_env_bool("LUCEBOX_HOST_HAS_AMD_GPU"), gpu_name=os.environ.get("LUCEBOX_HOST_GPU_NAME", ""), gpu_count=_env_int("LUCEBOX_HOST_GPU_COUNT"), vram_gb=_env_int("LUCEBOX_HOST_VRAM_GB"), gpu_sm=os.environ.get("LUCEBOX_HOST_GPU_SM", ""), driver_version=os.environ.get("LUCEBOX_HOST_DRIVER_VERSION", ""), driver_major=_env_int("LUCEBOX_HOST_DRIVER_MAJOR"), + rocm_version=os.environ.get("LUCEBOX_HOST_ROCM_VERSION", ""), + has_kfd=_env_bool("LUCEBOX_HOST_HAS_KFD"), + has_dri=_env_bool("LUCEBOX_HOST_HAS_DRI"), has_systemd=_env_bool("LUCEBOX_HOST_HAS_SYSTEMD"), is_wsl=_env_bool("LUCEBOX_HOST_IS_WSL"), has_docker=_env_bool("LUCEBOX_HOST_HAS_DOCKER"), diff --git a/lucebox/src/lucebox/types.py b/lucebox/src/lucebox/types.py index e1d3620d7..e9e2fa3ec 100644 --- a/lucebox/src/lucebox/types.py +++ b/lucebox/src/lucebox/types.py @@ -40,12 +40,20 @@ class HostFacts: nproc: int = 0 ram_gb: int = 0 gpu_vendor: GpuVendor = "none" + has_nvidia_gpu: bool = False + has_amd_gpu: bool = False gpu_name: str = "" gpu_count: int = 0 vram_gb: int = 0 - gpu_sm: str = "" # e.g. "120" — matches docker-bake arch lists - driver_version: str = "" # e.g. "595.71.05" + # NVIDIA compute capability without the dot ("120") or AMD gfx target + # ("gfx1151"). The historical ``gpu_sm`` name stays in the serialized + # contract for compatibility with HOST_INFO consumers. + gpu_sm: str = "" + driver_version: str = "" # NVIDIA driver, e.g. "595.71.05" driver_major: int = 0 + rocm_version: str = "" # AMD ROCm userspace, e.g. "7.2.4" + has_kfd: bool = False # /dev/kfd exists and is accessible to the user + has_dri: bool = False # at least one accessible /dev/dri/renderD* node has_systemd: bool = False is_wsl: bool = False has_docker: bool = False diff --git a/lucebox/tests/test_check.py b/lucebox/tests/test_check.py index 3fdd469d9..ba48b4d7d 100644 --- a/lucebox/tests/test_check.py +++ b/lucebox/tests/test_check.py @@ -116,3 +116,53 @@ def stub() -> HostFacts: assert result.exit_code == 1 # Host facts block still printed despite the failure. assert "Host facts" in result.stdout + + +def test_amd_rocm_host_passes_without_nvidia_ctk(monkeypatch: pytest.MonkeyPatch) -> None: + """ROCm readiness uses /dev/kfd + /dev/dri and never requires NVIDIA CTK.""" + + def stub() -> HostFacts: + return HostFacts( + nproc=32, + ram_gb=125, + gpu_vendor="amd", + has_amd_gpu=True, + gpu_name="AMD Radeon AI PRO R9700", + gpu_count=2, + vram_gb=31, + gpu_sm="gfx1201", + rocm_version="7.2.4", + has_kfd=True, + has_dri=True, + has_systemd=True, + has_docker=True, + docker_version="29.1.3", + ctk="none", + ) + + monkeypatch.setattr("lucebox.host_facts.from_env", stub) + monkeypatch.setattr("lucebox.cli.from_env", stub) + result = CliRunner().invoke(app, ["check"]) + assert result.exit_code == 0, result.stdout + assert "ROCm 7.2.4" in result.stdout + assert "gfx1201" in result.stdout + assert "/dev/kfd and /dev/dri are accessible" in result.stdout + assert "NVIDIA Container Toolkit" not in result.stdout + + +def test_amd_rocm_host_fails_when_devices_are_inaccessible() -> None: + host = HostFacts( + gpu_vendor="amd", + has_amd_gpu=True, + gpu_name="AMD Radeon Graphics", + vram_gb=125, + gpu_sm="gfx1151", + rocm_version="7.2.4", + has_docker=True, + has_kfd=False, + has_dri=True, + ) + results = host_check.run_checks(host) + devices = next(result for result in results if result.name == "devices") + assert devices.severity == "fail" + assert "/dev/kfd" in devices.message diff --git a/lucebox/tests/test_config.py b/lucebox/tests/test_config.py index 2ec795a12..9c2b913a3 100644 --- a/lucebox/tests/test_config.py +++ b/lucebox/tests/test_config.py @@ -176,6 +176,24 @@ def test_live_config_uses_recommend_preset_indirectly(tmp_path: Path) -> None: assert cfg.model.preset == "" +def test_live_config_selects_rocm_for_amd(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LUCEBOX_HOST_GPU_VENDOR", "amd") + monkeypatch.delenv("LUCEBOX_VARIANT", raising=False) + + cfg = config.live_config() + + assert cfg.variant == "rocm" + + +def test_live_config_variant_override_wins_on_amd(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LUCEBOX_HOST_GPU_VENDOR", "amd") + monkeypatch.setenv("LUCEBOX_VARIANT", "test-cuda12") + + cfg = config.live_config() + + assert cfg.variant == "test-cuda12" + + def test_seed_dflash_writes_heuristic_when_absent(tmp_path: Path) -> None: """First-time activate seeds the VRAM-tier heuristic into config.toml. diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index bb888514f..0f8999ea1 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -11,7 +11,7 @@ from pathlib import Path from lucebox.download import PRESETS -from lucebox.types import Config, DflashRuntime, ModelMeta +from lucebox.types import Config, DflashRuntime, HostFacts, ModelMeta from lucebox import docker_run @@ -63,6 +63,23 @@ def test_argv_flags_and_ordering() -> None: assert argv.index("--shm-size") < argv.index("img:tag") +def test_argv_amd_uses_rocm_device_contract() -> None: + spec = docker_run.DockerRunSpec(image="img:rocm", name="box", gpu_vendor="amd") + argv = spec.argv() + assert "--gpus" not in argv + assert ["--device", "/dev/kfd"] == argv[ + argv.index("--device") : argv.index("--device") + 2 + ] + assert "/dev/dri" in argv + assert ["--group-add", "video"] == argv[ + argv.index("--group-add") : argv.index("--group-add") + 2 + ] + assert "render" in argv + assert ["--security-opt", "seccomp=unconfined"] == argv[ + argv.index("--security-opt") : argv.index("--security-opt") + 2 + ] + + def test_printable_glues_value_taking_flags() -> None: spec = docker_run.DockerRunSpec( image="img:tag", @@ -148,6 +165,22 @@ def test_server_run_spec_top_level_shape(tmp_path: Path) -> None: assert (str(tmp_path), "/opt/lucebox-hub/server/models") in spec.volumes +def test_server_run_spec_rocm_uses_amd_devices_on_heterogeneous_host(tmp_path: Path) -> None: + cfg = Config( + variant="rocm", + models_dir=tmp_path, + # The generic probe may select NVIDIA by default on RTX + Strix. The + # explicit image variant must still control Docker's device contract. + host=HostFacts(gpu_vendor="nvidia", has_nvidia_gpu=True, has_amd_gpu=True), + ) + spec = docker_run.server_run_spec(cfg) + assert spec.gpu_vendor == "amd" + argv = spec.argv() + assert "--gpus" not in argv + assert "/dev/kfd" in argv + assert "/dev/dri" in argv + + def test_server_run_spec_always_emits_core_dflash_env(tmp_path: Path) -> None: cfg = Config(models_dir=tmp_path, dflash=DflashRuntime(budget=22, max_ctx=32768)) env = _env(docker_run.server_run_spec(cfg)) diff --git a/lucebox/tests/test_models_cli.py b/lucebox/tests/test_models_cli.py index f44583044..744f40c29 100644 --- a/lucebox/tests/test_models_cli.py +++ b/lucebox/tests/test_models_cli.py @@ -137,6 +137,9 @@ def test_installed_helpers_track_presence( target = cfg.models_dir / laguna.target_file target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(b"x" * (5 * 10**9)) + # Apparent size is what the CLI reports; a sparse file exercises the same + # stat path without allocating and writing a 5 GB Python byte string. + with target.open("wb") as sparse: + sparse.truncate(5 * 10**9) assert download_mod.installed_status(cfg, laguna) == "installed" assert download_mod.installed_size_gb(cfg, laguna) == pytest.approx(5.0, rel=0.01) diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index 3960760f6..f858690ba 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -32,6 +32,27 @@ if [ ! -f "$SCRIPT" ]; then exit 1 fi +# Make the whole suite safe on a contributor workstation. Several dispatch +# smoke tests intentionally invoke start/stop/install/uninstall/pull up to the +# missing-prerequisite boundary; without global shims those commands can touch +# a real user service or pull a multi-GB image on a fully provisioned host. +SUITE_SANDBOX=$(mktemp -d "${TMPDIR:-/tmp}/lucebox-shell-suite.XXXXXX") +SUITE_SHIMS="$SUITE_SANDBOX/shims" +mkdir -p "$SUITE_SHIMS" "$SUITE_SANDBOX/home" "$SUITE_SANDBOX/xdg" "$SUITE_SANDBOX/data" +for binname in docker systemctl journalctl loginctl; do + cat > "$SUITE_SHIMS/$binname" <<'STUB' +#!/usr/bin/env bash +exit 1 +STUB + chmod +x "$SUITE_SHIMS/$binname" +done +export HOME="$SUITE_SANDBOX/home" +export XDG_CONFIG_HOME="$SUITE_SANDBOX/xdg" +export XDG_DATA_HOME="$SUITE_SANDBOX/data" +export LUCEBOX_HOME="$SUITE_SANDBOX/home/.lucebox" +export PATH="$SUITE_SHIMS:$PATH" +trap 'rm -rf "$SUITE_SANDBOX"' EXIT + # entrypoint.sh ships with the docker-stack PR (#334). When it's absent # (e.g. on the lucebox-cli branch in isolation), skip the entire suite — # every section below either references $ENTRYPOINT in shellcheck targets, @@ -284,6 +305,7 @@ STUB LUCEBOX_HOST_HAS_DOCKER=1 \ LUCEBOX_HOST_HAS_CTK=runtime \ LUCEBOX_HOST_GPU_VENDOR=nvidia \ + LUCEBOX_HOST_HAS_NVIDIA_GPU=1 \ _LUCEBOX_HOST_PROBED=1 \ NO_COLOR=1 \ timeout 10 bash "$SCRIPT" install 2>&1 @@ -935,6 +957,7 @@ _run_wrapper_capture_docker() { LUCEBOX_HOST_HAS_DOCKER=1 \ LUCEBOX_HOST_HAS_CTK=runtime \ LUCEBOX_HOST_GPU_VENDOR=nvidia \ + LUCEBOX_HOST_HAS_NVIDIA_GPU=1 \ LUCEBOX_HOST_DRIVER_MAJOR=550 \ LUCEBOX_HOST_DRIVER_VERSION="550.00" \ LUCEBOX_HOST_GPU_NAME="Fake GPU" \ @@ -1071,6 +1094,125 @@ test_usage_mentions_exec_routing() { } test_usage_mentions_exec_routing "usage documents docker exec routing + --no-exec flag" +# ── cross-vendor GPU selection / Docker contract ────────────────────────── +test_amd_smi_parser() { + local label="$1" fn out + fn=$(awk '/^_parse_amd_smi_csv\(\) \{/,/^\}/' "$SCRIPT") + out=$(bash -c "$fn"$'\n''_parse_amd_smi_csv' <<'CSV' +gpu,market_name,vendor_id,vendor_name,subvendor_id,device_id,subsystem_id,rev_id,asic_serial,oam_id,num_compute_units,target_graphics_version,type,vendor,size,bit_width,max_bandwidth +0,AMD Radeon AI PRO R9700,0x1002,AMD,0xf111,0x7551,0x000a,0xc0,serial,N/A,64,gfx1201,GDDR6,SAMSUNG,32624,256,N/A +1,AMD Radeon Graphics,0x1002,AMD,0xf111,0x1586,0x000a,0xc1,serial,N/A,40,gfx1151,GDDR7,UNKNOWN,512,256,N/A +CSV +) + if ! grep -qF '0|AMD Radeon AI PRO R9700|gfx1201|32624' <<<"$out" \ + || ! grep -qF '1|AMD Radeon Graphics|gfx1151|512' <<<"$out"; then + report fail "$label" "unexpected normalized rows: $out" + return + fi + report ok "$label" +} +test_amd_smi_parser "amd-smi parser recognizes R9700 + Strix" + +test_variant_autoselection() { + local label="$1" fn common got + fn=$(awk '/^pick_variant\(\) \{/,/^\}/' "$SCRIPT") + common=$'_lucebox_config_get(){ :; }\nensure_probed(){ :; }\nLUCEBOX_VARIANT=""\n' + + got=$(bash -c "$fn"$'\n'"$common"$'LUCEBOX_HOST_HAS_NVIDIA_GPU=1\nLUCEBOX_HOST_HAS_AMD_GPU=1\npick_variant') + [ "$got" = "cuda12" ] || { report fail "$label" "mixed RTX + Strix chose $got"; return; } + + got=$(bash -c "$fn"$'\n'"$common"$'LUCEBOX_HOST_HAS_NVIDIA_GPU=0\nLUCEBOX_HOST_HAS_AMD_GPU=1\npick_variant') + [ "$got" = "rocm" ] || { report fail "$label" "R9700 + Strix chose $got"; return; } + + got=$(bash -c "$fn"$'\n'"$common"$'LUCEBOX_VARIANT=rocm\nLUCEBOX_HOST_HAS_NVIDIA_GPU=1\nLUCEBOX_HOST_HAS_AMD_GPU=1\npick_variant') + [ "$got" = "rocm" ] || { report fail "$label" "explicit override chose $got"; return; } + report ok "$label" +} +test_variant_autoselection "variant selection: RTX+Strix=cuda12, R9700+Strix=rocm" + +test_mixed_rtx_strix_probe_prefers_cuda() { + local label="$1" sandbox shim_dir out + sandbox=$(mktemp -d -t lucebox-mixed-gpu.XXXXXX) + shim_dir="$sandbox/bin" + mkdir -p "$shim_dir" + cat > "$shim_dir/nvidia-smi" <<'STUB' +#!/usr/bin/env bash +case "$*" in + *"name,memory.total,driver_version,compute_cap"*) + echo "NVIDIA GeForce RTX 3090, 24576, 550.00, 8.6" ;; + *"index,uuid,pci.bus_id,name,compute_cap,memory.total,power.limit"*) + echo "0, GPU-test, 00000000:01:00.0, NVIDIA GeForce RTX 3090, 8.6, 24576 MiB, 350.00 W" ;; + *"--query-gpu=name"*) echo "NVIDIA GeForce RTX 3090" ;; + "-L") echo "GPU 0: NVIDIA GeForce RTX 3090" ;; +esac +exit 0 +STUB + cat > "$shim_dir/amd-smi" <<'STUB' +#!/usr/bin/env bash +if [ "${1:-}" = "version" ]; then + echo "AMDSMI Tool: test | ROCm version: 7.2.4 | Platform: Linux Baremetal" + exit 0 +fi +cat <<'CSV' +gpu,market_name,vendor_id,vendor_name,subvendor_id,device_id,subsystem_id,rev_id,asic_serial,oam_id,num_compute_units,target_graphics_version,type,vendor,size,bit_width,max_bandwidth +0,AMD Radeon Graphics,0x1002,AMD,0xf111,0x1586,0x000a,0xc1,serial,N/A,40,gfx1151,GDDR7,UNKNOWN,512,256,N/A +CSV +STUB + cat > "$shim_dir/docker" <<'STUB' +#!/usr/bin/env bash +case "${1:-}" in + ps) exit 0 ;; + version) echo "29.1.3" ;; +esac +exit 0 +STUB + for binname in nvidia-container-runtime; do + printf '#!/usr/bin/env bash\nexit 0\n' > "$shim_dir/$binname" + done + chmod +x "$shim_dir"/* + out=$(HOME="$sandbox" LUCEBOX_HOME="$sandbox/.lucebox" \ + PATH="$shim_dir:$PATH" NO_COLOR=1 bash "$SCRIPT" check 2>&1) + rm -rf "$sandbox" + if ! grep -qF "NVIDIA GeForce RTX 3090" <<<"$out" \ + || ! grep -qF "AMD Radeon Graphics" <<<"$out" \ + || ! grep -qF ":cuda12 (NVIDIA selected)" <<<"$out"; then + report fail "$label" "mixed probe output: $(printf '%s' "$out" | head -20)" + return + fi + report ok "$label" +} +test_mixed_rtx_strix_probe_prefers_cuda "mixed RTX 3090 + Strix probe selects CUDA" + +test_cross_vendor_docker_args() { + local label="$1" helpers cuda_args rocm_args pinned_rocm_args + helpers=$(awk '/^_variant_is_rocm\(\) \{/,/^\}/' "$SCRIPT")$'\n' + helpers+=$(awk '/^_append_gpu_args\(\) \{/,/^\}/' "$SCRIPT") + cuda_args=$(bash -c "$helpers"$'\n''a=(); _append_gpu_args a cuda12; printf "%s\n" "${a[@]}"') + rocm_args=$(bash -c "$helpers"$'\n''a=(); _append_gpu_args a rocm; printf "%s\n" "${a[@]}"') + pinned_rocm_args=$(bash -c "$helpers"$'\n''a=(); _append_gpu_args a 0.3.0-rocm; printf "%s\n" "${a[@]}"') + if ! grep -qF -- '--gpus' <<<"$cuda_args" || ! grep -qF 'all' <<<"$cuda_args"; then + report fail "$label" "CUDA args missing --gpus all" + return + fi + for expected in /dev/kfd /dev/dri video render seccomp=unconfined; do + if ! grep -qF "$expected" <<<"$rocm_args"; then + report fail "$label" "ROCm args missing $expected" + return + fi + done + if grep -qF -- '--gpus' <<<"$rocm_args"; then + report fail "$label" "ROCm args incorrectly contain --gpus" + return + fi + if ! grep -qF '/dev/kfd' <<<"$pinned_rocm_args" \ + || grep -qF -- '--gpus' <<<"$pinned_rocm_args"; then + report fail "$label" "versioned ROCm tag did not select ROCm args" + return + fi + report ok "$label" +} +test_cross_vendor_docker_args "Docker args select CUDA or ROCm device contract" + # ── TTY flag selection. Regression guard for the process-substitution bug: # _set_tty_flags must run in the CALLER's scope so `[ -t 1 ]` inspects the # real terminal. If it is ever moved back behind `< <(...)` or `$(...)`, diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index cd0b5184e..78f96c60f 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -8,8 +8,8 @@ # Fallback path: a user runs the image directly (`docker run --gpus all # ghcr.io/luce-org/lucebox-hub:cuda12`) with no env-var prep. We then do a # minimal VRAM-tiered autotune — same tiers as `lucebox autotune`, kept in -# sync by hand. Anything more elaborate (driver-version probes, AMD paths, -# lspci fallbacks) belongs in the host CLI, not here. +# sync by hand. NVIDIA and AMD are both supported; elaborate driver/version +# diagnostics and heterogeneous-backend selection stay in the host CLI. set -euo pipefail @@ -227,13 +227,16 @@ _build_host_info_json() { printf '"kernel":%s,' "$(_json_str_or_null "${LUCEBOX_HOST_KERNEL:-}")" printf '"wsl_version":%s,' "$(_json_str_or_null "${LUCEBOX_HOST_WSL_VERSION:-}")" printf '"docker_version":%s,' "$(_json_str_or_null "${LUCEBOX_HOST_DOCKER_VERSION:-}")" + printf '"gpu_vendor":%s,' "$(_json_str_or_null "${LUCEBOX_HOST_GPU_VENDOR:-}")" printf '"nvidia_driver":%s,' "$(_json_str_or_null "${LUCEBOX_HOST_DRIVER_VERSION:-}")" printf '"nvidia_ctk_version":%s,' "$(_json_str_or_null "${LUCEBOX_HOST_NVIDIA_CTK_VERSION:-}")" + printf '"rocm_version":%s,' "$(_json_str_or_null "${LUCEBOX_HOST_ROCM_VERSION:-}")" printf '"cpu_model":%s,' "$(_json_str_or_null "${LUCEBOX_HOST_CPU_MODEL:-}")" printf '"nproc":%s,' "$(_json_int_or_null "${LUCEBOX_HOST_NPROC:-}")" printf '"ram_gb":%s,' "$(_json_int_or_null "${LUCEBOX_HOST_RAM_GB:-}")" printf '"gpus":%s,' "$(_emit_gpu_array)" printf '"cuda_visible_devices":%s,' "$(_json_str_or_null "${LUCEBOX_HOST_CUDA_VISIBLE_DEVICES:-}")" + printf '"hip_visible_devices":%s,' "$(_json_str_or_null "${LUCEBOX_HOST_HIP_VISIBLE_DEVICES:-}")" printf '"source":%s,' "$(_json_str_or_null "$source_tag")" printf '"collector":%s,' "$(_json_str_or_null "$collector_tag")" printf '"collected_at":%s' "$(_json_str_or_null "$collected_at")" @@ -243,17 +246,70 @@ _build_host_info_json() { write_host_info # ── detect ───────────────────────────────────────────────────────────────── -# nvidia-smi is always present here (--gpus all wires the driver in). +# The host wrapper wires either NVIDIA (--gpus all) or AMD (/dev/kfd + +# /dev/dri). Direct docker users get the same fallback detection here. GPU_VRAM_GB=0 +GPU_COUNT=0 if command -v nvidia-smi &>/dev/null; then if mem_mib=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null \ | head -1) && [ -n "$mem_mib" ]; then GPU_VRAM_GB=$((mem_mib / 1024)) fi -fi -GPU_COUNT=0 -if command -v nvidia-smi &>/dev/null; then GPU_COUNT=$(nvidia-smi -L 2>/dev/null | awk '/^GPU /{n++} END{print n+0}') || GPU_COUNT=0 +elif command -v amd-smi &>/dev/null; then + amd_stats=$(amd-smi static --asic --vram --csv 2>/dev/null | awk -F',' ' + NR == 1 { + for (i = 1; i <= NF; i++) { + key = $i; gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", key); col[key] = i + } + next + } + { + mem = $(col["size"]); arch = $(col["target_graphics_version"]) + gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", mem) + gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", arch) + if (mem ~ /^[0-9]+([.][0-9]+)?$/) { + n++ + if (mem > max) { max = mem; max_arch = arch } + } + } + END { if (n) printf "%d %d %s", max / 1024, n, max_arch } + ' || echo "") + if [ -n "$amd_stats" ]; then + read -r GPU_VRAM_GB GPU_COUNT GPU_ARCH <<<"$amd_stats" + fi +elif command -v rocm-smi &>/dev/null; then + amd_stats=$(rocm-smi --showproductname --showmeminfo vram --csv 2>/dev/null | awk -F',' ' + NR == 1 { + for (i = 1; i <= NF; i++) { + key = $i; gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", key); col[key] = i + } + next + } + { + bytes = $(col["VRAM Total Memory (B)"]); arch = $(col["GFX Version"]) + gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", bytes) + gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", arch) + if (bytes ~ /^[0-9]+$/) { + n++ + if (bytes > max) { max = bytes; max_arch = arch } + } + } + END { if (n) printf "%d %d %s", max / 1073741824, n, max_arch } + ' || echo "") + if [ -n "$amd_stats" ]; then + read -r GPU_VRAM_GB GPU_COUNT GPU_ARCH <<<"$amd_stats" + fi +fi + +# Strix Halo's unified memory is usable for model weights even when SMI only +# reports the small fixed VRAM carve-out. Mirror the host wrapper's effective +# capacity rule for direct `docker run` users. +if [ "${GPU_ARCH:-}" = "gfx1151" ] && [ "$GPU_VRAM_GB" -lt 12 ]; then + host_ram_gb=$(awk '/MemTotal/{printf "%.0f", $2/1024/1024}' /proc/meminfo 2>/dev/null || echo 0) + if [ "$host_ram_gb" -ge 32 ]; then + GPU_VRAM_GB=$host_ram_gb + fi fi # ── fallback autotune (only fills unset env) ─────────────────────────────── From 583694616be73efe70df8ac289682638feafd4bf Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 27 Jul 2026 13:12:29 +0200 Subject: [PATCH 03/25] fix(lucebox): harden config and target-only launches --- install.sh | 56 ++++++--------- lucebox.sh | 74 +++++++++++++++----- lucebox/src/lucebox/autotune.py | 5 +- lucebox/src/lucebox/cli.py | 19 ++--- lucebox/src/lucebox/config.py | 17 +++-- lucebox/src/lucebox/docker_run.py | 24 +++++-- lucebox/tests/test_config.py | 37 ++++++++++ lucebox/tests/test_config_cli.py | 13 ++++ lucebox/tests/test_docker_run.py | 33 ++++++++- scripts/test_lucebox_sh.sh | 112 +++++++++++++++++++++++++++--- server/scripts/entrypoint.sh | 6 +- 11 files changed, 305 insertions(+), 91 deletions(-) diff --git a/install.sh b/install.sh index c44f92712..c6f959bdc 100755 --- a/install.sh +++ b/install.sh @@ -36,43 +36,13 @@ die() { printf '%s[install] ✗%s %s\n' "$C_ERR" "$C_RST" "$*" >&2; exit 1; } command -v curl >/dev/null 2>&1 || die "curl is required (apt-get install curl)" -# ── fetch ───────────────────────────────────────────────────────────────── -tmp=$(mktemp -t lucebox.XXXXXX) || die "couldn't create temp file" -# shellcheck disable=SC2064 # we want $tmp expanded now, not at trap time -trap "rm -f '$tmp' '$tmp.bak'" EXIT -info "fetching $LUCEBOX_INSTALL_URL" -curl -fsSL "$LUCEBOX_INSTALL_URL" -o "$tmp" \ - || die "download failed from $LUCEBOX_INSTALL_URL" - -# ── sanity check ────────────────────────────────────────────────────────── -# Refuse to install something that isn't recognizably lucebox.sh. Catches -# 404 pages, redirects to HTML, and accidental URL typos. -head -1 "$tmp" | grep -q '^#!/usr/bin/env bash$' \ - || die "downloaded file does not look like a bash script (got: $(head -1 "$tmp"))" -grep -q '^VERSION=' "$tmp" \ - || die "downloaded file is missing VERSION marker — not lucebox.sh?" - # ── decide what gets baked in as the persisted channel ─────────────────── -# `lucebox update` reads LUCEBOX_INSTALLED_FROM from the installed copy and -# re-fetches from it. Persisting a SHA-pinned URL is a footgun — every -# future update would re-install the same frozen SHA forever, defeating -# the point of `update`. So: -# -# 1. If $LUCEBOX_INSTALL_CHANNEL is set, that's the persisted URL -# (caller takes responsibility for picking a real branch URL). -# 2. Else if LUCEBOX_INSTALL_URL has a 40-char hex SHA segment, refuse -# to persist it — tell the user to set LUCEBOX_INSTALL_CHANNEL. -# Common case: someone curl'd from /raw// to bypass a stale CDN -# cache during dev; they meant for updates to track the branch. -# 3. Else persist LUCEBOX_INSTALL_URL as-is (branch or canonical main). +# Do this before fetching so a SHA-pinned URL is refused for the intended +# reason even when that remote object is unavailable. channel_url="${LUCEBOX_INSTALL_CHANNEL:-}" if [ -z "$channel_url" ]; then - # Match a full 40-char hex SHA in the URL path, not the broader - # {7,40} range — a 7-39 char hex segment is more likely a branch - # name shaped like a short SHA (e.g. `feat/abc1234-hotfix`) than an - # actual SHA-pin. Keeping the gate at exactly 40 chars matches what - # `git rev-parse HEAD` emits and what `/raw//` URLs from - # GitHub's CDN actually carry. + # A full 40-char SHA is what `git rev-parse HEAD` and GitHub raw URLs use. + # Shorter hex-like path segments may be branch names, so don't reject them. if [[ "$LUCEBOX_INSTALL_URL" =~ /[0-9a-fA-F]{40}/[^/]+\.sh$ ]]; then die "$(cat < "$tmp.baked" mv "$tmp.baked" "$tmp" -grep -q "^LUCEBOX_INSTALLED_FROM=\"$escaped_url\"$" "$tmp" \ +grep -Fqx "LUCEBOX_INSTALLED_FROM=\"$channel_url\"" "$tmp" \ || die "failed to bake install source into the downloaded script" # ── install ─────────────────────────────────────────────────────────────── diff --git a/lucebox.sh b/lucebox.sh index ebdbb1139..d392064fd 100755 --- a/lucebox.sh +++ b/lucebox.sh @@ -101,7 +101,29 @@ _lucebox_config_get() { /=/ { if (current != want_section) next line = $0 - sub(/#.*$/, "", line) + # Strip a TOML comment only when # is outside a quoted string. + # Paths and image tags may legitimately contain #. + in_quote = 0 + escaped = 0 + for (i = 1; i <= length(line); i++) { + ch = substr(line, i, 1) + if (escaped) { + escaped = 0 + continue + } + if (in_quote && ch == "\\") { + escaped = 1 + continue + } + if (ch == "\"") { + in_quote = !in_quote + continue + } + if (ch == "#" && !in_quote) { + line = substr(line, 1, i - 1) + break + } + } eq = index(line, "=") if (eq == 0) next k = substr(line, 1, eq - 1) @@ -162,6 +184,7 @@ CONTAINER_NAME=$(_lucebox_resolve "${LUCEBOX_CONTAINER:-}" runtime.container_nam DEFAULT_PORT=$(_lucebox_resolve "${LUCEBOX_PORT:-}" runtime.port "8080") DEFAULT_MODELS_DIR=$(_lucebox_resolve "${LUCEBOX_MODELS:-}" paths.models "${XDG_DATA_HOME:-$HOME/.local/share}/lucebox/models") IMAGE_BASE=$(_lucebox_resolve "${LUCEBOX_IMAGE:-}" image.registry "$(_lucebox_derive_image "$LUCEBOX_INSTALLED_FROM")") +CONFIG_HOME="${LUCEBOX_HOME:-$HOME/.lucebox}" # ── LUCEBOX_HOST_* safe defaults (belt-and-suspenders) ──────────────────── # `set -u` makes any unbound LUCEBOX_HOST_* read fatal. Historically this has @@ -239,7 +262,7 @@ _parse_amd_smi_csv() { awk -F',' ' NR == 1 { for (i = 1; i <= NF; i++) { - key = $i + key = tolower($i) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", key) col[key] = i } @@ -266,7 +289,7 @@ _parse_rocm_smi_csv() { awk -F',' ' NR == 1 { for (i = 1; i <= NF; i++) { - key = $i + key = tolower($i) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", key) col[key] = i } @@ -274,9 +297,9 @@ _parse_rocm_smi_csv() { } { dev = $(col["device"]) - bytes = $(col["VRAM Total Memory (B)"]) - name = $(col["Card Series"]) - arch = $(col["GFX Version"]) + bytes = $(col["vram total memory (b)"]) + name = $(col["card series"]) + arch = $(col["gfx version"]) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", dev) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", bytes) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", name) @@ -559,9 +582,9 @@ pick_variant() { _variant_is_rocm() { # Variant names are not limited to the moving `rocm` tag. Releases and # CI produce tags such as `0.3.0-rocm` and `pr-335-rocm`; treat any tag - # containing the backend marker as ROCm. Spell out case-insensitivity - # instead of using Bash 4's `${value,,}` so host-only commands such as - # `check` keep working under the older Bash bundled with macOS too. + # containing the backend marker as ROCm. Spell out case-insensitivity so + # host-only commands such as `check` reach no Bash-4-only expansion. Full + # container dispatch still requires Bash 4.3+ for the argv namerefs below. case "$1" in *[Rr][Oo][Cc][Mm]*) return 0 ;; *) return 1 ;; @@ -686,6 +709,7 @@ _append_scalar_env() { _arr+=(-e "LUCEBOX_PORT=$DEFAULT_PORT") _arr+=(-e "LUCEBOX_CONTAINER=$CONTAINER_NAME") _arr+=(-e "LUCEBOX_MODELS=$DEFAULT_MODELS_DIR") + _arr+=(-e "LUCEBOX_HOME=$CONFIG_HOME") [ -n "${HF_TOKEN:-}" ] && _arr+=(-e "HF_TOKEN=$HF_TOKEN") return 0 } @@ -728,9 +752,9 @@ _set_tty_flags() { # usage: _set_tty_flags arrayname } build_orchestrator_argv() { - local variant="$1"; shift - local tty=() - _set_tty_flags tty + local variant="$1" caller_has_tty="$2"; shift 2 + local tty=(-i) + [ "$caller_has_tty" = "1" ] && tty=(-it) local argv=(docker run --rm "${tty[@]}") _append_gpu_args argv "$variant" argv+=(--name "${CONTAINER_NAME}-cli-$$") @@ -743,7 +767,11 @@ build_orchestrator_argv() { # when actually needed; pulling that mount when the host talks to # docker over TCP/SSH is fine. if [ -S "$DOCKER_SOCK_PATH" ]; then - argv+=(--group-add "$(stat -c '%g' "$DOCKER_SOCK_PATH")") + local socket_gid + socket_gid=$(stat -c '%g' "$DOCKER_SOCK_PATH" 2>/dev/null \ + || stat -f '%g' "$DOCKER_SOCK_PATH" 2>/dev/null \ + || echo "") + [ -n "$socket_gid" ] && argv+=(--group-add "$socket_gid") argv+=(-v "$DOCKER_SOCK_PATH:/var/run/docker.sock") fi argv+=(-v "$HOME:$HOME") @@ -753,7 +781,12 @@ build_orchestrator_argv() { # points XDG_DATA_HOME outside $HOME. mkdir -p "$DEFAULT_MODELS_DIR" argv+=(-v "$DEFAULT_MODELS_DIR:$DEFAULT_MODELS_DIR") - argv+=(-w "$PWD") + # A custom LUCEBOX_HOME may sit outside $HOME, so mount it explicitly. + # Use an image-owned working directory: callers may invoke lucebox from + # /tmp or another path that is not bind-mounted into the container. + mkdir -p "$CONFIG_HOME" + argv+=(-v "$CONFIG_HOME:$CONFIG_HOME") + argv+=(-w /opt/lucebox-hub) argv+=(-e "HOME=$HOME") # Host facts — Python side reads these instead of reprobing. _append_host_env argv @@ -836,7 +869,7 @@ cmd_serve() { esac local orch_argv server_argv - mapfile -t orch_argv < <(build_orchestrator_argv "$variant" print-serve-argv) + mapfile -t orch_argv < <(build_orchestrator_argv "$variant" 0 print-serve-argv) if mapfile -t server_argv < <("${orch_argv[@]}" 2>/dev/null) \ && [ "${#server_argv[@]}" -gt 0 ] \ @@ -859,6 +892,7 @@ cmd_serve() { --name "$CONTAINER_NAME" -p "$DEFAULT_PORT:8080" -v "$HOME:$HOME" + -v "$CONFIG_HOME:$CONFIG_HOME" -v "$fallback_models:/opt/lucebox-hub/server/models") _append_gpu_args fallback_argv "$variant" _append_host_env fallback_argv @@ -944,6 +978,7 @@ Environment=LUCEBOX_VARIANT=$variant Environment=LUCEBOX_PORT=$DEFAULT_PORT Environment=LUCEBOX_CONTAINER=$CONTAINER_NAME Environment=LUCEBOX_MODELS=$DEFAULT_MODELS_DIR +Environment=LUCEBOX_HOME=$CONFIG_HOME ExecStartPre=-$docker_bin rm -f $CONTAINER_NAME ExecStart=$SCRIPT_PATH serve ExecStop=$docker_bin stop -t 30 $CONTAINER_NAME @@ -1362,8 +1397,11 @@ cmd_in_container() { local variant variant=$(pick_variant) require_host_prereqs "$variant" - local argv - mapfile -t argv < <(build_orchestrator_argv "$variant" "$@") + local caller_has_tty=0 argv + if [ -t 0 ] && [ -t 1 ]; then + caller_has_tty=1 + fi + mapfile -t argv < <(build_orchestrator_argv "$variant" "$caller_has_tty" "$@") exec "${argv[@]}" } @@ -1404,7 +1442,7 @@ cmd_exec_in_container() { _set_tty_flags tty local argv=(docker exec "${tty[@]}") argv+=(--user "$(id -u):$(id -g)") - argv+=(-w "$PWD") + argv+=(-w /opt/lucebox-hub) argv+=(-e "HOME=$HOME") _append_host_env argv _append_scalar_env argv "$variant" diff --git a/lucebox/src/lucebox/autotune.py b/lucebox/src/lucebox/autotune.py index 5541cf271..0baf0fe91 100644 --- a/lucebox/src/lucebox/autotune.py +++ b/lucebox/src/lucebox/autotune.py @@ -26,8 +26,7 @@ def runtime_from_host(host: HostFacts) -> DflashRuntime: 22-31 — 24 GB-class consumer flagships (3090/4090/5090/5090-Laptop). 98 K with tq3_0 KV (~2 GB KV + ~18 GB model ≈ 20 GB). Confirmed on bragi (RTX 5090 Laptop, 23 GB VRAM) 2026-05-31. - 32-47 — RTX 6000 Ada / A100 40 GB. Full 128 K. - ≥48 — A100 80 GB / H100 / RTX 6000 Pro. Full 128 K. + ≥32 — RTX 6000 Ada / A100 / H100 class. Full 128 K. Prefix cache remains an explicit sweep tunable, but the automatic baseline keeps it off because tool prompts currently exercise a daemon @@ -75,6 +74,4 @@ def runtime_from_host(host: HostFacts) -> DflashRuntime: max_ctx=98304, cache_type_k="tq3_0", cache_type_v="tq3_0", ) - if host.vram_gb < 48: - return DflashRuntime(max_ctx=131072) return DflashRuntime(max_ctx=131072) diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index e3c6101bd..9c0e1809c 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -22,6 +22,7 @@ import typer from rich.console import Console +from rich.markup import escape from rich.table import Table import lucebox.config as config_mod @@ -102,7 +103,7 @@ def pull() -> None: """`docker pull` the image variant from config.toml.""" cfg = _load_or_build() tag = f"{cfg.image}:{cfg.variant}" - console.print(f"[bold]Pulling {tag}[/bold] (~14 GB; takes a while)…") + console.print(f"[bold]Pulling {escape(tag)}[/bold] (~14 GB; takes a while)…") rc = docker_run.docker_pull(tag) if rc != 0: raise typer.Exit(code=rc) @@ -145,10 +146,10 @@ def config_get_cmd( try: entries = config_get(key or None) except KeyError as exc: - console.print(f"[red]{exc}[/red]") + console.print(f"[red]{escape(str(exc))}[/red]") raise typer.Exit(code=2) from exc for k, (value, origin) in entries.items(): - console.print(f"{k} = {value!r} ([dim]from {origin}[/dim])") + console.print(f"{k} = {escape(repr(value))} ([dim]from {origin}[/dim])") @config_app.command("set") @@ -170,9 +171,9 @@ def config_set_cmd( try: config_set(key, value) except (KeyError, ValueError) as exc: - console.print(f"[red]{exc}[/red]") + console.print(f"[red]{escape(str(exc))}[/red]") raise typer.Exit(code=2) from exc - console.print(f"[green]Set[/green] {key} = {value}") + console.print(f"[green]Set[/green] {escape(key)} = {escape(value)}") @config_app.command("unset") @@ -183,12 +184,12 @@ def config_unset_cmd( try: changed = config_unset(key) except KeyError as exc: - console.print(f"[red]{exc}[/red]") + console.print(f"[red]{escape(str(exc))}[/red]") raise typer.Exit(code=2) from exc if changed: - console.print(f"[green]Unset[/green] {key}") + console.print(f"[green]Unset[/green] {escape(key)}") else: - console.print(f"[dim]{key} was not in config.toml; nothing to do[/dim]") + console.print(f"[dim]{escape(key)} was not in config.toml; nothing to do[/dim]") # ── models sub-app ───────────────────────────────────────────────────────── @@ -287,7 +288,7 @@ def models_download( try: pres = download_mod.resolve_preset(preset) except KeyError as exc: - console.print(f"[red]{exc}[/red]") + console.print(f"[red]{escape(str(exc))}[/red]") raise typer.Exit(code=2) from exc current = download_mod.status(cfg, pres) diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py index 6e3d113c8..7ca5d2445 100644 --- a/lucebox/src/lucebox/config.py +++ b/lucebox/src/lucebox/config.py @@ -26,7 +26,7 @@ from dataclasses import asdict, replace from datetime import UTC from pathlib import Path -from typing import Any +from typing import Any, Literal, cast import tomli_w @@ -55,11 +55,11 @@ def default_config_path() -> Path: # ── dotted-key registry ──────────────────────────────────────────────────── -def _cast_prefill_mode(v: Any) -> str: +def _cast_prefill_mode(v: Any) -> Literal["off", "auto", "always"]: s = str(v) if s not in {"off", "auto", "always"}: raise ValueError(f"prefill_mode must be off/auto/always, got {s!r}") - return s + return cast(Literal["off", "auto", "always"], s) def _cast_bool(v: Any) -> bool: @@ -236,12 +236,12 @@ def _from_dict(raw: dict[str, Any]) -> Config: dflash = DflashRuntime( budget=int(df.get("budget", 22)), max_ctx=int(df.get("max_ctx", 16384)), - lazy=bool(df.get("lazy", False)), + lazy=_cast_bool(df.get("lazy", False)), prefix_cache_slots=int(df.get("prefix_cache_slots", 0)), prefill_cache_slots=int(df.get("prefill_cache_slots", 0)), cache_type_k=str(df.get("cache_type_k", "")), cache_type_v=str(df.get("cache_type_v", "")), - prefill_mode=df.get("prefill_mode", "off"), + prefill_mode=_cast_prefill_mode(df.get("prefill_mode", "off")), prefill_keep_ratio=float(df.get("prefill_keep_ratio", 0.05)), prefill_threshold=int(df.get("prefill_threshold", 32000)), prefill_drafter=str(df.get("prefill_drafter", "")), @@ -249,7 +249,7 @@ def _from_dict(raw: dict[str, Any]) -> Config: fa_window=int(df.get("fa_window", 0)), think_soft_close_min_ratio=float( df.get("think_soft_close_min_ratio", 0.0)), - debug_thinking_logits=bool(df.get("debug_thinking_logits", False)), + debug_thinking_logits=_cast_bool(df.get("debug_thinking_logits", False)), ) host_raw = raw.get("host", {}) @@ -357,6 +357,11 @@ def seed_dflash_from_host(host: HostFacts, *, path: Path | None = None) -> bool: import lucebox.autotune as autotune_mod path = path or default_config_path() + # Preserve the normal load() migration contract even when this helper is + # called directly. Otherwise creating a fresh config.toml here would make + # an adjacent legacy config.env invisible to every future load. + if not path.exists() and path.with_suffix(".env").exists(): + load(path) doc = load_doc(path) if "dflash" in doc: return False diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index 8e61e355f..f8b688247 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -52,8 +52,8 @@ def _resolve_model_files(cfg: Config) -> tuple[str, str, str]: (a directory path) instead of the GGUF-file path, allowing the entrypoint to discover the safetensors file inside it. - Imported lazily to avoid the lucebox.types ↔ lucebox.download circular - import that surfaces when this module is imported from ``__init__``. + The preset registry is imported lazily so constructing a minimal Docker + spec does not import the Hugging Face download surface unnecessarily. """ target = cfg.model.target_file draft = cfg.model.draft_file @@ -75,12 +75,18 @@ def _resolve_model_files(cfg: Config) -> tuple[str, str, str]: def _runtime_volumes(cfg: Config) -> tuple[tuple[str, str], ...]: - """Mount models plus $HOME so absolute symlink targets remain valid.""" - home = str(Path.home()) + """Mount models, $HOME, and any config dir outside the home mount.""" + home_path = Path.home() + home = str(home_path) models = str(cfg.models_dir) volumes = [(models, "/opt/lucebox-hub/server/models")] if home != models: volumes.append((home, home)) + config_home = Path(os.environ.get("LUCEBOX_HOME", home_path / ".lucebox")).absolute() + # The same-path home mount normally covers config. If models_dir == HOME, + # that path is mounted at /opt/... instead, so add config explicitly too. + if home == models or not config_home.is_relative_to(home_path): + volumes.append((str(config_home), str(config_home))) return tuple(volumes) @@ -175,7 +181,11 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: # LUCEBOX_HOST_* first so they ride out front in the rendered argv, # making it obvious in `print-run` output what host facts get forwarded. env: list[tuple[str, str]] = list(_host_facts_env()) + config_home = str( + Path(os.environ.get("LUCEBOX_HOME", Path.home() / ".lucebox")).absolute() + ) env += [ + ("LUCEBOX_HOME", config_home), ("DFLASH_BUDGET", str(cfg.dflash.budget)), ("DFLASH_MAX_CTX", str(cfg.dflash.max_ctx)), ("DFLASH_PREFIX_CACHE_SLOTS", str(cfg.dflash.prefix_cache_slots)), @@ -198,6 +208,12 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: env.append(("DFLASH_DRAFT", f"/opt/lucebox-hub/server/models/draft/{draft_file}")) elif draft_dir: env.append(("DFLASH_DRAFT", f"/opt/lucebox-hub/server/models/draft/{draft_dir}")) + elif cfg.model.preset: + # An active target-only preset is an explicit choice, not permission + # for entrypoint.sh to scan models/draft and attach an unrelated stale + # draft left by a previously active model. The entrypoint treats this + # guaranteed-missing path as "run target-only". + env.append(("DFLASH_DRAFT", "/opt/lucebox-hub/server/models/.lucebox-no-draft")) if cfg.dflash.lazy: env.append(("DFLASH_LAZY", "1")) if cfg.dflash.cache_type_k: diff --git a/lucebox/tests/test_config.py b/lucebox/tests/test_config.py index 9c2b913a3..afedc7d91 100644 --- a/lucebox/tests/test_config.py +++ b/lucebox/tests/test_config.py @@ -35,6 +35,28 @@ def test_image_variant_round_trips_from_toml(tmp_path: Path) -> None: assert cfg.variant == "integration-props-uv-squared-clean-cuda12" +def test_toml_load_uses_strict_bool_and_prefill_casters(tmp_path: Path) -> None: + path = tmp_path / "config.toml" + path.write_text( + '[dflash]\nlazy = "false"\ndebug_thinking_logits = "false"\n' + 'prefill_mode = "auto"\n' + ) + + cfg = config._load_toml(path) + + assert cfg.dflash.lazy is False + assert cfg.dflash.debug_thinking_logits is False + assert cfg.dflash.prefill_mode == "auto" + + +def test_toml_load_rejects_invalid_prefill_mode(tmp_path: Path) -> None: + path = tmp_path / "config.toml" + path.write_text('[dflash]\nprefill_mode = "sometimes"\n') + + with pytest.raises(ValueError, match="prefill_mode"): + config._load_toml(path) + + def test_model_preset_round_trips_through_set_and_load(tmp_path: Path) -> None: """Setting model.preset writes a sparse TOML doc that loads back correctly.""" path = tmp_path / "config.toml" @@ -226,3 +248,18 @@ def test_seed_dflash_is_noop_when_dflash_present(tmp_path: Path) -> None: wrote = config.seed_dflash_from_host(HostFacts(vram_gb=80), path=path) assert wrote is False assert config.load(path).dflash.max_ctx == 4096 + + +def test_seed_dflash_migrates_legacy_config_before_writing(tmp_path: Path) -> None: + from lucebox.types import HostFacts + + path = tmp_path / "config.toml" + path.with_suffix(".env").write_text("DFLASH_PORT=9090\n") + + wrote = config.seed_dflash_from_host(HostFacts(vram_gb=24), path=path) + + assert wrote is True + loaded = config.load(path) + assert loaded is not None + assert loaded.port == 9090 + assert loaded.dflash.max_ctx == 98304 diff --git a/lucebox/tests/test_config_cli.py b/lucebox/tests/test_config_cli.py index 446ab41b6..6525ea319 100644 --- a/lucebox/tests/test_config_cli.py +++ b/lucebox/tests/test_config_cli.py @@ -76,6 +76,19 @@ def test_config_set_creates_file_when_missing( assert "port = 9090" in cfg_path.read_text() +def test_config_markup_characters_do_not_crash_output( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _set_config_path(tmp_path, monkeypatch) + + set_result = CliRunner().invoke(app, ["config", "set", "image=registry/[dev]"]) + get_result = CliRunner().invoke(app, ["config", "get", "image"]) + + assert set_result.exit_code == 0, set_result.stdout + assert get_result.exit_code == 0, get_result.stdout + assert "registry/[dev]" in get_result.stdout + + def test_load_or_build_env_overrides_persisted_config( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index 0f8999ea1..d4652cc5c 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -110,10 +110,26 @@ def test_runtime_volumes_mounts_models_and_home(tmp_path: Path) -> None: def test_runtime_volumes_dedupes_when_models_is_home(monkeypatch, tmp_path: Path) -> None: monkeypatch.setattr(Path, "home", staticmethod(lambda: tmp_path)) + monkeypatch.delenv("LUCEBOX_HOME", raising=False) cfg = Config(models_dir=tmp_path) vols = docker_run._runtime_volumes(cfg) - # models_dir == home → only the models mount, no duplicate home mount. - assert len(vols) == 1 + # models_dir == home is mounted at /opt/..., not at its host path, so the + # nested config directory still needs an explicit same-path mount. + assert len(vols) == 2 + assert (str(tmp_path / ".lucebox"), str(tmp_path / ".lucebox")) in vols + + +def test_runtime_volumes_mounts_custom_config_home( + monkeypatch, tmp_path: Path +) -> None: + home = tmp_path / "home" + config_home = tmp_path / "config-outside-home" + monkeypatch.setattr(Path, "home", staticmethod(lambda: home)) + monkeypatch.setenv("LUCEBOX_HOME", str(config_home)) + + vols = docker_run._runtime_volumes(Config(models_dir=tmp_path / "models")) + + assert (str(config_home), str(config_home)) in vols # ── _resolve_model_files ───────────────────────────────────────────────────── @@ -190,6 +206,19 @@ def test_server_run_spec_always_emits_core_dflash_env(tmp_path: Path) -> None: assert env["DFLASH_PREFILL_CACHE_SLOTS"] == "0" assert env["DFLASH_THINK_MAX"] == "15488" assert env["DFLASH_PORT"] == "8080" + assert env["LUCEBOX_HOME"] + + +def test_server_run_spec_target_only_preset_disables_stale_draft(tmp_path: Path) -> None: + cfg = Config( + models_dir=tmp_path, + model=ModelMeta(preset="qwen3.6-moe"), + ) + + env = _env(docker_run.server_run_spec(cfg)) + + assert env["DFLASH_TARGET"].endswith("Qwen3.6-35B-A3B-UD-Q4_K_M.gguf") + assert env["DFLASH_DRAFT"].endswith("/.lucebox-no-draft") def test_server_run_spec_optional_env_off_by_default(tmp_path: Path) -> None: diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index f858690ba..80125e8f0 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -85,8 +85,11 @@ report() { assert_runs() { local label="$1" cmd="$2" expect="${3:-}" local out rc - out=$(NO_COLOR=1 bash -c "$cmd" 2>&1) - rc=$? + if out=$(NO_COLOR=1 bash -c "$cmd" 2>&1); then + rc=0 + else + rc=$? + fi if [ "$rc" -ne 0 ]; then report fail "$label" "exit $rc; output: $(printf '%s' "$out" | head -3)" return @@ -337,6 +340,7 @@ STUB "Environment=LUCEBOX_VARIANT=" \ "Environment=LUCEBOX_PORT=" \ "Environment=LUCEBOX_MODELS=" \ + "Environment=LUCEBOX_HOME=" \ "SuccessExitStatus=143" \ ; do grep -qF "$needle" "$unit_path" || missing="$missing $needle" @@ -730,7 +734,7 @@ port = 9090 container_name = "luce-test" [paths] -models = "/srv/models" +models = "/srv/models#fast" # an actual TOML comment [dflash] budget = 22 @@ -745,7 +749,7 @@ TOML "|image.variant|cuda12|cuda13" "|runtime.port|8080|9090" "|runtime.container_name|lucebox|luce-test" - "|paths.models|/var/lib/lucebox|/srv/models" + "|paths.models|/var/lib/lucebox|/srv/models#fast" "OVERRIDE|image.registry|ghcr.io/luce-org/lucebox-hub|OVERRIDE" "|missing.key|fallback-default|fallback-default" ) @@ -822,14 +826,14 @@ test_cmd_start_already_active_shortcircuit "lucebox start has already-active + r # forever. With CHANNEL set, the bake-in uses the channel URL, not the # fetch URL. test_install_sha_pin_refusal_and_channel_override() { - local label="$1" tmp got rc + local label="$1" tmp got out rc tmp=$(mktemp -d -t lucebox-sha.XXXXXX) # Case 1: SHA-pinned URL without CHANNEL → must refuse - LUCEBOX_INSTALL_URL="https://raw.githubusercontent.com/easel/lucebox-hub/abc1234567/lucebox.sh" \ - LUCEBOX_INSTALL_DEST="$tmp/lucebox1" \ - NO_COLOR=1 \ - bash "$INSTALLER" >/dev/null 2>&1 && rc=0 || rc=$? + out=$(LUCEBOX_INSTALL_URL="https://raw.githubusercontent.com/easel/lucebox-hub/0123456789abcdef0123456789abcdef01234567/lucebox.sh" \ + LUCEBOX_INSTALL_DEST="$tmp/lucebox1" \ + NO_COLOR=1 \ + bash "$INSTALLER" 2>&1) && rc=0 || rc=$? if [ "$rc" -eq 0 ]; then rm -rf "$tmp" report fail "$label" "SHA-pinned URL without CHANNEL should have refused (rc=$rc, got success)" @@ -840,6 +844,11 @@ test_install_sha_pin_refusal_and_channel_override() { report fail "$label" "SHA-pinned URL refusal still wrote $tmp/lucebox1" return fi + if ! grep -qF 'is SHA-pinned' <<<"$out"; then + rm -rf "$tmp" + report fail "$label" "did not reach SHA-pin refusal branch: $(head -3 <<<"$out")" + return + fi # Case 2: SHA-pinned URL WITH CHANNEL → installs, bakes CHANNEL LUCEBOX_INSTALL_URL="file://$SCRIPT" \ @@ -952,7 +961,7 @@ _run_wrapper_capture_docker() { HOME="$sandbox" \ XDG_CONFIG_HOME="$sandbox/.config" \ XDG_DATA_HOME="$sandbox/.local/share" \ - LUCEBOX_HOME="$sandbox/.lucebox" \ + LUCEBOX_HOME="${TEST_LUCEBOX_HOME:-$sandbox/.lucebox}" \ PATH="$shim_dir:$PATH" \ LUCEBOX_HOST_HAS_DOCKER=1 \ LUCEBOX_HOST_HAS_CTK=runtime \ @@ -1001,6 +1010,10 @@ test_routes_to_exec_when_running() { report fail "$label" "exec argv missing 'LUCEBOX_IMAGE=' scalar env; got: $(head -3 <<<"$out")" return fi + if ! grep -q -- '-w /opt/lucebox-hub' <<<"$out"; then + report fail "$label" "exec argv uses a caller-dependent working directory: $(head -3 <<<"$out")" + return + fi report ok "$label" } test_routes_to_exec_when_running "config get routes to docker exec when container running" @@ -1023,6 +1036,85 @@ test_routes_to_run_when_not_running() { } test_routes_to_run_when_not_running "config get falls back to docker run when container not running" +test_custom_lucebox_home_is_mounted_and_forwarded() { + local label="$1" sandbox config_home out + sandbox=$(mktemp -d -t lucebox-route.XXXXXX) + config_home=$(mktemp -d -t lucebox-config.XXXXXX) + _make_docker_shim "$sandbox" 0 + out=$(TEST_LUCEBOX_HOME="$config_home" \ + _run_wrapper_capture_docker "$sandbox" config get model.preset || true) + rm -rf "$sandbox" "$config_home" + if ! grep -qF -- "-v $config_home:$config_home" <<<"$out"; then + report fail "$label" "custom config dir was not mounted: $(head -3 <<<"$out")" + return + fi + if ! grep -qF "LUCEBOX_HOME=$config_home" <<<"$out"; then + report fail "$label" "custom config dir was not forwarded: $(head -3 <<<"$out")" + return + fi + if ! grep -q -- '-w /opt/lucebox-hub' <<<"$out"; then + report fail "$label" "orchestrator uses a caller-dependent working directory" + return + fi + report ok "$label" +} +test_custom_lucebox_home_is_mounted_and_forwarded \ + "custom LUCEBOX_HOME is mounted + forwarded to orchestrator" + +test_run_route_preserves_tty() { + local label="$1" sandbox out + sandbox=$(mktemp -d -t lucebox-route-tty.XXXXXX) + _make_docker_shim "$sandbox" 0 + out=$(python3 - "$SCRIPT" "$sandbox" <<'PY' 2>/dev/null +import os +import pty +import sys + +script, sandbox = sys.argv[1:] +env = os.environ.copy() +env.update({ + "HOME": sandbox, + "XDG_CONFIG_HOME": sandbox + "/.config", + "XDG_DATA_HOME": sandbox + "/.local/share", + "LUCEBOX_HOME": sandbox + "/.lucebox", + "PATH": sandbox + "/bin:" + env["PATH"], + "LUCEBOX_HOST_HAS_DOCKER": "1", + "LUCEBOX_HOST_HAS_CTK": "runtime", + "LUCEBOX_HOST_GPU_VENDOR": "nvidia", + "LUCEBOX_HOST_HAS_NVIDIA_GPU": "1", + "LUCEBOX_HOST_DRIVER_MAJOR": "550", + "LUCEBOX_HOST_GPU_NAME": "Fake GPU", + "LUCEBOX_HOST_GPU_COUNT": "1", + "LUCEBOX_HOST_VRAM_GB": "24", + "LUCEBOX_HOST_GPU_SM": "89", + "_LUCEBOX_HOST_PROBED": "1", + "NO_COLOR": "1", +}) +pid, fd = pty.fork() +if pid == 0: + os.execve("/bin/bash", ["bash", script, "no-such-subcommand"], env) +buf = b"" +try: + while True: + chunk = os.read(fd, 4096) + if not chunk: + break + buf += chunk +except OSError: + pass +os.waitpid(pid, 0) +sys.stdout.write(buf.decode(errors="replace")) +PY +) + rm -rf "$sandbox" + if ! grep -qE '^DOCKER_INVOKED run .* -it( |$)' <<<"${out//$'\r'/}"; then + report fail "$label" "PTY route did not preserve docker -it: $(head -3 <<<"$out")" + return + fi + report ok "$label" +} +test_run_route_preserves_tty "docker-run route preserves caller TTY (-it)" + test_no_exec_flag_forces_run() { local label="$1" sandbox out sandbox=$(mktemp -d -t lucebox-route.XXXXXX) diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index 78f96c60f..95373aa1b 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -260,7 +260,7 @@ elif command -v amd-smi &>/dev/null; then amd_stats=$(amd-smi static --asic --vram --csv 2>/dev/null | awk -F',' ' NR == 1 { for (i = 1; i <= NF; i++) { - key = $i; gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", key); col[key] = i + key = tolower($i); gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", key); col[key] = i } next } @@ -282,12 +282,12 @@ elif command -v rocm-smi &>/dev/null; then amd_stats=$(rocm-smi --showproductname --showmeminfo vram --csv 2>/dev/null | awk -F',' ' NR == 1 { for (i = 1; i <= NF; i++) { - key = $i; gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", key); col[key] = i + key = tolower($i); gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", key); col[key] = i } next } { - bytes = $(col["VRAM Total Memory (B)"]); arch = $(col["GFX Version"]) + bytes = $(col["vram total memory (b)"]); arch = $(col["gfx version"]) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", bytes) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", arch) if (bytes ~ /^[0-9]+$/) { From 45431995679d20e48b59c0770512ce9a524c85aa Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 27 Jul 2026 13:25:17 +0200 Subject: [PATCH 04/25] fix(lucebox): align fallback config paths --- lucebox.sh | 4 +++- lucebox/src/lucebox/docker_run.py | 6 ++++-- lucebox/tests/test_docker_run.py | 12 ++++++++++++ scripts/test_lucebox_sh.sh | 31 ++++++++++++++++++++++++++++--- 4 files changed, 47 insertions(+), 6 deletions(-) diff --git a/lucebox.sh b/lucebox.sh index d392064fd..e03445131 100755 --- a/lucebox.sh +++ b/lucebox.sh @@ -893,9 +893,11 @@ cmd_serve() { -p "$DEFAULT_PORT:8080" -v "$HOME:$HOME" -v "$CONFIG_HOME:$CONFIG_HOME" - -v "$fallback_models:/opt/lucebox-hub/server/models") + -v "$fallback_models:/opt/lucebox-hub/server/models" + -e "HOME=$HOME") _append_gpu_args fallback_argv "$variant" _append_host_env fallback_argv + _append_scalar_env fallback_argv "$variant" fallback_argv+=("${IMAGE_BASE}:${variant}") _serve_and_track "${fallback_argv[@]}" } diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index f8b688247..ab34c837a 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -82,7 +82,9 @@ def _runtime_volumes(cfg: Config) -> tuple[tuple[str, str], ...]: volumes = [(models, "/opt/lucebox-hub/server/models")] if home != models: volumes.append((home, home)) - config_home = Path(os.environ.get("LUCEBOX_HOME", home_path / ".lucebox")).absolute() + config_home = Path( + os.environ.get("LUCEBOX_HOME") or home_path / ".lucebox" + ).absolute() # The same-path home mount normally covers config. If models_dir == HOME, # that path is mounted at /opt/... instead, so add config explicitly too. if home == models or not config_home.is_relative_to(home_path): @@ -182,7 +184,7 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: # making it obvious in `print-run` output what host facts get forwarded. env: list[tuple[str, str]] = list(_host_facts_env()) config_home = str( - Path(os.environ.get("LUCEBOX_HOME", Path.home() / ".lucebox")).absolute() + Path(os.environ.get("LUCEBOX_HOME") or Path.home() / ".lucebox").absolute() ) env += [ ("LUCEBOX_HOME", config_home), diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index d4652cc5c..f5a18ed59 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -132,6 +132,18 @@ def test_runtime_volumes_mounts_custom_config_home( assert (str(config_home), str(config_home)) in vols +def test_empty_lucebox_home_uses_default(monkeypatch, tmp_path: Path) -> None: + home = tmp_path / "home" + monkeypatch.setattr(Path, "home", staticmethod(lambda: home)) + monkeypatch.setenv("LUCEBOX_HOME", "") + + cfg = Config(models_dir=tmp_path / "models") + spec = docker_run.server_run_spec(cfg) + + assert _env(spec)["LUCEBOX_HOME"] == str(home / ".lucebox") + assert str(Path.cwd()) != _env(spec)["LUCEBOX_HOME"] + + # ── _resolve_model_files ───────────────────────────────────────────────────── diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index 80125e8f0..c9fc0577d 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -1061,11 +1061,29 @@ test_custom_lucebox_home_is_mounted_and_forwarded() { test_custom_lucebox_home_is_mounted_and_forwarded \ "custom LUCEBOX_HOME is mounted + forwarded to orchestrator" +test_serve_fallback_forwards_config_env() { + local label="$1" sandbox config_home out + sandbox=$(mktemp -d -t lucebox-serve-fallback.XXXXXX) + config_home=$(mktemp -d -t lucebox-config.XXXXXX) + _make_docker_shim "$sandbox" 0 + out=$(TEST_LUCEBOX_HOME="$config_home" \ + _run_wrapper_capture_docker "$sandbox" serve || true) + rm -rf "$sandbox" "$config_home" + if ! grep -qF "LUCEBOX_HOME=$config_home" <<<"$out" \ + || ! grep -qF "HOME=$sandbox" <<<"$out"; then + report fail "$label" "fallback server omitted HOME/config env: $(tail -3 <<<"$out")" + return + fi + report ok "$label" +} +test_serve_fallback_forwards_config_env \ + "serve fallback forwards HOME + LUCEBOX_HOME" + test_run_route_preserves_tty() { - local label="$1" sandbox out + local label="$1" sandbox out rc sandbox=$(mktemp -d -t lucebox-route-tty.XXXXXX) _make_docker_shim "$sandbox" 0 - out=$(python3 - "$SCRIPT" "$sandbox" <<'PY' 2>/dev/null + if out=$(timeout 15 python3 - "$SCRIPT" "$sandbox" <<'PY' 2>/dev/null import os import pty import sys @@ -1105,7 +1123,14 @@ except OSError: os.waitpid(pid, 0) sys.stdout.write(buf.decode(errors="replace")) PY -) + ); then + : + else + rc=$? + rm -rf "$sandbox" + report fail "$label" "PTY route timed out or failed (rc=$rc)" + return + fi rm -rf "$sandbox" if ! grep -qE '^DOCKER_INVOKED run .* -it( |$)' <<<"${out//$'\r'/}"; then report fail "$label" "PTY route did not preserve docker -it: $(head -3 <<<"$out")" From 40a04a3092c5b69c3ab5d47f849ec3528f57c343 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:35:13 +0200 Subject: [PATCH 05/25] fix(lucebox): harden the CLI package --- Dockerfile | 8 +- Dockerfile.rocm | 8 +- install.sh | 24 ++++ lucebox.sh | 155 ++++++++++++++++------- lucebox/README.md | 35 ++--- lucebox/pyproject.toml | 11 +- lucebox/src/lucebox/__main__.py | 6 +- lucebox/src/lucebox/cli.py | 99 ++++++++------- lucebox/src/lucebox/config.py | 117 +++++++++++++---- lucebox/src/lucebox/docker_run.py | 148 ++++++++++++++++++---- lucebox/src/lucebox/download.py | 44 +++++-- lucebox/src/lucebox/types.py | 33 ++++- lucebox/tests/test_cli.py | 10 +- lucebox/tests/test_config.py | 67 ++++++++++ lucebox/tests/test_config_cli.py | 13 ++ lucebox/tests/test_docker_run.py | 119 +++++++++++++++-- lucebox/tests/test_download.py | 18 +++ lucebox/tests/test_models_cli.py | 2 +- scripts/check_lucebox_wrapper_sandbox.sh | 13 +- scripts/test_lucebox_sh.sh | 148 ++++++++++++++++++++-- 20 files changed, 872 insertions(+), 206 deletions(-) diff --git a/Dockerfile b/Dockerfile index 6148f7ce5..20e93b2b4 100644 --- a/Dockerfile +++ b/Dockerfile @@ -126,7 +126,8 @@ COPY server/pyproject.toml server/README.md /src/server/ COPY server/scripts /src/server/scripts COPY optimizations/pflash /src/optimizations/pflash COPY optimizations/megakernel /src/optimizations/megakernel -COPY lucebox /src/lucebox +COPY lucebox/pyproject.toml lucebox/README.md /src/lucebox/ +COPY lucebox/src /src/lucebox/src # ─── Stage 2: runtime ─────────────────────────────────────────────────────── # Runtime image: ships nvidia driver libs but no nvcc / dev headers. Matches @@ -206,7 +207,12 @@ COPY --from=builder /src/server/build /opt/lucebox-hub/server/build # server/share/model_cards. The canonical copy also lives at # /opt/lucebox-hub/share/model_cards for any host-side tooling. COPY share/model_cards /opt/lucebox-hub/share/model_cards +# Dedicated targets for narrow, read-only binds when a selected model is a +# symlink to storage outside the main models directory. RUN mkdir -p /opt/lucebox-hub/server/share \ + /opt/lucebox-resolved/target \ + /opt/lucebox-resolved/draft \ + /opt/lucebox-resolved/draft-dir \ && ln -s /opt/lucebox-hub/share/model_cards \ /opt/lucebox-hub/server/share/model_cards diff --git a/Dockerfile.rocm b/Dockerfile.rocm index ba8b8bb7e..5cfd9b373 100644 --- a/Dockerfile.rocm +++ b/Dockerfile.rocm @@ -123,7 +123,8 @@ COPY server/pyproject.toml server/README.md /src/server/ COPY server/scripts /src/server/scripts COPY optimizations/pflash /src/optimizations/pflash COPY optimizations/megakernel /src/optimizations/megakernel -COPY lucebox /src/lucebox +COPY lucebox/pyproject.toml lucebox/README.md /src/lucebox/ +COPY lucebox/src /src/lucebox/src # ─── Stage 2: runtime ─────────────────────────────────────────────────────── # Runtime reuses the ROCm base so the HIP runtime libs (libamdhip64, @@ -182,7 +183,12 @@ COPY --from=builder /src/server/pyproject.toml /src/server/README.md \ COPY --from=builder /src/server/build /opt/lucebox-hub/server/build COPY share/model_cards /opt/lucebox-hub/share/model_cards +# Dedicated targets for narrow, read-only binds when a selected model is a +# symlink to storage outside the main models directory. RUN mkdir -p /opt/lucebox-hub/server/share \ + /opt/lucebox-resolved/target \ + /opt/lucebox-resolved/draft \ + /opt/lucebox-resolved/draft-dir \ && ln -s /opt/lucebox-hub/share/model_cards \ /opt/lucebox-hub/server/share/model_cards diff --git a/install.sh b/install.sh index c6f959bdc..d67dbd9b4 100755 --- a/install.sh +++ b/install.sh @@ -36,6 +36,16 @@ die() { printf '%s[install] ✗%s %s\n' "$C_ERR" "$C_RST" "$*" >&2; exit 1; } command -v curl >/dev/null 2>&1 || die "curl is required (apt-get install curl)" +sha256_file() { + if command -v sha256sum >/dev/null 2>&1; then + sha256sum "$1" | awk '{print $1}' + elif command -v shasum >/dev/null 2>&1; then + shasum -a 256 "$1" | awk '{print $1}' + else + die "checksum requested, but neither sha256sum nor shasum is installed" + fi +} + # ── decide what gets baked in as the persisted channel ─────────────────── # Do this before fetching so a SHA-pinned URL is refused for the intended # reason even when that remote object is unavailable. @@ -68,6 +78,20 @@ info "fetching $LUCEBOX_INSTALL_URL" curl --connect-timeout 10 --max-time 120 -fsSL "$LUCEBOX_INSTALL_URL" -o "$tmp" \ || die "download failed from $LUCEBOX_INSTALL_URL" +# Release automation can pin the exact wrapper payload. This is optional for +# branch-channel installs, where the URL intentionally moves over time. +expected_sha="${LUCEBOX_WRAPPER_SHA256:-}" +if [ -n "$expected_sha" ]; then + [[ "$expected_sha" =~ ^[0-9a-fA-F]{64}$ ]] \ + || die "LUCEBOX_WRAPPER_SHA256 must be exactly 64 hexadecimal characters" + actual_sha=$(sha256_file "$tmp") + actual_sha=$(printf '%s' "$actual_sha" | tr '[:upper:]' '[:lower:]') + expected_sha=$(printf '%s' "$expected_sha" | tr '[:upper:]' '[:lower:]') + [ "$actual_sha" = "$expected_sha" ] \ + || die "wrapper checksum mismatch (expected $expected_sha, got $actual_sha)" + ok "wrapper sha256 verified" +fi + # ── sanity check ────────────────────────────────────────────────────────── # Refuse to install something that isn't recognizably lucebox.sh. Catches # 404 pages, redirects to HTML, and accidental URL typos. diff --git a/lucebox.sh b/lucebox.sh index e03445131..273f9afaf 100755 --- a/lucebox.sh +++ b/lucebox.sh @@ -250,6 +250,19 @@ err() { printf '%b[ERROR]%b %s\n' "$C_ERR" "$C_RST" "$*" >&2; } hint() { printf ' %b%s%b\n' "$C_DIM" "$*" "$C_RST"; } die() { err "$*"; exit 1; } +sha256_file() { + local sum + if command -v sha256sum >/dev/null 2>&1; then + sum=$(sha256sum "$1") + elif command -v shasum >/dev/null 2>&1; then + sum=$(shasum -a 256 "$1") + else + die "checksum requested, but neither sha256sum nor shasum is installed" + fi + sum="${sum%% *}" + printf '%s' "$sum" | tr '[:upper:]' '[:lower:]' +} + # ── host probing ────────────────────────────────────────────────────────── # Sets the LUCEBOX_HOST_* variables consumed by the in-container Python CLI # (passed through with -e). The Python side trusts these and doesn't reprobe @@ -676,9 +689,9 @@ require_systemd() { # ── docker run construction ─────────────────────────────────────────────── # All the Python-CLI subcommands share the same docker run incantation: # mount the host docker socket (so the in-container CLI can spawn server / -# bench containers on the host daemon), mount $HOME at the same path (so -# paths look identical in and out), and pass host facts via env. The selected -# image gets its native accelerator contract: --gpus all for +# bench containers on the host daemon), mount only Lucebox's config/models +# state, and pass host facts via env. The selected image gets its native +# accelerator contract: --gpus all for # CUDA; /dev/kfd + /dev/dri and the render/video groups for ROCm. DOCKER_SOCK_PATH="${DOCKER_HOST:-/var/run/docker.sock}" @@ -774,20 +787,21 @@ build_orchestrator_argv() { [ -n "$socket_gid" ] && argv+=(--group-add "$socket_gid") argv+=(-v "$DOCKER_SOCK_PATH:/var/run/docker.sock") fi - argv+=(-v "$HOME:$HOME") - # Bind-mount the XDG models dir explicitly (host = container path) so - # paths line up in/out. The $HOME mount above already covers it when - # XDG_DATA_HOME is unset, but an explicit -v is required when the user - # points XDG_DATA_HOME outside $HOME. + # Bind only Lucebox-owned state. The orchestrator used to receive all of + # $HOME read-write to support model symlinks; the Python launch builder + # now handles selected symlinks with narrow read-only mounts instead. + # Keeping credentials and unrelated user files outside the container is + # the safer default for both buyer appliances and contributor machines. mkdir -p "$DEFAULT_MODELS_DIR" argv+=(-v "$DEFAULT_MODELS_DIR:$DEFAULT_MODELS_DIR") - # A custom LUCEBOX_HOME may sit outside $HOME, so mount it explicitly. + # A custom LUCEBOX_HOME may sit anywhere, so mount it explicitly and use + # it as the ephemeral CLI's HOME (its caches then remain app-scoped too). # Use an image-owned working directory: callers may invoke lucebox from # /tmp or another path that is not bind-mounted into the container. mkdir -p "$CONFIG_HOME" argv+=(-v "$CONFIG_HOME:$CONFIG_HOME") argv+=(-w /opt/lucebox-hub) - argv+=(-e "HOME=$HOME") + argv+=(-e "HOME=$CONFIG_HOME") # Host facts — Python side reads these instead of reprobing. _append_host_env argv # User overrides for image/port/container/models scalars + HF_TOKEN. @@ -868,17 +882,38 @@ cmd_serve() { ;; esac - local orch_argv server_argv + local orch_argv server_argv server_output orch_error_file orch_rc=0 mapfile -t orch_argv < <(build_orchestrator_argv "$variant" 0 print-serve-argv) - if mapfile -t server_argv < <("${orch_argv[@]}" 2>/dev/null) \ - && [ "${#server_argv[@]}" -gt 0 ] \ - && [ "${server_argv[0]}" = "docker" ]; then - info "Starting lucebox server (variant=$variant, from config.toml)" - _serve_and_track "${server_argv[@]}" - return $? + orch_error_file=$(mktemp -t lucebox-orchestrator.XXXXXX) \ + || die "couldn't create temporary orchestrator log" + if server_output=$("${orch_argv[@]}" 2>"$orch_error_file"); then + if [ -n "$server_output" ]; then + mapfile -t server_argv <<<"$server_output" + if [ "${#server_argv[@]}" -gt 0 ] \ + && [ "${server_argv[0]}" = "docker" ]; then + rm -f "$orch_error_file" + info "Starting lucebox server (variant=$variant, from config.toml)" + _serve_and_track "${server_argv[@]}" + return $? + fi + fi + else + orch_rc=$? fi + # Configuration/path errors must never turn into a successful launch with + # defaults: that could start the wrong model or ignore explicit tuning. + # Docker infrastructure failures (for example an image not pulled yet) + # retain the conservative fallback below. + if [ "$orch_rc" -eq 2 ] \ + || grep -qE 'Invalid configuration:|Cannot build server command:' "$orch_error_file"; then + [ ! -s "$orch_error_file" ] || cat "$orch_error_file" >&2 + rm -f "$orch_error_file" + die "refusing to ignore invalid Lucebox configuration" + fi + rm -f "$orch_error_file" + warn "Couldn't fetch server argv from container (image not pulled?) — using fallback" info "Starting lucebox server (variant=$variant, port=$DEFAULT_PORT, defaults only)" local fallback_models="$DEFAULT_MODELS_DIR" @@ -891,10 +926,9 @@ cmd_serve() { local fallback_argv=(docker run --rm --name "$CONTAINER_NAME" -p "$DEFAULT_PORT:8080" - -v "$HOME:$HOME" -v "$CONFIG_HOME:$CONFIG_HOME" -v "$fallback_models:/opt/lucebox-hub/server/models" - -e "HOME=$HOME") + -e "HOME=$CONFIG_HOME") _append_gpu_args fallback_argv "$variant" _append_host_env fallback_argv _append_scalar_env fallback_argv "$variant" @@ -1129,38 +1163,68 @@ cmd_pull() { } cmd_update() { - # Re-run the bootstrap installer against the channel we were installed - # from. The installer is the source of truth for "how do you install - # lucebox correctly" — chmod, atomic mv, validation, baking the source - # URL back into the new copy so the channel is preserved across - # upgrades. Keeping the logic in install.sh means it can evolve - # independently (sha verify, signature check, etc.) and the installed - # `lucebox update` picks those changes up on the next run. - # - # The installer URL is derived from LUCEBOX_INSTALLED_FROM by swapping - # `lucebox.sh` → `install.sh` in the same directory, so forks don't - # need a separate registration. Override the source channel via - # $LUCEBOX_INSTALL_URL (e.g. to switch from canonical to a dev fork). - local source_url installer_url target + # Download the wrapper itself as data; never execute a second remote + # installer. Override the persisted channel with LUCEBOX_INSTALL_URL. + local source_url target wrapper_tmp expected_sha actual_sha escaped_url source_url="${LUCEBOX_INSTALL_URL:-$LUCEBOX_INSTALLED_FROM}" if [[ "$source_url" != */lucebox.sh ]]; then die "LUCEBOX_INSTALLED_FROM doesn't end in /lucebox.sh: $source_url" fi - installer_url="${source_url%/lucebox.sh}/install.sh" + case "$source_url" in + *['"$`\']*|*$'\n'*|*$'\r'*) + die "update URL contains unsafe characters: $source_url" ;; + esac target=$(realpath "$SCRIPT_PATH") - info "Updating lucebox via $installer_url" + info "Updating lucebox" info " source: $source_url" info " target: $target" - # Pass the URLs through to install.sh via env. The installer reads - # $LUCEBOX_INSTALL_URL (which we set to source_url) and - # $LUCEBOX_INSTALL_DEST (the realpath of *this* file, so a symlinked - # install replaces the actual file behind the link). - LUCEBOX_INSTALL_URL="$source_url" \ - LUCEBOX_INSTALL_DEST="$target" \ - bash -c "$(curl -fsSL "$installer_url")" \ - || die "update failed (installer exited non-zero)" + # Create the temporary file next to the destination so the final rename + # is atomic even when /tmp and the install directory are different mounts. + wrapper_tmp=$(mktemp "${target}.update.XXXXXX") \ + || die "couldn't create temporary update file next to $target" + trap 'if [ -n "${wrapper_tmp:-}" ]; then rm -f "$wrapper_tmp" "$wrapper_tmp.baked"; fi' EXIT + curl --connect-timeout 10 --max-time 120 --retry 2 --retry-delay 1 \ + -fsSL "$source_url" -o "$wrapper_tmp" \ + || die "failed to download wrapper from $source_url" + + [ "$(head -1 "$wrapper_tmp")" = '#!/usr/bin/env bash' ] \ + || die "downloaded wrapper has an unexpected shebang" + grep -Fqx 'set -euo pipefail' "$wrapper_tmp" \ + || die "downloaded wrapper is missing strict shell mode" + grep -q '^VERSION=' "$wrapper_tmp" \ + || die "downloaded file is missing the Lucebox version marker" + grep -q '^LUCEBOX_INSTALLED_FROM=' "$wrapper_tmp" \ + || die "downloaded file is missing the Lucebox update-channel marker" + bash -n "$wrapper_tmp" \ + || die "downloaded wrapper does not parse as valid Bash" + + expected_sha="${LUCEBOX_WRAPPER_SHA256:-}" + if [ -n "$expected_sha" ]; then + [[ "$expected_sha" =~ ^[0-9a-fA-F]{64}$ ]] \ + || die "LUCEBOX_WRAPPER_SHA256 must be exactly 64 hexadecimal characters" + actual_sha=$(sha256_file "$wrapper_tmp") + expected_sha=$(printf '%s' "$expected_sha" | tr '[:upper:]' '[:lower:]') + [ "$actual_sha" = "$expected_sha" ] \ + || die "wrapper checksum mismatch (expected $expected_sha, got $actual_sha)" + ok "wrapper sha256 verified" + fi + + # Preserve the chosen branch/fork in the new copy. The URL was validated + # above for safe embedding in a Bash double-quoted string. + escaped_url=$(printf '%s' "$source_url" | sed 's/[&|]/\\&/g') + sed "s|^LUCEBOX_INSTALLED_FROM=.*|LUCEBOX_INSTALLED_FROM=\"$escaped_url\"|" \ + "$wrapper_tmp" > "$wrapper_tmp.baked" + mv "$wrapper_tmp.baked" "$wrapper_tmp" + grep -Fqx "LUCEBOX_INSTALLED_FROM=\"$source_url\"" "$wrapper_tmp" \ + || die "failed to preserve update channel in downloaded wrapper" + bash -n "$wrapper_tmp" || die "updated wrapper failed validation after channel rewrite" + + chmod +x "$wrapper_tmp" + mv "$wrapper_tmp" "$target" + trap - EXIT + ok "updated lucebox → $target" } cmd_completion() { @@ -1445,7 +1509,7 @@ cmd_exec_in_container() { local argv=(docker exec "${tty[@]}") argv+=(--user "$(id -u):$(id -g)") argv+=(-w /opt/lucebox-hub) - argv+=(-e "HOME=$HOME") + argv+=(-e "HOME=$CONFIG_HOME") _append_host_env argv _append_scalar_env argv "$variant" # The image has no top-level `lucebox` binary on PATH — that name only @@ -1536,7 +1600,10 @@ Environment overrides: LUCEBOX_VARIANT image tag override (default: cuda12 on NVIDIA, rocm on AMD) LUCEBOX_PORT host port for the server (default: 8080) LUCEBOX_CONTAINER server container name (default: lucebox) - LUCEBOX_MODELS host model directory (default: \$XDG_DATA_HOME/lucebox/models + LUCEBOX_MODELS host model directory (default: \$XDG_DATA_HOME/lucebox/models) + LUCEBOX_HOME config/state directory (default: \$HOME/.lucebox) + LUCEBOX_WRAPPER_SHA256 + optional 64-hex checksum pin for install/update LUCEBOX_NO_EXEC=1 force docker-run for in-container subcommands even when the container is up (equivalent to --no-exec) HF_TOKEN propagated to \`models download\` for gated HF repos diff --git a/lucebox/README.md b/lucebox/README.md index 6af3d2a02..f5bb7d635 100644 --- a/lucebox/README.md +++ b/lucebox/README.md @@ -1,21 +1,22 @@ -# lucebox — host CLI for the lucebox-hub container +# lucebox — CLI for the Lucebox inference appliance -This package ships *inside* the `ghcr.io/luce-org/lucebox-hub` Docker image -and is invoked from the host via the [`lucebox.sh`](../lucebox.sh) wrapper: +This Python package ships inside the `ghcr.io/luce-org/lucebox-hub` image. Most +users do not install it directly: they install the small +[`lucebox` host wrapper](https://github.com/Luce-Org/lucebox/blob/main/lucebox.sh), +which invokes the package in the appropriate container: - lucebox.sh check # `docker run … lucebox check` - lucebox.sh config get - lucebox.sh print-run + lucebox check + lucebox models list + lucebox models download qwen3.6-27b --activate + lucebox start -The wrapper is the only thing that runs on the host. It selects CUDA for a -working NVIDIA GPU (including RTX 3090 + Strix builds) and ROCm for AMD builds -(including R9700 + Strix), then applies the matching Docker device flags. -Everything else (host -checks, TOML config, docker daemon calls, model download) is Python in the -container. Host facts (driver/runtime, GPU, RAM, VRAM, systemd availability) are -passed in via `LUCEBOX_HOST_*` environment variables so the Python side -doesn't reprobe. The autotune sweep, profiling, and agent-client launchers -land in follow-up PRs. +The wrapper detects the host and selects CUDA for NVIDIA builds (including RTX +3090 + Strix) or ROCm for AMD builds (including R9700 + Strix). The package then +handles readiness checks, TOML configuration, model selection and download, +optimization settings, and construction of the final server command. Host +facts are passed through `LUCEBOX_HOST_*` environment variables so the +container never has to guess the host configuration. -Subcommands are defined in [`lucebox/cli.py`](src/lucebox/cli.py). See the -top-level [README.md](../README.md) for the user-facing flow. +See the [project README](https://github.com/Luce-Org/lucebox#readme) for the +installation and user flow. Contributors can find the CLI implementation in +[`src/lucebox`](https://github.com/Luce-Org/lucebox/tree/main/lucebox/src/lucebox). diff --git a/lucebox/pyproject.toml b/lucebox/pyproject.toml index 5277b268d..bb12d90d7 100644 --- a/lucebox/pyproject.toml +++ b/lucebox/pyproject.toml @@ -7,7 +7,7 @@ name = "lucebox" dynamic = ["version"] description = "Host-side CLI for the lucebox-hub container: launch, config, model download" readme = "README.md" -requires-python = ">=3.11" +requires-python = ">=3.12,<3.13" authors = [{ name = "Lucebox" }] license = { text = "Apache-2.0" } @@ -30,8 +30,13 @@ dependencies = [ # `uv pip install luce-bench` on the host running the scorer. ] +[project.urls] +Homepage = "https://github.com/Luce-Org/lucebox" +Repository = "https://github.com/Luce-Org/lucebox" +Issues = "https://github.com/Luce-Org/lucebox/issues" + [project.scripts] -lucebox = "lucebox.cli:app" +lucebox = "lucebox.cli:main" [build-system] requires = ["hatchling", "hatch-vcs"] @@ -42,7 +47,7 @@ source = "vcs" # Untagged checkouts (e.g. fresh clone before tagging lucebox-v0.2.1) # resolve to this rather than 0.0.0.dev0. fallback-version = "0.2.1.dev0" -raw-options.tag_regex = '''^lucebox-v(?P\d+\.\d+\.\d+)$''' +tag-pattern = '''^lucebox-v(?P\d+\.\d+\.\d+)$''' [tool.hatch.build.hooks.vcs] # Build hook writes the resolved version into src/lucebox/_version.py diff --git a/lucebox/src/lucebox/__main__.py b/lucebox/src/lucebox/__main__.py index 128e2ca87..5c9d49ba1 100644 --- a/lucebox/src/lucebox/__main__.py +++ b/lucebox/src/lucebox/__main__.py @@ -1,6 +1,6 @@ -"""Entry point for `python -m lucebox`.""" +"""Entry point for ``python -m lucebox``.""" -from lucebox.cli import app +from lucebox.cli import main if __name__ == "__main__": - app() + main() diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index 9c0e1809c..622e6de5a 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -14,10 +14,8 @@ from __future__ import annotations -import os import sys from dataclasses import replace -from pathlib import Path from typing import Annotated import typer @@ -32,57 +30,64 @@ from lucebox import __version__ from lucebox.config import config_get, config_set, config_unset, live_config from lucebox.host_facts import from_env +from lucebox.types import Config app = typer.Typer( name="lucebox", help="Host CLI for the lucebox-hub container. Invoked by lucebox.sh.", no_args_is_help=True, + invoke_without_command=True, add_completion=False, ) console = Console() +error_console = Console(stderr=True) + + +@app.callback() +def root_options( + version_flag: Annotated[ + bool, + typer.Option("--version", help="Print lucebox version and exit.", is_eager=True), + ] = False, +) -> None: + """Apply options shared by the top-level CLI.""" + if version_flag: + print(__version__) + raise typer.Exit() # ── helpers ──────────────────────────────────────────────────────────────── -def _load_or_build() -> config_mod.Config: # type: ignore[name-defined] +def _load_or_build() -> Config: """env > config.toml > dataclass defaults — the canonical precedence. - Without the env-overlay step below, `config_mod.load()` returned the - persisted config verbatim and `LUCEBOX_IMAGE` / `LUCEBOX_VARIANT` / - `LUCEBOX_PORT` / `LUCEBOX_CONTAINER` / `LUCEBOX_MODELS` from the - systemd unit's `Environment=` (or any one-shot shell export) were - silently dropped. That contradicted the precedence lucebox.sh - documents and applies — and bit sindri when its config.toml had - `[image]` without `registry`, so the dataclass default - `ghcr.io/luce-org/lucebox-hub` won over the unit's - `LUCEBOX_IMAGE=ghcr.io/easel/lucebox-hub`. - - Fix: overlay env on top of the loaded config (or the live_config - fallback when config.toml is absent). Only the five top-level - scalars have env hooks — dflash/host/model don't, by design. + Only the five documented top-level scalars have environment overrides; + dflash, host, and model settings intentionally remain config-driven. """ - cfg = config_mod.load() - if cfg is None: - cfg = live_config() - # Overlay live host facts. When ``config.toml`` exists without a - # ``[host]`` block (the common case — operators don't hand-edit - # host facts), ``cfg.host`` defaults to a zero-filled ``HostFacts`` - # and the DFLASH_* serve heuristic silently falls through to the - # "no VRAM signal" path. Re-probe from env so the wrapper-exported - # LUCEBOX_HOST_* facts always win over the persisted (possibly - # absent) snapshot. - live_host = from_env() - host = live_host if live_host.vram_gb > 0 or live_host.nproc > 0 else cfg.host - return replace( - cfg, - variant=os.environ.get("LUCEBOX_VARIANT", cfg.variant), - image=os.environ.get("LUCEBOX_IMAGE", cfg.image), - container_name=os.environ.get("LUCEBOX_CONTAINER", cfg.container_name), - port=int(os.environ.get("LUCEBOX_PORT", str(cfg.port))), - models_dir=Path(os.environ.get("LUCEBOX_MODELS", str(cfg.models_dir))), - host=host, - ) + try: + cfg = config_mod.load() + if cfg is None: + return live_config() + # Host facts exported by the wrapper take precedence over an absent or + # stale persisted snapshot. A zero-filled environment means the CLI was + # invoked directly, so retain any snapshot already in the config. + live_host = from_env() + host = live_host if live_host.vram_gb > 0 or live_host.nproc > 0 else cfg.host + return config_mod.overlay_env(replace(cfg, host=host)) + except (OSError, ValueError) as exc: + error_console.print(f"[red]Invalid configuration:[/red] {escape(str(exc))}") + raise typer.Exit(code=2) from exc + + +def _server_spec() -> docker_run.DockerRunSpec: + """Build the server command and turn user-path errors into clean CLI output.""" + cfg = _load_or_build() + try: + return docker_run.server_run_spec(cfg) + except (OSError, RuntimeError, ValueError) as exc: + error_console.print(f"[red]Cannot build server command:[/red] {escape(str(exc))}") + raise typer.Exit(code=2) from exc # ── subcommands ──────────────────────────────────────────────────────────── @@ -112,9 +117,7 @@ def pull() -> None: @app.command("print-run") def print_run() -> None: """Print the docker-run command for the server (copy-pasteable).""" - cfg = _load_or_build() - spec = docker_run.server_run_spec(cfg) - print(spec.printable()) + print(_server_spec().printable()) @app.command("print-serve-argv") @@ -125,9 +128,7 @@ def print_serve_argv() -> None: a separate command from `print-run` so the bash side has a guaranteed machine-readable contract that's independent of the pretty formatter. """ - cfg = _load_or_build() - spec = docker_run.server_run_spec(cfg) - for tok in spec.argv(): + for tok in _server_spec().argv(): print(tok) @@ -145,8 +146,8 @@ def config_get_cmd( """Print a single key (or every reachable key) with its origin annotation.""" try: entries = config_get(key or None) - except KeyError as exc: - console.print(f"[red]{escape(str(exc))}[/red]") + except (KeyError, OSError, ValueError) as exc: + error_console.print(f"[red]{escape(str(exc))}[/red]") raise typer.Exit(code=2) from exc for k, (value, origin) in entries.items(): console.print(f"{k} = {escape(repr(value))} ([dim]from {origin}[/dim])") @@ -170,8 +171,8 @@ def config_set_cmd( value = value.strip() try: config_set(key, value) - except (KeyError, ValueError) as exc: - console.print(f"[red]{escape(str(exc))}[/red]") + except (KeyError, OSError, ValueError) as exc: + error_console.print(f"[red]{escape(str(exc))}[/red]") raise typer.Exit(code=2) from exc console.print(f"[green]Set[/green] {escape(key)} = {escape(value)}") @@ -183,8 +184,8 @@ def config_unset_cmd( """Remove a key from config.toml. Next read uses the live default.""" try: changed = config_unset(key) - except KeyError as exc: - console.print(f"[red]{escape(str(exc))}[/red]") + except (KeyError, OSError, ValueError) as exc: + error_console.print(f"[red]{escape(str(exc))}[/red]") raise typer.Exit(code=2) from exc if changed: console.print(f"[green]Unset[/green] {escape(key)}") diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py index 7ca5d2445..75a601afb 100644 --- a/lucebox/src/lucebox/config.py +++ b/lucebox/src/lucebox/config.py @@ -9,7 +9,7 @@ The dotted-key surface area is small and flat: model.preset, model.target_file, model.draft_file port, models_dir, variant, image, container_name - dflash. for each of the 11 DflashRuntime knobs + think_max + dflash. for every registered DflashRuntime knob Load resolves the TOML file → ``Config`` object, with anything absent filled from ``Config()`` defaults. Save writes back only the keys that @@ -25,7 +25,7 @@ from collections.abc import Callable from dataclasses import asdict, replace from datetime import UTC -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import Any, Literal, cast import tomli_w @@ -44,8 +44,8 @@ def default_config_path() -> Path: """Where .lucebox/config.toml lives. Convention: under $LUCEBOX_HOME if set, otherwise $HOME/.lucebox. Lives in - the bind-mounted host home dir so the config survives container teardown - and is editable from the host. + an explicitly bind-mounted application directory so the config survives + container teardown without exposing the rest of the host home directory. """ base = os.environ.get("LUCEBOX_HOME") if base: @@ -85,6 +85,50 @@ def _cast_bool(v: Any) -> bool: raise ValueError(f"cannot parse boolean: {v!r}") +def _cast_port(v: Any) -> int: + value = int(v) + if not 1 <= value <= 65535: + raise ValueError(f"port must be in the interval [1, 65535], got {value!r}") + return value + + +def _cast_models_dir(v: Any) -> str: + value = str(v) + if not Path(value).is_absolute(): + raise ValueError(f"models_dir must be an absolute path, got {value!r}") + return value + + +def _cast_model_relative_path(v: Any) -> str: + value = str(v) + if not value: + return value + path = PurePosixPath(value) + if path.is_absolute() or ".." in path.parts or value == ".": + raise ValueError(f"model file must be below models_dir, got {value!r}") + return value + + +def _cast_prefill_keep_ratio(v: Any) -> float: + value = float(v) + if not 0.0 < value <= 1.0: + raise ValueError( + "prefill_keep_ratio must be in the interval (0.0, 1.0], " + f"got {value!r}" + ) + return value + + +def _cast_think_soft_close_min_ratio(v: Any) -> float: + value = float(v) + if not 0.0 <= value <= 1.0: + raise ValueError( + "think_soft_close_min_ratio must be in the interval [0.0, 1.0], " + f"got {value!r}" + ) + return value + + # Each entry: dotted-key → (toml_path, type_caster, default_getter). # ``toml_path`` is the (section, field) pair on disk; ``"_root"`` means the # key lives at the top level (no [section]). ``default_getter`` returns the @@ -93,11 +137,11 @@ def _cast_bool(v: Any) -> bool: "variant": (("image", "variant"), str), "image": (("image", "registry"), str), "container_name": (("runtime", "container_name"), str), - "port": (("runtime", "port"), int), - "models_dir": (("paths", "models"), str), + "port": (("runtime", "port"), _cast_port), + "models_dir": (("paths", "models"), _cast_models_dir), "model.preset": (("model", "preset"), str), - "model.target_file": (("model", "target_file"), str), - "model.draft_file": (("model", "draft_file"), str), + "model.target_file": (("model", "target_file"), _cast_model_relative_path), + "model.draft_file": (("model", "draft_file"), _cast_model_relative_path), "dflash.budget": (("dflash", "budget"), int), "dflash.max_ctx": (("dflash", "max_ctx"), int), "dflash.lazy": (("dflash", "lazy"), _cast_bool), @@ -106,15 +150,22 @@ def _cast_bool(v: Any) -> bool: "dflash.cache_type_k": (("dflash", "cache_type_k"), str), "dflash.cache_type_v": (("dflash", "cache_type_v"), str), "dflash.prefill_mode": (("dflash", "prefill_mode"), _cast_prefill_mode), - "dflash.prefill_keep_ratio": (("dflash", "prefill_keep_ratio"), float), + "dflash.prefill_keep_ratio": ( + ("dflash", "prefill_keep_ratio"), + _cast_prefill_keep_ratio, + ), "dflash.prefill_threshold": (("dflash", "prefill_threshold"), int), "dflash.prefill_drafter": (("dflash", "prefill_drafter"), str), "dflash.think_max": (("dflash", "think_max"), int), "dflash.fa_window": (("dflash", "fa_window"), int), "dflash.think_soft_close_min_ratio": ( - ("dflash", "think_soft_close_min_ratio"), float), + ("dflash", "think_soft_close_min_ratio"), + _cast_think_soft_close_min_ratio, + ), "dflash.debug_thinking_logits": ( - ("dflash", "debug_thinking_logits"), _cast_bool), + ("dflash", "debug_thinking_logits"), + _cast_bool, + ), } @@ -226,11 +277,11 @@ def _from_dict(raw: dict[str, Any]) -> Config: registry = img.get("registry", "ghcr.io/luce-org/lucebox-hub") runtime = raw.get("runtime", {}) - port = int(runtime.get("port", 8080)) + port = _cast_port(runtime.get("port", 8080)) container_name = str(runtime.get("container_name", "lucebox")) paths = raw.get("paths", {}) - models_dir = Path(paths.get("models", str(default_models_dir()))) + models_dir = Path(_cast_models_dir(paths.get("models", str(default_models_dir())))) df = raw.get("dflash", {}) dflash = DflashRuntime( @@ -242,13 +293,16 @@ def _from_dict(raw: dict[str, Any]) -> Config: cache_type_k=str(df.get("cache_type_k", "")), cache_type_v=str(df.get("cache_type_v", "")), prefill_mode=_cast_prefill_mode(df.get("prefill_mode", "off")), - prefill_keep_ratio=float(df.get("prefill_keep_ratio", 0.05)), + prefill_keep_ratio=_cast_prefill_keep_ratio( + df.get("prefill_keep_ratio", 0.05) + ), prefill_threshold=int(df.get("prefill_threshold", 32000)), prefill_drafter=str(df.get("prefill_drafter", "")), think_max=int(df.get("think_max", 15488)), fa_window=int(df.get("fa_window", 0)), - think_soft_close_min_ratio=float( - df.get("think_soft_close_min_ratio", 0.0)), + think_soft_close_min_ratio=_cast_think_soft_close_min_ratio( + df.get("think_soft_close_min_ratio", 0.0) + ), debug_thinking_logits=_cast_bool(df.get("debug_thinking_logits", False)), ) @@ -281,8 +335,8 @@ def _from_dict(raw: dict[str, Any]) -> Config: # them from the registry so users only have to write one key. mdl = raw.get("model", {}) preset_name = str(mdl.get("preset", "")) - target_file = str(mdl.get("target_file", "")) - draft_file = str(mdl.get("draft_file", "")) + target_file = _cast_model_relative_path(mdl.get("target_file", "")) + draft_file = _cast_model_relative_path(mdl.get("draft_file", "")) if preset_name and (not target_file or not draft_file): from lucebox.download import PRESETS @@ -466,6 +520,24 @@ def config_get(key: str | None = None, *, path: Path | None = None) -> dict[str, return out +def overlay_env(cfg: Config) -> Config: + """Apply supported process overrides to an existing config. + + Keeping this in one helper makes the documented ``env > TOML > default`` + precedence identical for both persisted and first-run configurations. + """ + return replace( + cfg, + variant=os.environ.get("LUCEBOX_VARIANT", cfg.variant), + image=os.environ.get("LUCEBOX_IMAGE", cfg.image), + container_name=os.environ.get("LUCEBOX_CONTAINER", cfg.container_name), + port=_cast_port(os.environ.get("LUCEBOX_PORT", str(cfg.port))), + models_dir=Path( + _cast_models_dir(os.environ.get("LUCEBOX_MODELS", str(cfg.models_dir))) + ), + ) + + def live_config() -> Config: """Build a fresh Config from current host facts + the DFLASH_* heuristic. @@ -481,13 +553,10 @@ def live_config() -> Config: host = from_env() default = Config() default_variant = "rocm" if host.gpu_vendor == "amd" else "cuda12" - return replace( + cfg = replace( default, - variant=os.environ.get("LUCEBOX_VARIANT", default_variant), - image=os.environ.get("LUCEBOX_IMAGE", default.image), - container_name=os.environ.get("LUCEBOX_CONTAINER", default.container_name), - port=int(os.environ.get("LUCEBOX_PORT", str(default.port))), - models_dir=Path(os.environ.get("LUCEBOX_MODELS", str(default.models_dir))), + variant=default_variant, dflash=autotune_mod.runtime_from_host(host), host=host, ) + return overlay_env(cfg) diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index ab34c837a..f9b618007 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -13,10 +13,42 @@ import shlex import subprocess from dataclasses import dataclass -from pathlib import Path +from pathlib import Path, PurePosixPath from lucebox.types import Config, GpuVendor +_CONTAINER_MODELS = "/opt/lucebox-hub/server/models" +_CONTAINER_RESOLVED_MODELS = "/opt/lucebox-resolved" + + +@dataclass(frozen=True, slots=True) +class BindMount: + """One explicit Docker bind mount. + + ``read_only`` is part of the type so sensitive compatibility mounts cannot + accidentally become writable during argv rendering. + """ + + source: str + target: str + read_only: bool = False + + def __post_init__(self) -> None: + if not Path(self.source).is_absolute(): + raise ValueError(f"bind-mount source must be absolute, got {self.source!r}") + if not PurePosixPath(self.target).is_absolute(): + raise ValueError(f"bind-mount target must be absolute, got {self.target!r}") + for label, value in (("source", self.source), ("target", self.target)): + if "," in value or any(char in value for char in "\n\r\0"): + raise ValueError( + f"bind-mount {label} contains a character unsupported by " + f"Docker --mount: {value!r}" + ) + + def argument(self) -> str: + option = f"type=bind,source={self.source},target={self.target}" + return f"{option},readonly" if self.read_only else option + def _host_facts_env() -> list[tuple[str, str]]: """Forward LUCEBOX_HOST_* from the orchestrator's env into the server. @@ -69,27 +101,68 @@ def _resolve_model_files(cfg: Config) -> tuple[str, str, str]: draft = pres.draft_file if not draft and pres.speculator_dir: spec_path = cfg.models_dir / "draft" / pres.speculator_dir - if spec_path.is_dir(): + if spec_path.is_dir() or spec_path.is_symlink(): draft_dir = pres.speculator_dir return target, draft, draft_dir -def _runtime_volumes(cfg: Config) -> tuple[tuple[str, str], ...]: - """Mount models, $HOME, and any config dir outside the home mount.""" - home_path = Path.home() - home = str(home_path) - models = str(cfg.models_dir) - volumes = [(models, "/opt/lucebox-hub/server/models")] - if home != models: - volumes.append((home, home)) +def _runtime_volumes(cfg: Config) -> tuple[BindMount, ...]: + """Mount only the writable application data needed by the server.""" + models = str(cfg.models_dir.absolute()) config_home = Path( - os.environ.get("LUCEBOX_HOME") or home_path / ".lucebox" + os.environ.get("LUCEBOX_HOME") or Path.home() / ".lucebox" ).absolute() - # The same-path home mount normally covers config. If models_dir == HOME, - # that path is mounted at /opt/... instead, so add config explicitly too. - if home == models or not config_home.is_relative_to(home_path): - volumes.append((str(config_home), str(config_home))) - return tuple(volumes) + return ( + BindMount(models, _CONTAINER_MODELS), + BindMount(str(config_home), str(config_home)), + ) + + +def _validate_model_relative_path(value: str, field: str) -> PurePosixPath: + """Validate a config-provided path intended to live below models_dir.""" + relative = PurePosixPath(value) + if relative.is_absolute() or ".." in relative.parts or value in {"", "."}: + raise ValueError(f"{field} must be a path below models_dir, got {value!r}") + return relative + + +def _selected_model_path( + cfg: Config, + value: str, + *, + field: str, + role: str, + under_draft: bool, + directory: bool, + mounts: list[BindMount], +) -> str: + """Return the container path for one explicitly selected model artifact. + + Normal files are already covered by the models-directory bind mount. For + a symlink, bind only its resolved file (or the selected directory for a + speculator) into a dedicated read-only location. This preserves symlinked + model workflows without exposing the user's home or adjacent model files. + """ + relative = _validate_model_relative_path(value, field) + base = cfg.models_dir / "draft" if under_draft else cfg.models_dir + host_path = base.joinpath(*relative.parts) + container_base = PurePosixPath(_CONTAINER_MODELS) + if under_draft: + container_base /= "draft" + canonical = str(container_base.joinpath(*relative.parts)) + lexical = host_path.absolute() + resolved = host_path.resolve(strict=False) + if resolved == lexical: + return canonical + + mount_target = f"{_CONTAINER_RESOLVED_MODELS}/{role}" + if directory: + mounts.append(BindMount(str(resolved), mount_target, read_only=True)) + return mount_target + + container_path = f"{mount_target}/{resolved.name}" + mounts.append(BindMount(str(resolved), container_path, read_only=True)) + return container_path @dataclass(frozen=True, slots=True) @@ -103,7 +176,7 @@ class DockerRunSpec: detach: bool = False remove: bool = True port_publish: tuple[int, int] | None = None # (host, container) - volumes: tuple[tuple[str, str], ...] = () + volumes: tuple[BindMount, ...] = () env: tuple[tuple[str, str], ...] = () entrypoint_args: tuple[str, ...] = () extra: tuple[str, ...] = () @@ -133,8 +206,8 @@ def argv(self) -> list[str]: if self.port_publish is not None: host, container = self.port_publish out += ["-p", f"{host}:{container}"] - for host_path, container_path in self.volumes: - out += ["-v", f"{host_path}:{container_path}"] + for mount in self.volumes: + out += ["--mount", mount.argument()] for k, v in self.env: out += ["-e", f"{k}={v}"] out += list(self.extra) @@ -161,6 +234,7 @@ def printable(self) -> str: "--gpus", "--env", "--volume", + "--mount", "--publish", "--entrypoint", "--device", @@ -204,12 +278,40 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: # draft_dir is a subdirectory of models/draft/ holding a safetensors speculator; # it takes effect only when draft_file is empty and the directory exists on disk. target_file, draft_file, draft_dir = _resolve_model_files(cfg) + volumes = list(_runtime_volumes(cfg)) if target_file: - env.append(("DFLASH_TARGET", f"/opt/lucebox-hub/server/models/{target_file}")) + target_path = _selected_model_path( + cfg, + target_file, + field="model.target_file", + role="target", + under_draft=False, + directory=False, + mounts=volumes, + ) + env.append(("DFLASH_TARGET", target_path)) if draft_file: - env.append(("DFLASH_DRAFT", f"/opt/lucebox-hub/server/models/draft/{draft_file}")) + draft_path = _selected_model_path( + cfg, + draft_file, + field="model.draft_file", + role="draft", + under_draft=True, + directory=False, + mounts=volumes, + ) + env.append(("DFLASH_DRAFT", draft_path)) elif draft_dir: - env.append(("DFLASH_DRAFT", f"/opt/lucebox-hub/server/models/draft/{draft_dir}")) + draft_path = _selected_model_path( + cfg, + draft_dir, + field="model.speculator_dir", + role="draft-dir", + under_draft=True, + directory=True, + mounts=volumes, + ) + env.append(("DFLASH_DRAFT", draft_path)) elif cfg.model.preset: # An active target-only preset is an explicit choice, not permission # for entrypoint.sh to scan models/draft and attach an unrelated stale @@ -267,7 +369,7 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: remove=True, detach=False, port_publish=(cfg.port, 8080), - volumes=_runtime_volumes(cfg), + volumes=tuple(volumes), env=tuple(env), ) diff --git a/lucebox/src/lucebox/download.py b/lucebox/src/lucebox/download.py index df1d5c96f..43aeb83ca 100644 --- a/lucebox/src/lucebox/download.py +++ b/lucebox/src/lucebox/download.py @@ -28,17 +28,8 @@ import time from dataclasses import dataclass from pathlib import Path +from typing import TYPE_CHECKING -# hf-xet (huggingface_hub ≥ 1.16) streams the entire file in one final -# burst — the polling-based progress bar sits at 0% for ~14 minutes -# then snaps to 100% on a 17 GB GGUF. Force the chunked Python -# downloader instead so bytes grow continuously and the Rich bar tracks -# reality. Set before importing hf_hub_download so the import picks -# the env up. `setdefault` lets a user override on the command line. -os.environ.setdefault("HF_HUB_DISABLE_XET", "1") - -from huggingface_hub import HfApi, hf_hub_download # noqa: E402 -from huggingface_hub._local_folder import get_local_download_paths # noqa: E402 from rich.console import Console from rich.progress import ( BarColumn, @@ -51,6 +42,27 @@ from lucebox.types import Config, HostFacts +if TYPE_CHECKING: + from huggingface_hub import HfApi + + +def _configure_huggingface_download() -> None: + """Select the progress-friendly downloader before importing HF Hub. + + hf-xet streams a multi-GB file in one final burst, so the polling-based + progress bar remains at zero until completion. Keep this environment + change local to download/status operations rather than mutating every + ``lucebox`` process merely because the CLI module was imported. + """ + os.environ.setdefault("HF_HUB_DISABLE_XET", "1") + + +def _new_hf_api() -> HfApi: + _configure_huggingface_download() + from huggingface_hub import HfApi + + return HfApi() + @dataclass(frozen=True, slots=True) class ModelPreset: @@ -244,6 +256,9 @@ def _incomplete_path_candidates(local_dir: Path, filename: str, etag: str | None when ``local_dir`` is set hf-hub always uses the local staging dir, so the two candidates above cover every code path we hit. """ + _configure_huggingface_download() + from huggingface_hub._local_folder import get_local_download_paths + paths = get_local_download_paths(local_dir, filename) candidates: list[Path] = [] if etag: @@ -302,6 +317,9 @@ def _download_with_progress( actual hf-xet staging path (a hashed filename under ``.cache/huggingface/download/``), not a guess. """ + _configure_huggingface_download() + from huggingface_hub import hf_hub_download + local_dir.mkdir(parents=True, exist_ok=True) target = local_dir / filename candidates = _incomplete_path_candidates(local_dir, filename, etag) @@ -384,7 +402,7 @@ def download_preset(cfg: Config, preset: ModelPreset | None = None) -> int: """ preset = preset or DEFAULT_PRESET console = Console() - api = HfApi() + api = _new_hf_api() models = cfg.models_dir models.mkdir(parents=True, exist_ok=True) draft = models / "draft" @@ -437,7 +455,7 @@ def installed_status(cfg: Config, preset: ModelPreset) -> str: def installed_size_gb(cfg: Config, preset: ModelPreset) -> float: - """Sum of on-disk byte sizes for the preset's files, in GB (binary 1e9).""" + """Sum of on-disk byte sizes for the preset's files, in decimal GB (1e9).""" total = 0 target = _local_target_path(cfg, preset) if target.exists(): @@ -478,7 +496,7 @@ def status(cfg: Config, preset: ModelPreset | None = None) -> dict[str, bool]: not a draft exists. """ preset = preset or DEFAULT_PRESET - api = HfApi() + api = _new_hf_api() out: dict[str, bool] = {} try: size, _ = _file_meta(api, preset.target_repo, preset.target_file) diff --git a/lucebox/src/lucebox/types.py b/lucebox/src/lucebox/types.py index e9e2fa3ec..d755083c0 100644 --- a/lucebox/src/lucebox/types.py +++ b/lucebox/src/lucebox/types.py @@ -60,6 +60,24 @@ class HostFacts: docker_version: str = "" ctk: CtkStatus = "none" + def __post_init__(self) -> None: + """Reject malformed persisted host snapshots. + + ``Literal`` annotations help static type-checkers, but TOML is runtime + input and can still contain arbitrary strings. Failing here keeps an + invalid snapshot from silently selecting the wrong accelerator path. + """ + if self.gpu_vendor not in {"nvidia", "amd", "none"}: + raise ValueError( + "gpu_vendor must be nvidia, amd, or none; " + f"got {self.gpu_vendor!r}" + ) + if self.ctk not in {"runtime", "cdi", "installed-unwired", "none"}: + raise ValueError( + "ctk must be runtime, cdi, installed-unwired, or none; " + f"got {self.ctk!r}" + ) + @dataclass(frozen=True, slots=True) class DflashRuntime: @@ -92,7 +110,7 @@ class DflashRuntime: # attention (server default). On gemma4's hybrid iSWA the full-attn # layers grow KV linearly with max_ctx; a sparse fa_window keeps # decode compute bounded on long prompts without changing the KV - # footprint. Q: passed through to the server's `--fa-window ` + # footprint. Passed through to the server's `--fa-window ` # flag (see server/src/server/server_main.cpp). fa_window: int = 0 # Soft-close thinking termination dial (PR #326 in lucebox-hub). @@ -112,6 +130,19 @@ class DflashRuntime: # in-flight requests); leave off in production. debug_thinking_logits: bool = False + def __post_init__(self) -> None: + """Validate the bounded tuning knobs before they reach the server.""" + if not 0.0 < self.prefill_keep_ratio <= 1.0: + raise ValueError( + "prefill_keep_ratio must be in the interval (0.0, 1.0]; " + f"got {self.prefill_keep_ratio!r}" + ) + if not 0.0 <= self.think_soft_close_min_ratio <= 1.0: + raise ValueError( + "think_soft_close_min_ratio must be in the interval [0.0, 1.0]; " + f"got {self.think_soft_close_min_ratio!r}" + ) + @dataclass(frozen=True, slots=True) class ModelMeta: diff --git a/lucebox/tests/test_cli.py b/lucebox/tests/test_cli.py index f7628e8b7..ee254d9bd 100644 --- a/lucebox/tests/test_cli.py +++ b/lucebox/tests/test_cli.py @@ -5,7 +5,7 @@ import os import pytest -from lucebox.cli import app +from lucebox.cli import __version__, app from typer.testing import CliRunner @@ -57,6 +57,14 @@ def test_core_verbs_present_in_app() -> None: assert verb in registered +@pytest.mark.parametrize("args", [["version"], ["--version"]]) +def test_version_command_and_option_match(args: list[str]) -> None: + result = CliRunner().invoke(app, args) + + assert result.exit_code == 0 + assert result.stdout.strip() == __version__ + + def test_legacy_subcommands_are_removed() -> None: """`configure` and `download-models` were folded into config/models.""" cfg = CliRunner().invoke(app, ["configure", "--help"]) diff --git a/lucebox/tests/test_config.py b/lucebox/tests/test_config.py index afedc7d91..a61f411f6 100644 --- a/lucebox/tests/test_config.py +++ b/lucebox/tests/test_config.py @@ -57,6 +57,56 @@ def test_toml_load_rejects_invalid_prefill_mode(tmp_path: Path) -> None: config._load_toml(path) +@pytest.mark.parametrize( + ("key", "value"), + [ + ("dflash.prefill_keep_ratio", "0"), + ("dflash.prefill_keep_ratio", "1.01"), + ("dflash.think_soft_close_min_ratio", "-0.01"), + ("dflash.think_soft_close_min_ratio", "1.01"), + ("dflash.think_soft_close_min_ratio", "nan"), + ], +) +def test_config_set_rejects_out_of_range_ratios( + tmp_path: Path, key: str, value: str +) -> None: + path = tmp_path / "config.toml" + + with pytest.raises(ValueError, match="interval"): + config_set(key, value, path=path) + + assert not path.exists() + + +def test_toml_load_rejects_out_of_range_ratios(tmp_path: Path) -> None: + path = tmp_path / "config.toml" + path.write_text( + "[dflash]\n" + "prefill_keep_ratio = 0.05\n" + "think_soft_close_min_ratio = 2.0\n" + ) + + with pytest.raises(ValueError, match="think_soft_close_min_ratio"): + config._load_toml(path) + + +@pytest.mark.parametrize( + "host_body", + [ + 'gpu_vendor = "intel"\n', + 'ctk = "maybe"\n', + ], +) +def test_toml_load_rejects_invalid_host_literals( + tmp_path: Path, host_body: str +) -> None: + path = tmp_path / "config.toml" + path.write_text(f"[host]\n{host_body}") + + with pytest.raises(ValueError): + config._load_toml(path) + + def test_model_preset_round_trips_through_set_and_load(tmp_path: Path) -> None: """Setting model.preset writes a sparse TOML doc that loads back correctly.""" path = tmp_path / "config.toml" @@ -171,6 +221,23 @@ def test_config_set_rejects_unknown_key(tmp_path: Path) -> None: config_set("not.a.key", 1, path=path) +@pytest.mark.parametrize( + ("key", "value", "message"), + [ + ("port", "0", "port"), + ("port", "65536", "port"), + ("models_dir", "relative/models", "absolute"), + ("model.target_file", "../secret.gguf", "below models_dir"), + ("model.draft_file", "/tmp/draft.gguf", "below models_dir"), + ], +) +def test_config_set_rejects_unsafe_runtime_values( + tmp_path: Path, key: str, value: str, message: str +) -> None: + with pytest.raises(ValueError, match=message): + config_set(key, value, path=tmp_path / "config.toml") + + def test_config_set_auto_creates_file(tmp_path: Path) -> None: """`config set` creates a missing config.toml on first write.""" path = tmp_path / "config.toml" diff --git a/lucebox/tests/test_config_cli.py b/lucebox/tests/test_config_cli.py index 6525ea319..8d36da760 100644 --- a/lucebox/tests/test_config_cli.py +++ b/lucebox/tests/test_config_cli.py @@ -138,3 +138,16 @@ def test_load_or_build_no_toml_env_overrides_defaults( monkeypatch.setenv("LUCEBOX_IMAGE", "ghcr.io/myfork/lucebox-hub") cfg = _load_or_build() assert cfg.image == "ghcr.io/myfork/lucebox-hub" + + +def test_print_run_reports_invalid_configuration_without_traceback( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _set_config_path(tmp_path, monkeypatch) + monkeypatch.setenv("LUCEBOX_PORT", "0") + + result = CliRunner().invoke(app, ["print-run"]) + + assert result.exit_code == 2 + assert "Invalid configuration" in result.output + assert "Traceback" not in result.output diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index f5a18ed59..9e693934f 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -10,6 +10,7 @@ from pathlib import Path +import pytest from lucebox.download import PRESETS from lucebox.types import Config, DflashRuntime, HostFacts, ModelMeta @@ -43,7 +44,9 @@ def test_argv_flags_and_ordering() -> None: detach=True, remove=False, port_publish=(8080, 8080), - volumes=(("/host/models", "/opt/lucebox-hub/server/models"),), + volumes=( + docker_run.BindMount("/host/models", "/opt/lucebox-hub/server/models"), + ), env=(("DFLASH_BUDGET", "22"),), entrypoint_args=("serve",), extra=("--shm-size", "1g"), @@ -53,8 +56,9 @@ def test_argv_flags_and_ordering() -> None: assert "-d" in argv # detach assert "--gpus" not in argv # gpus=False assert ["-p", "8080:8080"] == argv[argv.index("-p") : argv.index("-p") + 2] - assert ["-v", "/host/models:/opt/lucebox-hub/server/models"] == argv[ - argv.index("-v") : argv.index("-v") + 2 + mount_arg = "type=bind,source=/host/models,target=/opt/lucebox-hub/server/models" + assert ["--mount", mount_arg] == argv[ + argv.index("--mount") : argv.index("--mount") + 2 ] assert ["-e", "DFLASH_BUDGET=22"] == argv[argv.index("-e") : argv.index("-e") + 2] # extra flags precede the image; entrypoint_args follow it. @@ -63,6 +67,20 @@ def test_argv_flags_and_ordering() -> None: assert argv.index("--shm-size") < argv.index("img:tag") +@pytest.mark.parametrize( + ("source", "target"), + [ + ("relative", "/container"), + ("/host", "relative"), + ("/host,comma", "/container"), + ("/host", "/container\nnewline"), + ], +) +def test_bind_mount_rejects_ambiguous_paths(source: str, target: str) -> None: + with pytest.raises(ValueError, match="bind-mount"): + docker_run.BindMount(source, target) + + def test_argv_amd_uses_rocm_device_contract() -> None: spec = docker_run.DockerRunSpec(image="img:rocm", name="box", gpu_vendor="amd") argv = spec.argv() @@ -100,12 +118,21 @@ def test_printable_glues_value_taking_flags() -> None: # ── _runtime_volumes ───────────────────────────────────────────────────────── -def test_runtime_volumes_mounts_models_and_home(tmp_path: Path) -> None: +def test_runtime_volumes_mounts_only_models_and_config( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + home = tmp_path / "home" + monkeypatch.setattr(Path, "home", staticmethod(lambda: home)) + monkeypatch.delenv("LUCEBOX_HOME", raising=False) cfg = Config(models_dir=tmp_path / "models") vols = docker_run._runtime_volumes(cfg) - assert (str(tmp_path / "models"), "/opt/lucebox-hub/server/models") in vols - # $HOME is also mounted so absolute symlink targets resolve in-container. - assert any(host == str(Path.home()) for host, _ in vols) + assert docker_run.BindMount( + str(tmp_path / "models"), "/opt/lucebox-hub/server/models" + ) in vols + assert docker_run.BindMount( + str(home / ".lucebox"), str(home / ".lucebox") + ) in vols + assert all(mount.source != str(home) for mount in vols) def test_runtime_volumes_dedupes_when_models_is_home(monkeypatch, tmp_path: Path) -> None: @@ -113,10 +140,16 @@ def test_runtime_volumes_dedupes_when_models_is_home(monkeypatch, tmp_path: Path monkeypatch.delenv("LUCEBOX_HOME", raising=False) cfg = Config(models_dir=tmp_path) vols = docker_run._runtime_volumes(cfg) - # models_dir == home is mounted at /opt/..., not at its host path, so the - # nested config directory still needs an explicit same-path mount. + # models_dir is mounted at the image's canonical path, while config keeps + # its same-path mount. The parent home directory itself is never exposed. assert len(vols) == 2 - assert (str(tmp_path / ".lucebox"), str(tmp_path / ".lucebox")) in vols + assert docker_run.BindMount( + str(tmp_path / ".lucebox"), str(tmp_path / ".lucebox") + ) in vols + assert all( + mount.source != str(tmp_path) or mount.target != str(tmp_path) + for mount in vols + ) def test_runtime_volumes_mounts_custom_config_home( @@ -129,7 +162,7 @@ def test_runtime_volumes_mounts_custom_config_home( vols = docker_run._runtime_volumes(Config(models_dir=tmp_path / "models")) - assert (str(config_home), str(config_home)) in vols + assert docker_run.BindMount(str(config_home), str(config_home)) in vols def test_empty_lucebox_home_uses_default(monkeypatch, tmp_path: Path) -> None: @@ -190,7 +223,69 @@ def test_server_run_spec_top_level_shape(tmp_path: Path) -> None: assert spec.remove is True assert spec.detach is False assert spec.port_publish == (9000, 8080) - assert (str(tmp_path), "/opt/lucebox-hub/server/models") in spec.volumes + assert docker_run.BindMount( + str(tmp_path), "/opt/lucebox-hub/server/models" + ) in spec.volumes + + +def test_server_run_spec_mounts_selected_symlink_target_narrowly_read_only( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + home = tmp_path / "home" + models = tmp_path / "models" + external = home / "model-cache" + models.mkdir() + external.mkdir(parents=True) + target = external / "target.gguf" + target.write_bytes(b"model") + (models / "selected.gguf").symlink_to(target) + monkeypatch.setattr(Path, "home", staticmethod(lambda: home)) + + spec = docker_run.server_run_spec( + Config(models_dir=models, model=ModelMeta(target_file="selected.gguf")) + ) + + assert docker_run.BindMount( + str(target), + "/opt/lucebox-resolved/target/target.gguf", + read_only=True, + ) in spec.volumes + assert _env(spec)["DFLASH_TARGET"] == "/opt/lucebox-resolved/target/target.gguf" + assert all(mount.source != str(home) for mount in spec.volumes) + argv = spec.argv() + mount_args = [argv[i + 1] for i, token in enumerate(argv) if token == "--mount"] + assert any( + "target=/opt/lucebox-resolved/target/target.gguf,readonly" in argument + for argument in mount_args + ) + + +def test_server_run_spec_resolves_symlinked_model_parent(tmp_path: Path) -> None: + models = tmp_path / "models" + external = tmp_path / "external" + models.mkdir() + external.mkdir() + target = external / "nested.gguf" + target.write_bytes(b"model") + (models / "selected").symlink_to(external, target_is_directory=True) + + spec = docker_run.server_run_spec( + Config(models_dir=models, model=ModelMeta(target_file="selected/nested.gguf")) + ) + + assert docker_run.BindMount( + str(target), + "/opt/lucebox-resolved/target/nested.gguf", + read_only=True, + ) in spec.volumes + assert _env(spec)["DFLASH_TARGET"] == "/opt/lucebox-resolved/target/nested.gguf" + + +def test_server_run_spec_rejects_model_path_traversal(tmp_path: Path) -> None: + cfg = Config(models_dir=tmp_path, model=ModelMeta(target_file="../secret.gguf")) + + with pytest.raises(ValueError, match="below models_dir"): + docker_run.server_run_spec(cfg) def test_server_run_spec_rocm_uses_amd_devices_on_heterogeneous_host(tmp_path: Path) -> None: diff --git a/lucebox/tests/test_download.py b/lucebox/tests/test_download.py index 8b69e96b8..a4bf692e3 100644 --- a/lucebox/tests/test_download.py +++ b/lucebox/tests/test_download.py @@ -7,6 +7,8 @@ stays pinned without actually talking to the Hub. """ +import subprocess +import sys from pathlib import Path from types import SimpleNamespace @@ -23,6 +25,22 @@ from lucebox import download +def test_import_does_not_mutate_huggingface_environment() -> None: + code = ( + "import os; " + "os.environ.pop('HF_HUB_DISABLE_XET', None); " + "import lucebox.download; " + "assert 'HF_HUB_DISABLE_XET' not in os.environ" + ) + result = subprocess.run( + [sys.executable, "-c", code], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + def test_default_preset_uses_quantized_gguf_draft(): assert DEFAULT_PRESET.draft_repo == "spiritbuun/Qwen3.6-27B-DFlash-GGUF" assert DEFAULT_PRESET.draft_file == "dflash-draft-3.6-q4_k_m.gguf" diff --git a/lucebox/tests/test_models_cli.py b/lucebox/tests/test_models_cli.py index 744f40c29..dd4f244f5 100644 --- a/lucebox/tests/test_models_cli.py +++ b/lucebox/tests/test_models_cli.py @@ -44,7 +44,7 @@ def test_models_default_view_lists_only_installed( # No models on disk → default view says "no presets installed". result = CliRunner().invoke(app, ["models"]) assert result.exit_code == 0 - assert "No presets installed" in result.stdout or "Models dir" in result.stdout + assert "No presets installed" in result.stdout def test_models_download_recommends_when_empty( diff --git a/scripts/check_lucebox_wrapper_sandbox.sh b/scripts/check_lucebox_wrapper_sandbox.sh index df2b2b9bc..77f571c8b 100755 --- a/scripts/check_lucebox_wrapper_sandbox.sh +++ b/scripts/check_lucebox_wrapper_sandbox.sh @@ -95,8 +95,12 @@ run_logged_capture() { note "run: $* > $out" { printf '\n===== %s > %s =====\n' "$*" "$out" - "$@" - local rc=$? + local rc + if "$@"; then + rc=0 + else + rc=$? + fi printf '===== exit=%s =====\n' "$rc" return "$rc" } 2>&1 | tee "$out" | tee -a "$LOG" >/dev/null @@ -164,10 +168,9 @@ run_logged_capture "$ROOT/help.out" lucebox help assert_contains "$ROOT/help.out" "LUCEBOX_VARIANT" assert_contains "$ROOT/help.out" "LUCEBOX_IMAGE" -docker manifest inspect "${IMAGE}:${VARIANT}" >/dev/null -pass "image manifest exists: ${IMAGE}:${VARIANT}" - if [ "$RUN_PULL" = "1" ]; then + docker manifest inspect "${IMAGE}:${VARIANT}" >/dev/null + pass "image manifest exists: ${IMAGE}:${VARIANT}" run_logged_capture "$ROOT/pull.out" lucebox pull assert_contains "$ROOT/pull.out" "${IMAGE}:${VARIANT}" fi diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index c9fc0577d..d95513daf 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -26,6 +26,7 @@ ROOT="$(git rev-parse --show-toplevel 2>/dev/null || (cd "$(dirname "$0")/.." && SCRIPT="$ROOT/lucebox.sh" ENTRYPOINT="$ROOT/server/scripts/entrypoint.sh" INSTALLER="$ROOT/install.sh" +SELF_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/$(basename "${BASH_SOURCE[0]}")" if [ ! -f "$SCRIPT" ]; then echo "FAIL: lucebox.sh not found at $SCRIPT" >&2 @@ -46,6 +47,23 @@ exit 1 STUB chmod +x "$SUITE_SHIMS/$binname" done +if ! command -v timeout >/dev/null 2>&1; then + # macOS has no GNU `timeout`; keep the contributor test suite portable + # with a tiny subprocess wrapper. CI still uses the native coreutils tool. + cat > "$SUITE_SHIMS/timeout" <<'PYTHON_TIMEOUT' +#!/usr/bin/env python3 +import subprocess +import sys + +duration = float(sys.argv[1].removesuffix("s")) +try: + result = subprocess.run(sys.argv[2:], timeout=duration, check=False) +except subprocess.TimeoutExpired: + raise SystemExit(124) +raise SystemExit(result.returncode) +PYTHON_TIMEOUT + chmod +x "$SUITE_SHIMS/timeout" +fi export HOME="$SUITE_SANDBOX/home" export XDG_CONFIG_HOME="$SUITE_SANDBOX/xdg" export XDG_DATA_HOME="$SUITE_SANDBOX/data" @@ -142,10 +160,10 @@ SHELLCHECK_TARGETS=( ) # Add every scripts/*.sh except this one (don't recurse into our own tests). while IFS= read -r -d '' f; do - [ "$f" = "${BASH_SOURCE[0]}" ] && continue + [ "$f" = "$SELF_PATH" ] && continue SHELLCHECK_TARGETS+=("$f") done < <(find "$ROOT/scripts" -maxdepth 1 -name '*.sh' -type f -print0 2>/dev/null) -SHELLCHECK_TARGETS+=("${BASH_SOURCE[0]}") +SHELLCHECK_TARGETS+=("$SELF_PATH") if command -v shellcheck >/dev/null 2>&1; then sc_out=$(shellcheck --severity=error "${SHELLCHECK_TARGETS[@]}" 2>&1) || sc_rc=$? @@ -636,15 +654,35 @@ test_entrypoint_host_info_json "entrypoint HOST_INFO JSON shape (populated + unk # preserve the user's channel across upgrades. test_install_sh_bakes_source_url() { local label="$1" - local tmp dest_dir dest_path src_url out rc + local tmp dest_dir dest_path bad_dest src_url src_sha out rc tmp=$(mktemp -d -t lucebox-install.XXXXXX) # Use the real lucebox.sh as the "remote" file — `file://` works with # curl out of the box and exercises the same install.sh code path as # an https fetch would. src_url="file://$SCRIPT" + if command -v sha256sum >/dev/null 2>&1; then + src_sha=$(sha256sum "$SCRIPT") + else + src_sha=$(shasum -a 256 "$SCRIPT") + fi + src_sha="${src_sha%% *}" dest_dir="$tmp/bin" dest_path="$dest_dir/lucebox" + bad_dest="$dest_dir/bad-checksum" + + out=$(LUCEBOX_INSTALL_URL="$src_url" LUCEBOX_INSTALL_DEST="$bad_dest" \ + LUCEBOX_WRAPPER_SHA256="$(printf '0%.0s' {1..64})" \ + NO_COLOR=1 bash "$INSTALLER" 2>&1) && rc=0 || rc=$? + if [ "$rc" -eq 0 ] || [ -e "$bad_dest" ] \ + || ! grep -qF "wrapper checksum mismatch" <<<"$out"; then + rm -rf "$tmp" + report fail "$label" "bad wrapper checksum was not rejected" + return + fi + + rc=0 out=$(LUCEBOX_INSTALL_URL="$src_url" LUCEBOX_INSTALL_DEST="$dest_path" \ + LUCEBOX_WRAPPER_SHA256="$src_sha" \ NO_COLOR=1 bash "$INSTALLER" 2>&1) || rc=$? rc="${rc:-0}" if [ "$rc" -ne 0 ]; then @@ -662,15 +700,20 @@ test_install_sh_bakes_source_url() { report fail "$label" "LUCEBOX_INSTALLED_FROM not rewritten in installed copy" return fi + if ! grep -qF "wrapper sha256 verified" <<<"$out"; then + rm -rf "$tmp" + report fail "$label" "installer did not report checksum verification" + return + fi rm -rf "$tmp" report ok "$label" } test_install_sh_bakes_source_url "install.sh bakes LUCEBOX_INSTALLED_FROM into installed copy" # ── update dispatch ─────────────────────────────────────────────────────── -# `lucebox update` must dispatch to cmd_update — verify it's wired in the -# main case statement and appears in --help. We can't actually run the -# update (it'd curl + replace this very script) so the test is parse-level. +# `lucebox update` must dispatch to cmd_update — first verify it's wired in +# the main case statement and appears in --help, then exercise an isolated +# wrapper copy against a file:// channel below. test_update_subcommand_wired() { local label="$1" local out @@ -687,6 +730,65 @@ test_update_subcommand_wired() { } test_update_subcommand_wired "lucebox update subcommand is wired" +test_update_downloads_verifies_and_replaces_atomically() { + local label="$1" tmp installed source_url checksum out rc + tmp=$(mktemp -d -t lucebox-update.XXXXXX) + installed="$tmp/lucebox" + mkdir -p "$tmp/channel" + + # Bake a file:// channel into an isolated wrapper copy. The update must + # replace only that copy, never the repository script. + source_url="file://$tmp/channel/lucebox.sh" + sed "s|^LUCEBOX_INSTALLED_FROM=.*|LUCEBOX_INSTALLED_FROM=\"$source_url\"|" \ + "$SCRIPT" > "$installed" + chmod +x "$installed" + cat > "$tmp/channel/lucebox.sh" <<'WRAPPER' +#!/usr/bin/env bash +set -euo pipefail +VERSION="9.9.9" +LUCEBOX_INSTALLED_FROM="${LUCEBOX_INSTALLED_FROM:-https://example.invalid/lucebox.sh}" +printf 'updated %s from %s\n' "$VERSION" "$LUCEBOX_INSTALLED_FROM" +WRAPPER + + # A wrong explicit pin must leave the installed wrapper untouched. + out=$(LUCEBOX_WRAPPER_SHA256="$(printf '0%.0s' {1..64})" \ + NO_COLOR=1 bash "$installed" update 2>&1) && rc=0 || rc=$? + if [ "$rc" -eq 0 ] || ! grep -qF "wrapper checksum mismatch" <<<"$out" \ + || grep -qF 'VERSION="9.9.9"' "$installed"; then + rm -rf "$tmp" + report fail "$label" "bad checksum was not rejected: $(head -3 <<<"$out")" + return + fi + + if command -v sha256sum >/dev/null 2>&1; then + checksum=$(sha256sum "$tmp/channel/lucebox.sh") + else + checksum=$(shasum -a 256 "$tmp/channel/lucebox.sh") + fi + checksum="${checksum%% *}" + rc=0 + out=$(LUCEBOX_WRAPPER_SHA256="$checksum" \ + NO_COLOR=1 bash "$installed" update 2>&1) || rc=$? + rc="${rc:-0}" + if [ "$rc" -ne 0 ]; then + rm -rf "$tmp" + report fail "$label" "verified update exited $rc: $(head -3 <<<"$out")" + return + fi + if ! grep -qF 'VERSION="9.9.9"' "$installed" \ + || ! grep -Fqx "LUCEBOX_INSTALLED_FROM=\"$source_url\"" "$installed" \ + || ! grep -qF "wrapper sha256 verified" <<<"$out"; then + rm -rf "$tmp" + report fail "$label" "validated wrapper did not replace the isolated copy" + return + fi + + rm -rf "$tmp" + report ok "$label" +} +test_update_downloads_verifies_and_replaces_atomically \ + "lucebox update verifies and atomically replaces the wrapper" + # ── IMAGE_BASE derived from install source ──────────────────────────────── # Source lucebox.sh in a subshell with LUCEBOX_INSTALLED_FROM pointing at # various URLs, then check that IMAGE_BASE comes out right. Uses @@ -929,6 +1031,11 @@ case "\$1" in exit 0 ;; *) + if [ "\${DOCKER_FAKE_CONFIG_ERROR:-0}" = "1" ] \ + && [[ "\$*" == *"print-serve-argv"* ]]; then + echo "Invalid configuration: test fixture" >&2 + exit 2 + fi printf 'DOCKER_INVOKED' for a in "\$@"; do printf ' %q' "\$a"; done printf '\n' @@ -1070,14 +1177,39 @@ test_serve_fallback_forwards_config_env() { _run_wrapper_capture_docker "$sandbox" serve || true) rm -rf "$sandbox" "$config_home" if ! grep -qF "LUCEBOX_HOME=$config_home" <<<"$out" \ - || ! grep -qF "HOME=$sandbox" <<<"$out"; then + || ! grep -qF "HOME=$config_home" <<<"$out"; then report fail "$label" "fallback server omitted HOME/config env: $(tail -3 <<<"$out")" return fi + if grep -qF -- "-v $sandbox:$sandbox" <<<"$out"; then + report fail "$label" "fallback server exposed the full host HOME" + return + fi report ok "$label" } test_serve_fallback_forwards_config_env \ - "serve fallback forwards HOME + LUCEBOX_HOME" + "serve fallback isolates HOME and forwards LUCEBOX_HOME" + +test_serve_refuses_invalid_config_fallback() { + local label="$1" sandbox out + sandbox=$(mktemp -d -t lucebox-serve-invalid.XXXXXX) + _make_docker_shim "$sandbox" 0 + out=$(DOCKER_FAKE_CONFIG_ERROR=1 \ + _run_wrapper_capture_docker "$sandbox" serve || true) + rm -rf "$sandbox" + if ! grep -qF "Invalid configuration: test fixture" <<<"$out" \ + || ! grep -qF "refusing to ignore invalid Lucebox configuration" <<<"$out"; then + report fail "$label" "configuration error was not surfaced: $(tail -4 <<<"$out")" + return + fi + if grep -qF "using fallback" <<<"$out"; then + report fail "$label" "invalid configuration silently launched fallback defaults" + return + fi + report ok "$label" +} +test_serve_refuses_invalid_config_fallback \ + "serve refuses to replace invalid config with fallback defaults" test_run_route_preserves_tty() { local label="$1" sandbox out rc From f365f3d856a67bdf6ab9ef411f19a9e4fcf4d1e7 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:25:43 +0200 Subject: [PATCH 06/25] test(lucebox): validate speculator symlinks --- lucebox/src/lucebox/docker_run.py | 5 +++- lucebox/tests/test_docker_run.py | 44 +++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index f9b618007..c8c827f35 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -101,7 +101,10 @@ def _resolve_model_files(cfg: Config) -> tuple[str, str, str]: draft = pres.draft_file if not draft and pres.speculator_dir: spec_path = cfg.models_dir / "draft" / pres.speculator_dir - if spec_path.is_dir() or spec_path.is_symlink(): + # is_dir() follows valid directory symlinks. It also rejects + # dangling links and links to files, which cannot satisfy the + # speculator-directory contract. + if spec_path.is_dir(): draft_dir = pres.speculator_dir return target, draft, draft_dir diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index 9e693934f..257b369f8 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -205,6 +205,25 @@ def test_resolve_model_files_no_preset_no_override(tmp_path: Path) -> None: assert docker_run._resolve_model_files(cfg) == ("", "", "") +@pytest.mark.parametrize("invalid_target", ["file", "missing"]) +def test_resolve_model_files_ignores_invalid_speculator_symlink( + invalid_target: str, tmp_path: Path +) -> None: + draft_root = tmp_path / "draft" + draft_root.mkdir() + target = tmp_path / "external-speculator" + if invalid_target == "file": + target.write_bytes(b"not a directory") + (draft_root / "laguna-xs2-speculator").symlink_to( + target, target_is_directory=invalid_target == "missing" + ) + cfg = Config(models_dir=tmp_path, model=ModelMeta(preset="laguna-xs.2")) + + _, _, draft_dir = docker_run._resolve_model_files(cfg) + + assert draft_dir == "" + + # ── server_run_spec ────────────────────────────────────────────────────────── @@ -281,6 +300,31 @@ def test_server_run_spec_resolves_symlinked_model_parent(tmp_path: Path) -> None assert _env(spec)["DFLASH_TARGET"] == "/opt/lucebox-resolved/target/nested.gguf" +def test_server_run_spec_mounts_symlinked_speculator_directory_read_only( + tmp_path: Path, +) -> None: + models = tmp_path / "models" + draft_root = models / "draft" + external = tmp_path / "external-speculator" + draft_root.mkdir(parents=True) + external.mkdir() + (external / "model.safetensors").write_bytes(b"speculator") + (draft_root / "laguna-xs2-speculator").symlink_to( + external, target_is_directory=True + ) + + spec = docker_run.server_run_spec( + Config(models_dir=models, model=ModelMeta(preset="laguna-xs.2")) + ) + + assert docker_run.BindMount( + str(external), + "/opt/lucebox-resolved/draft-dir", + read_only=True, + ) in spec.volumes + assert _env(spec)["DFLASH_DRAFT"] == "/opt/lucebox-resolved/draft-dir" + + def test_server_run_spec_rejects_model_path_traversal(tmp_path: Path) -> None: cfg = Config(models_dir=tmp_path, model=ModelMeta(target_file="../secret.gguf")) From 1ff2515aaef1b5d71b1c2e32d7ba20d85f7088dc Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:04:38 +0200 Subject: [PATCH 07/25] refactor(lucebox): tighten package edge cases --- install.sh | 8 +++++--- lucebox/src/lucebox/docker_run.py | 12 +++++++++--- lucebox/tests/test_config_cli.py | 2 +- lucebox/tests/test_docker_run.py | 20 ++++++++++++++++++++ 4 files changed, 35 insertions(+), 7 deletions(-) diff --git a/install.sh b/install.sh index d67dbd9b4..545d64c8f 100755 --- a/install.sh +++ b/install.sh @@ -37,13 +37,16 @@ die() { printf '%s[install] ✗%s %s\n' "$C_ERR" "$C_RST" "$*" >&2; exit 1; } command -v curl >/dev/null 2>&1 || die "curl is required (apt-get install curl)" sha256_file() { + local sum if command -v sha256sum >/dev/null 2>&1; then - sha256sum "$1" | awk '{print $1}' + sum=$(sha256sum "$1") elif command -v shasum >/dev/null 2>&1; then - shasum -a 256 "$1" | awk '{print $1}' + sum=$(shasum -a 256 "$1") else die "checksum requested, but neither sha256sum nor shasum is installed" fi + sum="${sum%% *}" + printf '%s' "$sum" | tr '[:upper:]' '[:lower:]' } # ── decide what gets baked in as the persisted channel ─────────────────── @@ -85,7 +88,6 @@ if [ -n "$expected_sha" ]; then [[ "$expected_sha" =~ ^[0-9a-fA-F]{64}$ ]] \ || die "LUCEBOX_WRAPPER_SHA256 must be exactly 64 hexadecimal characters" actual_sha=$(sha256_file "$tmp") - actual_sha=$(printf '%s' "$actual_sha" | tr '[:upper:]' '[:lower:]') expected_sha=$(printf '%s' "$expected_sha" | tr '[:upper:]' '[:lower:]') [ "$actual_sha" = "$expected_sha" ] \ || die "wrapper checksum mismatch (expected $expected_sha, got $actual_sha)" diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index c8c827f35..1568fce02 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -153,9 +153,15 @@ def _selected_model_path( if under_draft: container_base /= "draft" canonical = str(container_base.joinpath(*relative.parts)) - lexical = host_path.absolute() resolved = host_path.resolve(strict=False) - if resolved == lexical: + # A symlink in models_dir or one of its ancestors is covered by the root bind. + # Resolve that root before comparing so only symlinks *within* the selected + # model path need a separate narrow mount. + expected = cfg.models_dir.resolve(strict=False) + if under_draft: + expected /= "draft" + expected = expected.joinpath(*relative.parts) + if resolved == expected: return canonical mount_target = f"{_CONTAINER_RESOLVED_MODELS}/{role}" @@ -308,7 +314,7 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: draft_path = _selected_model_path( cfg, draft_dir, - field="model.speculator_dir", + field="preset.speculator_dir", role="draft-dir", under_draft=True, directory=True, diff --git a/lucebox/tests/test_config_cli.py b/lucebox/tests/test_config_cli.py index 8d36da760..a9a727723 100644 --- a/lucebox/tests/test_config_cli.py +++ b/lucebox/tests/test_config_cli.py @@ -149,5 +149,5 @@ def test_print_run_reports_invalid_configuration_without_traceback( result = CliRunner().invoke(app, ["print-run"]) assert result.exit_code == 2 - assert "Invalid configuration" in result.output + assert "Invalid configuration" in result.stderr assert "Traceback" not in result.output diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index 257b369f8..7753b8e5a 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -300,6 +300,26 @@ def test_server_run_spec_resolves_symlinked_model_parent(tmp_path: Path) -> None assert _env(spec)["DFLASH_TARGET"] == "/opt/lucebox-resolved/target/nested.gguf" +def test_server_run_spec_does_not_remount_file_for_symlinked_models_dir( + tmp_path: Path, +) -> None: + actual_models = tmp_path / "actual-models" + models = tmp_path / "models" + actual_models.mkdir() + (actual_models / "target.gguf").write_bytes(b"model") + models.symlink_to(actual_models, target_is_directory=True) + + spec = docker_run.server_run_spec( + Config(models_dir=models, model=ModelMeta(target_file="target.gguf")) + ) + + assert _env(spec)["DFLASH_TARGET"] == "/opt/lucebox-hub/server/models/target.gguf" + assert all( + mount.target != "/opt/lucebox-resolved/target/target.gguf" + for mount in spec.volumes + ) + + def test_server_run_spec_mounts_symlinked_speculator_directory_read_only( tmp_path: Path, ) -> None: From 6f2a9810ab2d1beb67d85e18deccbf6d4f92d6e6 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:14:14 +0200 Subject: [PATCH 08/25] feat(lucebox): add guided inference menu --- README.md | 29 ++ install.sh | 12 +- lucebox.sh | 654 +++++++++++++++++++++++++++++-- lucebox/README.md | 10 +- lucebox/src/lucebox/cli.py | 212 +++++++++- lucebox/src/lucebox/config.py | 22 +- lucebox/tests/test_cli.py | 11 +- lucebox/tests/test_models_cli.py | 63 +++ scripts/test_lucebox_sh.sh | 73 +++- server/scripts/entrypoint.sh | 6 +- 10 files changed, 1029 insertions(+), 63 deletions(-) diff --git a/README.md b/README.md index 4ce49fd51..d7b3c1f08 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,35 @@ Reference target: **RTX 3090 (Ampere sm_86)** — all headline numbers. Other NV `server/` (DFlash) builds with CMake 3.18+ and vendors the required `ggml` sources directly; only `Block-Sparse-Attention` remains a git submodule. No PyTorch is needed for `server/`. `optimizations/megakernel/` is the only component requiring PyTorch 2.0+ (CUDAExtension links against torch C++ libs). Power-tune: `sudo nvidia-smi -pl 220` (3090 sweet spot, re-sweep for other cards). +## Lucebox CLI + +The easiest path for both Lucebox buyers and open-source users is one command: + +```bash +# Buyers receive this preinstalled. Other users install the small host wrapper: +curl -fsSL https://raw.githubusercontent.com/Luce-Org/lucebox/main/install.sh | bash + +lucebox +``` + +`lucebox` opens the branded menu. **Quick setup** detects NVIDIA CUDA or AMD +ROCm, lets you choose and download a model, applies safe hardware-aware +optimizations, and starts the inference service. Wi-Fi and device provisioning +are intentionally outside this inference-only CLI. + +Contributors can open the same menu from a repository checkout with +`./lucebox.sh`. Its **Developer tools** can build and run the native C++ engine +or launch Claude Code, Codex, OpenCode, Hermes, Pi, OpenClaw, and Open WebUI +harnesses. Every menu action also has a scriptable command, for example: + +```bash +lucebox models select +lucebox optimize --yes +lucebox start +./lucebox.sh build cuda +./lucebox.sh native cuda +``` + ## Quick Start On Harnesses [`harness/`](harness/) contains RTX 3090 client launchers and regression tests diff --git a/install.sh b/install.sh index 545d64c8f..1886bae18 100755 --- a/install.sh +++ b/install.sh @@ -3,7 +3,7 @@ # # Canonical install (Luce-Org main, stable channel): # -# curl -fsSL https://raw.githubusercontent.com/Luce-Org/lucebox-hub/main/install.sh | bash +# curl -fsSL https://raw.githubusercontent.com/Luce-Org/lucebox/main/install.sh | bash # # Install from a different fork / branch (dev channel). Note the env var # is on the `bash` side of the pipe — `VAR=val curl … | bash` would attach @@ -22,7 +22,7 @@ set -euo pipefail -LUCEBOX_INSTALL_URL="${LUCEBOX_INSTALL_URL:-https://raw.githubusercontent.com/Luce-Org/lucebox-hub/main/lucebox.sh}" +LUCEBOX_INSTALL_URL="${LUCEBOX_INSTALL_URL:-https://raw.githubusercontent.com/Luce-Org/lucebox/main/lucebox.sh}" DEST="${LUCEBOX_INSTALL_DEST:-$HOME/.local/bin/lucebox}" # ── helpers ─────────────────────────────────────────────────────────────── @@ -143,8 +143,10 @@ esac cat <&2; } hint() { printf ' %b%s%b\n' "$C_DIM" "$*" "$C_RST"; } die() { err "$*"; exit 1; } +print_logo() { + printf '%b' "$C_BRAND" + cat <<'EOF' + · ╱ + ·──✦──· █ █ █ ▄▀▀ █▀▀ █▀▀▄ ▄▀▀▄ █ █ + ╱ · █ █ █ █ █▀▀ █▀▀▄ █ █ █ + · ▀▀ ▀▀ ▀▀ ▀▀▀ ▀▀▀ ▀▀ ▀ ▀ +EOF + printf '%b local inference, made simple%b\n\n' "$C_DIM" "$C_RST" +} + +# Find a source checkout for contributor-only actions. An explicit path wins; +# otherwise inspect the current directory and the wrapper's own directory, +# walking upward until the repository markers are found. Buyer installs simply +# return no path and never see build/harness actions. +_find_repo_root() { + local candidate="${LUCEBOX_REPO:-}" dir + if [ -n "$candidate" ]; then + if [ -f "$candidate/server/CMakeLists.txt" ] && [ -d "$candidate/harness" ]; then + (cd "$candidate" && pwd) + return 0 + fi + return 1 + fi + + for candidate in "$PWD" "$(dirname "$SCRIPT_PATH")"; do + dir="$candidate" + while [ "$dir" != "/" ] && [ -n "$dir" ]; do + if [ -f "$dir/server/CMakeLists.txt" ] && [ -d "$dir/harness" ]; then + (cd "$dir" && pwd) + return 0 + fi + dir="$(dirname "$dir")" + done + done + return 1 +} + +_confirm() { + # usage: _confirm "question" [default_yes] + local question="$1" default_yes="${2:-1}" answer prompt + if [ "$default_yes" = "1" ]; then prompt="Y/n"; else prompt="y/N"; fi + printf '%s [%s] ' "$question" "$prompt" + IFS= read -r answer || return 1 + case "$answer" in + y|Y|yes|YES|Yes) return 0 ;; + n|N|no|NO|No) return 1 ;; + "") [ "$default_yes" = "1" ] ;; + *) return 1 ;; + esac +} + sha256_file() { local sum if command -v sha256sum >/dev/null 2>&1; then @@ -1141,13 +1200,46 @@ cmd_systemctl_passthrough() { } cmd_logs() { - require_systemd "logs" - # Pure passthrough: any flags the user wants (-f, -n, --since, ...) go - # straight to journalctl. Default is follow. - if [ $# -eq 0 ]; then - exec journalctl --user -u "$UNIT_NAME" -f + ensure_probed + if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ] && [ -f "$UNIT_PATH" ]; then + # Pure passthrough: any flags the user wants (-f, -n, --since, ...) + # go straight to journalctl. Default is follow. + if [ $# -eq 0 ]; then + exec journalctl --user -u "$UNIT_NAME" -f + fi + exec journalctl --user -u "$UNIT_NAME" "$@" fi - exec journalctl --user -u "$UNIT_NAME" "$@" + if _lucebox_container_running; then + if [ $# -eq 0 ]; then + exec docker logs -f "$CONTAINER_NAME" + fi + exec docker logs "$@" "$CONTAINER_NAME" + fi + die "the inference engine is not running — use '$SCRIPT_NAME start' or '$SCRIPT_NAME serve'" +} + +cmd_status() { + ensure_probed + if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ] && [ -f "$UNIT_PATH" ]; then + exec systemctl --user status "$UNIT_NAME" --no-pager + fi + if _lucebox_container_running; then + exec docker ps --filter "name=^${CONTAINER_NAME}\$" \ + --format 'table {{.Names}}\t{{.Status}}\t{{.Ports}}' + fi + info "Lucebox inference engine is stopped" + hint "Run '$SCRIPT_NAME setup' for first-time setup, or '$SCRIPT_NAME serve' in the foreground." +} + +cmd_stop() { + ensure_probed + if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ] && [ -f "$UNIT_PATH" ]; then + exec systemctl --user stop "$UNIT_NAME" + fi + if _lucebox_container_running; then + exec docker stop "$CONTAINER_NAME" + fi + ok "Lucebox inference engine is already stopped" } cmd_pull() { @@ -1162,6 +1254,251 @@ cmd_pull() { exec docker pull "${IMAGE_BASE}:${variant}" } +# ── contributor-native workflow ─────────────────────────────────────────── + +_native_backend() { + local requested="${1:-}" variant + case "$requested" in + cuda|cuda12|nvidia) printf 'cuda'; return ;; + rocm|hip|amd) printf 'rocm'; return ;; + "") + variant=$(pick_variant) + if _variant_is_rocm "$variant"; then printf 'rocm'; else printf 'cuda'; fi + return + ;; + *) die "unknown backend '$requested' — choose cuda or rocm" ;; + esac +} + +_native_build_dir() { + local repo="$1" backend="$2" + if [ -n "${LUCEBOX_BUILD_DIR:-}" ]; then + printf '%s' "$LUCEBOX_BUILD_DIR" + elif [ "$backend" = "rocm" ]; then + printf '%s/server/build-hip' "$repo" + else + printf '%s/server/build-cuda' "$repo" + fi +} + +_safe_model_relative_path() { + local value="$1" part parts=() + [ -n "$value" ] || return 1 + [[ "$value" != /* ]] || return 1 + IFS='/' read -r -a parts <<<"$value" + for part in "${parts[@]}"; do + [ "$part" != ".." ] || return 1 + done + return 0 +} + +_selected_model_paths() { + # Emit target path, draft path (or "none"), and model id on separate lines. + # `models select` persists the filenames, so preset mapping is only a + # compatibility fallback for hand-written config files. + local preset target_file draft_file target_path draft_path="none" + preset=$(_lucebox_config_get model.preset) + target_file=$(_lucebox_config_get model.target_file) + draft_file=$(_lucebox_config_get model.draft_file) + if [ -z "$target_file" ]; then + case "$preset" in + qwen3.6-27b) target_file="Qwen3.6-27B-Q4_K_M.gguf" ;; + gemma-4-26b) target_file="google_gemma-4-26B-A4B-it-Q4_K_M.gguf" ;; + gemma-4-31b) target_file="google_gemma-4-31B-it-Q4_K_M.gguf" ;; + laguna-xs.2) target_file="laguna-xs2-Q4_K_M.gguf" ;; + qwen3.6-moe) target_file="Qwen3.6-35B-A3B-UD-Q4_K_M.gguf" ;; + esac + fi + if [ -z "$draft_file" ]; then + case "$preset" in + qwen3.6-27b) draft_file="dflash-draft-3.6-q4_k_m.gguf" ;; + gemma-4-26b) draft_file="gemma-4-26B-A4B-it-DFlash-q8_0.gguf" ;; + gemma-4-31b) draft_file="gemma-4-31B-it-DFlash-q8_0.gguf" ;; + esac + fi + _safe_model_relative_path "$target_file" \ + || die "no valid model is selected — run '$SCRIPT_NAME models select' first" + target_path="$DEFAULT_MODELS_DIR/$target_file" + if [ -n "$draft_file" ]; then + _safe_model_relative_path "$draft_file" \ + || die "invalid model.draft_file in $(_lucebox_config_path)" + draft_path="$DEFAULT_MODELS_DIR/draft/$draft_file" + elif [ "$preset" = "laguna-xs.2" ] \ + && [ -d "$DEFAULT_MODELS_DIR/draft/laguna-xs2-speculator" ]; then + draft_path="$DEFAULT_MODELS_DIR/draft/laguna-xs2-speculator" + fi + printf '%s\n%s\n%s\n' "$target_path" "$draft_path" "${preset:-lucebox}" +} + +_export_native_config() { + local key env_name value + while IFS='|' read -r key env_name; do + [ -n "$key" ] || continue + value=$(_lucebox_config_get "$key") + [ -n "$value" ] || continue + case "$key" in + dflash.lazy|dflash.debug_thinking_logits) + case "$value" in + true|1|yes|on) value=1 ;; + *) value=0 ;; + esac + ;; + esac + export "$env_name=$value" + done <<'EOF' +dflash.budget|DFLASH_BUDGET +dflash.max_ctx|DFLASH_MAX_CTX +dflash.lazy|DFLASH_LAZY +dflash.prefix_cache_slots|DFLASH_PREFIX_CACHE_SLOTS +dflash.prefill_cache_slots|DFLASH_PREFILL_CACHE_SLOTS +dflash.cache_type_k|DFLASH_CACHE_TYPE_K +dflash.cache_type_v|DFLASH_CACHE_TYPE_V +dflash.prefill_mode|DFLASH_PREFILL_MODE +dflash.prefill_keep_ratio|DFLASH_PREFILL_KEEP +dflash.prefill_threshold|DFLASH_PREFILL_THRESHOLD +dflash.prefill_drafter|DFLASH_PREFILL_DRAFTER +dflash.think_max|DFLASH_THINK_MAX +dflash.fa_window|DFLASH_FA_WINDOW +dflash.think_soft_close_min_ratio|DFLASH_THINK_SOFT_CLOSE_MIN_RATIO +dflash.debug_thinking_logits|DFLASH_DEBUG_THINKING_LOGITS +EOF +} + +cmd_native_build() { + local repo backend build_dir jobs hip_wmma=OFF + repo=$(_find_repo_root) \ + || die "native build requires a lucebox repository checkout (cd into it or set LUCEBOX_REPO)" + ensure_probed + backend=$(_native_backend "${1:-}") + build_dir=$(_native_build_dir "$repo" "$backend") + command -v cmake >/dev/null 2>&1 || die "cmake is required to build the inference engine" + [ -f "$repo/server/deps/llama.cpp/ggml/CMakeLists.txt" ] \ + || die "git submodules are missing — run: git -C '$repo' submodule update --init --recursive" + + local configure=(cmake -S "$repo/server" -B "$build_dir" -DCMAKE_BUILD_TYPE=Release) + if [ "$backend" = "rocm" ]; then + [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ] \ + || die "ROCm native build selected but no AMD GPU was detected" + if [ -f /opt/rocm/include/rocwmma/rocwmma.hpp ] \ + || [ -f /usr/include/rocwmma/rocwmma.hpp ]; then + hip_wmma=ON + fi + configure+=( + -DDFLASH27B_GPU_BACKEND=hip + "-DDFLASH27B_HIP_ARCHITECTURES=${LUCEBOX_HOST_AMD_GPU_ARCH:-gfx1151}" + "-DDFLASH27B_HIP_SM80_EQUIV=$hip_wmma" + ) + else + [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ] \ + || die "CUDA native build selected but no NVIDIA GPU was detected" + configure+=(-DDFLASH27B_GPU_BACKEND=cuda) + if [[ "$LUCEBOX_HOST_GPU_SM" =~ ^[0-9]+$ ]]; then + configure+=("-DCMAKE_CUDA_ARCHITECTURES=$LUCEBOX_HOST_GPU_SM") + fi + fi + info "Configuring native $backend build in $build_dir" + "${configure[@]}" + jobs="${LUCEBOX_HOST_NPROC:-1}" + [ "$jobs" -gt 0 ] 2>/dev/null || jobs=1 + info "Building dflash_server ($jobs jobs)" + cmake --build "$build_dir" --target dflash_server -j "$jobs" + ok "Native engine ready: $build_dir/dflash_server" +} + +cmd_native_serve() { + local repo backend build_dir binary selected=() target draft model_id + repo=$(_find_repo_root) \ + || die "native run requires a lucebox repository checkout (cd into it or set LUCEBOX_REPO)" + ensure_probed + backend=$(_native_backend "${1:-}") + build_dir=$(_native_build_dir "$repo" "$backend") + binary="$build_dir/dflash_server" + [ -x "$binary" ] \ + || die "native engine is not built — run '$SCRIPT_NAME build $backend' first" + mapfile -t selected < <(_selected_model_paths) + target="${selected[0]:-}" + draft="${selected[1]:-none}" + model_id="${selected[2]:-lucebox}" + [ -f "$target" ] \ + || die "selected target is not installed: $target — run '$SCRIPT_NAME models select'" + if [ "$draft" != "none" ] && [ ! -e "$draft" ]; then + die "selected draft is not installed: $draft — run '$SCRIPT_NAME models select'" + fi + + _export_native_config + export DFLASH_DIR="$repo/server" + export DFLASH_SERVER_BIN="$binary" + export DFLASH_TARGET="$target" + if [ "$draft" = "none" ]; then + export DFLASH_DRAFT="$DEFAULT_MODELS_DIR/.lucebox-no-draft" + else + export DFLASH_DRAFT="$draft" + fi + export DFLASH_HOST="${LUCEBOX_NATIVE_HOST:-127.0.0.1}" + export DFLASH_PORT="$DEFAULT_PORT" + export DFLASH_MODEL_NAME="$model_id" + export LUCEBOX_NATIVE=1 + info "Starting native $backend engine at http://$DFLASH_HOST:$DFLASH_PORT" + exec "$repo/server/scripts/entrypoint.sh" serve +} + +cmd_harness() { + local repo name="${1:-}" backend build_dir binary script selected=() target draft + repo=$(_find_repo_root) \ + || die "harness launchers are contributor tools; run this command inside the lucebox repository" + if [ -z "$name" ]; then + cat <<'EOF' +Choose a harness: + 1 Claude Code + 2 Codex + 3 OpenCode + 4 Hermes + 5 Pi + 6 OpenClaw + 7 Open WebUI +EOF + printf 'Harness number: ' + IFS= read -r name || return 1 + fi + case "$name" in + 1|claude|claude-code) name=claude_code ;; + 2|codex) name=codex ;; + 3|opencode) name=opencode ;; + 4|hermes) name=hermes ;; + 5|pi) name=pi ;; + 6|openclaw) name=openclaw ;; + 7|openwebui|webui) name=openwebui ;; + *) die "unknown harness '$name'" ;; + esac + script="$repo/harness/clients/run_${name}.sh" + [ -x "$script" ] || die "harness launcher is missing: $script" + ensure_probed + backend=$(_native_backend "${LUCEBOX_HARNESS_BACKEND:-}") + build_dir=$(_native_build_dir "$repo" "$backend") + binary="$build_dir/dflash_server" + [ -x "$binary" ] \ + || die "native engine is not built — run '$SCRIPT_NAME build $backend' first" + mapfile -t selected < <(_selected_model_paths) + target="${selected[0]:-}" + draft="${selected[1]:-none}" + [ -f "$target" ] \ + || die "selected target is not installed: $target — run '$SCRIPT_NAME models select'" + + export REPO_DIR="$repo" + export DFLASH_SERVER_BIN="$binary" + export DFLASH_TARGET="$target" + export TARGET="$target" + export DFLASH_DRAFT="$draft" + export DRAFT="$draft" + local max_ctx budget + max_ctx=$(_lucebox_config_get dflash.max_ctx) + budget=$(_lucebox_config_get dflash.budget) + [ -z "$max_ctx" ] || export MAX_CTX="$max_ctx" + [ -z "$budget" ] || export BUDGET="$budget" + info "Launching $name with the selected Lucebox model" + exec "$script" +} + cmd_update() { # Download the wrapper itself as data; never execute a second remote # installer. Override the persisted channel with LUCEBOX_INSTALL_URL. @@ -1240,7 +1577,7 @@ cmd_completion() { # lucebox completion fish | source # # Keep this in sync with the dispatch table in main() and the sub-app - # verbs (config get/set/unset, models list/download). Adding a new + # verbs (config get/set/unset, models list/download/select). Adding a new # top-level command means adding it here too. local shell="${1:-}" case "$shell" in @@ -1253,11 +1590,11 @@ _lucebox_complete() { COMPREPLY=() cur="${COMP_WORDS[COMP_CWORD]}" prev="${COMP_WORDS[COMP_CWORD-1]}" - cmds="install uninstall start stop restart enable disable status logs \ - serve pull update check completion config models \ - print-run help version" + cmds="menu setup install uninstall start stop restart enable disable status logs \ + serve build native harness pull update check completion config models \ + optimize print-run help version" config_verbs="get set unset" - models_verbs="list download" + models_verbs="list download select" completion_shells="bash zsh fish" # Sub-app verbs / shell args. @@ -1295,14 +1632,14 @@ ZSH # lucebox fish completion. Source from ~/.config/fish/config.fish: # lucebox completion fish | source complete -c lucebox -f -set -l __lucebox_cmds install uninstall start stop restart enable disable \ - status logs serve pull update check completion config models \ - print-run help version +set -l __lucebox_cmds menu setup install uninstall start stop restart enable disable \ + status logs serve build native harness pull update check completion config models \ + optimize print-run help version for cmd in $__lucebox_cmds complete -c lucebox -n "not __fish_seen_subcommand_from $__lucebox_cmds" -a $cmd end complete -c lucebox -n "__fish_seen_subcommand_from config" -a "get set unset" -complete -c lucebox -n "__fish_seen_subcommand_from models" -a "list download" +complete -c lucebox -n "__fish_seen_subcommand_from models" -a "list download select" complete -c lucebox -n "__fish_seen_subcommand_from completion" -a "bash zsh fish" FISH ;; @@ -1534,7 +1871,7 @@ cmd_exec_in_container() { _lucebox_prefer_exec() { local cmd="$1"; shift case "$cmd" in - config|models|check|print-run|print-serve-argv) + config|models|optimize|check|print-run|print-serve-argv) return 0 ;; *) @@ -1567,9 +1904,250 @@ cmd_route_to_container() { cmd_in_container "$cmd" "$@" } +# ── interactive product surface ─────────────────────────────────────────── + +_engine_state() { + if command -v systemctl >/dev/null 2>&1 \ + && systemctl --user is-active --quiet "$UNIT_NAME" 2>/dev/null; then + printf 'running' + elif _lucebox_container_running; then + printf 'running' + else + printf 'stopped' + fi +} + +_menu_clear() { + if [ -t 1 ] && [ "${TERM:-dumb}" != "dumb" ] \ + && [ "${LUCEBOX_NO_CLEAR:-0}" != "1" ]; then + printf '\033[2J\033[H' + fi +} + +_menu_pause() { + [ -t 0 ] || return 0 + printf '\nPress Enter to return to the menu…' + IFS= read -r _ || true +} + +_menu_run() { + # Always invoke through bash: repository copies are not necessarily marked + # executable, while installed buyer copies are. This behaves identically + # in both places and lets exec-heavy subcommands return to the menu. + local rc + if bash "$SCRIPT_PATH" "$@"; then + return 0 + else + rc=$? + warn "Command failed (exit $rc)." + return "$rc" + fi +} + +_menu_start() { + ensure_probed + if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ] && [ -f "$UNIT_PATH" ]; then + _menu_run start + return + fi + warn "The background service is not installed." + if _confirm "Run the Docker engine in this terminal now?" 1; then + _menu_run serve + else + hint "Run '$SCRIPT_NAME setup' to install the background service." + fi +} + +_menu_restart_if_running() { + [ "$(_engine_state)" = "running" ] || return 0 + if ! _confirm "Restart the running engine to apply this change?" 1; then + warn "The change is saved and will apply at the next restart." + return 0 + fi + if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ] && [ -f "$UNIT_PATH" ]; then + _menu_run restart + else + _menu_run stop + warn "The foreground engine was stopped. Start it again from menu option 4." + fi +} + +cmd_setup() { + [ -t 0 ] && [ -t 1 ] \ + || die "guided setup needs a terminal — use '$SCRIPT_NAME --help' for non-interactive commands" + ensure_probed + _menu_clear + print_logo + printf '%bQuick setup%b\n\n' "$C_INFO" "$C_RST" + cmd_check + printf '\n' + + local variant + variant=$(pick_variant) + if [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ] \ + && [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ]; then + printf 'This build has NVIDIA and AMD graphics. Which accelerator should run Lucebox?\n' + printf ' 1 NVIDIA / CUDA %b(recommended for RTX + Strix builds)%b\n' "$C_DIM" "$C_RST" + printf ' 2 AMD / ROCm\n' + printf 'Choice [1]: ' + local backend_choice + IFS= read -r backend_choice || return 1 + case "$backend_choice" in + 2) variant=rocm ;; + *) variant=cuda12 ;; + esac + fi + info "Selected backend: $variant" + + if docker image inspect "${IMAGE_BASE}:${variant}" >/dev/null 2>&1; then + ok "Inference image is already installed (${IMAGE_BASE}:${variant})" + else + _confirm "Download the ${variant} inference image now?" 1 \ + || { warn "Setup stopped before the image download."; return 1; } + LUCEBOX_VARIANT="$variant" bash "$SCRIPT_PATH" pull || return $? + fi + + # Persist the accelerator choice only after its image is available; the + # config writer lives in that image on buyer installations. + LUCEBOX_VARIANT="$variant" bash "$SCRIPT_PATH" config set "variant=$variant" \ + || return $? + + printf '\n' + bash "$SCRIPT_PATH" models select || return $? + if [ -z "$(_lucebox_config_get model.preset)" ]; then + warn "No model was selected; setup stopped without starting the engine." + return 1 + fi + + printf '\n' + if _confirm "Use automatic GPU optimization?" 1; then + bash "$SCRIPT_PATH" optimize --yes || return $? + fi + + if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ]; then + if [ ! -f "$UNIT_PATH" ]; then + printf '\n' + if _confirm "Install Lucebox as a background service?" 1; then + bash "$SCRIPT_PATH" install || return $? + fi + fi + if [ -f "$UNIT_PATH" ] && _confirm "Start the inference engine now?" 1; then + if [ "$(_engine_state)" = "running" ]; then + bash "$SCRIPT_PATH" restart || return $? + else + bash "$SCRIPT_PATH" start || return $? + fi + fi + else + warn "Background services are unavailable on this host." + hint "Use '$SCRIPT_NAME serve' to run the engine in the foreground." + fi + + printf '\n' + ok "Lucebox is configured" + hint "API: http://127.0.0.1:$DEFAULT_PORT/v1" + hint "Menu: $SCRIPT_NAME" + hint "Status: $SCRIPT_NAME status" +} + +cmd_developer_menu() { + local repo choice + repo=$(_find_repo_root) \ + || { warn "Developer tools require a lucebox repository checkout."; return 1; } + while true; do + _menu_clear + print_logo + printf '%bDeveloper tools%b\n' "$C_INFO" "$C_RST" + printf 'Repository: %s\n\n' "$repo" + printf ' 1 Build the native inference engine\n' + printf ' 2 Run the native inference engine\n' + printf ' 3 Choose and run a client harness\n' + printf ' b Back\n\n' + printf 'Choose: ' + IFS= read -r choice || return 0 + case "$choice" in + 1) _menu_run build; _menu_pause ;; + 2) _menu_run native; _menu_pause ;; + 3) _menu_run harness; _menu_pause ;; + b|B|q|Q) return 0 ;; + *) warn "Choose 1–3 or b"; _menu_pause ;; + esac + done +} + +cmd_menu() { + local choice model variant state repo_hint + while true; do + ensure_probed + model=$(_lucebox_config_get model.preset) + model="${model:-not selected}" + variant=$(pick_variant) + state=$(_engine_state) + if _find_repo_root >/dev/null 2>&1; then + repo_hint="available" + else + repo_hint="not a source checkout" + fi + + _menu_clear + print_logo + printf ' GPU: %s\n' "${LUCEBOX_HOST_GPU_NAME:-${LUCEBOX_HOST_AMD_GPU_NAME:-not detected}}" + if [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ] \ + && [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ]; then + printf ' Other GPU: %s\n' "${LUCEBOX_HOST_AMD_GPU_NAME:-AMD GPU}" + fi + printf ' Backend: %s\n' "$variant" + printf ' Model: %s\n' "$model" + printf ' Optimization: automatic\n' + printf ' Engine: %s\n\n' "$state" + + printf ' 1 Quick setup\n' + printf ' 2 Choose or download a model\n' + printf ' 3 Apply automatic optimization\n' + printf ' 4 Start the inference engine\n' + printf ' 5 Stop the inference engine\n' + printf ' 6 Status\n' + printf ' 7 Recent logs\n' + printf ' d Developer tools %b(%s)%b\n' "$C_DIM" "$repo_hint" "$C_RST" + printf ' q Quit\n\n' + printf 'Choose: ' + IFS= read -r choice || return 0 + case "$choice" in + 1) cmd_setup; _menu_pause ;; + 2) + if _menu_run models select; then _menu_restart_if_running; fi + _menu_pause + ;; + 3) + if _menu_run optimize; then _menu_restart_if_running; fi + _menu_pause + ;; + 4) _menu_start; _menu_pause ;; + 5) _menu_run stop; _menu_pause ;; + 6) _menu_run status; _menu_pause ;; + 7) + if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ] && [ -f "$UNIT_PATH" ]; then + _menu_run logs -n 80 --no-pager + else + _menu_run logs --tail 80 + fi + _menu_pause + ;; + d|D) cmd_developer_menu ;; + q|Q|quit|exit) return 0 ;; + *) warn "Choose 1–7, d, or q"; _menu_pause ;; + esac + done +} + usage() { cat < print shell completion script (bash / zsh / fish) + models select numbered model picker; download + activate in one step models list / download / activate model presets + optimize apply safe hardware-aware inference defaults config read / write keys in .lucebox/config.toml print-run print the docker-run command for the server @@ -1633,22 +2219,42 @@ main() { *) args+=("$1"); shift ;; esac done - set -- "${args[@]}" + if [ "${#args[@]}" -gt 0 ]; then + set -- "${args[@]}" + else + set -- + fi - local cmd="${1:-help}" - [ $# -gt 0 ] && shift + local cmd + if [ $# -eq 0 ]; then + if [ -t 0 ] && [ -t 1 ]; then cmd=menu; else cmd=help; fi + else + cmd="$1" + shift + fi case "$cmd" in + # Branded interactive surface / guided first run. + menu) cmd_menu "$@" ;; + setup) cmd_setup "$@" ;; + # Systemd surface install) cmd_systemd_install "$@" ;; uninstall) cmd_systemd_uninstall "$@" ;; - start|stop|restart|enable|disable|status) + start|restart|enable|disable) cmd_systemctl_passthrough "$cmd" "$@" ;; + stop) cmd_stop "$@" ;; + status) cmd_status "$@" ;; logs) cmd_logs "$@" ;; # Direct server serve) cmd_serve "$@" ;; pull) cmd_pull "$@" ;; + # Native source-repository workflow. + build) cmd_native_build "$@" ;; + native) cmd_native_serve "$@" ;; + harness) cmd_harness "$@" ;; + # Self-update — re-runs the bootstrap installer against the channel # this script was installed from (LUCEBOX_INSTALLED_FROM). update) cmd_update "$@" ;; diff --git a/lucebox/README.md b/lucebox/README.md index f5bb7d635..6e5533fc6 100644 --- a/lucebox/README.md +++ b/lucebox/README.md @@ -5,9 +5,11 @@ users do not install it directly: they install the small [`lucebox` host wrapper](https://github.com/Luce-Org/lucebox/blob/main/lucebox.sh), which invokes the package in the appropriate container: + lucebox # branded interactive menu + lucebox setup # guided first run lucebox check - lucebox models list - lucebox models download qwen3.6-27b --activate + lucebox models select # numbered picker + download + activate + lucebox optimize # safe hardware-aware defaults lucebox start The wrapper detects the host and selects CUDA for NVIDIA builds (including RTX @@ -17,6 +19,10 @@ optimization settings, and construction of the final server command. Host facts are passed through `LUCEBOX_HOST_*` environment variables so the container never has to guess the host configuration. +Inside a source checkout, the same wrapper also exposes `lucebox build`, +`lucebox native`, and `lucebox harness` for contributors. Buyer installations +keep using the prebuilt container, so no compiler or source checkout is needed. + See the [project README](https://github.com/Luce-Org/lucebox#readme) for the installation and user flow. Contributors can find the CLI implementation in [`src/lucebox`](https://github.com/Luce-Org/lucebox/tree/main/lucebox/src/lucebox). diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index 622e6de5a..ff11ebf8c 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -4,8 +4,10 @@ doesn't intercept (everything outside the systemd surface) ends up here. Subcommand inventory: + (no command) — branded interactive menu check — readiness report config get/set/unset — read / write a single key in config.toml + optimize — apply the recommended hardware profile pull — docker pull the selected CUDA or ROCm image print-run — emit the docker-run command for the server print-serve-argv — same, raw argv lines (consumed by `lucebox serve`) @@ -23,6 +25,7 @@ from rich.markup import escape from rich.table import Table +import lucebox.autotune as autotune_mod import lucebox.config as config_mod import lucebox.docker_run as docker_run import lucebox.download as download_mod @@ -35,7 +38,7 @@ app = typer.Typer( name="lucebox", help="Host CLI for the lucebox-hub container. Invoked by lucebox.sh.", - no_args_is_help=True, + no_args_is_help=False, invoke_without_command=True, add_completion=False, ) @@ -45,6 +48,7 @@ @app.callback() def root_options( + ctx: typer.Context, version_flag: Annotated[ bool, typer.Option("--version", help="Print lucebox version and exit.", is_eager=True), @@ -54,11 +58,32 @@ def root_options( if version_flag: print(__version__) raise typer.Exit() + if ctx.invoked_subcommand is None: + if sys.stdin.isatty() and sys.stdout.isatty(): + _package_menu() + else: + _print_logo() + console.print(ctx.get_help()) + raise typer.Exit() # ── helpers ──────────────────────────────────────────────────────────────── +_LOGO = r""" + · ╱ + ·──✦──· █ █ █ ▄▀▀ █▀▀ █▀▀▄ ▄▀▀▄ █ █ + ╱ · █ █ █ █ █▀▀ █▀▀▄ █ █ █ + · ▀▀ ▀▀ ▀▀ ▀▀▀ ▀▀▀ ▀▀ ▀ ▀ +""".strip("\n") + + +def _print_logo() -> None: + """Render the compact Lucebox mark used by both interactive surfaces.""" + console.print(f"[bold gold1]{_LOGO}[/bold gold1]") + console.print("[dim] local inference, made simple[/dim]\n") + + def _load_or_build() -> Config: """env > config.toml > dataclass defaults — the canonical precedence. @@ -90,6 +115,43 @@ def _server_spec() -> docker_run.DockerRunSpec: raise typer.Exit(code=2) from exc +def _package_menu() -> None: + """Small in-package menu for direct installs and contributor workflows. + + Service lifecycle remains host-owned by ``lucebox.sh``. This menu covers + the package's own responsibilities and makes a direct ``python -m lucebox`` + invocation useful instead of dropping users into a wall of help text. + """ + while True: + _print_logo() + cfg = _load_or_build() + active = cfg.model.preset or "not selected" + console.print(f"Model: [bold]{escape(active)}[/bold]") + console.print("Optimization: [bold]Automatic[/bold] (safe defaults for this GPU)\n") + console.print(" [bold cyan]1[/bold cyan] Choose or download a model") + console.print(" [bold cyan]2[/bold cyan] Apply automatic optimization") + console.print(" [bold cyan]3[/bold cyan] Show configuration") + console.print(" [bold cyan]4[/bold cyan] Show Docker launch command") + console.print(" [bold cyan]q[/bold cyan] Quit") + try: + choice = typer.prompt("\nChoose", default="1").strip().lower() + except (EOFError, typer.Abort): + return + if choice in {"q", "quit", "exit"}: + return + if choice == "1": + models_select() + elif choice == "2": + optimize() + elif choice == "3": + config_get_cmd() + elif choice == "4": + print_run() + else: + console.print("[yellow]Choose 1–4 or q.[/yellow]") + console.print() + + # ── subcommands ──────────────────────────────────────────────────────────── @@ -223,6 +285,27 @@ def _print_installed_presets() -> None: console.print(f"[dim]Total disk usage: {total:.1f} GB[/dim]") +def _activate_preset(cfg: Config, preset: download_mod.ModelPreset) -> None: + """Persist one selected preset and seed first-run tuning.""" + config_set("model.preset", preset.name) + config_set("model.target_file", preset.target_file) + if preset.has_draft and preset.draft_file: + config_set("model.draft_file", preset.draft_file) + else: + # Drop any stale draft_file from a previous activation; the selected + # preset is explicitly target-only. + config_unset("model.draft_file") + console.print(f"[green]Activated:[/green] model.preset = {preset.name}") + # Never clobber a user-edited [dflash] section. Explicit profile resets + # go through `lucebox optimize` instead. + if config_mod.seed_dflash_from_host(cfg.host): + max_ctx = config_get("dflash.max_ctx")["dflash.max_ctx"][0] + console.print( + f"[green]Auto-tuned:[/green] VRAM-tier DFLASH_* defaults " + f"(max_ctx={max_ctx}) written to config.toml" + ) + + @models_app.callback(invoke_without_command=True) def models_default(ctx: typer.Context) -> None: """Default action: list installed presets, mark active with `*`.""" @@ -250,6 +333,82 @@ def models_list() -> None: console.print(table) +@models_app.command("select") +def models_select( + preset: Annotated[ + str, + typer.Argument(help="Preset name (omit for a numbered menu)."), + ] = "", + yes: Annotated[ + bool, + typer.Option("--yes", "-y", help="Download and activate without confirmation."), + ] = False, +) -> None: + """Choose, download, and activate a model in one guided step.""" + cfg = _load_or_build() + if not preset: + names = sorted(download_mod.PRESETS) + recommended = download_mod.recommend_preset(cfg.host) + default_name = cfg.model.preset or recommended or names[0] + default_index = names.index(default_name) + 1 if default_name in names else 1 + + table = Table(title="Choose a model", show_lines=False) + table.add_column("#", justify="right", style="cyan") + table.add_column("model") + table.add_column("download") + table.add_column("status") + table.add_column("notes") + for index, name in enumerate(names, start=1): + candidate = download_mod.PRESETS[name] + labels: list[str] = [] + if name == cfg.model.preset: + labels.append("active") + if name == recommended: + labels.append("recommended") + table.add_row( + str(index), + name, + f"~{candidate.approx_total_gb} GB", + download_mod.installed_status(cfg, candidate), + ", ".join(labels) or candidate.description, + ) + console.print(table) + try: + answer = typer.prompt( + "Model number or name", + default=str(default_index), + ).strip() + except (EOFError, typer.Abort): + console.print("[dim]No model changed.[/dim]") + return + if answer.isdigit() and 1 <= int(answer) <= len(names): + preset = names[int(answer) - 1] + else: + preset = answer + + try: + selected = download_mod.resolve_preset(preset) + except KeyError as exc: + error_console.print(f"[red]{escape(str(exc))}[/red]") + raise typer.Exit(code=2) from exc + + state = download_mod.installed_status(cfg, selected) + if state == "installed": + # Selection must work on a preloaded buyer appliance even when it is + # offline. Do not query Hugging Face merely to activate files that are + # already present locally. + _activate_preset(cfg, selected) + return + if state != "installed" and not yes: + if not typer.confirm( + f"Download about {selected.approx_total_gb} GB and activate {selected.name}?", + default=True, + ): + console.print("[dim]No model changed.[/dim]") + return + models_download(selected.name, activate=True) + + @models_app.command("download") def models_download( preset: Annotated[str, typer.Argument(help="Preset name (empty = recommend)")] = "", @@ -317,26 +476,37 @@ def models_download( console.print("[green]Done.[/green]") if activate: - config_set("model.preset", preset) - if pres.target_file: - config_set("model.target_file", pres.target_file) - if pres.has_draft and pres.draft_file: - config_set("model.draft_file", pres.draft_file) - else: - # Drop any stale draft_file from a previous activation; the - # active preset has no draft. - config_unset("model.draft_file") - console.print(f"[green]Activated:[/green] model.preset = {preset}") - # First-time setup: bake the VRAM-tier DFLASH_* heuristic into - # config.toml so `lucebox serve` is auto-tuned to this host instead - # of falling back to the conservative class defaults. Never clobbers - # an existing [dflash] section. - if config_mod.seed_dflash_from_host(cfg.host): - max_ctx = config_get("dflash.max_ctx")["dflash.max_ctx"][0] - console.print( - f"[green]Auto-tuned:[/green] VRAM-tier DFLASH_* defaults " - f"(max_ctx={max_ctx}) written to config.toml" - ) + _activate_preset(cfg, pres) + + +@app.command() +def optimize( + yes: Annotated[ + bool, + typer.Option("--yes", "-y", help="Apply without confirmation."), + ] = False, +) -> None: + """Apply the recommended hardware-aware inference profile. + + Stable CUDA/ROCm fast paths remain enabled by the engine itself. This + profile selects a safe context size, cache format, and DFlash budget from + detected VRAM while leaving experimental features off. + """ + cfg = _load_or_build() + recommended = autotune_mod.runtime_from_host(cfg.host) + console.print("[bold]Automatic optimization (recommended)[/bold]") + console.print(" DFlash decode automatic when the model has a matching draft") + console.print(" GPU fast paths enabled by the engine on CUDA and ROCm") + console.print(f" Maximum context {recommended.max_ctx:,} tokens") + cache = recommended.cache_type_k or "model default" + console.print(f" KV cache {cache}") + console.print(" Experimental paths off") + if not yes and not typer.confirm("Apply this profile?", default=True): + console.print("[dim]Optimization unchanged.[/dim]") + return + config_mod.seed_dflash_from_host(cfg.host, force=True) + console.print("[green]Automatic optimization applied.[/green]") + console.print("[dim]Advanced users can still use `lucebox config set dflash.KEY=VALUE`.[/dim]") @app.command() diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py index 75a601afb..778d81385 100644 --- a/lucebox/src/lucebox/config.py +++ b/lucebox/src/lucebox/config.py @@ -394,12 +394,19 @@ def save(cfg: Config, path: Path | None = None, *, doc: dict[str, Any] | None = return path -def seed_dflash_from_host(host: HostFacts, *, path: Path | None = None) -> bool: +def seed_dflash_from_host( + host: HostFacts, + *, + path: Path | None = None, + force: bool = False, +) -> bool: """Persist the VRAM-tier DFLASH_* heuristic to config.toml on first setup. - Returns True when it wrote. No-op (returns False) when a ``[dflash]`` - section already exists, so it never clobbers values a prior tune or the - user set. Called when a preset is first activated: without it a fresh + Returns True when it wrote. By default this is a no-op when a ``[dflash]`` + section already exists, so first-time setup never clobbers values a prior + tune or the user set. ``force=True`` intentionally replaces that section; + the interactive CLI uses it when the user explicitly selects the automatic + profile. Called when a preset is first activated: without it a fresh install serves at the conservative ``DflashRuntime`` class defaults (``load()`` returns those for a config.toml that has no ``[dflash]``), ignoring the host's VRAM tier. The ``live_config`` heuristic only fires @@ -417,8 +424,13 @@ def seed_dflash_from_host(host: HostFacts, *, path: Path | None = None) -> bool: if not path.exists() and path.with_suffix(".env").exists(): load(path) doc = load_doc(path) - if "dflash" in doc: + if "dflash" in doc and not force: return False + if force: + # Replace the whole section rather than updating known keys in place. + # That removes stale experimental fields and makes "Automatic" a real + # reset to the current hardware-derived defaults. + doc.pop("dflash", None) runtime = autotune_mod.runtime_from_host(host) for field, value in asdict(runtime).items(): _doc_set(doc, "dflash", field, _value_to_toml(value)) diff --git a/lucebox/tests/test_cli.py b/lucebox/tests/test_cli.py index ee254d9bd..7fb146ffd 100644 --- a/lucebox/tests/test_cli.py +++ b/lucebox/tests/test_cli.py @@ -53,10 +53,19 @@ def test_core_verbs_present_in_app() -> None: c.name or (c.callback.__name__ if c.callback else "") for c in app.registered_commands } - for verb in ("check", "pull", "print-run", "print-serve-argv", "version"): + for verb in ("check", "pull", "optimize", "print-run", "print-serve-argv", "version"): assert verb in registered +def test_no_args_has_branded_noninteractive_help() -> None: + """Pipes and CI get useful output instead of an input prompt.""" + result = CliRunner().invoke(app, []) + assert result.exit_code == 0 + assert "local inference, made simple" in result.output + assert "models" in result.output + assert "optimize" in result.output + + @pytest.mark.parametrize("args", [["version"], ["--version"]]) def test_version_command_and_option_match(args: list[str]) -> None: result = CliRunner().invoke(app, args) diff --git a/lucebox/tests/test_models_cli.py b/lucebox/tests/test_models_cli.py index dd4f244f5..97eaba8c7 100644 --- a/lucebox/tests/test_models_cli.py +++ b/lucebox/tests/test_models_cli.py @@ -122,6 +122,69 @@ def test_models_download_explicit_preset_with_activate( assert entries["model.preset"] == ("gemma-4-26b", "file") +def test_models_select_activates_preloaded_model_offline( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A factory-preloaded buyer can switch models without a network call.""" + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + preset = PRESETS["qwen3.6-27b"] + models = tmp_path / "models" + (models / preset.target_file).parent.mkdir(parents=True, exist_ok=True) + (models / preset.target_file).touch() + assert preset.draft_file is not None + (models / "draft").mkdir() + (models / "draft" / preset.draft_file).touch() + + def fail_network(*args: object, **kwargs: object) -> object: + raise AssertionError("preloaded selection must not contact Hugging Face") + + monkeypatch.setattr(download_mod, "status", fail_network) + monkeypatch.setattr(download_mod, "download_preset", fail_network) + + result = CliRunner().invoke(app, ["models", "select", preset.name, "--yes"]) + assert result.exit_code == 0 + assert "Activated" in result.output + entries = config_mod.config_get(path=cfg_path) + assert entries["model.preset"] == (preset.name, "file") + assert entries["model.target_file"] == (preset.target_file, "file") + assert entries["model.draft_file"] == (preset.draft_file, "file") + + +def test_models_select_numbered_picker_downloads_and_activates( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + monkeypatch.setattr(download_mod, "download_preset", lambda cfg, pres: 0) + monkeypatch.setattr( + download_mod, + "status", + lambda cfg, pres: {"target_present": False, "draft_present": False}, + ) + # qwen3.6-27b is the fourth entry in the stable alphabetical menu. + result = CliRunner().invoke(app, ["models", "select"], input="4\ny\n") + assert result.exit_code == 0 + assert "Choose a model" in result.output + entries = config_mod.config_get(path=cfg_path) + assert entries["model.preset"] == ("qwen3.6-27b", "file") + + +def test_optimize_resets_to_hardware_profile( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + config_mod.config_set("dflash.max_ctx", 4096, path=cfg_path) + + result = CliRunner().invoke(app, ["optimize", "--yes"]) + assert result.exit_code == 0 + assert "Automatic optimization applied" in result.output + entries = config_mod.config_get(path=cfg_path) + assert entries["dflash.max_ctx"] == (98304, "file") + assert entries["dflash.cache_type_k"] == ("tq3_0", "file") + + def test_installed_helpers_track_presence( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index d95513daf..46134abb3 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -184,9 +184,14 @@ if bash -n "$ENTRYPOINT"; then report ok "bash -n entrypoint.sh parses cleanly" else report fail "bash -n entrypoint.sh"; fi # ── 3. Trivial subcommands (zero-exit expected) ─────────────────────────── -assert_runs "help" "bash '$SCRIPT' help" "host-side wrapper" -assert_runs "--help" "bash '$SCRIPT' --help" "host-side wrapper" -assert_runs "-h" "bash '$SCRIPT' -h" "host-side wrapper" +assert_runs "help" "bash '$SCRIPT' help" "simple CLI for the Lucebox inference engine" +assert_runs "--help" "bash '$SCRIPT' --help" "simple CLI for the Lucebox inference engine" +assert_runs "-h" "bash '$SCRIPT' -h" "simple CLI for the Lucebox inference engine" +assert_runs "no args (non-interactive help)" \ + "bash '$SCRIPT' "$native_tmp/bin/cmake" <<'STUB' +#!/usr/bin/env bash +printf '%s\n' "$*" >> "${LUCEBOX_TEST_CMAKE_LOG:?}" +STUB +chmod +x "$native_tmp/bin/cmake" +native_log="$native_tmp/cmake.log" +native_out="" +native_rc=0 +native_out=$( + PATH="$native_tmp/bin:$PATH" \ + LUCEBOX_TEST_CMAKE_LOG="$native_log" \ + LUCEBOX_REPO="$ROOT" \ + LUCEBOX_BUILD_DIR="$native_tmp/build" \ + LUCEBOX_VARIANT=cuda12 \ + _LUCEBOX_HOST_PROBED=1 \ + LUCEBOX_HOST_HAS_NVIDIA_GPU=1 \ + LUCEBOX_HOST_GPU_SM=86 \ + LUCEBOX_HOST_NPROC=2 \ + bash "$SCRIPT" build cuda 2>&1 +) || native_rc=$? +if [ "$native_rc" -ne 0 ]; then + report fail "native CUDA build dispatch" "exit=$native_rc output=$(head -3 <<<"$native_out")" +elif ! grep -qF -- "-DDFLASH27B_GPU_BACKEND=cuda" "$native_log" \ + || ! grep -qF -- "-DCMAKE_CUDA_ARCHITECTURES=86" "$native_log" \ + || ! grep -qF -- "--target dflash_server -j 2" "$native_log"; then + report fail "native CUDA build dispatch" "unexpected cmake argv: $(tr '\n' ' ' < "$native_log")" +else + report ok "native CUDA build dispatch" +fi + +: > "$native_log" +native_rc=0 +native_out=$( + PATH="$native_tmp/bin:$PATH" \ + LUCEBOX_TEST_CMAKE_LOG="$native_log" \ + LUCEBOX_REPO="$ROOT" \ + LUCEBOX_BUILD_DIR="$native_tmp/build-hip" \ + LUCEBOX_VARIANT=rocm \ + _LUCEBOX_HOST_PROBED=1 \ + LUCEBOX_HOST_HAS_NVIDIA_GPU=0 \ + LUCEBOX_HOST_HAS_AMD_GPU=1 \ + LUCEBOX_HOST_AMD_GPU_ARCH=gfx1201 \ + LUCEBOX_HOST_NPROC=2 \ + bash "$SCRIPT" build rocm 2>&1 +) || native_rc=$? +if [ "$native_rc" -ne 0 ]; then + report fail "native ROCm build dispatch" "exit=$native_rc output=$(head -3 <<<"$native_out")" +elif ! grep -qF -- "-DDFLASH27B_GPU_BACKEND=hip" "$native_log" \ + || ! grep -qF -- "-DDFLASH27B_HIP_ARCHITECTURES=gfx1201" "$native_log"; then + report fail "native ROCm build dispatch" "unexpected cmake argv: $(tr '\n' ' ' < "$native_log")" +else + report ok "native ROCm build dispatch" +fi +rm -rf "$native_tmp" + # ── 7. Unknown subcommand → cmd_in_container fallback path. Same rule: # clean error, no raw bash leak. assert_no_set_u_leak "unknown subcommand dispatch" "$SCRIPT" no-such-subcommand @@ -797,6 +861,7 @@ test_image_base_derives_from_install_url() { local label="$1" url expected got for case in \ "https://raw.githubusercontent.com/easel/lucebox-hub/feat/lucebox-docker/lucebox.sh|ghcr.io/easel/lucebox-hub" \ + "https://raw.githubusercontent.com/Luce-Org/lucebox/main/lucebox.sh|ghcr.io/luce-org/lucebox-hub" \ "https://raw.githubusercontent.com/Luce-Org/lucebox-hub/main/lucebox.sh|ghcr.io/luce-org/lucebox-hub" \ "https://raw.githubusercontent.com/easel/lucebox-hub/601ab52/lucebox.sh|ghcr.io/easel/lucebox-hub" \ "https://example.com/bogus|ghcr.io/luce-org/lucebox-hub" @@ -817,7 +882,7 @@ test_image_base_derives_from_install_url() { done report ok "$label" } -test_image_base_derives_from_install_url "IMAGE_BASE derived from LUCEBOX_INSTALLED_FROM (4 URL shapes)" +test_image_base_derives_from_install_url "IMAGE_BASE derived from LUCEBOX_INSTALLED_FROM (5 URL shapes)" # ── config.toml reader + resolver ───────────────────────────────────────── # Drive _lucebox_config_get + _lucebox_resolve against a fixture diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index 95373aa1b..6187c8af8 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -610,7 +610,11 @@ if [ "$DFLASH_PREFILL_MODE" != "off" ]; then --prefill-drafter "$DFLASH_PREFILL_DRAFTER") fi -info "lucebox-hub container starting (target=$(basename "$DFLASH_TARGET"), max_ctx=$DFLASH_MAX_CTX, budget=$DFLASH_BUDGET, lazy=$DFLASH_LAZY)" +if [ "${LUCEBOX_NATIVE:-0}" = "1" ]; then + info "lucebox native server starting (target=$(basename "$DFLASH_TARGET"), max_ctx=$DFLASH_MAX_CTX, budget=$DFLASH_BUDGET, lazy=$DFLASH_LAZY)" +else + info "lucebox-hub container starting (target=$(basename "$DFLASH_TARGET"), max_ctx=$DFLASH_MAX_CTX, budget=$DFLASH_BUDGET, lazy=$DFLASH_LAZY)" +fi cd "$DFLASH_DIR" exec "${CMD[@]}" From 272cda3425424b3ee215e6548b64d6c62aaedab3 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:02:07 +0200 Subject: [PATCH 09/25] feat(lucebox): harden guided optimization setup Make the buyer and contributor CLI model-aware and hardware-aware across CUDA, ROCm, and heterogeneous Lucebox builds. Add explainable Automatic/Advanced DFlash, PFlash, KVFlash, and Spark profiles; safe artifact, memory, and cache validation; native harness forwarding; and R9700/Strix device pinning. --- README.md | 16 +- docs/specs/props-endpoint.md | 4 +- harness/clients/common.sh | 42 +++- lucebox.sh | 200 ++++++++++++++++-- lucebox/README.md | 20 +- lucebox/src/lucebox/autotune.py | 332 ++++++++++++++++++++++++------ lucebox/src/lucebox/cli.py | 278 ++++++++++++++++++++++--- lucebox/src/lucebox/config.py | 165 +++++++++++++-- lucebox/src/lucebox/docker_run.py | 52 +++-- lucebox/src/lucebox/download.py | 81 +++++++- lucebox/src/lucebox/types.py | 64 ++++-- lucebox/tests/test_autotune.py | 207 ++++++++++++++++++- lucebox/tests/test_config.py | 99 +++++++++ lucebox/tests/test_docker_run.py | 59 +++++- lucebox/tests/test_download.py | 12 +- lucebox/tests/test_models_cli.py | 106 +++++++++- scripts/test_lucebox_sh.sh | 93 ++++++++- server/scripts/entrypoint.sh | 49 ++++- 18 files changed, 1677 insertions(+), 202 deletions(-) diff --git a/README.md b/README.md index d7b3c1f08..72e696624 100644 --- a/README.md +++ b/README.md @@ -113,10 +113,21 @@ lucebox ``` `lucebox` opens the branded menu. **Quick setup** detects NVIDIA CUDA or AMD -ROCm, lets you choose and download a model, applies safe hardware-aware -optimizations, and starts the inference service. Wi-Fi and device provisioning +ROCm, lets you choose and download a model, applies a model- and hardware-aware +**Automatic** profile, and starts the inference service. Automatic explains its +DFlash, PFlash, KVFlash, and Spark choices and prefers exact, all-GPU execution +when it fits. It enables paging or expert offload only when the selected model +and primary GPU need them; Spark also requires enough host RAM for its cold +experts. **Review optimizations** offers an Advanced mode for users who want to +override those four choices. Wi-Fi and device provisioning are intentionally outside this inference-only CLI. +On heterogeneous builds the CLI selects one primary accelerator: CUDA on an +RTX 3090 + Strix machine, and the largest-VRAM AMD device on an R9700 + Strix +machine. Automatic tuning never adds VRAM across devices. The engine's advanced +layer-split paths remain contributor workflows; transparent multi-GPU sharding +is not part of this CLI release. + Contributors can open the same menu from a repository checkout with `./lucebox.sh`. Its **Developer tools** can build and run the native C++ engine or launch Claude Code, Codex, OpenCode, Hermes, Pi, OpenClaw, and Open WebUI @@ -125,6 +136,7 @@ harnesses. Every menu action also has a scriptable command, for example: ```bash lucebox models select lucebox optimize --yes +lucebox optimize --advanced lucebox start ./lucebox.sh build cuda ./lucebox.sh native cuda diff --git a/docs/specs/props-endpoint.md b/docs/specs/props-endpoint.md index e4238df3d..0aa604a9d 100644 --- a/docs/specs/props-endpoint.md +++ b/docs/specs/props-endpoint.md @@ -466,8 +466,8 @@ configuration drift between runs is possible. - `fa_window` — sliding-window attention window in tokens. - `kv_cache_k` / `kv_cache_v` — effective KV cache dtypes (e.g. `"q4_0"`, `"tq3_0"`, `"f16"`). Operator's CLI choice when set, - otherwise the binary's auto-default (`tq3_0` when - `max_ctx > 6144`, else `q4_0`, on CUDA). + otherwise the model-family default (Laguna uses `q8_0`; base families use + `q4_0`). Automatic profiles do not force `tq3_0`. - `lazy_draft` — whether the decode draft is parked when idle. - `target_sharding` — true when the target model is layer-split across multiple GPUs. diff --git a/harness/clients/common.sh b/harness/clients/common.sh index 3a5edf15f..ff35e6b7d 100755 --- a/harness/clients/common.sh +++ b/harness/clients/common.sh @@ -132,13 +132,22 @@ start_dflash_native_server() { echo " cmake --build $REPO_DIR/server/build --target dflash_server -j\$(nproc)" >&2 return 1 fi - if [[ ! -f "$TARGET" ]]; then + if [[ ! -s "$TARGET" ]]; then echo "target GGUF not found: $TARGET" >&2 echo "Set TARGET=/path/to/model.gguf or DFLASH_TARGET=/path/to/model.gguf, or download the default:" >&2 echo " hf download unsloth/Qwen3.6-27B-GGUF Qwen3.6-27B-Q4_K_M.gguf --local-dir $REPO_DIR/server/models/" >&2 return 1 fi - if draft_enabled && [[ ! -f "$DRAFT" ]]; then + if draft_enabled && [[ -d "$DRAFT" ]]; then + local resolved_draft="" + resolved_draft=$(find -L "$DRAFT" -maxdepth 4 -type f \ + \( -name 'model.safetensors' -o -name '*.safetensors' -o -name '*.gguf' \) \ + -print 2>/dev/null | sort | head -n 1) + if [[ -n "$resolved_draft" ]]; then + DRAFT="$resolved_draft" + fi + fi + if draft_enabled && [[ ! -s "$DRAFT" ]]; then echo "DFlash draft not found: $DRAFT" >&2 echo "Set DRAFT=/path/to/dflash-draft.gguf or DFLASH_DRAFT=/path/to/dflash-draft.gguf, or download the default:" >&2 echo " hf download Lucebox/Qwen3.6-27B-DFlash-GGUF dflash-draft-3.6-q4_k_m.gguf --local-dir $REPO_DIR/server/models/draft/" >&2 @@ -160,6 +169,34 @@ start_dflash_native_server() { if [[ -n "$FA_WINDOW" ]] && [[ "$FA_WINDOW" != "0" ]]; then fa_args=(--fa-window "$FA_WINDOW") fi + local optimization_args=() + if [[ -n "${DFLASH_PREFILL_DRAFTER:-}" ]]; then + if [[ ! -s "$DFLASH_PREFILL_DRAFTER" ]]; then + echo "PFlash/KVFlash scorer not found: $DFLASH_PREFILL_DRAFTER" >&2 + return 1 + fi + optimization_args+=(--prefill-drafter "$DFLASH_PREFILL_DRAFTER") + fi + if [[ -n "${DFLASH_PREFILL_MODE:-}" && "$DFLASH_PREFILL_MODE" != "off" ]]; then + optimization_args+=( + --prefill-compression "$DFLASH_PREFILL_MODE" + --prefill-keep-ratio "${DFLASH_PREFILL_KEEP:-0.10}" + --prefill-threshold "${DFLASH_PREFILL_THRESHOLD:-32768}" + ) + fi + if [[ -n "${DFLASH_KVFLASH:-}" && "$DFLASH_KVFLASH" != "off" ]]; then + optimization_args+=( + --kvflash "$DFLASH_KVFLASH" + --kvflash-policy "${DFLASH_KVFLASH_POLICY:-drafter}" + --kvflash-tau "${DFLASH_KVFLASH_TAU:-64}" + ) + fi + if [[ "${DFLASH_SPARK:-0}" == "1" ]]; then + optimization_args+=(--spark) + if [[ -n "${DFLASH_SPARK_VRAM_GB:-}" && "$DFLASH_SPARK_VRAM_GB" != "0" && "$DFLASH_SPARK_VRAM_GB" != "0.0" ]]; then + optimization_args+=(--spark-vram "$DFLASH_SPARK_VRAM_GB") + fi + fi # Export KV cache type env vars for the C++ server to pick up (only when # explicitly requested: the per-axis envs override family defaults). if [[ -n "$CACHE_TYPE_K" ]]; then export DFLASH27B_KV_K="$CACHE_TYPE_K"; fi @@ -173,6 +210,7 @@ start_dflash_native_server() { --model-name "$MODEL_ID" \ "${ddtree_args[@]}" \ "${fa_args[@]}" \ + "${optimization_args[@]}" \ "${extra_args[@]}" \ > "$SERVER_LOG" 2>&1 & SERVER_PID=$! diff --git a/lucebox.sh b/lucebox.sh index 9dc879d11..7f409f037 100755 --- a/lucebox.sh +++ b/lucebox.sh @@ -237,6 +237,7 @@ CONFIG_HOME="${LUCEBOX_HOME:-$HOME/.lucebox}" : "${LUCEBOX_HOST_GPU_LIST_CSV:=}" : "${LUCEBOX_HOST_CUDA_VISIBLE_DEVICES:=}" : "${LUCEBOX_HOST_HIP_VISIBLE_DEVICES:=}" +: "${LUCEBOX_HOST_ROCR_VISIBLE_DEVICES:=}" # Tracks whether probe_host has actually run; pieces of the code that need # fresh host facts (e.g. cmd_check, cmd_serve) gate on this. Default 0. : "${_LUCEBOX_HOST_PROBED:=0}" @@ -444,7 +445,7 @@ probe_host() { # NVIDIA GPU remains the default backend (RTX 3090 + Strix → cuda12), but # recording the AMD companion prevents the APU from confusing readiness # reporting and lets an explicit rocm variant remain possible. - local amd_csv="" amd_rows="" + local amd_csv="" amd_rows="" amd_primary_idx="" if command -v amd-smi &>/dev/null; then amd_csv=$(amd-smi static --asic --vram --csv 2>/dev/null || echo "") if [ -n "$amd_csv" ]; then @@ -480,6 +481,7 @@ probe_host() { amd_primary=$(printf '%s\n' "$amd_rows" | sort -t'|' -k4,4nr | head -1) local amd_idx amd_name amd_arch amd_mem_mib IFS='|' read -r amd_idx amd_name amd_arch amd_mem_mib <<<"$amd_primary" + amd_primary_idx="$amd_idx" LUCEBOX_HOST_AMD_GPU_NAME="$amd_name" LUCEBOX_HOST_AMD_GPU_ARCH="$amd_arch" LUCEBOX_HOST_AMD_VRAM_GB=$((amd_mem_mib / 1024)) @@ -528,7 +530,20 @@ probe_host() { done # CUDA_VISIBLE_DEVICES from the caller's env (empty default = "all GPUs"). LUCEBOX_HOST_CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-}" + # The native server currently uses one primary device unless an advanced + # user explicitly supplies a visibility list. On R9700 + Strix systems, + # pinning the largest-VRAM card keeps the runtime aligned with the GPU that + # the readiness screen and optimizer selected, even when ROCm enumerates + # the integrated GPU first. LUCEBOX_HOST_HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-}" + LUCEBOX_HOST_ROCR_VISIBLE_DEVICES="${ROCR_VISIBLE_DEVICES:-}" + if [ -z "$LUCEBOX_HOST_HIP_VISIBLE_DEVICES" ] \ + && [ -z "$LUCEBOX_HOST_ROCR_VISIBLE_DEVICES" ]; then + # AMD recommends ROCR_VISIBLE_DEVICES for Linux. Use one isolation + # layer, not both, so a physical device index is never filtered and + # then interpreted again in a renumbered HIP-visible set. + LUCEBOX_HOST_ROCR_VISIBLE_DEVICES="$amd_primary_idx" + fi # OS / kernel identity. /etc/os-release is the freedesktop spec for # "what distro is this?" and we keep PRETTY_NAME verbatim (it already @@ -614,7 +629,7 @@ probe_host() { export LUCEBOX_HOST_OS_PRETTY LUCEBOX_HOST_KERNEL LUCEBOX_HOST_WSL_VERSION export LUCEBOX_HOST_NVIDIA_CTK_VERSION LUCEBOX_HOST_CPU_MODEL export LUCEBOX_HOST_GPU_LIST_CSV LUCEBOX_HOST_CUDA_VISIBLE_DEVICES - export LUCEBOX_HOST_HIP_VISIBLE_DEVICES + export LUCEBOX_HOST_HIP_VISIBLE_DEVICES LUCEBOX_HOST_ROCR_VISIBLE_DEVICES _LUCEBOX_HOST_PROBED=1 } @@ -768,6 +783,26 @@ _append_host_env() { done } +# Override the generic primary-GPU facts when the user deliberately selects +# the AMD backend on a mixed NVIDIA + AMD machine. probe_host keeps NVIDIA as +# the mixed-build default, but Automatic must tune the accelerator that will +# actually run the model. +_append_selected_backend_facts() { # usage: arrayname variant + # shellcheck disable=SC2178 + local -n _facts_arr="$1" + local variant="$2" + if _variant_is_rocm "$variant" && [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ]; then + _facts_arr+=( + -e "LUCEBOX_HOST_GPU_VENDOR=amd" + -e "LUCEBOX_HOST_GPU_NAME=$LUCEBOX_HOST_AMD_GPU_NAME" + -e "LUCEBOX_HOST_GPU_COUNT=$LUCEBOX_HOST_AMD_GPU_COUNT" + -e "LUCEBOX_HOST_VRAM_GB=$LUCEBOX_HOST_AMD_VRAM_GB" + -e "LUCEBOX_HOST_GPU_SM=$LUCEBOX_HOST_AMD_GPU_ARCH" + -e "LUCEBOX_HOST_GPU_LIST_CSV=$LUCEBOX_HOST_AMD_GPU_LIST_CSV" + ) + fi +} + # Append the LUCEBOX_* scalar overrides (image/variant/port/container/models) # plus the optional HF_TOKEN guard onto the named docker-argv array. Shared # by the docker-run (build_orchestrator_argv) and docker-exec @@ -802,8 +837,18 @@ _append_gpu_args() { # usage: _append_gpu_args arrayname variant --group-add render --security-opt seccomp=unconfined ) + if [ -n "${LUCEBOX_HOST_ROCR_VISIBLE_DEVICES:-}" ]; then + _gpu_arr+=(-e "ROCR_VISIBLE_DEVICES=$LUCEBOX_HOST_ROCR_VISIBLE_DEVICES") + elif [ -n "${LUCEBOX_HOST_HIP_VISIBLE_DEVICES:-}" ]; then + _gpu_arr+=( + -e "HIP_VISIBLE_DEVICES=$LUCEBOX_HOST_HIP_VISIBLE_DEVICES" + ) + fi else _gpu_arr+=(--gpus all) + if [ -n "${LUCEBOX_HOST_CUDA_VISIBLE_DEVICES:-}" ]; then + _gpu_arr+=(-e "CUDA_VISIBLE_DEVICES=$LUCEBOX_HOST_CUDA_VISIBLE_DEVICES") + fi fi } @@ -863,6 +908,7 @@ build_orchestrator_argv() { argv+=(-e "HOME=$CONFIG_HOME") # Host facts — Python side reads these instead of reprobing. _append_host_env argv + _append_selected_backend_facts argv "$variant" # User overrides for image/port/container/models scalars + HF_TOKEN. # Always exports the resolved models dir so the in-container CLI sees # the same path the wrapper mounts (the XDG default flows through too). @@ -990,6 +1036,7 @@ cmd_serve() { -e "HOME=$CONFIG_HOME") _append_gpu_args fallback_argv "$variant" _append_host_env fallback_argv + _append_selected_backend_facts fallback_argv "$variant" _append_scalar_env fallback_argv "$variant" fallback_argv+=("${IMAGE_BASE}:${variant}") _serve_and_track "${fallback_argv[@]}" @@ -1292,11 +1339,26 @@ _safe_model_relative_path() { return 0 } +_model_artifact_ready() { + local path="$1" + if [ -f "$path" ]; then + [ -s "$path" ] + return + fi + if [ -d "$path" ]; then + [ -n "$(find -L "$path" -maxdepth 4 -type f \ + \( -name '*.gguf' -o -name '*.safetensors' \) \ + -size +0c -print -quit 2>/dev/null)" ] + return + fi + return 1 +} + _selected_model_paths() { # Emit target path, draft path (or "none"), and model id on separate lines. # `models select` persists the filenames, so preset mapping is only a # compatibility fallback for hand-written config files. - local preset target_file draft_file target_path draft_path="none" + local preset target_file draft_file target_path draft_path="none" speculative_decode preset=$(_lucebox_config_get model.preset) target_file=$(_lucebox_config_get model.target_file) draft_file=$(_lucebox_config_get model.draft_file) @@ -1316,6 +1378,10 @@ _selected_model_paths() { gemma-4-31b) draft_file="gemma-4-31B-it-DFlash-q8_0.gguf" ;; esac fi + speculative_decode=$(_lucebox_config_get dflash.speculative_decode) + case "$speculative_decode" in + false|0|no|off) draft_file="" ;; + esac _safe_model_relative_path "$target_file" \ || die "no valid model is selected — run '$SCRIPT_NAME models select' first" target_path="$DEFAULT_MODELS_DIR/$target_file" @@ -1324,6 +1390,10 @@ _selected_model_paths() { || die "invalid model.draft_file in $(_lucebox_config_path)" draft_path="$DEFAULT_MODELS_DIR/draft/$draft_file" elif [ "$preset" = "laguna-xs.2" ] \ + && [ "$speculative_decode" != "false" ] \ + && [ "$speculative_decode" != "0" ] \ + && [ "$speculative_decode" != "no" ] \ + && [ "$speculative_decode" != "off" ] \ && [ -d "$DEFAULT_MODELS_DIR/draft/laguna-xs2-speculator" ]; then draft_path="$DEFAULT_MODELS_DIR/draft/laguna-xs2-speculator" fi @@ -1332,22 +1402,40 @@ _selected_model_paths() { _export_native_config() { local key env_name value + if [ -n "${LUCEBOX_HOST_ROCR_VISIBLE_DEVICES:-}" ]; then + export ROCR_VISIBLE_DEVICES="$LUCEBOX_HOST_ROCR_VISIBLE_DEVICES" + unset HIP_VISIBLE_DEVICES + elif [ -n "${LUCEBOX_HOST_HIP_VISIBLE_DEVICES:-}" ]; then + export HIP_VISIBLE_DEVICES="$LUCEBOX_HOST_HIP_VISIBLE_DEVICES" + unset ROCR_VISIBLE_DEVICES + fi + if [ -n "${LUCEBOX_HOST_CUDA_VISIBLE_DEVICES:-}" ]; then + export CUDA_VISIBLE_DEVICES="$LUCEBOX_HOST_CUDA_VISIBLE_DEVICES" + fi while IFS='|' read -r key env_name; do [ -n "$key" ] || continue value=$(_lucebox_config_get "$key") [ -n "$value" ] || continue case "$key" in - dflash.lazy|dflash.debug_thinking_logits) + dflash.speculative_decode|dflash.lazy|dflash.spark|dflash.debug_thinking_logits) case "$value" in true|1|yes|on) value=1 ;; *) value=0 ;; esac ;; + dflash.prefill_drafter) + case "$value" in + /opt/lucebox-hub/server/models/*) + value="$DEFAULT_MODELS_DIR/${value#/opt/lucebox-hub/server/models/}" + ;; + esac + ;; esac export "$env_name=$value" done <<'EOF' dflash.budget|DFLASH_BUDGET dflash.max_ctx|DFLASH_MAX_CTX +dflash.speculative_decode|DFLASH_SPECULATIVE_DECODE dflash.lazy|DFLASH_LAZY dflash.prefix_cache_slots|DFLASH_PREFIX_CACHE_SLOTS dflash.prefill_cache_slots|DFLASH_PREFILL_CACHE_SLOTS @@ -1357,6 +1445,11 @@ dflash.prefill_mode|DFLASH_PREFILL_MODE dflash.prefill_keep_ratio|DFLASH_PREFILL_KEEP dflash.prefill_threshold|DFLASH_PREFILL_THRESHOLD dflash.prefill_drafter|DFLASH_PREFILL_DRAFTER +dflash.kvflash|DFLASH_KVFLASH +dflash.kvflash_policy|DFLASH_KVFLASH_POLICY +dflash.kvflash_tau|DFLASH_KVFLASH_TAU +dflash.spark|DFLASH_SPARK +dflash.spark_vram_gb|DFLASH_SPARK_VRAM_GB dflash.think_max|DFLASH_THINK_MAX dflash.fa_window|DFLASH_FA_WINDOW dflash.think_soft_close_min_ratio|DFLASH_THINK_SOFT_CLOSE_MIN_RATIO @@ -1419,9 +1512,9 @@ cmd_native_serve() { target="${selected[0]:-}" draft="${selected[1]:-none}" model_id="${selected[2]:-lucebox}" - [ -f "$target" ] \ + _model_artifact_ready "$target" \ || die "selected target is not installed: $target — run '$SCRIPT_NAME models select'" - if [ "$draft" != "none" ] && [ ! -e "$draft" ]; then + if [ "$draft" != "none" ] && ! _model_artifact_ready "$draft"; then die "selected draft is not installed: $draft — run '$SCRIPT_NAME models select'" fi @@ -1481,15 +1574,19 @@ EOF mapfile -t selected < <(_selected_model_paths) target="${selected[0]:-}" draft="${selected[1]:-none}" - [ -f "$target" ] \ + _model_artifact_ready "$target" \ || die "selected target is not installed: $target — run '$SCRIPT_NAME models select'" + _export_native_config export REPO_DIR="$repo" export DFLASH_SERVER_BIN="$binary" export DFLASH_TARGET="$target" export TARGET="$target" export DFLASH_DRAFT="$draft" export DRAFT="$draft" + if [ "$draft" != "none" ]; then + export VERIFY_MODE="${VERIFY_MODE:-ddtree}" + fi local max_ctx budget max_ctx=$(_lucebox_config_get dflash.max_ctx) budget=$(_lucebox_config_get dflash.budget) @@ -1726,6 +1823,11 @@ cmd_check() { _row warn "rocm" "GPU detected, userspace version unavailable" fi _row 1 "amd gpu" "${LUCEBOX_HOST_AMD_GPU_COUNT} device(s); primary ${LUCEBOX_HOST_AMD_GPU_NAME} (${LUCEBOX_HOST_AMD_GPU_ARCH}, ${LUCEBOX_HOST_AMD_VRAM_GB} GB effective)" + if [ "$LUCEBOX_HOST_AMD_GPU_COUNT" -gt 1 ]; then + local selected_device + selected_device="${LUCEBOX_HOST_ROCR_VISIBLE_DEVICES:-${LUCEBOX_HOST_HIP_VISIBLE_DEVICES:-0}}" + _row warn "gpu placement" "primary device ${selected_device}; multi-GPU sharding is not automatic" + fi local amd_line amd_device_idx amd_device_name amd_device_arch amd_device_mem while IFS= read -r amd_line; do [ -n "$amd_line" ] || continue @@ -1848,6 +1950,7 @@ cmd_exec_in_container() { argv+=(-w /opt/lucebox-hub) argv+=(-e "HOME=$CONFIG_HOME") _append_host_env argv + _append_selected_backend_facts argv "$variant" _append_scalar_env argv "$variant" # The image has no top-level `lucebox` binary on PATH — that name only # works as the first arg to /opt/lucebox-hub/server/scripts/entrypoint.sh, @@ -2020,8 +2123,10 @@ cmd_setup() { fi printf '\n' - if _confirm "Use automatic GPU optimization?" 1; then + if _confirm "Enable recommended optimizations (may add a ~1.2 GB shared scorer)?" 1; then bash "$SCRIPT_PATH" optimize --yes || return $? + else + hint "The safe base profile is active; optional scorer-based features remain off." fi if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ]; then @@ -2075,14 +2180,73 @@ cmd_developer_menu() { done } +_optimization_summary() { + local mode model decode prefill kvflash spark draft_file active="" + mode=$(_lucebox_config_get autotune.mode) + model=$(_lucebox_config_get model.preset) + decode=$(_lucebox_config_get dflash.speculative_decode) + prefill=$(_lucebox_config_get dflash.prefill_mode) + kvflash=$(_lucebox_config_get dflash.kvflash) + spark=$(_lucebox_config_get dflash.spark) + + if [ -z "$mode" ]; then + if [ -n "$(_lucebox_config_get dflash.max_ctx)" ]; then + mode="custom" + else + mode="not configured" + fi + fi + case "$decode" in false|0|no|off) ;; *) + case "$model" in + qwen3.6-27b|gemma-4-26b|gemma-4-31b) + draft_file=$(_lucebox_config_get model.draft_file) + if [ -z "$draft_file" ]; then + case "$model" in + qwen3.6-27b) draft_file="dflash-draft-3.6-q4_k_m.gguf" ;; + gemma-4-26b) draft_file="gemma-4-26B-A4B-it-DFlash-q8_0.gguf" ;; + gemma-4-31b) draft_file="gemma-4-31B-it-DFlash-q8_0.gguf" ;; + esac + fi + [ -s "$DEFAULT_MODELS_DIR/draft/$draft_file" ] && active="DFlash" + ;; + laguna-xs.2) + _model_artifact_ready "$DEFAULT_MODELS_DIR/draft/laguna-xs2-speculator" \ + && active="DFlash" + ;; + esac + ;; + esac + if [ -n "$prefill" ] && [ "$prefill" != "off" ]; then + active="${active:+$active, }PFlash" + fi + if [ -n "$kvflash" ] && [ "$kvflash" != "off" ]; then + active="${active:+$active, }KVFlash" + fi + case "$spark" in true|1|yes|on) active="${active:+$active, }Spark" ;; esac + printf '%s (%s)' "$mode" "${active:-standard engine}" +} + cmd_menu() { - local choice model variant state repo_hint + local choice model variant state repo_hint optimization gpu_name gpu_count other_gpu while true; do ensure_probed model=$(_lucebox_config_get model.preset) model="${model:-not selected}" variant=$(pick_variant) state=$(_engine_state) + optimization=$(_optimization_summary) + gpu_name="${LUCEBOX_HOST_GPU_NAME:-not detected}" + gpu_count="$LUCEBOX_HOST_GPU_COUNT" + other_gpu="" + if _variant_is_rocm "$variant" && [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ]; then + gpu_name="${LUCEBOX_HOST_AMD_GPU_NAME:-AMD GPU}" + gpu_count="$LUCEBOX_HOST_AMD_GPU_COUNT" + [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ] \ + && other_gpu="${LUCEBOX_HOST_GPU_NAME:-NVIDIA GPU}" + elif [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ] \ + && [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ]; then + other_gpu="${LUCEBOX_HOST_AMD_GPU_NAME:-AMD GPU}" + fi if _find_repo_root >/dev/null 2>&1; then repo_hint="available" else @@ -2091,19 +2255,23 @@ cmd_menu() { _menu_clear print_logo - printf ' GPU: %s\n' "${LUCEBOX_HOST_GPU_NAME:-${LUCEBOX_HOST_AMD_GPU_NAME:-not detected}}" - if [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ] \ - && [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ]; then - printf ' Other GPU: %s\n' "${LUCEBOX_HOST_AMD_GPU_NAME:-AMD GPU}" + if [ "$gpu_count" -gt 1 ]; then + printf ' GPU: %s (primary of %s)\n' \ + "$gpu_name" "$gpu_count" + else + printf ' GPU: %s\n' "$gpu_name" + fi + if [ -n "$other_gpu" ]; then + printf ' Other GPU: %s\n' "$other_gpu" fi printf ' Backend: %s\n' "$variant" printf ' Model: %s\n' "$model" - printf ' Optimization: automatic\n' + printf ' Optimization: %s\n' "$optimization" printf ' Engine: %s\n\n' "$state" printf ' 1 Quick setup\n' printf ' 2 Choose or download a model\n' - printf ' 3 Apply automatic optimization\n' + printf ' 3 Review optimizations\n' printf ' 4 Start the inference engine\n' printf ' 5 Stop the inference engine\n' printf ' 6 Status\n' @@ -2173,7 +2341,7 @@ Provisioning + workloads (delegated to the in-container Python CLI): completion print shell completion script (bash / zsh / fish) models select numbered model picker; download + activate in one step models list / download / activate model presets - optimize apply safe hardware-aware inference defaults + optimize automatic or guided DFlash/PFlash/KVFlash/Spark setup config read / write keys in .lucebox/config.toml print-run print the docker-run command for the server diff --git a/lucebox/README.md b/lucebox/README.md index 6e5533fc6..321f2706e 100644 --- a/lucebox/README.md +++ b/lucebox/README.md @@ -9,7 +9,8 @@ which invokes the package in the appropriate container: lucebox setup # guided first run lucebox check lucebox models select # numbered picker + download + activate - lucebox optimize # safe hardware-aware defaults + lucebox optimize # review/apply the Automatic profile + lucebox optimize --advanced # override individual optimizations lucebox start The wrapper detects the host and selects CUDA for NVIDIA builds (including RTX @@ -19,6 +20,23 @@ optimization settings, and construction of the final server command. Host facts are passed through `LUCEBOX_HOST_*` environment variables so the container never has to guess the host configuration. +Automatic mode considers the selected model and the primary GPU. It shows a +plain-language decision for DFlash, PFlash, KVFlash, and Spark, prefers exact +all-GPU execution when it fits, and enables memory-saving paths only under +pressure. Advanced mode lets experienced users override those four product +choices while the CLI retains safe context, cache, and DDTree defaults. A +factory-preloaded Lucebox can ship the models and shared optimizer scorer, so a +buyer does not need to download anything during first setup. + +Spark is enabled automatically only when an MoE model is under GPU-memory +pressure and the machine reports at least 32 GB of host RAM; its cold experts +live in system memory. Lower-memory or unknown hosts keep Spark off unless an +advanced user explicitly opts in. + +Heterogeneous machines use one detected primary accelerator by default. The +CLI pins ROCm to the largest-VRAM AMD device and never sums memory across GPUs; +automatic multi-GPU layer sharding is deliberately not claimed by this release. + Inside a source checkout, the same wrapper also exposes `lucebox build`, `lucebox native`, and `lucebox harness` for contributors. Buyer installations keep using the prebuilt container, so no compiler or source checkout is needed. diff --git a/lucebox/src/lucebox/autotune.py b/lucebox/src/lucebox/autotune.py index 0baf0fe91..553a4c5ba 100644 --- a/lucebox/src/lucebox/autotune.py +++ b/lucebox/src/lucebox/autotune.py @@ -1,77 +1,285 @@ -"""Heuristic autotune: VRAM tier → DflashRuntime defaults. +"""Safe, explainable optimization planning for a Lucebox machine. -The recommended runtime is computed from HostFacts (VRAM, is_wsl) — stateless: -it takes HostFacts in and returns a fresh DflashRuntime. ``config.live_config`` -applies it so ``lucebox print-serve-argv`` / ``docker_run`` bake conservative -DFLASH_* defaults into the serve command for the detected VRAM tier. - -The empirical sweep + per-workload profiles (``lucebox autotune --sweep``) -live in a follow-up PR; this module keeps only the host-derived heuristic -that the serve path depends on. +The planner combines host facts, the selected model preset, and locally +installed optimization assets. It intentionally does not claim to benchmark +the workload: every automatic decision is a conservative rule backed by a +known engine capability, and the CLI prints the reason for every on/off choice. """ from __future__ import annotations -from lucebox.types import DflashRuntime, HostFacts +from dataclasses import dataclass, replace +from typing import Literal + +from lucebox.types import Config, DflashRuntime, HostFacts + + +@dataclass(frozen=True, slots=True) +class OptimizationDecision: + """One user-visible optimization decision and its explanation.""" + + name: str + enabled: bool + available: bool + reason: str + + +@dataclass(frozen=True, slots=True) +class OptimizationPlan: + """Resolved runtime plus the four product-level optimization decisions.""" + + runtime: DflashRuntime + model_name: str + dflash: OptimizationDecision + pflash: OptimizationDecision + kvflash: OptimizationDecision + spark: OptimizationDecision + needs_optimizer_drafter: bool = False + + @property + def decisions(self) -> tuple[OptimizationDecision, ...]: + return (self.dflash, self.pflash, self.kvflash, self.spark) + + @property + def active_names(self) -> tuple[str, ...]: + return tuple(item.name for item in self.decisions if item.enabled) + + +def _budget_for_host(host: HostFacts) -> int: + """Use the documented per-architecture DDTree sweet spot when known.""" + if host.is_wsl and 22 <= host.vram_gb < 32: + return 16 + if host.gpu_sm == "gfx1100": + return 8 + if host.gpu_sm == "120": + return 40 + return 22 def runtime_from_host(host: HostFacts) -> DflashRuntime: - """Pick a conservative DflashRuntime that 'should work' on this VRAM tier. - - Tiers (selected CUDA/ROCm device, baseline = Qwen3.6-27B Q4_K_M ~18 GB total): - <12 GB — too small for 27B; pick min ctx as a floor so a fallback - start at least gets an error from the daemon rather than - a silent OOM. - 12-21 — fits but tight; cap ctx. - 22-31 — 24 GB-class consumer flagships (3090/4090/5090/5090-Laptop). - 98 K with tq3_0 KV (~2 GB KV + ~18 GB model ≈ 20 GB). - Confirmed on bragi (RTX 5090 Laptop, 23 GB VRAM) 2026-05-31. - ≥32 — RTX 6000 Ada / A100 / H100 class. Full 128 K. - - Prefix cache remains an explicit sweep tunable, but the automatic - baseline keeps it off because tool prompts currently exercise a daemon - snapshot path that is not reliable with prefix slots enabled. - Empirically confirmed on bragi 2026-05-31: prefix_cache_slots=32 - caused -19pp regression on agent_recorded (23.1% vs 42.3% baseline). - 5 previously-passing cases regressed; 0 new cases unlocked. See - docs/experiments/qwen3.6-27b-prefix-cache-regression-bragi-2026-05-31.md. - - On `lazy`: the C++ server requires `--prefill-drafter` (and `--draft`) - to be set for `--lazy-draft` to take effect, and silently ignores it - otherwise (`--lazy-draft ignored: requires both --prefill-drafter and - --draft`). Since the heuristic path does NOT set `prefill_drafter`, - we default `lazy=False` here — "what we say" matches "what runs". - Users who explicitly opt in via config.toml will be warned at server - startup that the flag is being dropped (see entrypoint.sh). + """Pick conservative context/cache defaults from the selected GPU. + + These tiers target the Qwen3.6-27B Q4_K_M stack. ``automatic_plan`` adds + model capability and installed-asset decisions on top of this baseline. + Only the selected/primary GPU is counted; multiple cards are never summed + until the engine has an explicit layer-split plan. """ + budget = _budget_for_host(host) if host.vram_gb <= 0: - return DflashRuntime() # no VRAM signal — stick with class defaults - + return DflashRuntime(budget=budget) if host.vram_gb < 12: - return DflashRuntime(max_ctx=4096) + return DflashRuntime(budget=budget, max_ctx=4096) if host.vram_gb < 22: - return DflashRuntime(max_ctx=32768) + return DflashRuntime(budget=budget, max_ctx=32768) if host.vram_gb < 32: - # 22-31 GB cards. tq3_0 KV is required at 98K: model (~18-19 GB) + - # q8_0 KV at 98K (~5-6 GB) = 24-25 GB → OOM, while tq3_0 KV (~2 GB) - # leaves ~3 GB headroom. Confirmed on bragi (RTX 5090 Laptop, 23 GB - # VRAM) 2026-05-30 — q8_0 timed out on every 98K cell; all tq3_0 cells - # passed. Preset-size-aware capping (large models → 32K) lives with the - # autotune sweep in a follow-up PR. - if host.is_wsl: - # Bumped from max_ctx=65536 → 98304 on 2026-05-30 after the - # coding-agent-loop sweep on sindri proved 98K serves real - # 90K-token agentic prompts with ~3 GB VRAM headroom and no - # CUDA VMM failures. See - # docs/experiments/gemma4-26b-coding-agent-loop-sweep-2026-05-30.md. - # The original 65K cap cited unverified VMM failures — - # bisect history showed no commit reproducing them. - return DflashRuntime( - budget=16, max_ctx=98304, - cache_type_k="tq3_0", cache_type_v="tq3_0", + # Cache types stay model-owned. In particular, tq3_0 is unsafe for + # Laguna and was removed from the engine's automatic defaults in + # July 2026. The CLI must not silently reintroduce it here. + return DflashRuntime(budget=budget, max_ctx=98304) + return DflashRuntime(budget=budget, max_ctx=131072) + + +def draft_available(cfg: Config, preset: object) -> bool: + """Whether the selected preset's speculative decoder is installed.""" + from lucebox.download import local_artifact_present + + draft_file = cfg.model.draft_file or getattr(preset, "draft_file", None) + if draft_file: + return local_artifact_present(cfg.models_dir / "draft" / str(draft_file)) + speculator_dir = getattr(preset, "speculator_dir", None) + if not speculator_dir: + return False + root = cfg.models_dir / "draft" / str(speculator_dir) + if not root.is_dir(): + return False + try: + return any( + local_artifact_present(candidate) + for pattern in ("*.gguf", "*.safetensors") + for candidate in root.rglob(pattern) + ) + except OSError: + return False + + +def _cap_exact_context(runtime: DflashRuntime, *, headroom_gb: int, arch: str) -> DflashRuntime: + """Keep exact-cache families inside conservative primary-GPU headroom. + + Qwen families have an automatic bounded-residency path below; other + families currently keep their exact cache in Automatic mode. Those models + get a smaller context when their weights/draft leave little working room. + """ + if arch in {"qwen35", "qwen35moe"}: + return runtime + if headroom_gb < 4: + return replace(runtime, max_ctx=min(runtime.max_ctx, 8192)) + if headroom_gb < 7: + return replace(runtime, max_ctx=min(runtime.max_ctx, 32768)) + return runtime + + +def automatic_plan( + cfg: Config, + *, + optimizer_drafter_available: bool | None = None, +) -> OptimizationPlan: + """Resolve the recommended profile for the selected model and primary GPU. + + Automatic mode favors exact/full-cache execution whenever it fits. PFlash + is enabled only for its validated Qwen3.6 dense profile. KVFlash is enabled + only under real memory pressure and only on qwen-family paths with a tuned + scorer. Spark is enabled only for supported MoE models that would otherwise + leave too little VRAM headroom. Advanced mode can override these choices. + """ + from lucebox import download as download_mod + + runtime = runtime_from_host(cfg.host) + preset = download_mod.PRESETS.get(cfg.model.preset) + if optimizer_drafter_available is None: + optimizer_drafter_available = download_mod.optimizer_drafter_installed(cfg) + + if preset is None: + dflash = OptimizationDecision( + "DFlash", + False, + False, + "choose a model first; it activates when that model has a matching draft", + ) + unavailable = "choose a model first" + return OptimizationPlan( + runtime=runtime, + model_name=cfg.model.preset or "not selected", + dflash=dflash, + pflash=OptimizationDecision("PFlash", False, False, unavailable), + kvflash=OptimizationDecision("KVFlash", False, False, unavailable), + spark=OptimizationDecision("Spark", False, False, unavailable), + ) + + runtime = replace(runtime, max_ctx=min(runtime.max_ctx, preset.native_context)) + vram_known = cfg.host.vram_gb > 0 + headroom = cfg.host.vram_gb - preset.approx_total_gb if vram_known else 0 + if vram_known: + runtime = _cap_exact_context( + runtime, + headroom_gb=headroom, + arch=preset.architecture, + ) + + has_draft = draft_available(cfg, preset) + runtime = replace(runtime, speculative_decode=has_draft) + if has_draft: + dflash_reason = "matching speculative draft is available for this model" + elif preset.speculator_dir: + dflash_reason = "optional model speculator is not installed" + elif preset.has_draft: + dflash_reason = "matching speculative draft is not installed" + else: + dflash_reason = "no compatible speculative draft is published for this model" + dflash = OptimizationDecision("DFlash", has_draft, has_draft, dflash_reason) + + long_context = runtime.max_ctx >= 32768 + + # PFlash automatic mode is deliberately narrow: this is the profile with + # end-to-end long-context quality validation. Other architectures remain + # unavailable until the engine exposes a compatible compression path. + pflash_supported = preset.architecture == "qwen35" + pflash_available = pflash_supported and long_context + pflash_profile = pflash_available and vram_known + pflash_enabled = pflash_profile and optimizer_drafter_available + if pflash_enabled: + pflash_reason = "long prompts use the installed scorer above 32K tokens" + runtime = replace( + runtime, + prefill_mode="auto", + prefill_keep_ratio=0.10, + prefill_threshold=32768, + prefill_drafter=download_mod.optimizer_drafter_container_path(), + ) + elif pflash_profile: + pflash_reason = "shared 1.2 GB scorer is not installed" + elif not pflash_supported: + pflash_reason = "this model architecture has no PFlash compression path" + elif not long_context: + pflash_reason = "the safe context for this GPU is below PFlash's long-prompt threshold" + else: + pflash_reason = "the conservative automatic profile keeps full prefill for this model" + pflash = OptimizationDecision( + "PFlash", + pflash_enabled, + pflash_available, + pflash_reason, + ) + + # Bounded KV residency changes long-context retrieval semantics, so it is + # not a blanket speed toggle. Use it only when the selected stack leaves + # less than 5 GB for KV/compute and a tuned qwen-family policy is available. + kvflash_profile = ( + vram_known + and long_context + and headroom < 5 + and preset.architecture in {"qwen35", "qwen35moe"} + ) + kvflash_enabled = kvflash_profile and ( + optimizer_drafter_available or preset.architecture == "qwen35" + ) + if kvflash_enabled: + policy: Literal["drafter", "qk"] = "drafter" if optimizer_drafter_available else "qk" + runtime = replace(runtime, kvflash="auto", kvflash_policy=policy) + if optimizer_drafter_available and not runtime.prefill_drafter: + runtime = replace( + runtime, + prefill_drafter=download_mod.optimizer_drafter_container_path(), ) - return DflashRuntime( - max_ctx=98304, - cache_type_k="tq3_0", cache_type_v="tq3_0", + kvflash_reason = f"only {headroom} GB VRAM headroom; {policy} policy bounds KV residency" + elif kvflash_profile: + kvflash_reason = "memory pressure detected, but no quality-safe scorer is installed" + elif not long_context: + kvflash_reason = "the selected context does not need bounded KV residency" + elif preset.architecture not in {"qwen35", "qwen35moe"}: + kvflash_reason = "cross-tokenizer residency remains an Advanced choice" + else: + kvflash_reason = "full KV cache fits; exact full-cache execution is preferred" + kvflash = OptimizationDecision( + "KVFlash", + kvflash_enabled, + long_context, + kvflash_reason, + ) + + spark_available = preset.architecture in {"laguna", "qwen35moe"} + spark_pressure = spark_available and vram_known and headroom < 6 + # Cold experts live in host RAM. A 20–22 GB preset plus the OS, runtime, + # KV spill, and working buffers is not a safe automatic fit below 32 GB. + # Advanced mode can still let an informed operator opt in. + spark_host_ready = cfg.host.ram_gb >= 32 + spark_enabled = spark_pressure and spark_host_ready + runtime = replace(runtime, spark=spark_enabled) + if spark_enabled: + spark_reason = f"MoE weights leave {headroom} GB headroom; expert residency self-tunes" + elif spark_pressure and cfg.host.ram_gb <= 0: + spark_reason = "GPU memory is tight, but host RAM is unknown; automatic offload stays off" + elif spark_pressure: + spark_reason = ( + f"GPU memory is tight, but {cfg.host.ram_gb} GB host RAM is below " + "the 32 GB automatic offload floor" ) - return DflashRuntime(max_ctx=131072) + elif spark_available and vram_known: + spark_reason = "the model fits the primary GPU; all-GPU execution is faster and simpler" + elif spark_available: + spark_reason = "GPU memory is unknown; automatic mode will not assume offload" + else: + spark_reason = "this is not a Spark-compatible MoE architecture" + spark = OptimizationDecision("Spark", spark_enabled, spark_available, spark_reason) + + needs_optimizer_drafter = not optimizer_drafter_available and ( + pflash_profile or (kvflash_profile and preset.architecture == "qwen35moe") + ) + return OptimizationPlan( + runtime=runtime, + model_name=preset.name, + dflash=dflash, + pflash=pflash, + kvflash=kvflash, + spark=spark, + needs_optimizer_drafter=needs_optimizer_drafter, + ) diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index ff11ebf8c..c940a157e 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -18,7 +18,7 @@ import sys from dataclasses import replace -from typing import Annotated +from typing import Annotated, Literal import typer from rich.console import Console @@ -33,7 +33,7 @@ from lucebox import __version__ from lucebox.config import config_get, config_set, config_unset, live_config from lucebox.host_facts import from_env -from lucebox.types import Config +from lucebox.types import Config, DflashRuntime, ModelMeta app = typer.Typer( name="lucebox", @@ -115,6 +115,36 @@ def _server_spec() -> docker_run.DockerRunSpec: raise typer.Exit(code=2) from exc +def _active_optimization_names(cfg: Config) -> list[str]: + """Return the product-level features represented by the loaded config.""" + active: list[str] = [] + preset = download_mod.PRESETS.get(cfg.model.preset) + if ( + cfg.dflash.speculative_decode + and preset is not None + and autotune_mod.draft_available(cfg, preset) + ): + active.append("DFlash") + if cfg.dflash.prefill_mode != "off": + active.append("PFlash") + if cfg.dflash.kvflash != "off": + active.append("KVFlash") + if cfg.dflash.spark: + active.append("Spark") + return active + + +def _optimization_label(cfg: Config) -> str: + mode = config_mod.optimization_mode() + mode_label = { + "automatic": "Automatic", + "custom": "Custom", + "unconfigured": "Not configured", + }[mode] + active = _active_optimization_names(cfg) + return f"{mode_label} ({', '.join(active) if active else 'standard engine'})" + + def _package_menu() -> None: """Small in-package menu for direct installs and contributor workflows. @@ -127,9 +157,9 @@ def _package_menu() -> None: cfg = _load_or_build() active = cfg.model.preset or "not selected" console.print(f"Model: [bold]{escape(active)}[/bold]") - console.print("Optimization: [bold]Automatic[/bold] (safe defaults for this GPU)\n") + console.print(f"Optimization: [bold]{escape(_optimization_label(cfg))}[/bold]\n") console.print(" [bold cyan]1[/bold cyan] Choose or download a model") - console.print(" [bold cyan]2[/bold cyan] Apply automatic optimization") + console.print(" [bold cyan]2[/bold cyan] Review optimizations") console.print(" [bold cyan]3[/bold cyan] Show configuration") console.print(" [bold cyan]4[/bold cyan] Show Docker launch command") console.print(" [bold cyan]q[/bold cyan] Quit") @@ -296,13 +326,28 @@ def _activate_preset(cfg: Config, preset: download_mod.ModelPreset) -> None: # preset is explicitly target-only. config_unset("model.draft_file") console.print(f"[green]Activated:[/green] model.preset = {preset.name}") - # Never clobber a user-edited [dflash] section. Explicit profile resets - # go through `lucebox optimize` instead. - if config_mod.seed_dflash_from_host(cfg.host): + selected_cfg = replace( + cfg, + model=ModelMeta( + preset=preset.name, + target_file=preset.target_file, + draft_file=preset.draft_file or "", + ), + ) + # Automatic profiles follow the selected model. A user-owned custom + # profile is preserved and called out instead of being silently reset. + if config_mod.seed_optimization_from_config(selected_cfg): max_ctx = config_get("dflash.max_ctx")["dflash.max_ctx"][0] + plan = autotune_mod.automatic_plan(selected_cfg) + active = ", ".join(plan.active_names) or "standard engine" + console.print( + f"[green]Optimized:[/green] Automatic profile for {preset.name} " + f"({active}; max_ctx={max_ctx})" + ) + else: console.print( - f"[green]Auto-tuned:[/green] VRAM-tier DFLASH_* defaults " - f"(max_ctx={max_ctx}) written to config.toml" + "[yellow]Custom optimization profile kept.[/yellow] " + "Run `lucebox optimize` to review it for this model." ) @@ -479,34 +524,215 @@ def models_download( _activate_preset(cfg, pres) +def _render_optimization_plan( + plan: autotune_mod.OptimizationPlan, + *, + title: str = "Automatic optimization (recommended)", +) -> None: + console.print(f"[bold]{title}[/bold]") + console.print(f"Model: {escape(plan.model_name)}") + table = Table(show_header=True, box=None, pad_edge=False) + table.add_column("Optimization", style="bold") + table.add_column("State") + table.add_column("Why") + for decision in plan.decisions: + if decision.enabled: + state = "[green]ON[/green]" + elif decision.available: + state = "[dim]off[/dim]" + else: + state = "[dim]not available[/dim]" + table.add_row(decision.name, state, escape(decision.reason)) + console.print(table) + cache = plan.runtime.cache_type_k or "model default" + console.print( + f"Context: [bold]{plan.runtime.max_ctx:,}[/bold] tokens · " + f"KV format: [bold]{cache}[/bold] · " + f"DFlash budget: [bold]{plan.runtime.budget}[/bold]" + ) + + +def _ensure_optimizer_drafter(cfg: Config, *, assume_yes: bool = False) -> bool: + """Offer the one shared scorer asset and degrade cleanly when offline.""" + if download_mod.optimizer_drafter_installed(cfg): + return True + question = ( + f"Install the shared PFlash/KVFlash scorer " + f"(~{download_mod.OPTIMIZER_DRAFTER_APPROX_GB:g} GB)?" + ) + if not assume_yes and not typer.confirm(question, default=True): + return False + console.print("[bold]Installing the shared long-context scorer…[/bold]") + if download_mod.download_optimizer_drafter(cfg) != 0: + console.print( + "[yellow]Continuing without the scorer; the rest of the automatic " + "profile is still safe.[/yellow]" + ) + return False + if not download_mod.optimizer_drafter_installed(cfg): + console.print("[yellow]The scorer download completed but its file is missing.[/yellow]") + return False + console.print("[green]Shared optimizer installed.[/green]") + return True + + +def _customize_runtime( + cfg: Config, + plan: autotune_mod.OptimizationPlan, +) -> DflashRuntime: + """Prompt only for product-level choices; keep low-level knobs automatic.""" + preset = download_mod.PRESETS.get(cfg.model.preset) + runtime = plan.runtime + console.print("\n[bold]Customize optimizations[/bold]") + console.print("[dim]Context, cache format, and budgets remain hardware-tuned.[/dim]\n") + + if plan.dflash.available: + speculative_decode = typer.confirm( + "Enable DFlash speculative decode?", default=plan.dflash.enabled + ) + else: + speculative_decode = False + console.print(f"DFlash: [dim]unavailable — {escape(plan.dflash.reason)}[/dim]") + + pflash = False + if plan.pflash.available: + pflash = typer.confirm( + "Enable PFlash automatically for long prompts?", default=plan.pflash.enabled + ) + if pflash and not _ensure_optimizer_drafter(cfg): + pflash = False + console.print("[yellow]PFlash left off because its scorer is unavailable.[/yellow]") + else: + console.print(f"PFlash: [dim]unavailable — {escape(plan.pflash.reason)}[/dim]") + + kvflash = False + kvflash_policy: Literal["drafter", "lru", "qk"] = "drafter" + if preset is not None and plan.kvflash.available: + kvflash = typer.confirm( + "Enable KVFlash bounded long-context memory?", default=plan.kvflash.enabled + ) + else: + console.print(f"KVFlash: [dim]unavailable — {escape(plan.kvflash.reason)}[/dim]") + if kvflash: + has_scorer = download_mod.optimizer_drafter_installed(cfg) + if has_scorer: + policy_choices = "drafter" + if preset is not None and preset.architecture == "qwen35": + policy_choices += "/qk/lru" + else: + policy_choices += "/lru" + answer = typer.prompt("KVFlash policy", default="drafter").strip().lower() + if answer not in policy_choices.split("/"): + fallback_policy = policy_choices.split("/")[0] + console.print(f"[yellow]Unknown policy; using {fallback_policy}.[/yellow]") + answer = fallback_policy + if answer == "qk": + kvflash_policy = "qk" + elif answer == "lru": + kvflash_policy = "lru" + else: + kvflash_policy = "drafter" + elif preset is not None and preset.architecture == "qwen35": + kvflash_policy = "qk" + console.print("KVFlash policy: [bold]qk[/bold] (no extra scorer required)") + else: + console.print( + "[yellow]KVFlash would use recency-only LRU without the shared scorer; " + "older context can be evicted.[/yellow]" + ) + if typer.confirm("Install the scorer instead?", default=True): + if _ensure_optimizer_drafter(cfg): + kvflash_policy = "drafter" + else: + kvflash_policy = "lru" + else: + kvflash_policy = "lru" + + spark = False + if plan.spark.available: + spark = typer.confirm( + "Enable Spark self-tuning MoE expert offload?", default=plan.spark.enabled + ) + else: + console.print(f"Spark: [dim]unavailable — {escape(plan.spark.reason)}[/dim]") + + uses_scorer = pflash or (kvflash and kvflash_policy == "drafter") + prefill_drafter = download_mod.optimizer_drafter_container_path() if uses_scorer else "" + if kvflash and runtime.fa_window > 0: + console.print("[dim]KVFlash selected: disabling the incompatible FA window.[/dim]") + + return replace( + runtime, + speculative_decode=speculative_decode, + lazy=False if not speculative_decode else runtime.lazy, + prefill_mode="auto" if pflash else "off", + prefill_keep_ratio=0.10 if pflash else runtime.prefill_keep_ratio, + prefill_threshold=32768, + prefill_drafter=prefill_drafter, + kvflash="auto" if kvflash else "off", + kvflash_policy=kvflash_policy, + spark=spark, + spark_vram_gb=0.0, + fa_window=0 if kvflash else runtime.fa_window, + ) + + @app.command() def optimize( yes: Annotated[ bool, - typer.Option("--yes", "-y", help="Apply without confirmation."), + typer.Option("--yes", "-y", help="Apply Automatic without prompts."), + ] = False, + advanced: Annotated[ + bool, + typer.Option("--advanced", help="Review DFlash/PFlash/KVFlash/Spark choices."), ] = False, ) -> None: - """Apply the recommended hardware-aware inference profile. + """Choose and apply a model-aware, hardware-aware optimization profile.""" + if yes and advanced: + error_console.print( + "[red]--advanced is interactive and cannot be combined with --yes[/red]" + ) + raise typer.Exit(code=2) - Stable CUDA/ROCm fast paths remain enabled by the engine itself. This - profile selects a safe context size, cache format, and DFlash budget from - detected VRAM while leaving experimental features off. - """ cfg = _load_or_build() - recommended = autotune_mod.runtime_from_host(cfg.host) - console.print("[bold]Automatic optimization (recommended)[/bold]") - console.print(" DFlash decode automatic when the model has a matching draft") - console.print(" GPU fast paths enabled by the engine on CUDA and ROCm") - console.print(f" Maximum context {recommended.max_ctx:,} tokens") - cache = recommended.cache_type_k or "model default" - console.print(f" KV cache {cache}") - console.print(" Experimental paths off") - if not yes and not typer.confirm("Apply this profile?", default=True): + plan = autotune_mod.automatic_plan(cfg) + _render_optimization_plan(plan) + + custom = advanced + if not yes and not advanced: + console.print("\n [bold cyan]1[/bold cyan] Apply Automatic (recommended)") + console.print(" [bold cyan]2[/bold cyan] Customize") + console.print(" [bold cyan]q[/bold cyan] Cancel") + choice = typer.prompt("\nChoose", default="1").strip().lower() + if choice in {"q", "quit", "cancel"}: + console.print("[dim]Optimization unchanged.[/dim]") + return + if choice not in {"1", "2"}: + error_console.print("[red]Choose 1, 2, or q.[/red]") + raise typer.Exit(code=2) + custom = choice == "2" + + if custom: + runtime = _customize_runtime(cfg, plan) + if not typer.confirm("Apply this custom profile?", default=True): + console.print("[dim]Optimization unchanged.[/dim]") + return + config_mod.write_optimization_runtime(runtime, mode="custom", source="guided") + console.print("[green]Custom optimization profile applied.[/green]") + return + + if plan.needs_optimizer_drafter and _ensure_optimizer_drafter(cfg, assume_yes=yes): + plan = autotune_mod.automatic_plan(cfg) + _render_optimization_plan(plan, title="Updated automatic plan") + if not yes and not typer.confirm("Apply this automatic profile?", default=True): console.print("[dim]Optimization unchanged.[/dim]") return - config_mod.seed_dflash_from_host(cfg.host, force=True) + config_mod.write_optimization_runtime(plan.runtime) console.print("[green]Automatic optimization applied.[/green]") - console.print("[dim]Advanced users can still use `lucebox config set dflash.KEY=VALUE`.[/dim]") + console.print( + "[dim]Run `lucebox optimize --advanced` anytime to review individual features.[/dim]" + ) @app.command() diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py index 778d81385..e51faeffb 100644 --- a/lucebox/src/lucebox/config.py +++ b/lucebox/src/lucebox/config.py @@ -55,6 +55,7 @@ def default_config_path() -> Path: # ── dotted-key registry ──────────────────────────────────────────────────── + def _cast_prefill_mode(v: Any) -> Literal["off", "auto", "always"]: s = str(v) if s not in {"off", "auto", "always"}: @@ -62,6 +63,42 @@ def _cast_prefill_mode(v: Any) -> Literal["off", "auto", "always"]: return cast(Literal["off", "auto", "always"], s) +def _cast_kvflash(v: Any) -> str: + value = str(v).strip().lower() + if value in {"off", "auto"}: + return value + try: + pool_tokens = int(value) + except ValueError as exc: + raise ValueError( + f"kvflash must be off, auto, or a positive token count, got {value!r}" + ) from exc + if pool_tokens <= 0: + raise ValueError(f"kvflash token count must be positive, got {value!r}") + return str(pool_tokens) + + +def _cast_kvflash_policy(v: Any) -> Literal["drafter", "lru", "qk"]: + value = str(v).strip().lower() + if value not in {"drafter", "lru", "qk"}: + raise ValueError(f"kvflash_policy must be drafter/lru/qk, got {value!r}") + return cast(Literal["drafter", "lru", "qk"], value) + + +def _cast_positive_int(v: Any) -> int: + value = int(v) + if value <= 0: + raise ValueError(f"value must be positive, got {value!r}") + return value + + +def _cast_nonnegative_float(v: Any) -> float: + value = float(v) + if value < 0.0: + raise ValueError(f"value must be zero or positive, got {value!r}") + return value + + def _cast_bool(v: Any) -> bool: """Strict-ish boolean coercion for config values. @@ -112,10 +149,7 @@ def _cast_model_relative_path(v: Any) -> str: def _cast_prefill_keep_ratio(v: Any) -> float: value = float(v) if not 0.0 < value <= 1.0: - raise ValueError( - "prefill_keep_ratio must be in the interval (0.0, 1.0], " - f"got {value!r}" - ) + raise ValueError(f"prefill_keep_ratio must be in the interval (0.0, 1.0], got {value!r}") return value @@ -123,8 +157,7 @@ def _cast_think_soft_close_min_ratio(v: Any) -> float: value = float(v) if not 0.0 <= value <= 1.0: raise ValueError( - "think_soft_close_min_ratio must be in the interval [0.0, 1.0], " - f"got {value!r}" + f"think_soft_close_min_ratio must be in the interval [0.0, 1.0], got {value!r}" ) return value @@ -142,6 +175,7 @@ def _cast_think_soft_close_min_ratio(v: Any) -> float: "model.preset": (("model", "preset"), str), "model.target_file": (("model", "target_file"), _cast_model_relative_path), "model.draft_file": (("model", "draft_file"), _cast_model_relative_path), + "dflash.speculative_decode": (("dflash", "speculative_decode"), _cast_bool), "dflash.budget": (("dflash", "budget"), int), "dflash.max_ctx": (("dflash", "max_ctx"), int), "dflash.lazy": (("dflash", "lazy"), _cast_bool), @@ -156,6 +190,11 @@ def _cast_think_soft_close_min_ratio(v: Any) -> float: ), "dflash.prefill_threshold": (("dflash", "prefill_threshold"), int), "dflash.prefill_drafter": (("dflash", "prefill_drafter"), str), + "dflash.kvflash": (("dflash", "kvflash"), _cast_kvflash), + "dflash.kvflash_policy": (("dflash", "kvflash_policy"), _cast_kvflash_policy), + "dflash.kvflash_tau": (("dflash", "kvflash_tau"), _cast_positive_int), + "dflash.spark": (("dflash", "spark"), _cast_bool), + "dflash.spark_vram_gb": (("dflash", "spark_vram_gb"), _cast_nonnegative_float), "dflash.think_max": (("dflash", "think_max"), int), "dflash.fa_window": (("dflash", "fa_window"), int), "dflash.think_soft_close_min_ratio": ( @@ -239,9 +278,14 @@ def load_doc(path: Path | None = None) -> dict[str, Any]: _LEGACY_KEY_MAP: dict[str, tuple[str, str, Callable[[str], Any]]] = { "DFLASH_BUDGET": ("dflash", "budget", int), "DFLASH_MAX_CTX": ("dflash", "max_ctx", int), - "DFLASH_LAZY": ("dflash", "lazy", - lambda v: str(v).strip().lower() in ("1", "true", "yes", "on")), + "DFLASH_LAZY": ( + "dflash", + "lazy", + lambda v: str(v).strip().lower() in ("1", "true", "yes", "on"), + ), "DFLASH_PREFIX_CACHE_SLOTS": ("dflash", "prefix_cache_slots", int), + "DFLASH_KVFLASH": ("dflash", "kvflash", _cast_kvflash), + "DFLASH_SPARK": ("dflash", "spark", _cast_bool), "DFLASH_PORT": ("runtime", "port", int), "LUCEBOX_VARIANT": ("image", "variant", str), "LUCEBOX_IMAGE": ("image", "registry", str), @@ -285,6 +329,7 @@ def _from_dict(raw: dict[str, Any]) -> Config: df = raw.get("dflash", {}) dflash = DflashRuntime( + speculative_decode=_cast_bool(df.get("speculative_decode", True)), budget=int(df.get("budget", 22)), max_ctx=int(df.get("max_ctx", 16384)), lazy=_cast_bool(df.get("lazy", False)), @@ -293,11 +338,14 @@ def _from_dict(raw: dict[str, Any]) -> Config: cache_type_k=str(df.get("cache_type_k", "")), cache_type_v=str(df.get("cache_type_v", "")), prefill_mode=_cast_prefill_mode(df.get("prefill_mode", "off")), - prefill_keep_ratio=_cast_prefill_keep_ratio( - df.get("prefill_keep_ratio", 0.05) - ), + prefill_keep_ratio=_cast_prefill_keep_ratio(df.get("prefill_keep_ratio", 0.05)), prefill_threshold=int(df.get("prefill_threshold", 32000)), prefill_drafter=str(df.get("prefill_drafter", "")), + kvflash=_cast_kvflash(df.get("kvflash", "off")), + kvflash_policy=_cast_kvflash_policy(df.get("kvflash_policy", "drafter")), + kvflash_tau=_cast_positive_int(df.get("kvflash_tau", 64)), + spark=_cast_bool(df.get("spark", False)), + spark_vram_gb=_cast_nonnegative_float(df.get("spark_vram_gb", 0.0)), think_max=int(df.get("think_max", 15488)), fa_window=int(df.get("fa_window", 0)), think_soft_close_min_ratio=_cast_think_soft_close_min_ratio( @@ -413,8 +461,6 @@ def seed_dflash_from_host( when there is no config.toml at all, which stops being true the moment a model is activated — so the heuristic is persisted here instead. """ - from datetime import datetime - import lucebox.autotune as autotune_mod path = path or default_config_path() @@ -432,15 +478,79 @@ def seed_dflash_from_host( # reset to the current hardware-derived defaults. doc.pop("dflash", None) runtime = autotune_mod.runtime_from_host(host) + _write_runtime_doc(doc, runtime, mode="automatic", source="heuristic") + path.parent.mkdir(parents=True, exist_ok=True) + _atomic_write_doc(path, doc) + return True + + +def _write_runtime_doc( + doc: dict[str, Any], + runtime: DflashRuntime, + *, + mode: str, + source: str, +) -> None: + """Replace ``[dflash]`` with one coherent resolved profile.""" + from datetime import datetime + + doc.pop("dflash", None) for field, value in asdict(runtime).items(): _doc_set(doc, "dflash", field, _value_to_toml(value)) - _doc_set(doc, "autotune", "source", "heuristic") + _doc_set(doc, "autotune", "mode", mode) + _doc_set(doc, "autotune", "source", source) _doc_set( - doc, "autotune", "timestamp", + doc, + "autotune", + "timestamp", datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"), ) + + +def write_optimization_runtime( + runtime: DflashRuntime, + *, + path: Path | None = None, + mode: str = "automatic", + source: str = "model+hardware", +) -> None: + """Atomically persist a complete automatic or custom optimization profile.""" + if mode not in {"automatic", "custom"}: + raise ValueError(f"optimization mode must be automatic or custom, got {mode!r}") + path = path or default_config_path() + doc = load_doc(path) + _write_runtime_doc(doc, runtime, mode=mode, source=source) path.parent.mkdir(parents=True, exist_ok=True) _atomic_write_doc(path, doc) + + +def optimization_mode(*, path: Path | None = None) -> str: + """Return automatic/custom/unconfigured for menu and model-switch behavior.""" + doc = load_doc(path) + auto = doc.get("autotune", {}) + if isinstance(auto, dict) and auto.get("mode") in {"automatic", "custom"}: + return str(auto["mode"]) + if "dflash" not in doc: + return "unconfigured" + # Existing profiles predate mode metadata and may contain hand tuning. + return "custom" + + +def seed_optimization_from_config( + cfg: Config, + *, + path: Path | None = None, + force: bool = False, +) -> bool: + """Apply a model+hardware plan unless the user owns a custom profile.""" + import lucebox.autotune as autotune_mod + + path = path or default_config_path() + mode = optimization_mode(path=path) + if not force and mode == "custom": + return False + plan = autotune_mod.automatic_plan(cfg) + write_optimization_runtime(plan.runtime, path=path) return True @@ -485,6 +595,20 @@ def config_set(key: str, value: Any, *, path: Path | None = None) -> None: path = path or default_config_path() doc = load_doc(path) if path.exists() else {} _doc_set(doc, section, field, _value_to_toml(cast_value)) + if section == "dflash": + # Validate the complete profile before replacing the file. Individual + # fields can be valid while their combination is not (for example, + # KVFlash and a finite FA window are mutually exclusive). + try: + _from_dict({"dflash": doc.get("dflash", {})}) + except (TypeError, ValueError) as exc: + raise ValueError(f"invalid dflash profile after setting {key}: {exc}") from exc + # A direct low-level edit transfers ownership from the automatic + # planner to the user. Model changes will preserve it until they pick + # Automatic again in ``lucebox optimize``. + _doc_set(doc, "autotune", "mode", "custom") + _doc_set(doc, "autotune", "source", "manual") + _doc_unset(doc, "autotune", "timestamp") path.parent.mkdir(parents=True, exist_ok=True) _atomic_write_doc(path, doc) @@ -501,6 +625,13 @@ def config_unset(key: str, *, path: Path | None = None) -> bool: doc = load_doc(path) changed = _doc_unset(doc, section, field) if changed: + if section == "dflash": + if "dflash" in doc: + _doc_set(doc, "autotune", "mode", "custom") + _doc_set(doc, "autotune", "source", "manual") + _doc_unset(doc, "autotune", "timestamp") + else: + doc.pop("autotune", None) # Leave the file in place even when empty — `config set` will # repopulate; deleting would surprise users who expect their # config dir to exist. @@ -544,9 +675,7 @@ def overlay_env(cfg: Config) -> Config: image=os.environ.get("LUCEBOX_IMAGE", cfg.image), container_name=os.environ.get("LUCEBOX_CONTAINER", cfg.container_name), port=_cast_port(os.environ.get("LUCEBOX_PORT", str(cfg.port))), - models_dir=Path( - _cast_models_dir(os.environ.get("LUCEBOX_MODELS", str(cfg.models_dir))) - ), + models_dir=Path(_cast_models_dir(os.environ.get("LUCEBOX_MODELS", str(cfg.models_dir)))), ) diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index 1568fce02..7e866fb78 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -112,9 +112,7 @@ def _resolve_model_files(cfg: Config) -> tuple[str, str, str]: def _runtime_volumes(cfg: Config) -> tuple[BindMount, ...]: """Mount only the writable application data needed by the server.""" models = str(cfg.models_dir.absolute()) - config_home = Path( - os.environ.get("LUCEBOX_HOME") or Path.home() / ".lucebox" - ).absolute() + config_home = Path(os.environ.get("LUCEBOX_HOME") or Path.home() / ".lucebox").absolute() return ( BindMount(models, _CONTAINER_MODELS), BindMount(str(config_home), str(config_home)), @@ -266,9 +264,7 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: # LUCEBOX_HOST_* first so they ride out front in the rendered argv, # making it obvious in `print-run` output what host facts get forwarded. env: list[tuple[str, str]] = list(_host_facts_env()) - config_home = str( - Path(os.environ.get("LUCEBOX_HOME") or Path.home() / ".lucebox").absolute() - ) + config_home = str(Path(os.environ.get("LUCEBOX_HOME") or Path.home() / ".lucebox").absolute()) env += [ ("LUCEBOX_HOME", config_home), ("DFLASH_BUDGET", str(cfg.dflash.budget)), @@ -299,7 +295,9 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: mounts=volumes, ) env.append(("DFLASH_TARGET", target_path)) - if draft_file: + if not cfg.dflash.speculative_decode: + env.append(("DFLASH_DRAFT", "/opt/lucebox-hub/server/models/.lucebox-no-draft")) + elif draft_file: draft_path = _selected_model_path( cfg, draft_file, @@ -333,14 +331,24 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: env.append(("DFLASH_CACHE_TYPE_K", cfg.dflash.cache_type_k)) if cfg.dflash.cache_type_v: env.append(("DFLASH_CACHE_TYPE_V", cfg.dflash.cache_type_v)) + if cfg.dflash.prefill_drafter: + env.append(("DFLASH_PREFILL_DRAFTER", cfg.dflash.prefill_drafter)) if cfg.dflash.prefill_mode != "off": env += [ ("DFLASH_PREFILL_MODE", cfg.dflash.prefill_mode), ("DFLASH_PREFILL_KEEP", str(cfg.dflash.prefill_keep_ratio)), ("DFLASH_PREFILL_THRESHOLD", str(cfg.dflash.prefill_threshold)), ] - if cfg.dflash.prefill_drafter: - env.append(("DFLASH_PREFILL_DRAFTER", cfg.dflash.prefill_drafter)) + if cfg.dflash.kvflash != "off": + env += [ + ("DFLASH_KVFLASH", cfg.dflash.kvflash), + ("DFLASH_KVFLASH_POLICY", cfg.dflash.kvflash_policy), + ("DFLASH_KVFLASH_TAU", str(cfg.dflash.kvflash_tau)), + ] + if cfg.dflash.spark: + env.append(("DFLASH_SPARK", "1")) + if cfg.dflash.spark_vram_gb > 0.0: + env.append(("DFLASH_SPARK_VRAM_GB", f"{cfg.dflash.spark_vram_gb:g}")) # fa_window=0 is the server's own default (full attention); only emit # the env when the operator has selected a sparse decode window. The # entrypoint mirrors this guard so an unset env reproduces the @@ -351,10 +359,12 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: # to pre-PR-#326 behavior). Emit only when nonzero to keep the # docker env minimal and mirror the entrypoint's `case` guard. if cfg.dflash.think_soft_close_min_ratio > 0.0: - env.append(( - "DFLASH_THINK_SOFT_CLOSE_MIN_RATIO", - f"{cfg.dflash.think_soft_close_min_ratio:g}", - )) + env.append( + ( + "DFLASH_THINK_SOFT_CLOSE_MIN_RATIO", + f"{cfg.dflash.think_soft_close_min_ratio:g}", + ) + ) if cfg.dflash.debug_thinking_logits: env.append(("DFLASH_DEBUG_THINKING_LOGITS", "1")) @@ -370,6 +380,22 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: else: gpu_vendor = cfg.host.gpu_vendor if cfg.host.gpu_vendor != "none" else "nvidia" + # Keep execution on the same primary device used by host detection and + # automatic tuning. This matters on an R9700 + Strix build where ROCm may + # enumerate the integrated GPU before the larger discrete card. Explicit + # CUDA/ROCR/HIP visibility supplied by an advanced user remains authoritative. + if gpu_vendor == "amd": + rocr_visible = os.environ.get("LUCEBOX_HOST_ROCR_VISIBLE_DEVICES", "").strip() + hip_visible = os.environ.get("LUCEBOX_HOST_HIP_VISIBLE_DEVICES", "").strip() + if rocr_visible: + env.append(("ROCR_VISIBLE_DEVICES", rocr_visible)) + elif hip_visible: + env.append(("HIP_VISIBLE_DEVICES", hip_visible)) + elif gpu_vendor == "nvidia": + visible = os.environ.get("LUCEBOX_HOST_CUDA_VISIBLE_DEVICES", "").strip() + if visible: + env.append(("CUDA_VISIBLE_DEVICES", visible)) + return DockerRunSpec( image=f"{cfg.image}:{cfg.variant}", name=cfg.container_name, diff --git a/lucebox/src/lucebox/download.py b/lucebox/src/lucebox/download.py index 43aeb83ca..b39bd728a 100644 --- a/lucebox/src/lucebox/download.py +++ b/lucebox/src/lucebox/download.py @@ -87,6 +87,8 @@ class ModelPreset: draft_repo: str | None draft_file: str | None approx_total_gb: int + architecture: str + native_context: int description: str = "" speculator_dir: str | None = None @@ -107,6 +109,8 @@ def has_draft(self) -> bool: draft_repo="spiritbuun/Qwen3.6-27B-DFlash-GGUF", draft_file="dflash-draft-3.6-q4_k_m.gguf", approx_total_gb=17, + architecture="qwen35", + native_context=262144, description="Qwen3.6 27B dense (Q4_K_M) + Qwen3.6 DFlash draft. Lucebox default.", ), "gemma-4-26b": ModelPreset( @@ -116,6 +120,8 @@ def has_draft(self) -> bool: draft_repo="Lucebox/gemma-4-26B-A4B-it-DFlash-GGUF", draft_file="gemma-4-26B-A4B-it-DFlash-q8_0.gguf", approx_total_gb=18, + architecture="gemma4", + native_context=262144, description="Gemma 4 26B-A4B IT MoE (Q4_K_M) + Lucebox DFlash q8_0 draft.", ), "gemma-4-31b": ModelPreset( @@ -125,6 +131,8 @@ def has_draft(self) -> bool: draft_repo="Lucebox/gemma-4-31B-it-DFlash-GGUF", draft_file="gemma-4-31B-it-DFlash-q8_0.gguf", approx_total_gb=21, + architecture="gemma4", + native_context=262144, description="Gemma 4 31B IT dense (Q4_K_M) + Lucebox DFlash q8_0 draft.", ), "laguna-xs.2": ModelPreset( @@ -140,6 +148,8 @@ def has_draft(self) -> bool: draft_file=None, speculator_dir="laguna-xs2-speculator", approx_total_gb=20, + architecture="laguna", + native_context=4096, description=( "Laguna-XS.2 MoE code model (Q4_K_M). " "DFlash safetensors speculator in draft/laguna-xs2-speculator/ " @@ -162,6 +172,8 @@ def has_draft(self) -> bool: draft_repo=None, draft_file=None, approx_total_gb=22, + architecture="qwen35moe", + native_context=262144, description=( "Qwen3.6 35B-A3B MoE (3B active per token), Q4_K_M unsloth " "dynamic quant. Target-only — no DFlash MoE draft published " @@ -172,6 +184,13 @@ def has_draft(self) -> bool: DEFAULT_PRESET = PRESETS["qwen3.6-27b"] +# Shared scorer used by PFlash and by KVFlash's drafter residency policy. +# It is deliberately separate from each target preset: one ~1.2 GB file is +# reused by every installed model and can be factory-preloaded on a Lucebox. +OPTIMIZER_DRAFTER_REPO = "unsloth/Qwen3-0.6B-GGUF" +OPTIMIZER_DRAFTER_FILE = "Qwen3-0.6B-BF16.gguf" +OPTIMIZER_DRAFTER_APPROX_GB = 1.2 + def resolve_preset(name: str | None) -> ModelPreset: """Look up a preset by name, with a friendly error on typos. @@ -184,7 +203,7 @@ def resolve_preset(name: str | None) -> ModelPreset: if name in PRESETS: return PRESETS[name] # Build a suggestion list — show every known preset; the user's - # search space is small (4 entries today) so listing them all is + # search space is small, so listing them all is # cheaper and clearer than a fuzzy-match heuristic. known = ", ".join(sorted(PRESETS.keys())) raise KeyError(f"unknown preset {name!r}. Known presets: {known}") @@ -425,6 +444,47 @@ def download_preset(cfg: Config, preset: ModelPreset | None = None) -> int: return 0 +def optimizer_drafter_path(cfg: Config) -> Path: + """Host path of the scorer shared by PFlash and KVFlash.""" + return cfg.models_dir / "drafter" / OPTIMIZER_DRAFTER_FILE + + +def optimizer_drafter_container_path() -> str: + """Path seen by the native server inside the runtime image.""" + return f"/opt/lucebox-hub/server/models/drafter/{OPTIMIZER_DRAFTER_FILE}" + + +def optimizer_drafter_installed(cfg: Config) -> bool: + """Presence-only check suitable for an offline, factory-preloaded box.""" + return local_artifact_present(optimizer_drafter_path(cfg)) + + +def local_artifact_present(path: Path) -> bool: + """Reject directories, placeholders, and interrupted zero-byte copies.""" + try: + return path.is_file() and path.stat().st_size > 0 + except OSError: + return False + + +def download_optimizer_drafter(cfg: Config) -> int: + """Download and verify the shared PFlash/KVFlash scorer.""" + console = Console() + api = _new_hf_api() + try: + _fetch( + api, + OPTIMIZER_DRAFTER_REPO, + OPTIMIZER_DRAFTER_FILE, + cfg.models_dir / "drafter", + console, + ) + except Exception as exc: + console.print(f"[red]optimizer download failed:[/red] {exc}") + return 1 + return 0 + + def _local_target_path(cfg: Config, preset: ModelPreset) -> Path: return cfg.models_dir / preset.target_file @@ -438,15 +498,16 @@ def _local_draft_path(cfg: Config, preset: ModelPreset) -> Path | None: def installed_status(cfg: Config, preset: ModelPreset) -> str: """Return ``"installed"`` / ``"partial"`` / ``"absent"`` for a preset. - Size-only — doesn't hash. ``"installed"`` requires the target (and - draft when one is published) to exist on disk; ``"partial"`` means - at least one of the two is present but the set is incomplete. + Presence-only — doesn't query the network or hash. ``"installed"`` + requires a non-empty target (and draft when one is published); + ``"partial"`` means at least one of the two is present but the set is + incomplete. """ - target_exists = _local_target_path(cfg, preset).exists() + target_exists = local_artifact_present(_local_target_path(cfg, preset)) draft_path = _local_draft_path(cfg, preset) if draft_path is None: return "installed" if target_exists else "absent" - draft_exists = draft_path.exists() + draft_exists = local_artifact_present(draft_path) if target_exists and draft_exists: return "installed" if target_exists or draft_exists: @@ -522,12 +583,12 @@ def recommend_preset(host: HostFacts) -> str | None: """Pick a default preset for first-run install. None = ask the user. Tiers follow the model size catalog: 22 GB+ → Qwen3.6-27B (the - Lucebox default), 16-21 GB → Laguna-XS.2 (small target-only). Below - 16 GB we punt and let the user pick explicitly — the registered - presets all need at least 16 GB to run usefully. + Lucebox default); 16-21 GB plus at least 32 GB host RAM → Laguna-XS.2 + with Spark expert offload. Otherwise ask explicitly instead of proposing + a model that may not have enough GPU/host memory to start safely. """ if host.vram_gb >= 22: return "qwen3.6-27b" - if host.vram_gb >= 16: + if host.vram_gb >= 16 and host.ram_gb >= 32: return "laguna-xs.2" return None diff --git a/lucebox/src/lucebox/types.py b/lucebox/src/lucebox/types.py index d755083c0..6fe873c4c 100644 --- a/lucebox/src/lucebox/types.py +++ b/lucebox/src/lucebox/types.py @@ -7,6 +7,7 @@ from __future__ import annotations +import math import os from dataclasses import dataclass, field from pathlib import Path @@ -68,29 +69,25 @@ def __post_init__(self) -> None: invalid snapshot from silently selecting the wrong accelerator path. """ if self.gpu_vendor not in {"nvidia", "amd", "none"}: - raise ValueError( - "gpu_vendor must be nvidia, amd, or none; " - f"got {self.gpu_vendor!r}" - ) + raise ValueError(f"gpu_vendor must be nvidia, amd, or none; got {self.gpu_vendor!r}") if self.ctk not in {"runtime", "cdi", "installed-unwired", "none"}: raise ValueError( - "ctk must be runtime, cdi, installed-unwired, or none; " - f"got {self.ctk!r}" + f"ctk must be runtime, cdi, installed-unwired, or none; got {self.ctk!r}" ) @dataclass(frozen=True, slots=True) class DflashRuntime: - """The DFLASH_* knobs as typed values. Serialized under [dflash] in TOML - and emitted as -e DFLASH_FOO=bar args to docker run. - - The 11 fields below (budget through prefill_drafter) form the strict - allowlist mirrored by lucebench's snapshot config.json — keep both - in lockstep. ``think_max`` is a separate phase-1 thinking cap that - isn't part of the runtime snapshot allowlist (it's per-request, not - per-server). + """Typed inference and optimization settings persisted under ``[dflash]``. + + The name is historical: the same native server owns DFlash speculative + decode, PFlash prefill compression, KVFlash bounded residency, and Spark + MoE offload. Keeping their launch settings together gives the optimizer a + single atomic profile to write and the container one explicit contract to + consume. """ + speculative_decode: bool = True budget: int = 22 max_ctx: int = 16384 lazy: bool = False @@ -102,6 +99,11 @@ class DflashRuntime: prefill_keep_ratio: float = 0.05 prefill_threshold: int = 32000 prefill_drafter: str = "" + kvflash: str = "off" + kvflash_policy: Literal["drafter", "lru", "qk"] = "drafter" + kvflash_tau: int = 64 + spark: bool = False + spark_vram_gb: float = 0.0 # Phase-1 (thinking) cap when a request opts into thinking. Default mirrors # antirez/ds4 ds4_eval.c: think_max_tokens = max_tokens - hard_limit_reply # budget = 16000 - 512 = 15488. The server's own hardcoded default is 10000. @@ -132,6 +134,12 @@ class DflashRuntime: def __post_init__(self) -> None: """Validate the bounded tuning knobs before they reach the server.""" + if self.budget < 0: + raise ValueError(f"budget must be zero or positive; got {self.budget!r}") + if self.max_ctx <= 0: + raise ValueError(f"max_ctx must be positive; got {self.max_ctx!r}") + if self.prefix_cache_slots < 0 or self.prefill_cache_slots < 0: + raise ValueError("cache slot counts must be zero or positive") if not 0.0 < self.prefill_keep_ratio <= 1.0: raise ValueError( "prefill_keep_ratio must be in the interval (0.0, 1.0]; " @@ -142,6 +150,34 @@ def __post_init__(self) -> None: "think_soft_close_min_ratio must be in the interval [0.0, 1.0]; " f"got {self.think_soft_close_min_ratio!r}" ) + if self.kvflash not in {"off", "auto"}: + try: + pool_tokens = int(self.kvflash) + except ValueError as exc: + raise ValueError( + f"kvflash must be off, auto, or a positive token count; got {self.kvflash!r}" + ) from exc + if pool_tokens <= 0: + raise ValueError(f"kvflash token count must be positive; got {self.kvflash!r}") + if self.kvflash_policy not in {"drafter", "lru", "qk"}: + raise ValueError( + f"kvflash_policy must be drafter, lru, or qk; got {self.kvflash_policy!r}" + ) + if self.kvflash_tau <= 0: + raise ValueError(f"kvflash_tau must be positive; got {self.kvflash_tau!r}") + if not math.isfinite(self.spark_vram_gb) or self.spark_vram_gb < 0.0: + raise ValueError( + "spark_vram_gb must be finite and zero (automatic) or positive; " + f"got {self.spark_vram_gb!r}" + ) + if self.prefill_threshold <= 0: + raise ValueError(f"prefill_threshold must be positive; got {self.prefill_threshold!r}") + if self.think_max < 0: + raise ValueError(f"think_max must be zero or positive; got {self.think_max!r}") + if self.fa_window < 0: + raise ValueError(f"fa_window must be zero or positive; got {self.fa_window!r}") + if self.kvflash != "off" and self.fa_window > 0: + raise ValueError("kvflash and fa_window are mutually exclusive") @dataclass(frozen=True, slots=True) diff --git a/lucebox/tests/test_autotune.py b/lucebox/tests/test_autotune.py index c47c17c9e..41187af01 100644 --- a/lucebox/tests/test_autotune.py +++ b/lucebox/tests/test_autotune.py @@ -1,5 +1,9 @@ -from lucebox.autotune import runtime_from_host -from lucebox.types import HostFacts +from pathlib import Path + +from lucebox.autotune import automatic_plan, runtime_from_host +from lucebox.types import Config, HostFacts, ModelMeta + +from lucebox import download def test_wsl_24gb_defaults_leave_cuda_headroom() -> None: @@ -46,3 +50,202 @@ def test_no_heuristic_tier_sets_lazy_without_prefill_drafter() -> None: f"vram={vram} is_wsl={is_wsl}: lazy=True without " f"prefill_drafter → silent no-op on the C++ server" ) + + +def _cfg(tmp_path: Path, preset: str, host: HostFacts) -> Config: + return Config(models_dir=tmp_path / "models", host=host, model=ModelMeta(preset=preset)) + + +def _install_optimizer_drafter(cfg: Config) -> None: + path = download.optimizer_drafter_path(cfg) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"test-scorer") + + +def _install_decode_draft(cfg: Config) -> None: + preset = download.PRESETS[cfg.model.preset] + assert preset.draft_file is not None + path = cfg.models_dir / "draft" / preset.draft_file + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"test-draft") + + +def test_unconfigured_plan_does_not_claim_an_active_optimization(tmp_path: Path) -> None: + cfg = Config( + models_dir=tmp_path / "models", + host=HostFacts(gpu_vendor="amd", vram_gb=31, gpu_sm="gfx1201"), + ) + + plan = automatic_plan(cfg) + + assert plan.active_names == () + assert plan.dflash.available is False + assert "choose a model" in plan.dflash.reason + + +def test_qwen_24gb_enables_dflash_and_pflash_but_prefers_full_kv( + tmp_path: Path, +) -> None: + cfg = _cfg( + tmp_path, + "qwen3.6-27b", + HostFacts(gpu_vendor="nvidia", vram_gb=24, gpu_sm="86"), + ) + _install_optimizer_drafter(cfg) + _install_decode_draft(cfg) + + plan = automatic_plan(cfg) + + assert plan.dflash.enabled is True + assert plan.pflash.enabled is True + assert plan.kvflash.enabled is False + assert plan.spark.enabled is False + assert plan.runtime.prefill_mode == "auto" + assert plan.runtime.prefill_keep_ratio == 0.10 + assert plan.runtime.max_ctx == 98304 + assert plan.active_names == ("DFlash", "PFlash") + + +def test_published_draft_is_not_claimed_until_it_is_installed(tmp_path: Path) -> None: + cfg = _cfg( + tmp_path, + "qwen3.6-27b", + HostFacts(gpu_vendor="nvidia", vram_gb=24, gpu_sm="86"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.dflash.enabled is False + assert plan.dflash.available is False + assert "not installed" in plan.dflash.reason + + +def test_tight_qwen_uses_drafter_free_qk_kvflash_until_scorer_is_installed( + tmp_path: Path, +) -> None: + cfg = _cfg( + tmp_path, + "qwen3.6-27b", + HostFacts(gpu_vendor="nvidia", vram_gb=20, gpu_sm="86"), + ) + + plan = automatic_plan(cfg) + + assert plan.pflash.enabled is False + assert plan.kvflash.enabled is True + assert plan.runtime.kvflash == "auto" + assert plan.runtime.kvflash_policy == "qk" + assert plan.needs_optimizer_drafter is True + + +def test_constrained_qwen_moe_enables_spark_and_scored_kvflash( + tmp_path: Path, +) -> None: + cfg = _cfg( + tmp_path, + "qwen3.6-moe", + HostFacts(gpu_vendor="nvidia", vram_gb=24, ram_gb=64, gpu_sm="86"), + ) + _install_optimizer_drafter(cfg) + + plan = automatic_plan(cfg) + + assert plan.dflash.enabled is False + assert plan.pflash.enabled is False + assert plan.kvflash.enabled is True + assert plan.runtime.kvflash_policy == "drafter" + assert plan.spark.enabled is True + + +def test_r9700_qwen_moe_keeps_exact_all_gpu_path(tmp_path: Path) -> None: + cfg = _cfg( + tmp_path, + "qwen3.6-moe", + HostFacts(gpu_vendor="amd", vram_gb=31, gpu_sm="gfx1201"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.kvflash.enabled is False + assert plan.spark.enabled is False + assert "fits" in plan.spark.reason + + +def test_laguna_uses_native_context_and_spark_under_24gb_pressure( + tmp_path: Path, +) -> None: + cfg = _cfg( + tmp_path, + "laguna-xs.2", + HostFacts(gpu_vendor="nvidia", vram_gb=24, ram_gb=64, gpu_sm="86"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.runtime.max_ctx == 4096 + assert plan.dflash.enabled is False + assert plan.pflash.available is False + assert plan.kvflash.enabled is False + assert plan.spark.enabled is True + + +def test_spark_stays_off_when_host_ram_cannot_hold_cold_experts( + tmp_path: Path, +) -> None: + cfg = _cfg( + tmp_path, + "laguna-xs.2", + HostFacts(gpu_vendor="nvidia", vram_gb=24, ram_gb=16, gpu_sm="86"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.spark.available is True + assert plan.spark.enabled is False + assert plan.runtime.spark is False + assert "16 GB host RAM" in plan.spark.reason + + +def test_gemma_does_not_offer_unsupported_pflash(tmp_path: Path) -> None: + cfg = _cfg( + tmp_path, + "gemma-4-31b", + HostFacts(gpu_vendor="amd", vram_gb=31, gpu_sm="gfx1201"), + ) + _install_optimizer_drafter(cfg) + + plan = automatic_plan(cfg) + + assert plan.pflash.available is False + assert plan.pflash.enabled is False + assert plan.runtime.prefill_mode == "off" + assert "no PFlash" in plan.pflash.reason + + +def test_gemma_context_is_capped_when_exact_cache_has_little_headroom( + tmp_path: Path, +) -> None: + cfg = _cfg( + tmp_path, + "gemma-4-31b", + HostFacts(gpu_vendor="nvidia", vram_gb=24, gpu_sm="86"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.runtime.max_ctx == 8192 + assert plan.runtime.cache_type_k == "" + assert plan.runtime.cache_type_v == "" + + +def test_automatic_hardware_tiers_never_force_quality_risky_kv_types() -> None: + for vram_gb in (0, 8, 16, 24, 32, 80): + runtime = runtime_from_host(HostFacts(vram_gb=vram_gb)) + assert runtime.cache_type_k == "" + assert runtime.cache_type_v == "" + + +def test_known_gpu_specific_ddtree_budgets() -> None: + assert runtime_from_host(HostFacts(vram_gb=24, gpu_sm="gfx1100")).budget == 8 + assert runtime_from_host(HostFacts(vram_gb=32, gpu_sm="120")).budget == 40 + assert runtime_from_host(HostFacts(vram_gb=31, gpu_sm="gfx1201")).budget == 22 diff --git a/lucebox/tests/test_config.py b/lucebox/tests/test_config.py index a61f411f6..5d73d39ef 100644 --- a/lucebox/tests/test_config.py +++ b/lucebox/tests/test_config.py @@ -330,3 +330,102 @@ def test_seed_dflash_migrates_legacy_config_before_writing(tmp_path: Path) -> No assert loaded is not None assert loaded.port == 9090 assert loaded.dflash.max_ctx == 98304 + + +def test_optimization_fields_round_trip_and_validate(tmp_path: Path) -> None: + path = tmp_path / "config.toml" + path.write_text( + "[dflash]\n" + "speculative_decode = false\n" + 'kvflash = "auto"\n' + 'kvflash_policy = "qk"\n' + "kvflash_tau = 96\n" + "spark = true\n" + "spark_vram_gb = 14.5\n" + ) + + loaded = config.load(path) + assert loaded is not None + assert loaded.dflash.speculative_decode is False + assert loaded.dflash.kvflash == "auto" + assert loaded.dflash.kvflash_policy == "qk" + assert loaded.dflash.kvflash_tau == 96 + assert loaded.dflash.spark is True + assert loaded.dflash.spark_vram_gb == 14.5 + + +@pytest.mark.parametrize( + ("key", "value"), + [ + ("dflash.kvflash", "banana"), + ("dflash.kvflash", "0"), + ("dflash.kvflash_policy", "random"), + ("dflash.kvflash_tau", "0"), + ("dflash.spark_vram_gb", "-1"), + ], +) +def test_config_rejects_invalid_optimization_values( + tmp_path: Path, key: str, value: str +) -> None: + with pytest.raises(ValueError): + config_set(key, value, path=tmp_path / "config.toml") + + +def test_direct_optimization_edit_marks_profile_custom(tmp_path: Path) -> None: + path = tmp_path / "config.toml" + + config_set("dflash.spark", "true", path=path) + + assert config.optimization_mode(path=path) == "custom" + assert config.load_doc(path)["autotune"]["source"] == "manual" + + +def test_manual_edit_rejects_incompatible_kvflash_and_fa_window( + tmp_path: Path, +) -> None: + path = tmp_path / "config.toml" + config_set("dflash.fa_window", 512, path=path) + + with pytest.raises(ValueError, match="mutually exclusive"): + config_set("dflash.kvflash", "auto", path=path) + + loaded = config.load(path) + assert loaded is not None + assert loaded.dflash.fa_window == 512 + assert loaded.dflash.kvflash == "off" + + +def test_automatic_profile_replans_on_model_switch(tmp_path: Path) -> None: + from lucebox.types import Config, HostFacts, ModelMeta + + path = tmp_path / "config.toml" + config.seed_dflash_from_host(HostFacts(vram_gb=24), path=path) + cfg = Config( + models_dir=tmp_path / "models", + host=HostFacts(gpu_vendor="nvidia", vram_gb=24, ram_gb=64, gpu_sm="86"), + model=ModelMeta(preset="qwen3.6-moe"), + ) + + assert config.seed_optimization_from_config(cfg, path=path) is True + loaded = config.load(path) + assert loaded is not None + assert loaded.dflash.speculative_decode is False + assert loaded.dflash.spark is True + assert config.optimization_mode(path=path) == "automatic" + + +def test_model_switch_preserves_custom_profile(tmp_path: Path) -> None: + from lucebox.types import Config, HostFacts, ModelMeta + + path = tmp_path / "config.toml" + config_set("dflash.max_ctx", 8192, path=path) + cfg = Config( + models_dir=tmp_path / "models", + host=HostFacts(vram_gb=24), + model=ModelMeta(preset="qwen3.6-moe"), + ) + + assert config.seed_optimization_from_config(cfg, path=path) is False + loaded = config.load(path) + assert loaded is not None + assert loaded.dflash.max_ctx == 8192 diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index 7753b8e5a..053046fd9 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -399,6 +399,12 @@ def test_server_run_spec_optional_env_off_by_default(tmp_path: Path) -> None: "DFLASH_CACHE_TYPE_K", "DFLASH_CACHE_TYPE_V", "DFLASH_PREFILL_MODE", + "DFLASH_PREFILL_DRAFTER", + "DFLASH_KVFLASH", + "DFLASH_KVFLASH_POLICY", + "DFLASH_KVFLASH_TAU", + "DFLASH_SPARK", + "DFLASH_SPARK_VRAM_GB", "DFLASH_FA_WINDOW", "DFLASH_THINK_SOFT_CLOSE_MIN_RATIO", "DFLASH_DEBUG_THINKING_LOGITS", @@ -419,7 +425,11 @@ def test_server_run_spec_optional_env_emitted_when_set(tmp_path: Path) -> None: prefill_keep_ratio=0.1, prefill_threshold=20000, prefill_drafter="drafter.gguf", - fa_window=512, + kvflash="auto", + kvflash_policy="qk", + kvflash_tau=96, + spark=True, + spark_vram_gb=14.5, think_soft_close_min_ratio=0.5, debug_thinking_logits=True, ), @@ -432,11 +442,34 @@ def test_server_run_spec_optional_env_emitted_when_set(tmp_path: Path) -> None: assert env["DFLASH_PREFILL_KEEP"] == "0.1" assert env["DFLASH_PREFILL_THRESHOLD"] == "20000" assert env["DFLASH_PREFILL_DRAFTER"] == "drafter.gguf" - assert env["DFLASH_FA_WINDOW"] == "512" + assert env["DFLASH_KVFLASH"] == "auto" + assert env["DFLASH_KVFLASH_POLICY"] == "qk" + assert env["DFLASH_KVFLASH_TAU"] == "96" + assert env["DFLASH_SPARK"] == "1" + assert env["DFLASH_SPARK_VRAM_GB"] == "14.5" assert env["DFLASH_THINK_SOFT_CLOSE_MIN_RATIO"] == "0.5" assert env["DFLASH_DEBUG_THINKING_LOGITS"] == "1" +def test_server_run_spec_rejects_kvflash_with_fa_window() -> None: + with pytest.raises(ValueError, match="mutually exclusive"): + DflashRuntime(kvflash="auto", fa_window=512) + + +def test_server_run_spec_forwards_primary_rocm_device( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("LUCEBOX_HOST_ROCR_VISIBLE_DEVICES", "1") + monkeypatch.delenv("LUCEBOX_HOST_HIP_VISIBLE_DEVICES", raising=False) + cfg = Config(variant="rocm", models_dir=tmp_path) + + env = _env(docker_run.server_run_spec(cfg)) + + assert env["ROCR_VISIBLE_DEVICES"] == "1" + assert "HIP_VISIBLE_DEVICES" not in env + + def test_server_run_spec_resolves_target_and_draft_paths(tmp_path: Path) -> None: pres = PRESETS["qwen3.6-27b"] cfg = Config(models_dir=tmp_path, model=ModelMeta(preset="qwen3.6-27b")) @@ -448,6 +481,18 @@ def test_server_run_spec_resolves_target_and_draft_paths(tmp_path: Path) -> None ) +def test_server_run_spec_can_disable_preset_dflash_draft(tmp_path: Path) -> None: + cfg = Config( + models_dir=tmp_path, + model=ModelMeta(preset="qwen3.6-27b"), + dflash=DflashRuntime(speculative_decode=False), + ) + + env = _env(docker_run.server_run_spec(cfg)) + + assert env["DFLASH_DRAFT"].endswith("/.lucebox-no-draft") + + def test_server_run_spec_forwards_host_env(monkeypatch, tmp_path: Path) -> None: monkeypatch.setenv("LUCEBOX_HOST_OS_PRETTY", "Ubuntu 22.04") monkeypatch.setenv("LUCEBOX_HOST_GPU_NAME", "RTX 5090") @@ -457,15 +502,7 @@ def test_server_run_spec_forwards_host_env(monkeypatch, tmp_path: Path) -> None: def test_large_preset_serves_at_safe_default_ctx(tmp_path: Path) -> None: - """Regression guard for the preset-cap analysis (#5). - - Activating a preset writes only [model], never [dflash], so a loaded - Config keeps the conservative DflashRuntime() floor (max_ctx=16384). - The VRAM-tier heuristic's higher caps only apply via `autotune --apply` - (which threads cfg.model.preset and is a separate PR). This test pins - that a large preset does NOT silently serve at a high, OOM-prone ctx - through the default serve path. - """ + """A bare low-level Config retains the conservative 16K context floor.""" cfg = Config(models_dir=tmp_path, model=ModelMeta(preset="qwen3.6-27b")) env = _env(docker_run.server_run_spec(cfg)) assert env["DFLASH_MAX_CTX"] == "16384" diff --git a/lucebox/tests/test_download.py b/lucebox/tests/test_download.py index a4bf692e3..bb5b1e4b9 100644 --- a/lucebox/tests/test_download.py +++ b/lucebox/tests/test_download.py @@ -327,15 +327,15 @@ def _meta(_api, repo_id: str, filename: str) -> tuple[int, None]: def test_recommend_preset_tiers() -> None: - """First-run preset recommendation is a pure VRAM-tier function. + """First-run recommendations account for GPU and host-memory capacity. - 22 GB+ → the Lucebox default (qwen3.6-27b); 16-21 GB → laguna-xs.2; - below 16 GB → None (the registered presets need ≥16 GB, so we punt to - an explicit choice rather than recommend something that can't run). + 22 GB+ → the Lucebox default (qwen3.6-27b); 16-21 GB plus enough host + RAM for Spark offload → laguna-xs.2; otherwise ask explicitly. """ assert recommend_preset(HostFacts(vram_gb=24)) == "qwen3.6-27b" assert recommend_preset(HostFacts(vram_gb=22)) == "qwen3.6-27b" - assert recommend_preset(HostFacts(vram_gb=20)) == "laguna-xs.2" - assert recommend_preset(HostFacts(vram_gb=16)) == "laguna-xs.2" + assert recommend_preset(HostFacts(vram_gb=20, ram_gb=64)) == "laguna-xs.2" + assert recommend_preset(HostFacts(vram_gb=16, ram_gb=32)) == "laguna-xs.2" + assert recommend_preset(HostFacts(vram_gb=20, ram_gb=16)) is None assert recommend_preset(HostFacts(vram_gb=12)) is None assert recommend_preset(HostFacts(vram_gb=0)) is None diff --git a/lucebox/tests/test_models_cli.py b/lucebox/tests/test_models_cli.py index 97eaba8c7..00b0c98e6 100644 --- a/lucebox/tests/test_models_cli.py +++ b/lucebox/tests/test_models_cli.py @@ -21,8 +21,9 @@ def _set_config_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: def _stub_host(monkeypatch: pytest.MonkeyPatch, vram_gb: int) -> None: - monkeypatch.setattr("lucebox.host_facts.from_env", lambda: HostFacts(vram_gb=vram_gb)) - monkeypatch.setattr("lucebox.cli.from_env", lambda: HostFacts(vram_gb=vram_gb)) + host = HostFacts(vram_gb=vram_gb, ram_gb=64) + monkeypatch.setattr("lucebox.host_facts.from_env", lambda: host) + monkeypatch.setattr("lucebox.cli.from_env", lambda: host) def test_models_list_shows_every_registered_preset( @@ -131,10 +132,10 @@ def test_models_select_activates_preloaded_model_offline( preset = PRESETS["qwen3.6-27b"] models = tmp_path / "models" (models / preset.target_file).parent.mkdir(parents=True, exist_ok=True) - (models / preset.target_file).touch() + (models / preset.target_file).write_bytes(b"preloaded-target") assert preset.draft_file is not None (models / "draft").mkdir() - (models / "draft" / preset.draft_file).touch() + (models / "draft" / preset.draft_file).write_bytes(b"preloaded-draft") def fail_network(*args: object, **kwargs: object) -> object: raise AssertionError("preloaded selection must not contact Hugging Face") @@ -182,7 +183,100 @@ def test_optimize_resets_to_hardware_profile( assert "Automatic optimization applied" in result.output entries = config_mod.config_get(path=cfg_path) assert entries["dflash.max_ctx"] == (98304, "file") - assert entries["dflash.cache_type_k"] == ("tq3_0", "file") + assert entries["dflash.cache_type_k"] == ("", "file") + + +def test_optimize_applies_model_aware_qwen_stack( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + config_mod.config_set("model.preset", "qwen3.6-27b", path=cfg_path) + cfg = config_mod.load(cfg_path) + assert cfg is not None + cfg = config_mod.overlay_env(cfg) + scorer = download_mod.optimizer_drafter_path(cfg) + scorer.parent.mkdir(parents=True, exist_ok=True) + scorer.write_bytes(b"preloaded") + preset = PRESETS["qwen3.6-27b"] + assert preset.draft_file is not None + draft = cfg.models_dir / "draft" / preset.draft_file + draft.parent.mkdir(parents=True, exist_ok=True) + draft.write_bytes(b"preloaded-draft") + + result = CliRunner().invoke(app, ["optimize", "--yes"]) + + assert result.exit_code == 0 + assert "DFlash" in result.output + assert "PFlash" in result.output + loaded = config_mod.load(cfg_path) + assert loaded is not None + assert loaded.dflash.speculative_decode is True + assert loaded.dflash.prefill_mode == "auto" + assert loaded.dflash.kvflash == "off" + assert config_mod.optimization_mode(path=cfg_path) == "automatic" + + +def test_optimize_installs_shared_scorer_for_constrained_moe( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + config_mod.config_set("model.preset", "qwen3.6-moe", path=cfg_path) + + def fake_download(cfg: object) -> int: + assert isinstance(cfg, config_mod.Config) + path = download_mod.optimizer_drafter_path(cfg) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"downloaded") + return 0 + + monkeypatch.setattr(download_mod, "download_optimizer_drafter", fake_download) + + result = CliRunner().invoke(app, ["optimize", "--yes"]) + + assert result.exit_code == 0 + assert "Shared optimizer installed" in result.output + loaded = config_mod.load(cfg_path) + assert loaded is not None + assert loaded.dflash.kvflash == "auto" + assert loaded.dflash.kvflash_policy == "drafter" + assert loaded.dflash.spark is True + + +def test_optimize_advanced_writes_a_custom_profile( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + config_mod.config_set("model.preset", "qwen3.6-27b", path=cfg_path) + cfg = config_mod.load(cfg_path) + assert cfg is not None + cfg = config_mod.overlay_env(cfg) + scorer = download_mod.optimizer_drafter_path(cfg) + scorer.parent.mkdir(parents=True, exist_ok=True) + scorer.write_bytes(b"preloaded") + preset = PRESETS["qwen3.6-27b"] + assert preset.draft_file is not None + draft = cfg.models_dir / "draft" / preset.draft_file + draft.parent.mkdir(parents=True, exist_ok=True) + draft.write_bytes(b"preloaded-draft") + + # DFlash yes, PFlash no, KVFlash yes, qk policy, apply yes. + result = CliRunner().invoke( + app, + ["optimize", "--advanced"], + input="y\nn\ny\nqk\ny\n", + ) + + assert result.exit_code == 0 + loaded = config_mod.load(cfg_path) + assert loaded is not None + assert loaded.dflash.speculative_decode is True + assert loaded.dflash.prefill_mode == "off" + assert loaded.dflash.kvflash == "auto" + assert loaded.dflash.kvflash_policy == "qk" + assert config_mod.optimization_mode(path=cfg_path) == "custom" def test_installed_helpers_track_presence( @@ -200,6 +294,8 @@ def test_installed_helpers_track_presence( target = cfg.models_dir / laguna.target_file target.parent.mkdir(parents=True, exist_ok=True) + target.touch() + assert download_mod.installed_status(cfg, laguna) == "absent" # Apparent size is what the CLI reports; a sparse file exercises the same # stat path without allocating and writing a 5 GB Python byte string. with target.open("wb") as sparse: diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index 46134abb3..def9acbd8 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -554,6 +554,66 @@ test_entrypoint_serve_path "entrypoint serve: gemma-4-31b family match" \ test_entrypoint_serve_path "entrypoint serve: generic fallback" \ "Mystery-Model-7B.gguf" "model.gguf" +test_entrypoint_optimization_flags() { + local label="$1" sandbox draft_dir models_dir bin_dir shim_dir out + _make_entrypoint_sandbox + touch "$models_dir/Qwen3.6-27B-Q4_K_M.gguf" + touch "$draft_dir/dflash-draft-3.6-test.gguf" + touch "$models_dir/Qwen3-0.6B-BF16.gguf" + out=$(PATH="$shim_dir:$PATH" \ + DFLASH_DIR="$sandbox" \ + DFLASH_SERVER_BIN="$bin_dir/dflash_server" \ + DFLASH_TARGET="$models_dir/Qwen3.6-27B-Q4_K_M.gguf" \ + DFLASH_DRAFT="$draft_dir" \ + DFLASH_PREFILL_MODE=auto \ + DFLASH_PREFILL_DRAFTER="$models_dir/Qwen3-0.6B-BF16.gguf" \ + DFLASH_KVFLASH=auto \ + DFLASH_KVFLASH_POLICY=qk \ + DFLASH_KVFLASH_TAU=96 \ + DFLASH_SPARK=1 \ + DFLASH_SPARK_VRAM_GB=14 \ + timeout 10 bash "$ENTRYPOINT" serve 2>&1 || true) + rm -rf "$sandbox" + local required missing="" + for required in \ + "--prefill-drafter" "--prefill-compression auto" \ + "--kvflash auto" "--kvflash-policy qk" "--kvflash-tau 96" \ + "--spark" "--spark-vram 14"; do + [[ "$out" == *"$required"* ]] || missing+=" $required" + done + if [ -n "$missing" ]; then + report fail "$label" "missing argv:$missing; output: $(tail -3 <<<"$out")" + elif [ "$(grep -o -- '--prefill-drafter' <<<"$out" | wc -l | tr -d ' ')" != "1" ]; then + report fail "$label" "--prefill-drafter was emitted more than once" + else + report ok "$label" + fi +} +test_entrypoint_optimization_flags \ + "entrypoint forwards PFlash, KVFlash, and Spark exactly once" + +test_entrypoint_keeps_family_kv_default() { + local label="$1" sandbox draft_dir models_dir bin_dir shim_dir out + _make_entrypoint_sandbox + touch "$models_dir/Qwen3.6-27B-Q4_K_M.gguf" + out=$(PATH="$shim_dir:$PATH" \ + DFLASH_DIR="$sandbox" \ + DFLASH_SERVER_BIN="$bin_dir/dflash_server" \ + DFLASH_TARGET="$models_dir/Qwen3.6-27B-Q4_K_M.gguf" \ + DFLASH_DRAFT="$models_dir/.lucebox-no-draft" \ + timeout 10 bash "$ENTRYPOINT" serve 2>&1 || true) + rm -rf "$sandbox" + if [[ "$out" == *"--cache-type-k"* || "$out" == *"--cache-type-v"* ]]; then + report fail "$label" "24 GB fallback forced a quality-risky KV type: $(tail -3 <<<"$out")" + elif [[ "$out" != *"[shim] dflash_server"* ]]; then + report fail "$label" "server shim did not run: $(tail -3 <<<"$out")" + else + report ok "$label" + fi +} +test_entrypoint_keeps_family_kv_default \ + "entrypoint preserves model-family KV defaults on 24 GB GPUs" + # ── 11. entrypoint.sh serve-path with MULTIPLE target-sized GGUFs in # models/. The single-candidate fixture in test 10 doesn't exercise the # auto-detect path that picks "first alphabetically" when more than one @@ -1498,12 +1558,13 @@ STUB test_mixed_rtx_strix_probe_prefers_cuda "mixed RTX 3090 + Strix probe selects CUDA" test_cross_vendor_docker_args() { - local label="$1" helpers cuda_args rocm_args pinned_rocm_args + local label="$1" helpers cuda_args rocm_args pinned_rocm_args selected_rocm_args helpers=$(awk '/^_variant_is_rocm\(\) \{/,/^\}/' "$SCRIPT")$'\n' helpers+=$(awk '/^_append_gpu_args\(\) \{/,/^\}/' "$SCRIPT") cuda_args=$(bash -c "$helpers"$'\n''a=(); _append_gpu_args a cuda12; printf "%s\n" "${a[@]}"') rocm_args=$(bash -c "$helpers"$'\n''a=(); _append_gpu_args a rocm; printf "%s\n" "${a[@]}"') pinned_rocm_args=$(bash -c "$helpers"$'\n''a=(); _append_gpu_args a 0.3.0-rocm; printf "%s\n" "${a[@]}"') + selected_rocm_args=$(bash -c "$helpers"$'\n''LUCEBOX_HOST_ROCR_VISIBLE_DEVICES=1; a=(); _append_gpu_args a rocm; printf "%s\n" "${a[@]}"') if ! grep -qF -- '--gpus' <<<"$cuda_args" || ! grep -qF 'all' <<<"$cuda_args"; then report fail "$label" "CUDA args missing --gpus all" return @@ -1523,10 +1584,40 @@ test_cross_vendor_docker_args() { report fail "$label" "versioned ROCm tag did not select ROCm args" return fi + if ! grep -qF 'ROCR_VISIBLE_DEVICES=1' <<<"$selected_rocm_args" \ + || grep -qF 'HIP_VISIBLE_DEVICES=' <<<"$selected_rocm_args"; then + report fail "$label" "ROCm primary device was not pinned" + return + fi report ok "$label" } test_cross_vendor_docker_args "Docker args select CUDA or ROCm device contract" +test_selected_backend_facts() { + local label="$1" helpers out + helpers=$(awk '/^_variant_is_rocm\(\) \{/,/^\}/' "$SCRIPT")$'\n' + helpers+=$(awk '/^_append_selected_backend_facts\(\) \{/,/^\}/' "$SCRIPT") + out=$(bash -c "$helpers"$'\n'' + LUCEBOX_HOST_HAS_AMD_GPU=1 + LUCEBOX_HOST_AMD_GPU_NAME="AMD Strix" + LUCEBOX_HOST_AMD_GPU_COUNT=1 + LUCEBOX_HOST_AMD_VRAM_GB=64 + LUCEBOX_HOST_AMD_GPU_ARCH=gfx1151 + LUCEBOX_HOST_AMD_GPU_LIST_CSV="0, , , AMD Strix, gfx1151, 65536 MiB," + a=() + _append_selected_backend_facts a rocm + printf "%s\n" "${a[@]}" + ') + if ! grep -qF 'LUCEBOX_HOST_GPU_VENDOR=amd' <<<"$out" \ + || ! grep -qF 'LUCEBOX_HOST_VRAM_GB=64' <<<"$out" \ + || ! grep -qF 'LUCEBOX_HOST_GPU_SM=gfx1151' <<<"$out"; then + report fail "$label" "selected ROCm facts were not forwarded: $out" + return + fi + report ok "$label" +} +test_selected_backend_facts "mixed-build autotune facts follow the selected ROCm backend" + # ── TTY flag selection. Regression guard for the process-substitution bug: # _set_tty_flags must run in the CALLER's scope so `[ -t 1 ]` inspects the # real terminal. If it is ever moved back behind `< <(...)` or `$(...)`, diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index 6187c8af8..eb789c82e 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -324,25 +324,22 @@ if [ "$GPU_VRAM_GB" -gt 0 ]; then IS_WSL=1 fi if [ "$GPU_VRAM_GB" -lt 12 ]; then - : "${DFLASH_LAZY:=1}" : "${DFLASH_MAX_CTX:=4096}" warn "VRAM ${GPU_VRAM_GB} GB < 12 GB — 27B target unlikely to fit" elif [ "$GPU_VRAM_GB" -lt 22 ]; then - : "${DFLASH_LAZY:=1}" : "${DFLASH_MAX_CTX:=32768}" elif [ "$GPU_VRAM_GB" -lt 32 ]; then - : "${DFLASH_LAZY:=1}" if [ "$IS_WSL" = "1" ]; then : "${DFLASH_BUDGET:=16}" - : "${DFLASH_MAX_CTX:=65536}" - else - : "${DFLASH_MAX_CTX:=98304}" fi + : "${DFLASH_MAX_CTX:=98304}" else : "${DFLASH_MAX_CTX:=131072}" fi fi +[ "${GPU_ARCH:-}" = "gfx1100" ] && : "${DFLASH_BUDGET:=8}" + : "${DFLASH_BIN:=$DFLASH_DIR/build/test_dflash}" : "${DFLASH_SERVER_BIN:=$DFLASH_DIR/build/dflash_server}" : "${DFLASH_HOST:=0.0.0.0}" @@ -359,6 +356,11 @@ fi : "${DFLASH_PREFILL_KEEP:=0.05}" : "${DFLASH_PREFILL_THRESHOLD:=32000}" : "${DFLASH_PREFILL_DRAFTER:=}" +: "${DFLASH_KVFLASH:=off}" +: "${DFLASH_KVFLASH_POLICY:=drafter}" +: "${DFLASH_KVFLASH_TAU:=64}" +: "${DFLASH_SPARK:=0}" +: "${DFLASH_SPARK_VRAM_GB:=0}" # Optional server default for requests that omit max_tokens. When unset, # the C++ server uses the model-card default. : "${DFLASH_DEFAULT_MAX_TOKENS:=}" @@ -404,7 +406,10 @@ fi if [ -z "$DFLASH_TARGET" ] && [ -d "$DFLASH_DIR/models" ]; then # Collect candidates: .gguf files ≥5 GB (target-sized), excluding # anything under models/draft/. Sort alphabetically for determinism. - mapfile -t TARGET_CANDIDATES < <( + TARGET_CANDIDATES=() + while IFS= read -r candidate; do + [ -n "$candidate" ] && TARGET_CANDIDATES+=("$candidate") + done < <( find -L "$DFLASH_DIR/models" -maxdepth 4 -type f -name '*.gguf' \ -size +5G \ -not -path '*/draft/*' \ @@ -489,6 +494,10 @@ if [ -d "$DFLASH_DRAFT" ]; then # then generic dflash-draft-*.gguf legacy, then last-resort *.gguf. # The 31B match in the Lucebox repo uses capital B in the filename — # -iname handles that without needing to enumerate every case form. + # Bash 3.2 + `set -u` treats an expansion of an empty array as an + # unbound variable. Keep a non-empty sentinel for an unknown family and + # build the combined search list explicitly below. The production image + # uses a newer Bash, but this also keeps host-side smoke tests portable. case "$(echo "$TARGET_BASENAME" | tr 'A-Z' 'a-z')" in *gemma-4-26b*|*gemma4-26b*) FAMILY_GLOBS=('*gemma*4*26b*dflash*.gguf' '*dflash*gemma*4*26b*.gguf') ;; @@ -499,7 +508,7 @@ if [ -d "$DFLASH_DRAFT" ]; then *qwen3.6*|*qwen36*) FAMILY_GLOBS=('dflash-draft-3.6-*.gguf' '*qwen*3.6*dflash*.gguf') ;; *) - FAMILY_GLOBS=() ;; + FAMILY_GLOBS=('') ;; esac DRAFT_FILE="" @@ -514,8 +523,14 @@ if [ -d "$DFLASH_DRAFT" ]; then # `*.gguf` / safetensors fallbacks. GENERIC_GLOBS=('dflash-draft-*.gguf' '*dflash*.gguf' '*.gguf' 'model.safetensors' '*.safetensors') family_count="${#FAMILY_GLOBS[@]}" + if [ -z "${FAMILY_GLOBS[0]}" ]; then + family_count=0 + ALL_DRAFT_GLOBS=("${GENERIC_GLOBS[@]}") + else + ALL_DRAFT_GLOBS=("${FAMILY_GLOBS[@]}" "${GENERIC_GLOBS[@]}") + fi i=0 - for pattern in "${FAMILY_GLOBS[@]}" "${GENERIC_GLOBS[@]}"; do + for pattern in "${ALL_DRAFT_GLOBS[@]}"; do # Sort matches lexicographically so the pick is deterministic across # filesystems (find's traversal order is filesystem-dependent without # an explicit sort). First lexicographic match wins. @@ -575,6 +590,7 @@ CMD=("$DFLASH_SERVER_BIN" "$DFLASH_TARGET" [ -n "$DRAFT_ARG" ] && CMD+=(--ddtree --ddtree-budget "$DFLASH_BUDGET") [ -n "$DFLASH_DEFAULT_MAX_TOKENS" ] && CMD+=(--default-max-tokens "$DFLASH_DEFAULT_MAX_TOKENS") [ -n "$DFLASH_MODEL_NAME" ] && CMD+=(--model-name "$DFLASH_MODEL_NAME") +[ -n "$DFLASH_PREFILL_DRAFTER" ] && CMD+=(--prefill-drafter "$DFLASH_PREFILL_DRAFTER") # `--lazy-draft` is silently dropped by the C++ server unless both # `--prefill-drafter` and `--draft` are present (look for the runtime # warning `--lazy-draft ignored: requires both --prefill-drafter and @@ -591,6 +607,18 @@ fi [ -n "$DFLASH_CACHE_TYPE_K" ] && CMD+=(--cache-type-k "$DFLASH_CACHE_TYPE_K") [ -n "$DFLASH_CACHE_TYPE_V" ] && CMD+=(--cache-type-v "$DFLASH_CACHE_TYPE_V") [ "$DFLASH_FA_WINDOW" -gt 0 ] 2>/dev/null && CMD+=(--fa-window "$DFLASH_FA_WINDOW") +if [ "$DFLASH_KVFLASH" != "off" ] && [ -n "$DFLASH_KVFLASH" ]; then + CMD+=(--kvflash "$DFLASH_KVFLASH" + --kvflash-policy "$DFLASH_KVFLASH_POLICY" + --kvflash-tau "$DFLASH_KVFLASH_TAU") +fi +if [ "$DFLASH_SPARK" = "1" ]; then + CMD+=(--spark) + case "$DFLASH_SPARK_VRAM_GB" in + 0|0.0|0.00|"") ;; + *) CMD+=(--spark-vram "$DFLASH_SPARK_VRAM_GB") ;; + esac +fi # Soft-close ratio: emit only when nonzero. The default-string compare # guards against the floating-point quirks of `[` numeric tests for # values like 0.0/0/0.00 — anything non-"0.0" passes through to the @@ -606,8 +634,7 @@ if [ "$DFLASH_PREFILL_MODE" != "off" ]; then [ -f "$DFLASH_PREFILL_DRAFTER" ] || die "Prefill drafter not found at $DFLASH_PREFILL_DRAFTER" CMD+=(--prefill-compression "$DFLASH_PREFILL_MODE" --prefill-keep-ratio "$DFLASH_PREFILL_KEEP" - --prefill-threshold "$DFLASH_PREFILL_THRESHOLD" - --prefill-drafter "$DFLASH_PREFILL_DRAFTER") + --prefill-threshold "$DFLASH_PREFILL_THRESHOLD") fi if [ "${LUCEBOX_NATIVE:-0}" = "1" ]; then From e90528c7b9bed94709b56f336e31d4538043673e Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:37:49 +0200 Subject: [PATCH 10/25] fix(lucebox): resolve final CLI review findings --- harness/clients/common.sh | 28 +++++++++++------ lucebox.sh | 39 +++++++++++++++-------- lucebox/src/lucebox/autotune.py | 6 +++- lucebox/src/lucebox/cli.py | 33 +++++++++++++++++++- lucebox/src/lucebox/config.py | 5 +-- lucebox/src/lucebox/types.py | 10 ++++-- lucebox/tests/test_autotune.py | 10 +++--- lucebox/tests/test_config.py | 42 ++++++++++++++----------- lucebox/tests/test_models_cli.py | 6 ++-- scripts/test_lucebox_sh.sh | 53 +++++++++++++++++++++++++++++--- server/scripts/entrypoint.sh | 8 ++++- 11 files changed, 179 insertions(+), 61 deletions(-) diff --git a/harness/clients/common.sh b/harness/clients/common.sh index ff35e6b7d..d8dd28d8a 100755 --- a/harness/clients/common.sh +++ b/harness/clients/common.sh @@ -132,22 +132,30 @@ start_dflash_native_server() { echo " cmake --build $REPO_DIR/server/build --target dflash_server -j\$(nproc)" >&2 return 1 fi - if [[ ! -s "$TARGET" ]]; then + if [[ ! -f "$TARGET" || ! -s "$TARGET" ]]; then echo "target GGUF not found: $TARGET" >&2 echo "Set TARGET=/path/to/model.gguf or DFLASH_TARGET=/path/to/model.gguf, or download the default:" >&2 echo " hf download unsloth/Qwen3.6-27B-GGUF Qwen3.6-27B-Q4_K_M.gguf --local-dir $REPO_DIR/server/models/" >&2 return 1 fi if draft_enabled && [[ -d "$DRAFT" ]]; then - local resolved_draft="" - resolved_draft=$(find -L "$DRAFT" -maxdepth 4 -type f \ - \( -name 'model.safetensors' -o -name '*.safetensors' -o -name '*.gguf' \) \ - -print 2>/dev/null | sort | head -n 1) - if [[ -n "$resolved_draft" ]]; then - DRAFT="$resolved_draft" - fi + local draft_candidates=() candidate + while IFS= read -r candidate; do + [[ -n "$candidate" ]] && draft_candidates+=("$candidate") + done < <(find -L "$DRAFT" -maxdepth 4 -type f \ + \( -name '*.safetensors' -o -name '*.gguf' \) \ + -size +0c -print 2>/dev/null | sort) + case "${#draft_candidates[@]}" in + 0) ;; + 1) DRAFT="${draft_candidates[0]}" ;; + *) + echo "multiple DFlash draft candidates in $DRAFT; choose one file explicitly:" >&2 + printf ' %s\n' "${draft_candidates[@]}" >&2 + return 1 + ;; + esac fi - if draft_enabled && [[ ! -s "$DRAFT" ]]; then + if draft_enabled && [[ ! -f "$DRAFT" || ! -s "$DRAFT" ]]; then echo "DFlash draft not found: $DRAFT" >&2 echo "Set DRAFT=/path/to/dflash-draft.gguf or DFLASH_DRAFT=/path/to/dflash-draft.gguf, or download the default:" >&2 echo " hf download Lucebox/Qwen3.6-27B-DFlash-GGUF dflash-draft-3.6-q4_k_m.gguf --local-dir $REPO_DIR/server/models/draft/" >&2 @@ -171,7 +179,7 @@ start_dflash_native_server() { fi local optimization_args=() if [[ -n "${DFLASH_PREFILL_DRAFTER:-}" ]]; then - if [[ ! -s "$DFLASH_PREFILL_DRAFTER" ]]; then + if [[ ! -f "$DFLASH_PREFILL_DRAFTER" || ! -s "$DFLASH_PREFILL_DRAFTER" ]]; then echo "PFlash/KVFlash scorer not found: $DFLASH_PREFILL_DRAFTER" >&2 return 1 fi diff --git a/lucebox.sh b/lucebox.sh index 7f409f037..4aba877a8 100755 --- a/lucebox.sh +++ b/lucebox.sh @@ -329,8 +329,10 @@ sha256_file() { # — it can't see the host's /proc anyway, only the container's. # Normalize ``amd-smi static --asic --vram --csv`` into the compact internal -# form ``index|name|gfx_arch|vram_mib``. Header lookup keeps this resilient to -# columns being added or reordered by future amd-smi releases. +# form ``index|name|gfx_arch|vram_mib|rocr_selector``. Discrete cards use +# their stable GPU UUID instead of assuming amd-smi and ROCr enumerate devices +# in the same order; devices without a usable ASIC serial fall back to index. +# Header lookup keeps this resilient to columns being added or reordered. _parse_amd_smi_csv() { awk -F',' ' NR == 1 { @@ -346,12 +348,19 @@ _parse_amd_smi_csv() { name = $(col["market_name"]) arch = $(col["target_graphics_version"]) mem = $(col["size"]) + serial = $(col["asic_serial"]) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", idx) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", name) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", arch) gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", mem) + gsub(/^[[:space:]]+|[[:space:]\r]+$/, "", serial) + selector = idx + serial = tolower(serial) + sub(/^0x/, "", serial) + if (serial ~ /^[0-9a-f]+$/ && serial !~ /^0+$/) + selector = "GPU-" serial if (idx ~ /^[0-9]+$/ && arch ~ /^gfx[0-9a-z]+$/ && mem ~ /^[0-9]+([.][0-9]+)?$/) - printf "%s|%s|%s|%d\n", idx, name, arch, mem + printf "%s|%s|%s|%d|%s\n", idx, name, arch, mem, selector } ' } @@ -380,7 +389,7 @@ _parse_rocm_smi_csv() { idx = dev sub(/^card/, "", idx) if (idx ~ /^[0-9]+$/ && arch ~ /^gfx[0-9a-z]+$/ && bytes ~ /^[0-9]+$/) - printf "%s|%s|%s|%d\n", idx, name, arch, bytes / 1048576 + printf "%s|%s|%s|%d|%s\n", idx, name, arch, bytes / 1048576, idx } ' } @@ -445,7 +454,7 @@ probe_host() { # NVIDIA GPU remains the default backend (RTX 3090 + Strix → cuda12), but # recording the AMD companion prevents the APU from confusing readiness # reporting and lets an explicit rocm variant remain possible. - local amd_csv="" amd_rows="" amd_primary_idx="" + local amd_csv="" amd_rows="" amd_primary_selector="" if command -v amd-smi &>/dev/null; then amd_csv=$(amd-smi static --asic --vram --csv 2>/dev/null || echo "") if [ -n "$amd_csv" ]; then @@ -467,7 +476,7 @@ probe_host() { if [ -n "$roc_arches" ]; then local amd_idx=0 arch while IFS= read -r arch; do - amd_rows+="${amd_idx}|AMD GPU|${arch}|0"$'\n' + amd_rows+="${amd_idx}|AMD GPU|${arch}|0|${amd_idx}"$'\n' amd_idx=$((amd_idx + 1)) done <<<"$roc_arches" amd_rows=${amd_rows%$'\n'} @@ -479,9 +488,9 @@ probe_host() { LUCEBOX_HOST_AMD_GPU_COUNT=$(printf '%s\n' "$amd_rows" | awk 'NF{n++} END{print n+0}') local amd_primary amd_primary=$(printf '%s\n' "$amd_rows" | sort -t'|' -k4,4nr | head -1) - local amd_idx amd_name amd_arch amd_mem_mib - IFS='|' read -r amd_idx amd_name amd_arch amd_mem_mib <<<"$amd_primary" - amd_primary_idx="$amd_idx" + local amd_idx amd_name amd_arch amd_mem_mib amd_selector + IFS='|' read -r amd_idx amd_name amd_arch amd_mem_mib amd_selector <<<"$amd_primary" + amd_primary_selector="${amd_selector:-$amd_idx}" LUCEBOX_HOST_AMD_GPU_NAME="$amd_name" LUCEBOX_HOST_AMD_GPU_ARCH="$amd_arch" LUCEBOX_HOST_AMD_VRAM_GB=$((amd_mem_mib / 1024)) @@ -539,10 +548,11 @@ probe_host() { LUCEBOX_HOST_ROCR_VISIBLE_DEVICES="${ROCR_VISIBLE_DEVICES:-}" if [ -z "$LUCEBOX_HOST_HIP_VISIBLE_DEVICES" ] \ && [ -z "$LUCEBOX_HOST_ROCR_VISIBLE_DEVICES" ]; then - # AMD recommends ROCR_VISIBLE_DEVICES for Linux. Use one isolation - # layer, not both, so a physical device index is never filtered and - # then interpreted again in a renumbered HIP-visible set. - LUCEBOX_HOST_ROCR_VISIBLE_DEVICES="$amd_primary_idx" + # AMD recommends ROCR_VISIBLE_DEVICES for Linux. Prefer the physical + # GPU UUID derived from amd-smi's ASIC serial; it remains stable even + # if SMI and ROCr enumeration orders differ. Use one isolation layer, + # not both, so an index fallback is never interpreted twice. + LUCEBOX_HOST_ROCR_VISIBLE_DEVICES="$amd_primary_selector" fi # OS / kernel identity. /etc/os-release is the freedesktop spec for @@ -1585,6 +1595,9 @@ EOF export DFLASH_DRAFT="$draft" export DRAFT="$draft" if [ "$draft" != "none" ]; then + # Match `lucebox native` and the container entrypoint: a selected + # DFlash draft uses DDTree by default. An explicit contributor + # VERIFY_MODE remains authoritative for comparison experiments. export VERIFY_MODE="${VERIFY_MODE:-ddtree}" fi local max_ctx budget diff --git a/lucebox/src/lucebox/autotune.py b/lucebox/src/lucebox/autotune.py index 553a4c5ba..e3c77610a 100644 --- a/lucebox/src/lucebox/autotune.py +++ b/lucebox/src/lucebox/autotune.py @@ -75,7 +75,11 @@ def runtime_from_host(host: HostFacts) -> DflashRuntime: # Cache types stay model-owned. In particular, tq3_0 is unsafe for # Laguna and was removed from the engine's automatic defaults in # July 2026. The CLI must not silently reintroduce it here. - return DflashRuntime(budget=budget, max_ctx=98304) + # The earlier 98K WSL result depended on a blanket tq3_0 cache. + # Automatic now preserves model-family cache types, so keep extra + # headroom for WSL's virtualization overhead instead. + max_ctx = 65536 if host.is_wsl else 98304 + return DflashRuntime(budget=budget, max_ctx=max_ctx) return DflashRuntime(budget=budget, max_ctx=131072) diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index c940a157e..fa045e0ea 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -552,6 +552,32 @@ def _render_optimization_plan( ) +def _render_custom_runtime(runtime: DflashRuntime) -> None: + """Show the exact product choices before a custom profile is committed.""" + console.print("\n[bold]Your custom profile[/bold]") + table = Table(show_header=True, box=None, pad_edge=False) + table.add_column("Optimization", style="bold") + table.add_column("Selected") + selections = ( + ("DFlash", runtime.speculative_decode), + ("PFlash", runtime.prefill_mode != "off"), + ("KVFlash", runtime.kvflash != "off"), + ("Spark", runtime.spark), + ) + for name, enabled in selections: + table.add_row(name, "[green]ON[/green]" if enabled else "[dim]off[/dim]") + console.print(table) + if runtime.cache_type_k == runtime.cache_type_v: + cache = runtime.cache_type_k or "model default" + else: + cache = f"K={runtime.cache_type_k or 'default'}, V={runtime.cache_type_v or 'default'}" + console.print( + f"Context: [bold]{runtime.max_ctx:,}[/bold] tokens · " + f"KV format: [bold]{cache}[/bold] · " + f"DFlash budget: [bold]{runtime.budget}[/bold]" + ) + + def _ensure_optimizer_drafter(cfg: Config, *, assume_yes: bool = False) -> bool: """Offer the one shared scorer asset and degrade cleanly when offline.""" if download_mod.optimizer_drafter_installed(cfg): @@ -621,7 +647,11 @@ def _customize_runtime( policy_choices += "/qk/lru" else: policy_choices += "/lru" - answer = typer.prompt("KVFlash policy", default="drafter").strip().lower() + answer = ( + typer.prompt(f"KVFlash policy ({policy_choices})", default="drafter") + .strip() + .lower() + ) if answer not in policy_choices.split("/"): fallback_policy = policy_choices.split("/")[0] console.print(f"[yellow]Unknown policy; using {fallback_policy}.[/yellow]") @@ -715,6 +745,7 @@ def optimize( if custom: runtime = _customize_runtime(cfg, plan) + _render_custom_runtime(runtime) if not typer.confirm("Apply this custom profile?", default=True): console.print("[dim]Optimization unchanged.[/dim]") return diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py index e51faeffb..762816811 100644 --- a/lucebox/src/lucebox/config.py +++ b/lucebox/src/lucebox/config.py @@ -19,6 +19,7 @@ from __future__ import annotations +import math import os import re import tomllib @@ -94,8 +95,8 @@ def _cast_positive_int(v: Any) -> int: def _cast_nonnegative_float(v: Any) -> float: value = float(v) - if value < 0.0: - raise ValueError(f"value must be zero or positive, got {value!r}") + if not math.isfinite(value) or value < 0.0: + raise ValueError(f"value must be finite and zero or positive, got {value!r}") return value diff --git a/lucebox/src/lucebox/types.py b/lucebox/src/lucebox/types.py index 6fe873c4c..6164fc6c7 100644 --- a/lucebox/src/lucebox/types.py +++ b/lucebox/src/lucebox/types.py @@ -138,8 +138,14 @@ def __post_init__(self) -> None: raise ValueError(f"budget must be zero or positive; got {self.budget!r}") if self.max_ctx <= 0: raise ValueError(f"max_ctx must be positive; got {self.max_ctx!r}") - if self.prefix_cache_slots < 0 or self.prefill_cache_slots < 0: - raise ValueError("cache slot counts must be zero or positive") + if self.prefix_cache_slots < 0: + raise ValueError( + f"prefix_cache_slots must be zero or positive; got {self.prefix_cache_slots!r}" + ) + if self.prefill_cache_slots < 0: + raise ValueError( + f"prefill_cache_slots must be zero or positive; got {self.prefill_cache_slots!r}" + ) if not 0.0 < self.prefill_keep_ratio <= 1.0: raise ValueError( "prefill_keep_ratio must be in the interval (0.0, 1.0]; " diff --git a/lucebox/tests/test_autotune.py b/lucebox/tests/test_autotune.py index 41187af01..8b481c1a0 100644 --- a/lucebox/tests/test_autotune.py +++ b/lucebox/tests/test_autotune.py @@ -10,12 +10,10 @@ def test_wsl_24gb_defaults_leave_cuda_headroom() -> None: runtime = runtime_from_host(HostFacts(vram_gb=24, is_wsl=True)) assert runtime.budget == 16 - # Bumped 65536 → 98304 on 2026-05-30 after the gemma4-26b coding- - # agent-loop sweep proved 98K serves 90K-token agentic prompts - # with ~3 GB VRAM headroom and no CUDA VMM failures on the 3090 Ti - # WSL configuration (see - # docs/experiments/gemma4-26b-coding-agent-loop-sweep-2026-05-30.md). - assert runtime.max_ctx == 98304 + # The prior 98K result relied on forcing tq3_0. Automatic profiles now + # preserve the model family's quality-safe cache type, so WSL keeps a + # conservative context cap for virtualization overhead. + assert runtime.max_ctx == 65536 # lazy is False because the heuristic path does NOT set prefill_drafter, # and the C++ server silently ignores --lazy-draft without it. Flipping # to False makes the host config match runtime behaviour. See the diff --git a/lucebox/tests/test_config.py b/lucebox/tests/test_config.py index 5d73d39ef..a1ec51043 100644 --- a/lucebox/tests/test_config.py +++ b/lucebox/tests/test_config.py @@ -38,8 +38,7 @@ def test_image_variant_round_trips_from_toml(tmp_path: Path) -> None: def test_toml_load_uses_strict_bool_and_prefill_casters(tmp_path: Path) -> None: path = tmp_path / "config.toml" path.write_text( - '[dflash]\nlazy = "false"\ndebug_thinking_logits = "false"\n' - 'prefill_mode = "auto"\n' + '[dflash]\nlazy = "false"\ndebug_thinking_logits = "false"\nprefill_mode = "auto"\n' ) cfg = config._load_toml(path) @@ -67,9 +66,7 @@ def test_toml_load_rejects_invalid_prefill_mode(tmp_path: Path) -> None: ("dflash.think_soft_close_min_ratio", "nan"), ], ) -def test_config_set_rejects_out_of_range_ratios( - tmp_path: Path, key: str, value: str -) -> None: +def test_config_set_rejects_out_of_range_ratios(tmp_path: Path, key: str, value: str) -> None: path = tmp_path / "config.toml" with pytest.raises(ValueError, match="interval"): @@ -80,11 +77,7 @@ def test_config_set_rejects_out_of_range_ratios( def test_toml_load_rejects_out_of_range_ratios(tmp_path: Path) -> None: path = tmp_path / "config.toml" - path.write_text( - "[dflash]\n" - "prefill_keep_ratio = 0.05\n" - "think_soft_close_min_ratio = 2.0\n" - ) + path.write_text("[dflash]\nprefill_keep_ratio = 0.05\nthink_soft_close_min_ratio = 2.0\n") with pytest.raises(ValueError, match="think_soft_close_min_ratio"): config._load_toml(path) @@ -97,9 +90,7 @@ def test_toml_load_rejects_out_of_range_ratios(tmp_path: Path) -> None: 'ctk = "maybe"\n', ], ) -def test_toml_load_rejects_invalid_host_literals( - tmp_path: Path, host_body: str -) -> None: +def test_toml_load_rejects_invalid_host_literals(tmp_path: Path, host_body: str) -> None: path = tmp_path / "config.toml" path.write_text(f"[host]\n{host_body}") @@ -298,8 +289,9 @@ def test_seed_dflash_writes_heuristic_when_absent(tmp_path: Path) -> None: assert wrote is True loaded = config.load(path) assert loaded is not None - # 24 GB tier heuristic → 98304, not the 16384 class default. - assert loaded.dflash.max_ctx == 98304 + # A 24 GB WSL host keeps extra virtualization headroom while still + # replacing the 16K class default. + assert loaded.dflash.max_ctx == 65536 # Provenance recorded; [model] preserved. doc = config.load_doc(path) assert doc["autotune"]["source"] == "heuristic" @@ -362,15 +354,29 @@ def test_optimization_fields_round_trip_and_validate(tmp_path: Path) -> None: ("dflash.kvflash_policy", "random"), ("dflash.kvflash_tau", "0"), ("dflash.spark_vram_gb", "-1"), + ("dflash.spark_vram_gb", "nan"), + ("dflash.spark_vram_gb", "inf"), ], ) -def test_config_rejects_invalid_optimization_values( - tmp_path: Path, key: str, value: str -) -> None: +def test_config_rejects_invalid_optimization_values(tmp_path: Path, key: str, value: str) -> None: with pytest.raises(ValueError): config_set(key, value, path=tmp_path / "config.toml") +@pytest.mark.parametrize( + ("field", "message"), + [ + ("prefix_cache_slots", "prefix_cache_slots"), + ("prefill_cache_slots", "prefill_cache_slots"), + ], +) +def test_cache_slot_validation_names_the_invalid_field(field: str, message: str) -> None: + from lucebox.types import DflashRuntime + + with pytest.raises(ValueError, match=message): + DflashRuntime(**{field: -1}) + + def test_direct_optimization_edit_marks_profile_custom(tmp_path: Path) -> None: path = tmp_path / "config.toml" diff --git a/lucebox/tests/test_models_cli.py b/lucebox/tests/test_models_cli.py index 00b0c98e6..a9f86e81c 100644 --- a/lucebox/tests/test_models_cli.py +++ b/lucebox/tests/test_models_cli.py @@ -270,6 +270,8 @@ def test_optimize_advanced_writes_a_custom_profile( ) assert result.exit_code == 0 + assert "KVFlash policy (drafter/qk/lru)" in result.output + assert "Your custom profile" in result.output loaded = config_mod.load(cfg_path) assert loaded is not None assert loaded.dflash.speculative_decode is True @@ -279,9 +281,7 @@ def test_optimize_advanced_writes_a_custom_profile( assert config_mod.optimization_mode(path=cfg_path) == "custom" -def test_installed_helpers_track_presence( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_installed_helpers_track_presence(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """``installed_status`` / ``installed_size_gb`` reflect on-disk byte counts.""" _set_config_path(tmp_path, monkeypatch) _stub_host(monkeypatch, vram_gb=24) diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index def9acbd8..77f5ba5ad 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -26,6 +26,7 @@ ROOT="$(git rev-parse --show-toplevel 2>/dev/null || (cd "$(dirname "$0")/.." && SCRIPT="$ROOT/lucebox.sh" ENTRYPOINT="$ROOT/server/scripts/entrypoint.sh" INSTALLER="$ROOT/install.sh" +HARNESS_COMMON="$ROOT/harness/clients/common.sh" SELF_PATH="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/$(basename "${BASH_SOURCE[0]}")" if [ ! -f "$SCRIPT" ]; then @@ -157,6 +158,7 @@ SHELLCHECK_TARGETS=( "$SCRIPT" "$ENTRYPOINT" "$INSTALLER" + "$HARNESS_COMMON" ) # Add every scripts/*.sh except this one (don't recurse into our own tests). while IFS= read -r -d '' f; do @@ -182,6 +184,8 @@ if bash -n "$SCRIPT"; then report ok "bash -n lucebox.sh parses cleanly" else report fail "bash -n lucebox.sh"; fi if bash -n "$ENTRYPOINT"; then report ok "bash -n entrypoint.sh parses cleanly" else report fail "bash -n entrypoint.sh"; fi +if bash -n "$HARNESS_COMMON"; then report ok "bash -n harness common.sh parses cleanly" +else report fail "bash -n harness common.sh"; fi # ── 3. Trivial subcommands (zero-exit expected) ─────────────────────────── assert_runs "help" "bash '$SCRIPT' help" "simple CLI for the Lucebox inference engine" @@ -592,6 +596,47 @@ test_entrypoint_optimization_flags() { test_entrypoint_optimization_flags \ "entrypoint forwards PFlash, KVFlash, and Spark exactly once" +test_harness_draft_directory_validation() { + local label="$1" sandbox target server ambiguous empty out rc out_empty rc_empty + sandbox=$(mktemp -d -t lucebox-harness-draft.XXXXXX) + target="$sandbox/target.gguf" + server="$sandbox/dflash_server" + ambiguous="$sandbox/ambiguous" + empty="$sandbox/empty" + mkdir -p "$ambiguous" "$empty" "$sandbox/work" + printf 'target' > "$target" + printf 'draft-a' > "$ambiguous/a.gguf" + printf 'draft-b' > "$ambiguous/b.safetensors" + printf '#!/usr/bin/env bash\nexit 0\n' > "$server" + chmod +x "$server" + + if out=$(REPO_DIR="$ROOT" CLIENT_WORK_DIR="$sandbox/work" \ + TARGET="$target" DRAFT="$ambiguous" DFLASH_SERVER_BIN="$server" \ + bash -c 'source "$1"; start_dflash_native_server' _ "$HARNESS_COMMON" 2>&1); then + rc=0 + else + rc=$? + fi + if out_empty=$(REPO_DIR="$ROOT" CLIENT_WORK_DIR="$sandbox/work" \ + TARGET="$target" DRAFT="$empty" DFLASH_SERVER_BIN="$server" \ + bash -c 'source "$1"; start_dflash_native_server' _ "$HARNESS_COMMON" 2>&1); then + rc_empty=0 + else + rc_empty=$? + fi + rm -rf "$sandbox" + + if [ "$rc" -eq 0 ] || [[ "$out" != *"multiple DFlash draft candidates"* ]]; then + report fail "$label" "ambiguous directory was accepted: $out" + elif [ "$rc_empty" -eq 0 ] || [[ "$out_empty" != *"DFlash draft not found"* ]]; then + report fail "$label" "empty directory was accepted: $out_empty" + else + report ok "$label" + fi +} +test_harness_draft_directory_validation \ + "native harness rejects empty or ambiguous draft directories" + test_entrypoint_keeps_family_kv_default() { local label="$1" sandbox draft_dir models_dir bin_dir shim_dir out _make_entrypoint_sandbox @@ -1474,12 +1519,12 @@ test_amd_smi_parser() { fn=$(awk '/^_parse_amd_smi_csv\(\) \{/,/^\}/' "$SCRIPT") out=$(bash -c "$fn"$'\n''_parse_amd_smi_csv' <<'CSV' gpu,market_name,vendor_id,vendor_name,subvendor_id,device_id,subsystem_id,rev_id,asic_serial,oam_id,num_compute_units,target_graphics_version,type,vendor,size,bit_width,max_bandwidth -0,AMD Radeon AI PRO R9700,0x1002,AMD,0xf111,0x7551,0x000a,0xc0,serial,N/A,64,gfx1201,GDDR6,SAMSUNG,32624,256,N/A -1,AMD Radeon Graphics,0x1002,AMD,0xf111,0x1586,0x000a,0xc1,serial,N/A,40,gfx1151,GDDR7,UNKNOWN,512,256,N/A +0,AMD Radeon AI PRO R9700,0x1002,AMD,0xf111,0x7551,0x000a,0xc0,0xE099917E6553AFAA,N/A,64,gfx1201,GDDR6,SAMSUNG,32624,256,N/A +1,AMD Radeon Graphics,0x1002,AMD,0xf111,0x1586,0x000a,0xc1,0x0000000000000000,N/A,40,gfx1151,GDDR7,UNKNOWN,512,256,N/A CSV ) - if ! grep -qF '0|AMD Radeon AI PRO R9700|gfx1201|32624' <<<"$out" \ - || ! grep -qF '1|AMD Radeon Graphics|gfx1151|512' <<<"$out"; then + if ! grep -qF '0|AMD Radeon AI PRO R9700|gfx1201|32624|GPU-e099917e6553afaa' <<<"$out" \ + || ! grep -qF '1|AMD Radeon Graphics|gfx1151|512|1' <<<"$out"; then report fail "$label" "unexpected normalized rows: $out" return fi diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index eb789c82e..500f03f06 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -331,13 +331,19 @@ if [ "$GPU_VRAM_GB" -gt 0 ]; then elif [ "$GPU_VRAM_GB" -lt 32 ]; then if [ "$IS_WSL" = "1" ]; then : "${DFLASH_BUDGET:=16}" + : "${DFLASH_MAX_CTX:=65536}" + else + : "${DFLASH_MAX_CTX:=98304}" fi - : "${DFLASH_MAX_CTX:=98304}" else : "${DFLASH_MAX_CTX:=131072}" fi fi +# Do not synthesize DFLASH_LAZY here. dflash_server ignores --lazy-draft +# unless both a decode draft and a prefill scorer are configured, so setting +# it from VRAM alone would claim an optimization the runtime cannot apply. + [ "${GPU_ARCH:-}" = "gfx1100" ] && : "${DFLASH_BUDGET:=8}" : "${DFLASH_BIN:=$DFLASH_DIR/build/test_dflash}" From d15fc1f893d97e3ba72d958f8b5fd61b5d5df8a6 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:55:23 +0200 Subject: [PATCH 11/25] feat(lucebox): connect installed clients safely --- README.md | 22 +- harness/clients/run_hermes.sh | 7 +- harness/clients/run_openclaw.sh | 3 +- lucebox.sh | 584 +++++++++++++++++++++++++++++- lucebox/src/lucebox/__init__.py | 8 +- lucebox/src/lucebox/docker_run.py | 5 + lucebox/tests/test_docker_run.py | 2 + scripts/test_lucebox_sh.sh | 226 ++++++++++++ 8 files changed, 841 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 72e696624..d34a281c3 100644 --- a/README.md +++ b/README.md @@ -128,10 +128,28 @@ machine. Automatic tuning never adds VRAM across devices. The engine's advanced layer-split paths remain contributor workflows; transparent multi-GPU sharding is not part of this CLI release. +**Connect or open your harness** links the running local API to a client the +engine user already installed: Claude Code, Codex, OpenCode, Hermes, Pi, +OpenClaw, or Open WebUI. It is available from the same main menu in an installed +Lucebox and a source checkout. Lucebox does not install the client and does not +replace its normal cloud configuration. It launches the installed binary with +an invocation-only override or an additive, Lucebox-owned profile, remembers +the selection, and shows it on the main menu. For example: + +```bash +lucebox connect codex # link and open Codex +lucebox connect opencode --no-launch +``` + +The connector verifies the local server before opening the client. Starting +the same client normally, without `lucebox connect`, continues to use the +user's existing provider and account. + Contributors can open the same menu from a repository checkout with `./lucebox.sh`. Its **Developer tools** can build and run the native C++ engine -or launch Claude Code, Codex, OpenCode, Hermes, Pi, OpenClaw, and Open WebUI -harnesses. Every menu action also has a scriptable command, for example: +or run isolated compatibility harnesses for Claude Code, Codex, OpenCode, +Hermes, Pi, OpenClaw, and Open WebUI. Every menu action also has a scriptable +command, for example: ```bash lucebox models select diff --git a/harness/clients/run_hermes.sh b/harness/clients/run_hermes.sh index d5207d3b9..50a403b27 100755 --- a/harness/clients/run_hermes.sh +++ b/harness/clients/run_hermes.sh @@ -18,7 +18,7 @@ mkdir -p "$HOME_DIR" cat > "$HOME_DIR/config.yaml" < "$HOME_DIR/.env" < "$CONFIG_PATCH" </dev/null + return + fi + command -v "$command_name" 2>/dev/null +} + +_connector_selected() { + local selected + [ -f "$CONNECTOR_SELECTION_FILE" ] || return 0 + IFS= read -r selected < "$CONNECTOR_SELECTION_FILE" || return 0 + _connector_normalize "$selected" 2>/dev/null || true +} + +_connector_write_file() { + # Atomically replace a Lucebox-owned connector file with private mode. + # Content is read from stdin so callers can use a quoted or expanded + # heredoc without lossy shell escaping. + local target="$1" mode="${2:-600}" parent tmp + parent=$(dirname "$target") + mkdir -p "$parent" + tmp=$(mktemp "${target}.tmp.XXXXXX") \ + || die "couldn't create a temporary connector file next to $target" + if ! cat > "$tmp"; then + rm -f "$tmp" + die "couldn't write connector file: $target" + fi + chmod "$mode" "$tmp" + mv "$tmp" "$target" +} + +_connector_private_dir() { + local path="$1" + mkdir -p "$path" + chmod 700 "$path" +} + +_connector_remember() { + local client="$1" + _connector_private_dir "$CONNECTOR_STATE_DIR" + _connector_write_file "$CONNECTOR_SELECTION_FILE" 600 </dev/null 2>&1 || return 1 + response=$(curl -fsS --connect-timeout 2 --max-time 4 \ + "$api_root/v1/models" 2>/dev/null) || return 1 + # A generic 200 response from another process on the configured port is + # not sufficient. Require an OpenAI-compatible model catalog shape. + [[ "$response" == *'"data"'* || "$response" == *'"models"'* ]] +} + +_connector_ensure_api() { + local api_root="$1" i state + if ! command -v curl >/dev/null 2>&1; then + err "curl is required to verify the local Lucebox API" + return 1 + fi + if _connector_api_ready "$api_root"; then + return 0 + fi + ensure_probed + state=$(_engine_state) + if [ "$state" = "stopped" ] && [ -t 0 ] \ + && [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ] \ + && [ -f "$UNIT_PATH" ] \ + && _confirm "The inference engine is stopped. Start it now?" 1; then + bash "$SCRIPT_PATH" start || return $? + state=running + for i in 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15; do + _connector_api_ready "$api_root" && return 0 + sleep 1 + done + fi + if [ "$state" = "running" ]; then + err "the inference engine is running, but its API is not ready at $api_root" + hint "Wait for model loading to finish, or inspect: $SCRIPT_NAME logs" + else + err "the Lucebox API is not reachable at $api_root" + hint "Start it from the menu or run: $SCRIPT_NAME start" + fi + return 1 +} + +_connector_print_argv() { + local arg + printf ' ' + for arg in "$@"; do + printf '%q ' "$arg" + done + printf '\n' +} + +_connector_choose() { + local client label selected status number selected_number="" + selected=$(_connector_selected) + printf 'Choose an installed harness to connect:\n' + number=0 + for client in claude_code codex opencode hermes pi openclaw openwebui; do + number=$((number + 1)) + label=$(_connector_label "$client") + if _connector_binary "$client" >/dev/null 2>&1; then + status="installed" + [ "$client" = "$selected" ] && selected_number="$number" + else + status="not found" + fi + [ "$client" = "$selected" ] && status="$status, selected" + printf ' %s %-12s %b(%s)%b\n' "$number" "$label" "$C_DIM" "$status" "$C_RST" + done + printf ' b Back\n\n' + if [ -n "$selected_number" ]; then + printf 'Harness number [%s]: ' "$selected_number" + else + printf 'Harness number: ' + fi + IFS= read -r client || return 1 + case "$client" in b|B|q|Q) return 1 ;; esac + [ -n "$client" ] || client="$selected_number" + [ -n "$client" ] || return 1 + CONNECTOR_CHOICE=$(_connector_normalize "$client") \ + || { warn "Choose 1–7 or b"; return 1; } +} + +# The prepare helpers generate only Lucebox-owned files or additive named +# profiles. They populate CONNECTOR_SETUP_ARGV (optional) and +# CONNECTOR_LAUNCH_ARGV; the caller applies setup before remembering the choice. +_connector_prepare_claude() { + local binary="$1" api_root="$2" model="$3" + CONNECTOR_LAUNCH_ARGV=( + env + -u CLAUDE_CODE_USE_BEDROCK + -u CLAUDE_CODE_USE_VERTEX + -u CLAUDE_CODE_USE_FOUNDRY + -u ANTHROPIC_API_KEY + -u CLAUDE_CODE_OAUTH_TOKEN + "ANTHROPIC_BASE_URL=$api_root" + "ANTHROPIC_AUTH_TOKEN=lucebox-local" + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1" + "$binary" --model "$model" + ) + CONNECTOR_NOTE="Only this Claude Code session uses Lucebox; your normal Claude configuration is unchanged." +} + +_connector_prepare_codex() { + local binary="$1" api_root="$2" model="$3" codex_home profile + codex_home="${CODEX_HOME:-$HOME/.codex}" + profile="$codex_home/lucebox-local.config.toml" + mkdir -p "$codex_home" + if [ -e "$profile" ] \ + && ! grep -Fqx '# Generated by Lucebox. Normal Codex configuration is not modified.' "$profile"; then + die "refusing to replace an unowned Codex profile: $profile" + fi + _connector_write_file "$profile" 600 </dev/null | head -n 1 || true) + if [[ "$version" =~ (^|[^0-9])([0-9]+)\.[0-9]+ ]]; then + major="${BASH_REMATCH[2]}" + fi + if [ "$major" -ge 2 ]; then + _connector_write_file "$config" 600 </dev/null 2>&1 \ + || connector_label="$connector_label (not found)" + else + connector_label="not selected" + fi _menu_clear print_logo @@ -2280,7 +2849,8 @@ cmd_menu() { printf ' Backend: %s\n' "$variant" printf ' Model: %s\n' "$model" printf ' Optimization: %s\n' "$optimization" - printf ' Engine: %s\n\n' "$state" + printf ' Engine: %s\n' "$state" + printf ' Harness: %s\n\n' "$connector_label" printf ' 1 Quick setup\n' printf ' 2 Choose or download a model\n' @@ -2289,6 +2859,7 @@ cmd_menu() { printf ' 5 Stop the inference engine\n' printf ' 6 Status\n' printf ' 7 Recent logs\n' + printf ' 8 Connect or open your harness\n' printf ' d Developer tools %b(%s)%b\n' "$C_DIM" "$repo_hint" "$C_RST" printf ' q Quit\n\n' printf 'Choose: ' @@ -2314,9 +2885,10 @@ cmd_menu() { fi _menu_pause ;; + 8) _menu_run connect; _menu_pause ;; d|D) cmd_developer_menu ;; q|Q|quit|exit) return 0 ;; - *) warn "Choose 1–7, d, or q"; _menu_pause ;; + *) warn "Choose 1–8, d, or q"; _menu_pause ;; esac done } @@ -2329,6 +2901,7 @@ Interactive: $SCRIPT_NAME open the branded menu (when run in a terminal) menu open the menu explicitly setup guided backend, model, optimization, and service setup + connect [harness] link an installed AI client to the local Lucebox API Service management (via user systemd): install install user systemd unit @@ -2417,6 +2990,7 @@ main() { # Branded interactive surface / guided first run. menu) cmd_menu "$@" ;; setup) cmd_setup "$@" ;; + connect) cmd_connect "$@" ;; # Systemd surface install) cmd_systemd_install "$@" ;; diff --git a/lucebox/src/lucebox/__init__.py b/lucebox/src/lucebox/__init__.py index 8ca821024..79e8f7a8e 100644 --- a/lucebox/src/lucebox/__init__.py +++ b/lucebox/src/lucebox/__init__.py @@ -1,10 +1,10 @@ """lucebox — host-side CLI for the lucebox-hub container. Runs inside the container; the host wrapper at ../lucebox.sh handles `docker -run` plumbing and systemd integration. This package owns: TOML config, the -host-derived DFLASH_* serve heuristic, docker daemon calls (via the mounted -socket), and model download. The empirical autotune sweep, profiling, and -agent-client launchers land in follow-up PRs. +run` plumbing, systemd integration, and isolated profiles for already-installed +AI clients. This package owns: TOML config, the host-derived DFLASH_* serve +heuristic, docker daemon calls (via the mounted socket), and model download. +The empirical autotune sweep and profiling remain separate workflows. """ # Version is generated by hatch-vcs at build time into _version.py. diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index 7e866fb78..853128369 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -274,6 +274,11 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: ("DFLASH_THINK_MAX", str(cfg.dflash.think_max)), ("DFLASH_PORT", "8080"), ] + # Keep the API model catalog aligned with the preset selected by the host + # CLI. Client connectors use this id for explicit model selection and + # clients such as Codex also discover it through ``/v1/models``. + if cfg.model.preset: + env.append(("DFLASH_MODEL_NAME", cfg.model.preset)) # Resolve target/draft GGUFs in priority order: # 1. cfg.model.target_file / draft_file (explicit override in config.toml) # 2. PRESETS[cfg.model.preset].target_file / draft_file / speculator_dir (registry) diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index 053046fd9..251a6028c 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -390,6 +390,7 @@ def test_server_run_spec_target_only_preset_disables_stale_draft(tmp_path: Path) assert env["DFLASH_TARGET"].endswith("Qwen3.6-35B-A3B-UD-Q4_K_M.gguf") assert env["DFLASH_DRAFT"].endswith("/.lucebox-no-draft") + assert env["DFLASH_MODEL_NAME"] == "qwen3.6-moe" def test_server_run_spec_optional_env_off_by_default(tmp_path: Path) -> None: @@ -410,6 +411,7 @@ def test_server_run_spec_optional_env_off_by_default(tmp_path: Path) -> None: "DFLASH_DEBUG_THINKING_LOGITS", "DFLASH_TARGET", "DFLASH_DRAFT", + "DFLASH_MODEL_NAME", ): assert absent not in env diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index 77f5ba5ad..a2edb094d 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -289,6 +289,232 @@ else fi rm -rf "$native_tmp" +# Engine client connectors: use already-installed binaries, preserve every +# normal client config, and point a Lucebox-owned profile at the local API. +# Fake binaries make this an execution-path test without installing or +# launching any real client. +test_engine_client_connectors() { + local label="engine client connectors are isolated and executable" + local sandbox bin_dir state_dir home_dir log out rc client command_name + sandbox=$(mktemp -d "${TMPDIR:-/tmp}/lucebox-connectors.XXXXXX") + bin_dir="$sandbox/bin" + state_dir="$sandbox/lucebox-home" + home_dir="$sandbox/home" + log="$sandbox/client.log" + mkdir -p "$bin_dir" "$state_dir" "$home_dir/.codex" "$home_dir/.config/opencode" + cat > "$state_dir/config.toml" <<'TOML' +[runtime] +port = 18080 + +[model] +preset = "qwen3.6-27b" + +[dflash] +max_ctx = 98304 +TOML + printf 'personal-codex-config\n' > "$home_dir/.codex/config.toml" + printf 'personal-opencode-config\n' > "$home_dir/.config/opencode/opencode.json" + + cat > "$bin_dir/curl" <<'STUB' +#!/usr/bin/env bash +printf '%s\n' '{"object":"list","data":[{"id":"qwen3.6-27b"}]}' +exit 0 +STUB + chmod +x "$bin_dir/curl" + for command_name in claude codex opencode hermes pi openclaw open-webui; do + cat > "$bin_dir/$command_name" <<'STUB' +#!/usr/bin/env bash +connector_context="${ANTHROPIC_BASE_URL:-}" +[ -n "$connector_context" ] || connector_context="${OPENCODE_CONFIG:-}" +[ -n "$connector_context" ] || connector_context="${HERMES_HOME:-}" +[ -n "$connector_context" ] || connector_context="${PI_CODING_AGENT_DIR:-}" +[ -n "$connector_context" ] || connector_context="${OPENCLAW_STATE_DIR:-}" +[ -n "$connector_context" ] || connector_context="${OPENAI_API_BASE_URL:-}" +printf '%s|%s|%s|%s|%s|%s\n' \ + "$(basename "$0")" "$*" "$connector_context" \ + "${OPENAI_BASE_URL:-}" "${OPENCLAW_WORKSPACE_DIR:-}" "${WEBUI_AUTH:-}" \ + >> "${LUCEBOX_TEST_CONNECTOR_LOG:?}" +if [ "$(basename "$0")" = "claude" ] \ + && { [ -n "${ANTHROPIC_API_KEY:-}" ] || [ -n "${CLAUDE_CODE_OAUTH_TOKEN:-}" ]; }; then + exit 42 +fi +if [ "$(basename "$0")" = "opencode" ] && [ "${1:-}" = "--version" ]; then + printf '2.0.0\n' +fi +if [ "$(basename "$0")" = "openclaw" ] \ + && [ "${1:-}" = "config" ] && [ "${2:-}" = "patch" ]; then + mkdir -p "$(dirname "${OPENCLAW_CONFIG_PATH:?}")" + cp "${4:?}" "$OPENCLAW_CONFIG_PATH" +fi +exit 0 +STUB + chmod +x "$bin_dir/$command_name" + done + + rc=0 + for client in claude codex opencode hermes pi openclaw openwebui; do + if ! out=$(HOME="$home_dir" XDG_CONFIG_HOME="$home_dir/.config" \ + PATH="$bin_dir:$PATH" LUCEBOX_HOME="$state_dir" \ + LUCEBOX_TEST_CONNECTOR_LOG="$log" NO_COLOR=1 \ + ANTHROPIC_API_KEY=personal-key CLAUDE_CODE_OAUTH_TOKEN=personal-oauth \ + bash "$SCRIPT" connect "$client" 2>&1); then + rc=1 + break + fi + if [[ "$out" != *"is linked to Lucebox"* ]]; then + rc=1 + break + fi + done + + local missing="" + [ "$rc" -eq 0 ] || missing+=" launch-$client" + grep -qF 'personal-codex-config' "$home_dir/.codex/config.toml" \ + || missing+=" codex-default-overwritten" + grep -qF 'personal-opencode-config' "$home_dir/.config/opencode/opencode.json" \ + || missing+=" opencode-default-overwritten" + grep -qF 'base_url = "http://127.0.0.1:18080/v1"' \ + "$home_dir/.codex/lucebox-local.config.toml" \ + || missing+=" codex-profile" + grep -qF '"providers"' "$state_dir/connectors/opencode/opencode.json" \ + || missing+=" opencode-v2-profile" + grep -qF 'http://127.0.0.1:18080/v1' "$state_dir/connectors/hermes/config.yaml" \ + || missing+=" hermes-profile" + grep -qF 'provider: "custom"' "$state_dir/connectors/hermes/config.yaml" \ + || missing+=" hermes-custom-provider" + grep -qF '"baseUrl": "http://127.0.0.1:18080/v1"' \ + "$state_dir/connectors/pi/models.json" \ + || missing+=" pi-profile" + grep -qF '"api": "openai-completions"' "$state_dir/connectors/pi/models.json" \ + || missing+=" pi-chat-completions" + grep -qF '"primary": "lucebox/qwen3.6-27b"' \ + "$state_dir/connectors/openclaw/lucebox.patch.json" \ + || missing+=" openclaw-profile" + cmp -s "$state_dir/connectors/openclaw/lucebox.patch.json" \ + "$state_dir/connectors/openclaw/state/openclaw.json" \ + || missing+=" openclaw-isolation" + python3 -m json.tool "$state_dir/connectors/opencode/opencode.json" >/dev/null \ + || missing+=" opencode-invalid-json" + python3 -m json.tool "$state_dir/connectors/pi/models.json" >/dev/null \ + || missing+=" pi-invalid-json" + python3 -m json.tool "$state_dir/connectors/openclaw/lucebox.patch.json" >/dev/null \ + || missing+=" openclaw-invalid-json" + python3 -c 'import sys, tomllib; tomllib.load(open(sys.argv[1], "rb"))' \ + "$home_dir/.codex/lucebox-local.config.toml" \ + || missing+=" codex-invalid-toml" + grep -qF 'open-webui|serve --host 127.0.0.1 --port 3000' "$log" \ + || missing+=" openwebui-launch" + grep -qF 'open-webui|serve --host 127.0.0.1 --port 3000|http://127.0.0.1:18080/v1|||True' "$log" \ + || missing+=" openwebui-auth" + grep -qF 'claude|--model qwen3.6-27b|http://127.0.0.1:18080' "$log" \ + || missing+=" claude-endpoint" + grep -qF 'codex|--profile lucebox-local --model qwen3.6-27b' "$log" \ + || missing+=" codex-launch" + [ "$(cat "$state_dir/connectors/selected")" = "openwebui" ] \ + || missing+=" remembered-selection" + [ -z "$(find "$state_dir/connectors" -type f ! -perm 600 -print -quit)" ] \ + || missing+=" connector-file-mode" + [ -z "$(find "$state_dir/connectors" -type d ! -perm 700 -print -quit)" ] \ + || missing+=" connector-directory-mode" + [ "$(find "$home_dir/.codex/lucebox-local.config.toml" -perm 600 -print)" ] \ + || missing+=" codex-profile-mode" + + if [ -n "$missing" ]; then + report fail "$label" "missing:$missing; output=${out:-}; log=$(tail -10 "$log" 2>/dev/null)" + else + report ok "$label" + fi + + out=$(HOME="$home_dir" PATH="$bin_dir:$PATH" LUCEBOX_HOME="$state_dir" \ + LUCEBOX_CODEX_BIN="$sandbox/does-not-exist/codex" NO_COLOR=1 \ + bash "$SCRIPT" connect codex --no-launch 2>&1 || true) + if [[ "$out" == *"does not install harnesses"* ]]; then + report ok "missing harness is reported without installing it" + else + report fail "missing harness is reported without installing it" "output: $out" + fi + + local unowned_home="$sandbox/unowned-home" + mkdir -p "$unowned_home/.codex" + printf 'user-owned-lucebox-profile\n' > "$unowned_home/.codex/lucebox-local.config.toml" + out=$(HOME="$unowned_home" PATH="$bin_dir:$PATH" LUCEBOX_HOME="$state_dir" \ + LUCEBOX_TEST_CONNECTOR_LOG="$log" NO_COLOR=1 \ + bash "$SCRIPT" connect codex --no-launch 2>&1 || true) + if [[ "$out" == *"refusing to replace an unowned Codex profile"* ]] \ + && grep -qF 'user-owned-lucebox-profile' \ + "$unowned_home/.codex/lucebox-local.config.toml"; then + report ok "Codex connector refuses to overwrite an unowned profile" + else + report fail "Codex connector refuses to overwrite an unowned profile" "output: $out" + fi + + local offline_bin="$sandbox/offline-bin" before_lines after_lines + mkdir -p "$offline_bin" + cp "$bin_dir/codex" "$offline_bin/codex" + cat > "$offline_bin/curl" <<'STUB' +#!/usr/bin/env bash +exit 1 +STUB + chmod +x "$offline_bin/curl" + before_lines=$(wc -l < "$log" | tr -d ' ') + out=$(HOME="$home_dir" PATH="$offline_bin:$PATH" LUCEBOX_HOME="$state_dir" \ + LUCEBOX_TEST_CONNECTOR_LOG="$log" NO_COLOR=1 \ + bash "$SCRIPT" connect codex 2>&1 || true) + after_lines=$(wc -l < "$log" | tr -d ' ') + if [[ "$out" == *"Lucebox API is not reachable"* ]] \ + && [ "$before_lines" = "$after_lines" ]; then + report ok "connector does not launch a client when the API is unavailable" + else + report fail "connector does not launch a client when the API is unavailable" \ + "before=$before_lines after=$after_lines output=$out" + fi + + local small_state="$sandbox/small-state" + mkdir -p "$small_state" + cat > "$small_state/config.toml" <<'TOML' +[runtime] +port = 18080 + +[model] +preset = "qwen3.6-27b" + +[dflash] +max_ctx = 4096 +TOML + out=$(HOME="$home_dir" PATH="$bin_dir:$PATH" LUCEBOX_HOME="$small_state" \ + LUCEBOX_TEST_CONNECTOR_LOG="$log" NO_COLOR=1 \ + bash "$SCRIPT" connect pi --no-launch 2>&1 || true) + if grep -qF '"contextWindow": 4096' "$small_state/connectors/pi/models.json" \ + && grep -qF '"maxTokens": 2048' "$small_state/connectors/pi/models.json"; then + report ok "connector output allowance fits a small context profile" + else + report fail "connector output allowance fits a small context profile" "output: $out" + fi + + out=$(HOME="$home_dir" PATH="$bin_dir:$PATH" LUCEBOX_HOME="$small_state" \ + LUCEBOX_TEST_CONNECTOR_LOG="$log" NO_COLOR=1 \ + bash "$SCRIPT" connect hermes --no-launch 2>&1 || true) + if [[ "$out" == *"Hermes requires at least a 65536-token context"* ]]; then + report ok "Hermes connector rejects an unsupported small context" + else + report fail "Hermes connector rejects an unsupported small context" "output: $out" + fi + + out=$(HOME="$home_dir" XDG_CONFIG_HOME="$home_dir/.config" \ + PATH="$bin_dir:$PATH" LUCEBOX_HOME="$state_dir" \ + LUCEBOX_TEST_CONNECTOR_LOG="$log" LUCEBOX_NO_CLEAR=1 NO_COLOR=1 \ + bash -c "printf 'q\\n' | bash '$SCRIPT' menu" 2>&1 || true) + if [[ "$out" == *"Harness: Codex"* ]] \ + && [[ "$out" == *"Connect or open your harness"* ]]; then + report ok "main menu shows the selected engine harness" + else + report fail "main menu shows the selected engine harness" "output: $(tail -20 <<<"$out")" + fi + + rm -rf "$sandbox" +} +test_engine_client_connectors + # ── 7. Unknown subcommand → cmd_in_container fallback path. Same rule: # clean error, no raw bash leak. assert_no_set_u_leak "unknown subcommand dispatch" "$SCRIPT" no-such-subcommand From e6d2ae5ae25ef020d44747f02a49350991705877 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:45:56 +0200 Subject: [PATCH 12/25] fix(ci): include CMake test discovery in images --- scripts/test_lucebox_sh.sh | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index a2edb094d..22eef4ea8 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -187,6 +187,20 @@ else report fail "bash -n entrypoint.sh"; fi if bash -n "$HARNESS_COMMON"; then report ok "bash -n harness common.sh parses cleanly" else report fail "bash -n harness common.sh"; fi +# The CMake unit-test target runs a post-link discovery helper. Dockerfiles +# copy build inputs selectively for cache efficiency, so omitting server/cmake +# does not fail until the final link has already spent several minutes. Keep +# the CUDA and ROCm build contexts aligned with that CMake contract. +for dockerfile in Dockerfile Dockerfile.rocm; do + if [ -f "$ROOT/server/cmake/DiscoverCppUnitTests.cmake" ] \ + && grep -Eq '^COPY[[:space:]]+server/cmake([[:space:]]|/)' "$ROOT/$dockerfile"; then + report ok "$dockerfile includes CMake discovery helpers" + else + report fail "$dockerfile includes CMake discovery helpers" \ + "server/cmake/DiscoverCppUnitTests.cmake would be missing from the image build context" + fi +done + # ── 3. Trivial subcommands (zero-exit expected) ─────────────────────────── assert_runs "help" "bash '$SCRIPT' help" "simple CLI for the Lucebox inference engine" assert_runs "--help" "bash '$SCRIPT' --help" "simple CLI for the Lucebox inference engine" From f3da3e1f483ce3f91980c1d9c9cce914a9165441 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:46:24 +0200 Subject: [PATCH 13/25] feat(lucebox): automate heterogeneous inference placement --- Dockerfile | 9 +- Dockerfile.rocm | 9 +- README.md | 23 +- lucebox.sh | 570 +++++++++++++++--- lucebox/README.md | 46 +- lucebox/pyproject.toml | 2 +- lucebox/src/lucebox/autotune.py | 73 ++- lucebox/src/lucebox/cli.py | 142 +++-- lucebox/src/lucebox/config.py | 139 ++++- lucebox/src/lucebox/docker_run.py | 39 +- lucebox/src/lucebox/download.py | 14 +- lucebox/src/lucebox/host_facts.py | 11 + lucebox/src/lucebox/placement.py | 434 +++++++++++++ lucebox/src/lucebox/topology.py | 230 +++++++ lucebox/src/lucebox/types.py | 145 +++++ lucebox/tests/test_check.py | 4 + lucebox/tests/test_cli.py | 3 +- lucebox/tests/test_config.py | 97 +++ lucebox/tests/test_config_cli.py | 23 +- lucebox/tests/test_docker_run.py | 174 ++++-- lucebox/tests/test_placement.py | 292 +++++++++ scripts/test_lucebox_sh.sh | 344 ++++++++++- server/scripts/entrypoint.sh | 57 +- .../src/common/dflash_layer_split_runtime.h | 30 + server/src/common/peer_access.cpp | 73 ++- server/src/common/peer_access.h | 2 +- .../src/gemma4/gemma4_layer_split_adapter.cpp | 18 +- server/src/ipc/backend_ipc_main.cpp | 1 + .../src/laguna/laguna_layer_split_adapter.cpp | 9 +- server/src/qwen35/layer_split_forward.cpp | 6 +- 30 files changed, 2763 insertions(+), 256 deletions(-) create mode 100644 lucebox/src/lucebox/placement.py create mode 100644 lucebox/src/lucebox/topology.py create mode 100644 lucebox/tests/test_placement.py diff --git a/Dockerfile b/Dockerfile index 20e93b2b4..cfeb6f29d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -100,16 +100,18 @@ RUN cmake -S /src/server -B /src/server/build \ -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ -DDFLASH27B_USER_CUDA_ARCHITECTURES="${DFLASH_CUDA_ARCHES}" \ -DCMAKE_CUDA_ARCHITECTURES="${DFLASH_CUDA_ARCHES}" \ - && cmake --build /src/server/build --target test_dflash dflash_server test_server_unit --parallel + && cmake --build /src/server/build \ + --target test_dflash dflash_server backend_ipc_daemon test_server_unit --parallel # Prune the build tree to only what the runtime stage needs: the native server, -# test_dflash, test_server_unit, and the ggml shared libs their embedded rpath +# backend IPC companion, tests, and the ggml shared libs their embedded rpath # ($ORIGIN/deps/...) looks up. Drops ~1 GB per image of CMakeFiles/, # libdflash27b.a (statically linked into the binaries), ninja state, # compile_commands.json, and the template-instance .o tree from ggml-cuda. RUN cd /src/server/build \ && find . -mindepth 1 -maxdepth 1 \ - ! -name test_dflash ! -name dflash_server ! -name test_server_unit ! -name deps -exec rm -rf {} + \ + ! -name test_dflash ! -name dflash_server ! -name backend_ipc_daemon \ + ! -name test_server_unit ! -name deps -exec rm -rf {} + \ && find deps -mindepth 1 -type f ! -name 'lib*.so*' -delete \ && find deps -depth -type d -empty -delete @@ -218,6 +220,7 @@ RUN mkdir -p /opt/lucebox-hub/server/share \ RUN test -x /opt/lucebox-hub/server/build/test_dflash \ && test -x /opt/lucebox-hub/server/build/dflash_server \ + && test -x /opt/lucebox-hub/server/build/backend_ipc_daemon \ && test -x /opt/lucebox-hub/server/build/test_server_unit \ && test -f /opt/lucebox-hub/server/share/model_cards/qwen3.6-27b.json \ && chmod +x /opt/lucebox-hub/server/scripts/entrypoint.sh diff --git a/Dockerfile.rocm b/Dockerfile.rocm index 5cfd9b373..007f805a1 100644 --- a/Dockerfile.rocm +++ b/Dockerfile.rocm @@ -106,14 +106,16 @@ RUN cmake -S /src/server -B /src/server/build \ -DDFLASH27B_HIP_ARCHITECTURES="${DFLASH_HIP_ARCHES}" \ -DDFLASH27B_FA_ALL_QUANTS=OFF \ -DDFLASH27B_ENABLE_BSA=OFF \ - && cmake --build /src/server/build --target test_dflash dflash_server test_server_unit --parallel + && cmake --build /src/server/build \ + --target test_dflash dflash_server backend_ipc_daemon test_server_unit --parallel # Prune the build tree to only what the runtime stage needs: the native server, -# test_dflash, test_server_unit, and the ggml shared libs their embedded rpath +# backend IPC companion, tests, and the ggml shared libs their embedded rpath # ($ORIGIN/deps/...) looks up. RUN cd /src/server/build \ && find . -mindepth 1 -maxdepth 1 \ - ! -name test_dflash ! -name dflash_server ! -name test_server_unit ! -name deps -exec rm -rf {} + \ + ! -name test_dflash ! -name dflash_server ! -name backend_ipc_daemon \ + ! -name test_server_unit ! -name deps -exec rm -rf {} + \ && find deps -mindepth 1 -type f ! -name 'lib*.so*' -delete \ && find deps -depth -type d -empty -delete @@ -194,6 +196,7 @@ RUN mkdir -p /opt/lucebox-hub/server/share \ RUN test -x /opt/lucebox-hub/server/build/test_dflash \ && test -x /opt/lucebox-hub/server/build/dflash_server \ + && test -x /opt/lucebox-hub/server/build/backend_ipc_daemon \ && test -x /opt/lucebox-hub/server/build/test_server_unit \ && test -f /opt/lucebox-hub/server/share/model_cards/qwen3.6-27b.json \ && chmod +x /opt/lucebox-hub/server/scripts/entrypoint.sh diff --git a/README.md b/README.md index d34a281c3..c8fe80116 100644 --- a/README.md +++ b/README.md @@ -122,11 +122,22 @@ experts. **Review optimizations** offers an Advanced mode for users who want to override those four choices. Wi-Fi and device provisioning are intentionally outside this inference-only CLI. -On heterogeneous builds the CLI selects one primary accelerator: CUDA on an -RTX 3090 + Strix machine, and the largest-VRAM AMD device on an R9700 + Strix -machine. Automatic tuning never adds VRAM across devices. The engine's advanced -layer-split paths remain contributor workflows; transparent multi-GPU sharding -is not part of this CLI release. +Automatic inventories every NVIDIA and AMD accelerator, then writes one +validated execution plan together with the optimization profile. A model that +fits stays on the faster primary GPU. When capacity requires it, supported +architectures are split across same-backend GPUs; DFlash or PFlash work can be +placed on a companion GPU when the target fits but the complete stack does not. +Spark-compatible MoE models can place cold experts on a companion accelerator. +Cross-GPU copies use the correctness-first staged path unless an operator +explicitly opts into topology-specific peer access. + +RTX 3090 + Strix builds use CUDA for the main server and can use the Strix HIP +companion through the paired native runtime. R9700 + Strix builds use ROCm and +prefer the discrete R9700, while retaining Strix as a capacity/optimizer +device. Factory Lucebox systems can ship that runtime preinstalled; a source +checkout can create the same layout with `./lucebox.sh build hybrid` followed +by `./lucebox.sh package-runtime`. A normal single-GPU NVIDIA or AMD machine +keeps using its prebuilt Docker image and requires no native compilation. **Connect or open your harness** links the running local API to a client the engine user already installed: Claude Code, Codex, OpenCode, Hermes, Pi, @@ -158,6 +169,8 @@ lucebox optimize --advanced lucebox start ./lucebox.sh build cuda ./lucebox.sh native cuda +./lucebox.sh build hybrid # CUDA + HIP contributor build +./lucebox.sh package-runtime # factory-installable paired runtime ``` ## Quick Start On Harnesses diff --git a/lucebox.sh b/lucebox.sh index 849a4051c..11addb8bb 100755 --- a/lucebox.sh +++ b/lucebox.sh @@ -79,8 +79,12 @@ _lucebox_config_path() { # key = "string" # surrounding double-quotes are stripped # key = 8080 # bare scalars passed through verbatim # key = true # same -# Inline `# comment` is honored. Arrays / inline tables / multi-line -# strings aren't written by the Python persister, so we don't parse them. +# key = [ # simple multi-line arrays written by tomli_w +# "hip:0", +# "hip:1", +# ] +# Inline `# comment` is honored. Inline tables and multi-line strings are not +# part of the host wrapper's config contract. _lucebox_config_get() { local dotted="$1" cfg cfg="$(_lucebox_config_path)" @@ -89,24 +93,11 @@ _lucebox_config_get() { local key="${dotted##*.}" [ "$section" = "$dotted" ] && section="" awk -v want_section="$section" -v want_key="$key" ' - BEGIN { current = "" } - /^[[:space:]]*\[/ { - t = $0 - sub(/^[[:space:]]*\[[[:space:]]*/, "", t) - sub(/[[:space:]]*\][[:space:]]*$/, "", t) - current = t - next - } - /^[[:space:]]*#/ { next } - /=/ { - if (current != want_section) next - line = $0 - # Strip a TOML comment only when # is outside a quoted string. - # Paths and image tags may legitimately contain #. + function strip_comment(text, i, ch, in_quote, escaped) { in_quote = 0 escaped = 0 - for (i = 1; i <= length(line); i++) { - ch = substr(line, i, 1) + for (i = 1; i <= length(text); i++) { + ch = substr(text, i, 1) if (escaped) { escaped = 0 continue @@ -119,11 +110,23 @@ _lucebox_config_get() { in_quote = !in_quote continue } - if (ch == "#" && !in_quote) { - line = substr(line, 1, i - 1) - break - } + if (ch == "#" && !in_quote) + return substr(text, 1, i - 1) } + return text + } + BEGIN { current = "" } + /^[[:space:]]*\[/ { + t = $0 + sub(/^[[:space:]]*\[[[:space:]]*/, "", t) + sub(/[[:space:]]*\][[:space:]]*$/, "", t) + current = t + next + } + /^[[:space:]]*#/ { next } + /=/ { + if (current != want_section) next + line = strip_comment($0) eq = index(line, "=") if (eq == 0) next k = substr(line, 1, eq - 1) @@ -131,6 +134,14 @@ _lucebox_config_get() { gsub(/^[[:space:]]+|[[:space:]]+$/, "", k) gsub(/^[[:space:]]+|[[:space:]]+$/, "", v) if (k != want_key) next + if (substr(v, 1, 1) == "[" && index(v, "]") == 0) { + while ((getline continuation) > 0) { + continuation = strip_comment(continuation) + gsub(/^[[:space:]]+|[[:space:]]+$/, "", continuation) + v = v continuation + if (index(continuation, "]") > 0) break + } + } if (length(v) >= 2 && substr(v, 1, 1) == "\"" && substr(v, length(v), 1) == "\"") v = substr(v, 2, length(v) - 2) print v @@ -214,6 +225,11 @@ CONNECTOR_SELECTION_FILE="$CONNECTOR_STATE_DIR/selected" : "${LUCEBOX_HOST_GPU_SM:=}" : "${LUCEBOX_HOST_DRIVER_VERSION:=}" : "${LUCEBOX_HOST_DRIVER_MAJOR:=0}" +: "${LUCEBOX_HOST_NVIDIA_GPU_NAME:=}" +: "${LUCEBOX_HOST_NVIDIA_GPU_COUNT:=0}" +: "${LUCEBOX_HOST_NVIDIA_VRAM_GB:=0}" +: "${LUCEBOX_HOST_NVIDIA_GPU_ARCH:=}" +: "${LUCEBOX_HOST_NVIDIA_GPU_LIST_CSV:=}" : "${LUCEBOX_HOST_ROCM_VERSION:=}" : "${LUCEBOX_HOST_HAS_KFD:=0}" : "${LUCEBOX_HOST_HAS_DRI:=0}" @@ -240,6 +256,11 @@ CONNECTOR_SELECTION_FILE="$CONNECTOR_STATE_DIR/selected" : "${LUCEBOX_HOST_CUDA_VISIBLE_DEVICES:=}" : "${LUCEBOX_HOST_HIP_VISIBLE_DEVICES:=}" : "${LUCEBOX_HOST_ROCR_VISIBLE_DEVICES:=}" +: "${LUCEBOX_HOST_HAS_HYBRID_RUNTIME:=0}" +: "${LUCEBOX_HOST_HYBRID_SERVER_BIN:=}" +: "${LUCEBOX_HOST_HYBRID_IPC_BIN:=}" +: "${LUCEBOX_HOST_HYBRID_DFLASH_DIR:=}" +: "${LUCEBOX_HOST_HYBRID_ENTRYPOINT:=}" # Tracks whether probe_host has actually run; pieces of the code that need # fresh host facts (e.g. cmd_check, cmd_serve) gate on this. Default 0. : "${_LUCEBOX_HOST_PROBED:=0}" @@ -298,6 +319,14 @@ _find_repo_root() { return 1 } +_native_binary_ready() { + local binary="$1" dependencies + [ -x "$binary" ] || return 1 + command -v ldd >/dev/null 2>&1 || return 0 + dependencies=$(ldd "$binary" 2>&1 || true) + [[ "$dependencies" != *"not found"* ]] +} + _confirm() { # usage: _confirm "question" [default_yes] local question="$1" default_yes="${2:-1}" answer prompt @@ -416,6 +445,11 @@ probe_host() { LUCEBOX_HOST_DRIVER_VERSION="" LUCEBOX_HOST_DRIVER_MAJOR=0 LUCEBOX_HOST_GPU_LIST_CSV="" + LUCEBOX_HOST_NVIDIA_GPU_NAME="" + LUCEBOX_HOST_NVIDIA_GPU_COUNT=0 + LUCEBOX_HOST_NVIDIA_VRAM_GB=0 + LUCEBOX_HOST_NVIDIA_GPU_ARCH="" + LUCEBOX_HOST_NVIDIA_GPU_LIST_CSV="" LUCEBOX_HOST_ROCM_VERSION="" LUCEBOX_HOST_HAS_KFD=0 LUCEBOX_HOST_HAS_DRI=0 @@ -424,6 +458,11 @@ probe_host() { LUCEBOX_HOST_AMD_VRAM_GB=0 LUCEBOX_HOST_AMD_GPU_ARCH="" LUCEBOX_HOST_AMD_GPU_LIST_CSV="" + LUCEBOX_HOST_HAS_HYBRID_RUNTIME=0 + LUCEBOX_HOST_HYBRID_SERVER_BIN="" + LUCEBOX_HOST_HYBRID_IPC_BIN="" + LUCEBOX_HOST_HYBRID_DFLASH_DIR="" + LUCEBOX_HOST_HYBRID_ENTRYPOINT="" if command -v nvidia-smi &>/dev/null; then local q @@ -441,6 +480,10 @@ probe_host() { cc=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $4}') LUCEBOX_HOST_GPU_SM="${cc//./}" LUCEBOX_HOST_GPU_COUNT=$(printf '%s\n' "$q" | wc -l) + LUCEBOX_HOST_NVIDIA_GPU_NAME="$LUCEBOX_HOST_GPU_NAME" + LUCEBOX_HOST_NVIDIA_GPU_COUNT="$LUCEBOX_HOST_GPU_COUNT" + LUCEBOX_HOST_NVIDIA_VRAM_GB="$LUCEBOX_HOST_VRAM_GB" + LUCEBOX_HOST_NVIDIA_GPU_ARCH="$LUCEBOX_HOST_GPU_SM" fi # Multi-GPU enumeration for /props.host. The single-GPU vars # above (GPU_NAME / GPU_SM / VRAM_GB / DRIVER_VERSION) keep @@ -450,6 +493,7 @@ probe_host() { LUCEBOX_HOST_GPU_LIST_CSV=$(nvidia-smi \ --query-gpu=index,uuid,pci.bus_id,name,compute_cap,memory.total,power.limit \ --format=csv,noheader 2>/dev/null || echo "") + LUCEBOX_HOST_NVIDIA_GPU_LIST_CSV="$LUCEBOX_HOST_GPU_LIST_CSV" fi # Probe AMD independently even on mixed NVIDIA + Strix systems. A working @@ -489,7 +533,20 @@ probe_host() { LUCEBOX_HOST_HAS_AMD_GPU=1 LUCEBOX_HOST_AMD_GPU_COUNT=$(printf '%s\n' "$amd_rows" | awk 'NF{n++} END{print n+0}') local amd_primary - amd_primary=$(printf '%s\n' "$amd_rows" | sort -t'|' -k4,4nr | head -1) + # Prefer a discrete accelerator over Strix Halo UMA even when a newer + # firmware reports the APU's large shared-memory aperture as VRAM. + # Within the same memory class, use the device with most physical + # memory. This keeps R9700 + Strix builds on the faster R9700 while a + # Strix-only machine still selects its integrated GPU. + amd_primary=$(printf '%s\n' "$amd_rows" \ + | awk -F'|' '$3 != "gfx1151"' \ + | sort -t'|' -k4,4nr \ + | head -1) + if [ -z "$amd_primary" ]; then + amd_primary=$(printf '%s\n' "$amd_rows" \ + | sort -t'|' -k4,4nr \ + | head -1) + fi local amd_idx amd_name amd_arch amd_mem_mib amd_selector IFS='|' read -r amd_idx amd_name amd_arch amd_mem_mib amd_selector <<<"$amd_primary" amd_primary_selector="${amd_selector:-$amd_idx}" @@ -501,7 +558,7 @@ probe_host() { # Strix Halo exposes most memory as unified system RAM, while SMI may # report only a 512 MiB carve-out. Use host RAM as the effective model # capacity on a Strix-only build; a discrete R9700 remains primary on - # the R9700 + Strix build because it has the largest physical VRAM. + # the R9700 + Strix build by the policy above. if [ "$amd_arch" = "gfx1151" ] \ && [ "$LUCEBOX_HOST_AMD_VRAM_GB" -lt 12 ] \ && [ "$LUCEBOX_HOST_RAM_GB" -ge 32 ]; then @@ -541,11 +598,10 @@ probe_host() { done # CUDA_VISIBLE_DEVICES from the caller's env (empty default = "all GPUs"). LUCEBOX_HOST_CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-}" - # The native server currently uses one primary device unless an advanced - # user explicitly supplies a visibility list. On R9700 + Strix systems, - # pinning the largest-VRAM card keeps the runtime aligned with the GPU that - # the readiness screen and optimizer selected, even when ROCm enumerates - # the integrated GPU first. + # Legacy profiles use one isolated primary device. A resolved placement + # later clears this mask so the engine can address every named device. On + # R9700 + Strix systems, the legacy pin remains aligned with the discrete + # primary even when ROCm enumerates the integrated GPU first. LUCEBOX_HOST_HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-}" LUCEBOX_HOST_ROCR_VISIBLE_DEVICES="${ROCR_VISIBLE_DEVICES:-}" if [ -z "$LUCEBOX_HOST_HIP_VISIBLE_DEVICES" ] \ @@ -557,6 +613,52 @@ probe_host() { LUCEBOX_HOST_ROCR_VISIBLE_DEVICES="$amd_primary_selector" fi + # A mixed CUDA/HIP process cannot be assembled from a single-backend + # Docker image. Buyers may receive the paired native runtime under + # /opt/lucebox/runtime; contributors get the same contract after building + # both backends in their checkout. Record it only when both executables + # exist so Automatic never emits an unlaunchable cross-vendor plan. + if [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ] \ + && [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ]; then + local hybrid_repo="" hybrid_server="" hybrid_ipc="" + local hybrid_dir="" hybrid_entrypoint="" + hybrid_server="${LUCEBOX_HYBRID_SERVER_BIN:-}" + hybrid_ipc="${LUCEBOX_HYBRID_IPC_BIN:-}" + hybrid_dir="${LUCEBOX_HYBRID_DFLASH_DIR:-}" + hybrid_entrypoint="${LUCEBOX_HYBRID_ENTRYPOINT:-}" + if [ -z "$hybrid_entrypoint" ] && [ -n "$hybrid_dir" ]; then + hybrid_entrypoint="$hybrid_dir/scripts/entrypoint.sh" + fi + if ! _native_binary_ready "$hybrid_server" \ + || ! _native_binary_ready "$hybrid_ipc" \ + || [ ! -f "$hybrid_entrypoint" ]; then + hybrid_repo=$(_find_repo_root 2>/dev/null || true) + if [ -n "$hybrid_repo" ]; then + hybrid_server="$hybrid_repo/server/build-cuda/dflash_server" + hybrid_ipc="$hybrid_repo/server/build-hip/backend_ipc_daemon" + hybrid_dir="$hybrid_repo/server" + hybrid_entrypoint="$hybrid_dir/scripts/entrypoint.sh" + fi + fi + if ! _native_binary_ready "$hybrid_server" \ + || ! _native_binary_ready "$hybrid_ipc" \ + || [ ! -f "$hybrid_entrypoint" ]; then + hybrid_server="/opt/lucebox/runtime/cuda/dflash_server" + hybrid_ipc="/opt/lucebox/runtime/hip/backend_ipc_daemon" + hybrid_dir="/opt/lucebox/runtime/server" + hybrid_entrypoint="$hybrid_dir/scripts/entrypoint.sh" + fi + if _native_binary_ready "$hybrid_server" \ + && _native_binary_ready "$hybrid_ipc" \ + && [ -f "$hybrid_entrypoint" ]; then + LUCEBOX_HOST_HAS_HYBRID_RUNTIME=1 + LUCEBOX_HOST_HYBRID_SERVER_BIN="$hybrid_server" + LUCEBOX_HOST_HYBRID_IPC_BIN="$hybrid_ipc" + LUCEBOX_HOST_HYBRID_DFLASH_DIR="$hybrid_dir" + LUCEBOX_HOST_HYBRID_ENTRYPOINT="$hybrid_entrypoint" + fi + fi + # OS / kernel identity. /etc/os-release is the freedesktop spec for # "what distro is this?" and we keep PRETTY_NAME verbatim (it already # includes the version, e.g. "Ubuntu 22.04.3 LTS"). @@ -642,6 +744,12 @@ probe_host() { export LUCEBOX_HOST_NVIDIA_CTK_VERSION LUCEBOX_HOST_CPU_MODEL export LUCEBOX_HOST_GPU_LIST_CSV LUCEBOX_HOST_CUDA_VISIBLE_DEVICES export LUCEBOX_HOST_HIP_VISIBLE_DEVICES LUCEBOX_HOST_ROCR_VISIBLE_DEVICES + export LUCEBOX_HOST_NVIDIA_GPU_NAME LUCEBOX_HOST_NVIDIA_GPU_COUNT + export LUCEBOX_HOST_NVIDIA_VRAM_GB LUCEBOX_HOST_NVIDIA_GPU_ARCH + export LUCEBOX_HOST_NVIDIA_GPU_LIST_CSV + export LUCEBOX_HOST_HAS_HYBRID_RUNTIME + export LUCEBOX_HOST_HYBRID_SERVER_BIN LUCEBOX_HOST_HYBRID_IPC_BIN + export LUCEBOX_HOST_HYBRID_DFLASH_DIR LUCEBOX_HOST_HYBRID_ENTRYPOINT _LUCEBOX_HOST_PROBED=1 } @@ -935,6 +1043,127 @@ build_orchestrator_argv() { # ── subcommand implementations ──────────────────────────────────────────── +_config_requires_hybrid_runtime() { + local value target target_devices remote target_backend remote_backend + for value in \ + "$(_lucebox_config_get placement.remote_draft)" \ + "$(_lucebox_config_get placement.remote_target_shard)"; do + case "$value" in true|1|yes|on) return 0 ;; esac + done + + remote=$(_lucebox_config_get placement.remote_expert_device) + [ -n "$remote" ] || return 1 + target=$(_lucebox_config_get placement.target_device) + if [ -z "$target" ]; then + target_devices=$(_toml_array_to_csv \ + "$(_lucebox_config_get placement.target_devices)") + target="${target_devices%%,*}" + fi + [ -n "$target" ] || return 1 + target_backend="${target%%:*}" + remote_backend="${remote%%:*}" + [ "$target_backend" != "$remote_backend" ] +} + +_validate_hybrid_profile() { + local target target_devices draft remote_expert value device + local remote_draft=0 remote_target=0 seen_remote=0 + local hybrid_targets=() + target=$(_lucebox_config_get placement.target_device) + target_devices=$(_toml_array_to_csv \ + "$(_lucebox_config_get placement.target_devices)") + if [ -z "$target" ]; then + target="${target_devices%%,*}" + fi + [[ "$target" =~ ^cuda:[0-9]+$ ]] \ + || die "the installed hybrid runtime requires a CUDA target server (got '${target:-none}')" + + value=$(_lucebox_config_get placement.remote_draft) + case "$value" in true|1|yes|on) remote_draft=1 ;; esac + value=$(_lucebox_config_get placement.remote_target_shard) + case "$value" in true|1|yes|on) remote_target=1 ;; esac + + if [ "$remote_draft" = "1" ]; then + draft=$(_lucebox_config_get placement.draft_device) + [[ "$draft" =~ ^hip:[0-9]+$ ]] \ + || die "the installed hybrid runtime requires the remote draft on HIP (got '${draft:-none}')" + fi + + if [ "$remote_target" = "1" ]; then + [ -n "$target_devices" ] \ + || die "remote target sharding requires placement.target_devices" + IFS=',' read -r -a hybrid_targets <<<"$target_devices" + for device in "${hybrid_targets[@]}"; do + if [[ "$device" =~ ^cuda:[0-9]+$ ]]; then + [ "$seen_remote" = "0" ] \ + || die "hybrid target devices must list CUDA shards before HIP shards" + elif [[ "$device" =~ ^hip:[0-9]+$ ]]; then + seen_remote=1 + else + die "bad hybrid target device '$device'" + fi + done + [ "$seen_remote" = "1" ] \ + || die "remote target sharding requires at least one HIP shard" + fi + + remote_expert=$(_lucebox_config_get placement.remote_expert_device) + if [ -n "$remote_expert" ] \ + && [ "${remote_expert%%:*}" != "${target%%:*}" ]; then + [[ "$remote_expert" =~ ^hip:[0-9]+$ ]] \ + || die "the installed hybrid runtime requires remote Spark experts on HIP (got '$remote_expert')" + fi +} + +cmd_hybrid_serve() { + ensure_probed + [ "$LUCEBOX_HOST_HAS_HYBRID_RUNTIME" = "1" ] \ + || die "this profile needs the paired CUDA + HIP runtime, but it is not installed" + [ -x "$LUCEBOX_HOST_HYBRID_SERVER_BIN" ] \ + || die "hybrid server is missing: $LUCEBOX_HOST_HYBRID_SERVER_BIN" + [ -x "$LUCEBOX_HOST_HYBRID_IPC_BIN" ] \ + || die "hybrid backend daemon is missing: $LUCEBOX_HOST_HYBRID_IPC_BIN" + [ -f "$LUCEBOX_HOST_HYBRID_ENTRYPOINT" ] \ + || die "hybrid runtime entrypoint is missing: $LUCEBOX_HOST_HYBRID_ENTRYPOINT" + _validate_hybrid_profile + + if [ -z "${INVOCATION_ID:-}" ] \ + && systemctl --user is-active --quiet "$UNIT_NAME" 2>/dev/null; then + die "$UNIT_NAME is already running; use '$SCRIPT_NAME restart' or '$SCRIPT_NAME logs'" + fi + if _lucebox_container_running; then + die "container '$CONTAINER_NAME' is already running; stop it before starting the hybrid runtime" + fi + + local selected=() target draft model_id + mapfile -t selected < <(_selected_model_paths) + target="${selected[0]:-}" + draft="${selected[1]:-none}" + model_id="${selected[2]:-lucebox}" + _model_artifact_ready "$target" \ + || die "selected target is not installed: $target — run '$SCRIPT_NAME models select'" + if [ "$draft" != "none" ] && ! _model_artifact_ready "$draft"; then + die "selected draft is not installed: $draft — run '$SCRIPT_NAME models select'" + fi + + _export_native_config + export DFLASH_DIR="$LUCEBOX_HOST_HYBRID_DFLASH_DIR" + export DFLASH_SERVER_BIN="$LUCEBOX_HOST_HYBRID_SERVER_BIN" + export DFLASH_BACKEND_IPC_BIN="$LUCEBOX_HOST_HYBRID_IPC_BIN" + export DFLASH_TARGET="$target" + if [ "$draft" = "none" ]; then + export DFLASH_DRAFT="$DEFAULT_MODELS_DIR/.lucebox-no-draft" + else + export DFLASH_DRAFT="$draft" + fi + export DFLASH_HOST="${LUCEBOX_NATIVE_HOST:-127.0.0.1}" + export DFLASH_PORT="$DEFAULT_PORT" + export DFLASH_MODEL_NAME="$model_id" + export LUCEBOX_NATIVE=1 + info "Starting topology-aware CUDA + HIP engine at http://$DFLASH_HOST:$DFLASH_PORT" + exec bash "$LUCEBOX_HOST_HYBRID_ENTRYPOINT" serve +} + cmd_serve() { # Long-running foreground server. Also what systemd's ExecStart= calls. # @@ -950,6 +1179,10 @@ cmd_serve() { ensure_probed local variant variant=$(pick_variant) + if _config_requires_hybrid_runtime; then + cmd_hybrid_serve + return $? + fi require_host_prereqs "$variant" require_ctk "$variant" @@ -1091,12 +1324,33 @@ _serve_and_track() { cmd_systemd_install() { ensure_probed - local variant + local variant unit_after unit_wants unit_pre="" unit_stop="" unit_hybrid_env="" variant=$(pick_variant) - require_host_prereqs "$variant" require_systemd "service install" - local docker_bin - docker_bin=$(command -v docker) + unit_after="network-online.target" + unit_wants="network-online.target" + if _config_requires_hybrid_runtime; then + [ "$LUCEBOX_HOST_HAS_HYBRID_RUNTIME" = "1" ] \ + || die "install the paired CUDA + HIP runtime before enabling this profile" + _validate_hybrid_profile + # A contributor build may live in a checkout that systemd cannot + # rediscover from its working directory. Persist the already-validated + # executable contract in the unit; factory installs use the same lines + # with their stable /opt/lucebox/runtime paths. + unit_hybrid_env=$(printf '%s\n%s\n%s\n%s' \ + "Environment=LUCEBOX_HYBRID_SERVER_BIN=$LUCEBOX_HOST_HYBRID_SERVER_BIN" \ + "Environment=LUCEBOX_HYBRID_IPC_BIN=$LUCEBOX_HOST_HYBRID_IPC_BIN" \ + "Environment=LUCEBOX_HYBRID_DFLASH_DIR=$LUCEBOX_HOST_HYBRID_DFLASH_DIR" \ + "Environment=LUCEBOX_HYBRID_ENTRYPOINT=$LUCEBOX_HOST_HYBRID_ENTRYPOINT") + else + require_host_prereqs "$variant" + local docker_bin + docker_bin=$(command -v docker) + unit_after+=" docker.service" + unit_wants+=" docker.service" + unit_pre="ExecStartPre=-$docker_bin rm -f $CONTAINER_NAME" + unit_stop="ExecStop=$docker_bin stop -t 30 $CONTAINER_NAME" + fi mkdir -p "$(dirname "$UNIT_PATH")" # Capture the user's resolved env at install time so the unit launches @@ -1106,25 +1360,22 @@ cmd_systemd_install() { # in-script defaults and silently pick a different image or models # directory than the user's interactive session uses. # - # ExecStartPre cleans up any orphaned container with the target name - # left behind by a previous crash (docker's `--rm` only fires on clean - # exit — a SIGKILL or daemon restart leaves the name claimed, and the - # next ExecStart would die with "name already in use" while systemd - # reports a useless "exit code 125"). + # Docker profiles add an ExecStartPre cleanup for an orphaned container + # name. Native hybrid profiles intentionally leave both Docker directives + # empty because their server process is owned directly by systemd. cat > "$UNIT_PATH" </dev/null 2>&1 || die "cmake is required to build the inference engine" [ -f "$repo/server/deps/llama.cpp/ggml/CMakeLists.txt" ] \ || die "git submodules are missing — run: git -C '$repo' submodule update --init --recursive" - local configure=(cmake -S "$repo/server" -B "$build_dir" -DCMAKE_BUILD_TYPE=Release) + local configure=( + cmake -S "$repo/server" -B "$build_dir" + -DCMAKE_BUILD_TYPE=Release + -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON + ) if [ "$backend" = "rocm" ]; then [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ] \ || die "ROCm native build selected but no AMD GPU was detected" @@ -2037,33 +2366,99 @@ cmd_native_build() { || [ -f /usr/include/rocwmma/rocwmma.hpp ]; then hip_wmma=ON fi + build_arches=$(_amd_build_arches) + build_arches="${build_arches:-${LUCEBOX_HOST_AMD_GPU_ARCH:-gfx1151}}" configure+=( -DDFLASH27B_GPU_BACKEND=hip - "-DDFLASH27B_HIP_ARCHITECTURES=${LUCEBOX_HOST_AMD_GPU_ARCH:-gfx1151}" + "-DDFLASH27B_HIP_ARCHITECTURES=$build_arches" "-DDFLASH27B_HIP_SM80_EQUIV=$hip_wmma" ) else [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ] \ || die "CUDA native build selected but no NVIDIA GPU was detected" configure+=(-DDFLASH27B_GPU_BACKEND=cuda) - if [[ "$LUCEBOX_HOST_GPU_SM" =~ ^[0-9]+$ ]]; then - configure+=("-DCMAKE_CUDA_ARCHITECTURES=$LUCEBOX_HOST_GPU_SM") + build_arches=$(_nvidia_build_arches) + build_arches="${build_arches:-${LUCEBOX_HOST_NVIDIA_GPU_ARCH:-$LUCEBOX_HOST_GPU_SM}}" + if [[ "$build_arches" =~ ^[0-9]+([;][0-9]+)*$ ]]; then + configure+=("-DCMAKE_CUDA_ARCHITECTURES=$build_arches") fi fi info "Configuring native $backend build in $build_dir" "${configure[@]}" jobs="${LUCEBOX_HOST_NPROC:-1}" [ "$jobs" -gt 0 ] 2>/dev/null || jobs=1 - info "Building dflash_server ($jobs jobs)" - cmake --build "$build_dir" --target dflash_server -j "$jobs" + info "Building dflash_server + backend_ipc_daemon ($jobs jobs)" + cmake --build "$build_dir" \ + --target dflash_server backend_ipc_daemon -j "$jobs" ok "Native engine ready: $build_dir/dflash_server" + ok "Backend companion ready: $build_dir/backend_ipc_daemon" +} + +cmd_package_runtime() { + local repo destination stage runtime_dir backend build_dir package_backend + repo=$(_find_repo_root) \ + || die "runtime packaging requires a lucebox repository checkout" + destination="${1:-$repo/dist/lucebox-runtime}" + case "$destination" in + ""|/|"$repo") die "refusing unsafe runtime package destination: $destination" ;; + esac + [ ! -e "$destination" ] \ + || die "runtime package destination already exists: $destination" + + for backend in cuda rocm; do + build_dir=$(_native_build_dir "$repo" "$backend") + [ -x "$build_dir/dflash_server" ] \ + || die "missing $backend server — run '$SCRIPT_NAME build hybrid' first" + [ -x "$build_dir/backend_ipc_daemon" ] \ + || die "missing $backend companion — run '$SCRIPT_NAME build hybrid' first" + done + + stage=$(mktemp -d -t lucebox-runtime.XXXXXX) \ + || die "could not create runtime staging directory" + trap 'if [ -n "${stage:-}" ] && [ -d "$stage" ]; then rm -rf -- "$stage"; fi' EXIT + runtime_dir="$stage/lucebox-runtime" + mkdir -p "$runtime_dir/cuda" "$runtime_dir/hip" \ + "$runtime_dir/server/scripts" "$runtime_dir/server/share" + + for backend in cuda rocm; do + build_dir=$(_native_build_dir "$repo" "$backend") + package_backend="$backend" + [ "$backend" = "rocm" ] && package_backend=hip + cp "$build_dir/dflash_server" "$runtime_dir/$package_backend/" + cp "$build_dir/backend_ipc_daemon" "$runtime_dir/$package_backend/" + if [ -d "$build_dir/deps" ]; then + cp -R "$build_dir/deps" "$runtime_dir/$package_backend/deps" + find "$runtime_dir/$package_backend/deps" -type f \ + ! -name 'lib*.so*' -delete + find "$runtime_dir/$package_backend/deps" -depth -type d \ + -empty -delete + fi + done + cp "$repo/server/scripts/entrypoint.sh" "$runtime_dir/server/scripts/" + cp -R "$repo/server/share/." "$runtime_dir/server/share/" + printf 'source_commit=%s\ncreated_at=%s\n' \ + "$(git -C "$repo" rev-parse HEAD 2>/dev/null || printf unknown)" \ + "$(date -u +%FT%TZ 2>/dev/null || printf unknown)" \ + > "$runtime_dir/MANIFEST" + + mkdir -p "$(dirname "$destination")" + mv "$runtime_dir" "$destination" + rmdir "$stage" 2>/dev/null || true + stage="" + trap - EXIT + ok "Paired native runtime packaged: $destination" + hint "Factory install location: /opt/lucebox/runtime" } cmd_native_serve() { local repo backend build_dir binary selected=() target draft model_id + ensure_probed + if _config_requires_hybrid_runtime; then + cmd_hybrid_serve + return $? + fi repo=$(_find_repo_root) \ || die "native run requires a lucebox repository checkout (cd into it or set LUCEBOX_REPO)" - ensure_probed backend=$(_native_backend "${1:-}") build_dir=$(_native_build_dir "$repo" "$backend") binary="$build_dir/dflash_server" @@ -2082,6 +2477,7 @@ cmd_native_serve() { _export_native_config export DFLASH_DIR="$repo/server" export DFLASH_SERVER_BIN="$binary" + export DFLASH_BACKEND_IPC_BIN="$build_dir/backend_ipc_daemon" export DFLASH_TARGET="$target" if [ "$draft" = "none" ]; then export DFLASH_DRAFT="$DEFAULT_MODELS_DIR/.lucebox-no-draft" @@ -2127,6 +2523,9 @@ EOF script="$repo/harness/clients/run_${name}.sh" [ -x "$script" ] || die "harness launcher is missing: $script" ensure_probed + if _config_requires_hybrid_runtime; then + die "this profile uses the paired CUDA + HIP runtime; start it with '$SCRIPT_NAME native', then connect the harness with '$SCRIPT_NAME connect $name'" + fi backend=$(_native_backend "${LUCEBOX_HARNESS_BACKEND:-}") build_dir=$(_native_build_dir "$repo" "$backend") binary="$build_dir/dflash_server" @@ -2252,7 +2651,7 @@ _lucebox_complete() { cur="${COMP_WORDS[COMP_CWORD]}" prev="${COMP_WORDS[COMP_CWORD-1]}" cmds="menu setup connect install uninstall start stop restart enable disable status logs \ - serve build native harness pull update check completion config models \ + serve build package-runtime native harness pull update check completion config models \ optimize print-run help version" config_verbs="get set unset" models_verbs="list download select" @@ -2296,7 +2695,7 @@ ZSH # lucebox completion fish | source complete -c lucebox -f set -l __lucebox_cmds menu setup connect install uninstall start stop restart enable disable \ - status logs serve build native harness pull update check completion config models \ + status logs serve build package-runtime native harness pull update check completion config models \ optimize print-run help version for cmd in $__lucebox_cmds complete -c lucebox -n "not __fish_seen_subcommand_from $__lucebox_cmds" -a $cmd @@ -2393,7 +2792,7 @@ cmd_check() { if [ "$LUCEBOX_HOST_AMD_GPU_COUNT" -gt 1 ]; then local selected_device selected_device="${LUCEBOX_HOST_ROCR_VISIBLE_DEVICES:-${LUCEBOX_HOST_HIP_VISIBLE_DEVICES:-0}}" - _row warn "gpu placement" "primary device ${selected_device}; multi-GPU sharding is not automatic" + _row 1 "gpu placement" "Automatic resolves single- or multi-GPU execution after model selection (primary ${selected_device})" fi local amd_line amd_device_idx amd_device_name amd_device_arch amd_device_mem while IFS= read -r amd_line; do @@ -2740,6 +3139,7 @@ cmd_developer_menu() { printf ' 1 Build the native inference engine\n' printf ' 2 Run the native inference engine\n' printf ' 3 Choose and run a client harness\n' + printf ' 4 Package the CUDA + HIP buyer runtime\n' printf ' b Back\n\n' printf 'Choose: ' IFS= read -r choice || return 0 @@ -2747,8 +3147,9 @@ cmd_developer_menu() { 1) _menu_run build; _menu_pause ;; 2) _menu_run native; _menu_pause ;; 3) _menu_run harness; _menu_pause ;; + 4) _menu_run package-runtime; _menu_pause ;; b|B|q|Q) return 0 ;; - *) warn "Choose 1–3 or b"; _menu_pause ;; + *) warn "Choose 1–4 or b"; _menu_pause ;; esac done } @@ -2799,8 +3200,29 @@ _optimization_summary() { printf '%s (%s)' "$mode" "${active:-standard engine}" } +_placement_summary() { + local mode target targets draft remote_expert + mode=$(_lucebox_config_get placement.mode) + target=$(_lucebox_config_get placement.target_device) + targets=$(_toml_array_to_csv \ + "$(_lucebox_config_get placement.target_devices)") + draft=$(_lucebox_config_get placement.draft_device) + remote_expert=$(_lucebox_config_get placement.remote_expert_device) + if [ -n "$remote_expert" ]; then + printf '%s + %s Spark experts' "${target:-$targets}" "$remote_expert" + elif [ -n "$draft" ]; then + printf '%s + %s draft/scorer' "${target:-$targets}" "$draft" + elif [ -n "$targets" ]; then + printf '%s target split' "$targets" + elif [ -n "$target" ]; then + printf '%s' "$target" + else + printf '%s' "${mode:-server default}" + fi +} + cmd_menu() { - local choice model variant state repo_hint optimization gpu_name gpu_count other_gpu + local choice model variant state repo_hint optimization placement gpu_name gpu_count other_gpu local connector connector_label while true; do ensure_probed @@ -2809,6 +3231,7 @@ cmd_menu() { variant=$(pick_variant) state=$(_engine_state) optimization=$(_optimization_summary) + placement=$(_placement_summary) gpu_name="${LUCEBOX_HOST_GPU_NAME:-not detected}" gpu_count="$LUCEBOX_HOST_GPU_COUNT" other_gpu="" @@ -2849,6 +3272,7 @@ cmd_menu() { printf ' Backend: %s\n' "$variant" printf ' Model: %s\n' "$model" printf ' Optimization: %s\n' "$optimization" + printf ' Execution: %s\n' "$placement" printf ' Engine: %s\n' "$state" printf ' Harness: %s\n\n' "$connector_label" @@ -2915,7 +3339,10 @@ Direct server invocation (foreground, no systemd): serve docker run the server in the foreground Contributor workflow (inside a repository checkout): - build [cuda|rocm] build the native dflash_server + build [cuda|rocm] build one native backend and its IPC companion + build hybrid build the paired CUDA + HIP runtime + package-runtime [dir] + stage a relocatable factory runtime (default: dist/) native [cuda|rocm] run the selected model with the native engine harness [name] run Claude, Codex, OpenCode, Hermes, Pi, OpenClaw, or Open WebUI against the native engine @@ -3007,6 +3434,7 @@ main() { # Native source-repository workflow. build) cmd_native_build "$@" ;; + package-runtime) cmd_package_runtime "$@" ;; native) cmd_native_serve "$@" ;; harness) cmd_harness "$@" ;; diff --git a/lucebox/README.md b/lucebox/README.md index 321f2706e..c37dc2953 100644 --- a/lucebox/README.md +++ b/lucebox/README.md @@ -13,33 +13,41 @@ which invokes the package in the appropriate container: lucebox optimize --advanced # override individual optimizations lucebox start -The wrapper detects the host and selects CUDA for NVIDIA builds (including RTX -3090 + Strix) or ROCm for AMD builds (including R9700 + Strix). The package then -handles readiness checks, TOML configuration, model selection and download, -optimization settings, and construction of the final server command. Host -facts are passed through `LUCEBOX_HOST_*` environment variables so the -container never has to guess the host configuration. - -Automatic mode considers the selected model and the primary GPU. It shows a -plain-language decision for DFlash, PFlash, KVFlash, and Spark, prefers exact -all-GPU execution when it fits, and enables memory-saving paths only under -pressure. Advanced mode lets experienced users override those four product -choices while the CLI retains safe context, cache, and DDTree defaults. A -factory-preloaded Lucebox can ship the models and shared optimizer scorer, so a -buyer does not need to download anything during first setup. +The wrapper inventories the host and selects CUDA for NVIDIA builds (including +RTX 3090 + Strix) or ROCm for AMD builds (including R9700 + Strix). The package +then handles readiness checks, TOML configuration, model selection and +download, optimization and placement settings, and construction of the final +server command. Host facts are passed through `LUCEBOX_HOST_*` environment +variables so the package never has to guess the host configuration. + +Automatic mode considers the selected model and the complete accelerator +topology. It shows plain-language decisions for DFlash, PFlash, KVFlash, Spark, +and execution placement. It keeps fitting targets on the faster primary GPU, +uses supported secondary-device workloads under memory pressure, and enables +memory-saving paths only when needed. Advanced mode lets experienced users +override the four optimization choices while placement invariants, context, +cache, and DDTree defaults remain guarded. A factory-preloaded Lucebox can ship +the models, shared optimizer scorer, and paired runtime, so a buyer does not +need to download or compile anything during first setup. Spark is enabled automatically only when an MoE model is under GPU-memory pressure and the machine reports at least 32 GB of host RAM; its cold experts live in system memory. Lower-memory or unknown hosts keep Spark off unless an advanced user explicitly opts in. -Heterogeneous machines use one detected primary accelerator by default. The -CLI pins ROCm to the largest-VRAM AMD device and never sums memory across GPUs; -automatic multi-GPU layer sharding is deliberately not claimed by this release. +On a same-backend multi-GPU host, supported models can use a layer split when +the target does not fit one device. On R9700 + Strix, ROCm prefers the discrete +R9700 and can use Strix as a companion. RTX + Strix cross-vendor execution uses +a validated CUDA-main/HIP-companion native package; it is selected only when +that paired runtime is actually installed. Cross-GPU transfers default to safe +host staging, with peer access left as an explicit expert opt-in. Inside a source checkout, the same wrapper also exposes `lucebox build`, -`lucebox native`, and `lucebox harness` for contributors. Buyer installations -keep using the prebuilt container, so no compiler or source checkout is needed. +`lucebox native`, and `lucebox harness` for contributors. `lucebox build hybrid` +builds both native backends and `lucebox package-runtime` stages the factory +layout. Single-backend buyer installations keep using the prebuilt container; +heterogeneous buyers can receive the paired runtime preinstalled. Neither path +requires a source checkout during normal use. See the [project README](https://github.com/Luce-Org/lucebox#readme) for the installation and user flow. Contributors can find the CLI implementation in diff --git a/lucebox/pyproject.toml b/lucebox/pyproject.toml index bb12d90d7..8c873e3a1 100644 --- a/lucebox/pyproject.toml +++ b/lucebox/pyproject.toml @@ -5,7 +5,7 @@ name = "lucebox" # tag get a `.devN+g` suffix so dev installs are visibly distinct # from releases. Single source of truth: the git tag. dynamic = ["version"] -description = "Host-side CLI for the lucebox-hub container: launch, config, model download" +description = "Guided CLI for Lucebox inference: setup, models, optimization, and launch" readme = "README.md" requires-python = ">=3.12,<3.13" authors = [{ name = "Lucebox" }] diff --git a/lucebox/src/lucebox/autotune.py b/lucebox/src/lucebox/autotune.py index e3c77610a..18b43f293 100644 --- a/lucebox/src/lucebox/autotune.py +++ b/lucebox/src/lucebox/autotune.py @@ -11,6 +11,7 @@ from dataclasses import dataclass, replace from typing import Literal +from lucebox.placement import PlacementPlan, automatic_placement from lucebox.types import Config, DflashRuntime, HostFacts @@ -26,10 +27,11 @@ class OptimizationDecision: @dataclass(frozen=True, slots=True) class OptimizationPlan: - """Resolved runtime plus the four product-level optimization decisions.""" + """Resolved optimization and accelerator-placement decisions.""" runtime: DflashRuntime model_name: str + placement: PlacementPlan dflash: OptimizationDecision pflash: OptimizationDecision kvflash: OptimizationDecision @@ -106,7 +108,34 @@ def draft_available(cfg: Config, preset: object) -> bool: return False -def _cap_exact_context(runtime: DflashRuntime, *, headroom_gb: int, arch: str) -> DflashRuntime: +def placement_for_runtime(cfg: Config, runtime: DflashRuntime) -> PlacementPlan: + """Resolve placement after an Advanced-mode optimization edit.""" + from lucebox import download as download_mod + + preset = download_mod.PRESETS.get(cfg.model.preset) + if preset is None: + return automatic_placement( + cfg, + runtime, + object(), + has_draft=False, + optimizer_drafter_available=False, + ) + return automatic_placement( + cfg, + runtime, + preset, + has_draft=runtime.speculative_decode and draft_available(cfg, preset), + optimizer_drafter_available=download_mod.optimizer_drafter_installed(cfg), + ) + + +def _cap_exact_context( + runtime: DflashRuntime, + *, + headroom_gb: float, + arch: str, +) -> DflashRuntime: """Keep exact-cache families inside conservative primary-GPU headroom. Qwen families have an automatic bounded-residency path below; other @@ -150,9 +179,17 @@ def automatic_plan( "choose a model first; it activates when that model has a matching draft", ) unavailable = "choose a model first" + placement = automatic_placement( + cfg, + runtime, + object(), + has_draft=False, + optimizer_drafter_available=optimizer_drafter_available, + ) return OptimizationPlan( runtime=runtime, model_name=cfg.model.preset or "not selected", + placement=placement, dflash=dflash, pflash=OptimizationDecision("PFlash", False, False, unavailable), kvflash=OptimizationDecision("KVFlash", False, False, unavailable), @@ -278,9 +315,39 @@ def automatic_plan( needs_optimizer_drafter = not optimizer_drafter_available and ( pflash_profile or (kvflash_profile and preset.architecture == "qwen35moe") ) + placement = automatic_placement( + cfg, + runtime, + preset, + has_draft=has_draft, + optimizer_drafter_available=optimizer_drafter_available, + ) + adjusted = placement.optimization_runtime + if dflash.enabled and not adjusted.speculative_decode: + dflash = OptimizationDecision( + "DFlash", + False, + dflash.available, + "disabled because the selected target placement has no compatible draft path", + ) + if pflash.enabled and adjusted.prefill_mode == "off": + pflash = OptimizationDecision( + "PFlash", + False, + pflash.available, + "disabled because the selected target placement has no compatible compression path", + ) + if spark.enabled and not adjusted.spark: + spark = OptimizationDecision( + "Spark", + False, + spark.available, + "disabled because Automatic does not compose Spark with target layer splitting", + ) return OptimizationPlan( - runtime=runtime, + runtime=adjusted, model_name=preset.name, + placement=placement, dflash=dflash, pflash=pflash, kvflash=kvflash, diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index fa045e0ea..a5471d41a 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -33,6 +33,7 @@ from lucebox import __version__ from lucebox.config import config_get, config_set, config_unset, live_config from lucebox.host_facts import from_env +from lucebox.placement import PlacementPlan from lucebox.types import Config, DflashRuntime, ModelMeta app = typer.Typer( @@ -87,8 +88,9 @@ def _print_logo() -> None: def _load_or_build() -> Config: """env > config.toml > dataclass defaults — the canonical precedence. - Only the five documented top-level scalars have environment overrides; - dflash, host, and model settings intentionally remain config-driven. + Only the five documented top-level scalars have environment overrides. + Optimization, placement, and model settings remain config-driven, while + fresh host facts replace an absent or stale persisted snapshot. """ try: cfg = config_mod.load() @@ -145,6 +147,18 @@ def _optimization_label(cfg: Config) -> str: return f"{mode_label} ({', '.join(active) if active else 'standard engine'})" +def _placement_label(cfg: Config) -> str: + placement = cfg.placement + target = placement.target_device or ", ".join(placement.target_devices) + if placement.remote_expert_device: + return f"{target} target + {placement.remote_expert_device} Spark experts" + if placement.draft_device: + return f"{target} target + {placement.draft_device} draft/scorer" + if placement.target_devices: + return f"{target} target split" + return target or "server default" + + def _package_menu() -> None: """Small in-package menu for direct installs and contributor workflows. @@ -157,7 +171,8 @@ def _package_menu() -> None: cfg = _load_or_build() active = cfg.model.preset or "not selected" console.print(f"Model: [bold]{escape(active)}[/bold]") - console.print(f"Optimization: [bold]{escape(_optimization_label(cfg))}[/bold]\n") + console.print(f"Optimization: [bold]{escape(_optimization_label(cfg))}[/bold]") + console.print(f"Execution: [bold]{escape(_placement_label(cfg))}[/bold]\n") console.print(" [bold cyan]1[/bold cyan] Choose or download a model") console.print(" [bold cyan]2[/bold cyan] Review optimizations") console.print(" [bold cyan]3[/bold cyan] Show configuration") @@ -316,39 +331,66 @@ def _print_installed_presets() -> None: def _activate_preset(cfg: Config, preset: download_mod.ModelPreset) -> None: - """Persist one selected preset and seed first-run tuning.""" - config_set("model.preset", preset.name) - config_set("model.target_file", preset.target_file) - if preset.has_draft and preset.draft_file: - config_set("model.draft_file", preset.draft_file) - else: - # Drop any stale draft_file from a previous activation; the selected - # preset is explicitly target-only. - config_unset("model.draft_file") - console.print(f"[green]Activated:[/green] model.preset = {preset.name}") + """Activate one preset together with a runnable execution profile.""" + selected_model = ModelMeta( + preset=preset.name, + target_file=preset.target_file, + draft_file=preset.draft_file if preset.has_draft and preset.draft_file else "", + ) selected_cfg = replace( cfg, - model=ModelMeta( - preset=preset.name, - target_file=preset.target_file, - draft_file=preset.draft_file or "", - ), + model=selected_model, ) - # Automatic profiles follow the selected model. A user-owned custom - # profile is preserved and called out instead of being silently reset. - if config_mod.seed_optimization_from_config(selected_cfg): - max_ctx = config_get("dflash.max_ctx")["dflash.max_ctx"][0] - plan = autotune_mod.automatic_plan(selected_cfg) - active = ", ".join(plan.active_names) or "standard engine" - console.print( - f"[green]Optimized:[/green] Automatic profile for {preset.name} " - f"({active}; max_ctx={max_ctx})" + + mode = config_mod.optimization_mode() + if mode == "custom": + placement = autotune_mod.placement_for_runtime(selected_cfg, cfg.dflash) + if not placement.runnable: + error_console.print( + "[red]Model not activated: the custom profile has no runnable " + f"placement on this machine.[/red]\n[dim]{escape(placement.reason)}[/dim]" + ) + raise typer.Exit(code=2) + runtime = placement.optimization_runtime + config_mod.write_model_profile( + selected_model, + runtime, + placement.runtime, + mode="custom", + source="model-switch", ) - else: - console.print( - "[yellow]Custom optimization profile kept.[/yellow] " - "Run `lucebox optimize` to review it for this model." + console.print(f"[green]Activated:[/green] model.preset = {preset.name}") + if runtime != cfg.dflash: + console.print( + "[yellow]Custom profile adjusted.[/yellow] Features incompatible " + "with the new placement were disabled; run `lucebox optimize " + "--advanced` to review it." + ) + else: + console.print( + "[yellow]Custom optimization profile kept.[/yellow] " + "Execution placement was recalculated for this model." + ) + return + + plan = autotune_mod.automatic_plan(selected_cfg) + if not plan.placement.runnable: + error_console.print( + "[red]Model not activated: Automatic found no runnable placement " + f"on this machine.[/red]\n[dim]{escape(plan.placement.reason)}[/dim]" ) + raise typer.Exit(code=2) + config_mod.write_model_profile( + selected_model, + plan.runtime, + plan.placement.runtime, + ) + console.print(f"[green]Activated:[/green] model.preset = {preset.name}") + active = ", ".join(plan.active_names) or "standard engine" + console.print( + f"[green]Optimized:[/green] Automatic profile for {preset.name} " + f"({active}; max_ctx={plan.runtime.max_ctx})" + ) @models_app.callback(invoke_without_command=True) @@ -531,6 +573,9 @@ def _render_optimization_plan( ) -> None: console.print(f"[bold]{title}[/bold]") console.print(f"Model: {escape(plan.model_name)}") + placement_state = "[green]ready[/green]" if plan.placement.runnable else "[red]blocked[/red]" + console.print(f"Execution: [bold]{escape(plan.placement.summary)}[/bold] ({placement_state})") + console.print(f"[dim]{escape(plan.placement.reason)}[/dim]") table = Table(show_header=True, box=None, pad_edge=False) table.add_column("Optimization", style="bold") table.add_column("State") @@ -552,7 +597,10 @@ def _render_optimization_plan( ) -def _render_custom_runtime(runtime: DflashRuntime) -> None: +def _render_custom_runtime( + runtime: DflashRuntime, + placement: PlacementPlan | None = None, +) -> None: """Show the exact product choices before a custom profile is committed.""" console.print("\n[bold]Your custom profile[/bold]") table = Table(show_header=True, box=None, pad_edge=False) @@ -576,6 +624,10 @@ def _render_custom_runtime(runtime: DflashRuntime) -> None: f"KV format: [bold]{cache}[/bold] · " f"DFlash budget: [bold]{runtime.budget}[/bold]" ) + if placement is not None: + console.print( + f"Execution: [bold]{escape(placement.summary)}[/bold] · {escape(placement.reason)}" + ) def _ensure_optimizer_drafter(cfg: Config, *, assume_yes: bool = False) -> bool: @@ -745,11 +797,23 @@ def optimize( if custom: runtime = _customize_runtime(cfg, plan) - _render_custom_runtime(runtime) + placement = autotune_mod.placement_for_runtime(cfg, runtime) + runtime = placement.optimization_runtime + _render_custom_runtime(runtime, placement) + if not placement.runnable: + error_console.print( + "[red]This profile has no runnable placement on the detected hardware.[/red]" + ) + raise typer.Exit(code=2) if not typer.confirm("Apply this custom profile?", default=True): console.print("[dim]Optimization unchanged.[/dim]") return - config_mod.write_optimization_runtime(runtime, mode="custom", source="guided") + config_mod.write_optimization_runtime( + runtime, + placement=placement.runtime, + mode="custom", + source="guided", + ) console.print("[green]Custom optimization profile applied.[/green]") return @@ -759,7 +823,15 @@ def optimize( if not yes and not typer.confirm("Apply this automatic profile?", default=True): console.print("[dim]Optimization unchanged.[/dim]") return - config_mod.write_optimization_runtime(plan.runtime) + if not plan.placement.runnable: + error_console.print( + "[red]Automatic cannot produce a runnable placement for this model and machine.[/red]" + ) + raise typer.Exit(code=2) + config_mod.write_optimization_runtime( + plan.runtime, + placement=plan.placement.runtime, + ) console.print("[green]Automatic optimization applied.[/green]") console.print( "[dim]Run `lucebox optimize --advanced` anytime to review individual features.[/dim]" diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py index 762816811..aaedcc3c3 100644 --- a/lucebox/src/lucebox/config.py +++ b/lucebox/src/lucebox/config.py @@ -6,11 +6,16 @@ stay implicit, so `config.toml` reads like a diff against live defaults and upgrades that add new fields don't gratuitously rewrite every file. -The dotted-key surface area is small and flat: +The user-editable dotted-key surface area is small and flat: model.preset, model.target_file, model.draft_file port, models_dir, variant, image, container_name dflash. for every registered DflashRuntime knob +The resolved ``[placement]`` section is deliberately written only as part of +an atomic optimization plan. Its fields have cross-field invariants, so +exposing them through one-at-a-time ``config set`` operations would make it +easy to persist an intermediate profile that cannot be loaded. + Load resolves the TOML file → ``Config`` object, with anything absent filled from ``Config()`` defaults. Save writes back only the keys that appear in the TOML doc (tracked on ``Config._user_set``). The TOML doc @@ -36,6 +41,8 @@ DflashRuntime, HostFacts, ModelMeta, + PlacementMode, + PlacementRuntime, Variant, default_models_dir, ) @@ -163,6 +170,24 @@ def _cast_think_soft_close_min_ratio(v: Any) -> float: return value +def _cast_placement_mode(v: Any) -> PlacementMode: + value = str(v).strip().lower() + allowed = {"single", "draft-offload", "layer-split", "heterogeneous"} + if value not in allowed: + raise ValueError(f"placement mode must be one of {sorted(allowed)}, got {value!r}") + return cast(PlacementMode, value) + + +def _cast_string_tuple(v: Any) -> tuple[str, ...]: + values = v if isinstance(v, (list, tuple)) else str(v).split(",") + return tuple(str(value).strip() for value in values if str(value).strip()) + + +def _cast_float_tuple(v: Any) -> tuple[float, ...]: + values = v if isinstance(v, (list, tuple)) else str(v).split(",") + return tuple(float(value) for value in values if str(value).strip()) + + # Each entry: dotted-key → (toml_path, type_caster, default_getter). # ``toml_path`` is the (section, field) pair on disk; ``"_root"`` means the # key lives at the top level (no [section]). ``default_getter`` returns the @@ -355,13 +380,26 @@ def _from_dict(raw: dict[str, Any]) -> Config: debug_thinking_logits=_cast_bool(df.get("debug_thinking_logits", False)), ) + placement_raw = raw.get("placement", {}) + placement = PlacementRuntime( + mode=_cast_placement_mode(placement_raw.get("mode", "single")), + target_device=str(placement_raw.get("target_device", "")), + target_devices=_cast_string_tuple(placement_raw.get("target_devices", ())), + target_layer_split=_cast_float_tuple(placement_raw.get("target_layer_split", ())), + draft_device=str(placement_raw.get("draft_device", "")), + remote_draft=_cast_bool(placement_raw.get("remote_draft", False)), + remote_target_shard=_cast_bool(placement_raw.get("remote_target_shard", False)), + peer_access=_cast_bool(placement_raw.get("peer_access", False)), + remote_expert_device=str(placement_raw.get("remote_expert_device", "")), + ) + host_raw = raw.get("host", {}) host = HostFacts( nproc=int(host_raw.get("nproc", 0)), ram_gb=int(host_raw.get("ram_gb", 0)), gpu_vendor=host_raw.get("gpu_vendor", "none"), - has_nvidia_gpu=bool(host_raw.get("has_nvidia_gpu", False)), - has_amd_gpu=bool(host_raw.get("has_amd_gpu", False)), + has_nvidia_gpu=_cast_bool(host_raw.get("has_nvidia_gpu", False)), + has_amd_gpu=_cast_bool(host_raw.get("has_amd_gpu", False)), gpu_name=str(host_raw.get("gpu_name", "")), gpu_count=int(host_raw.get("gpu_count", 0)), vram_gb=int(host_raw.get("vram_gb", 0)), @@ -369,13 +407,24 @@ def _from_dict(raw: dict[str, Any]) -> Config: driver_version=str(host_raw.get("driver_version", "")), driver_major=int(host_raw.get("driver_major", 0)), rocm_version=str(host_raw.get("rocm_version", "")), - has_kfd=bool(host_raw.get("has_kfd", False)), - has_dri=bool(host_raw.get("has_dri", False)), - has_systemd=bool(host_raw.get("has_systemd", False)), - is_wsl=bool(host_raw.get("is_wsl", False)), - has_docker=bool(host_raw.get("has_docker", False)), + has_kfd=_cast_bool(host_raw.get("has_kfd", False)), + has_dri=_cast_bool(host_raw.get("has_dri", False)), + has_systemd=_cast_bool(host_raw.get("has_systemd", False)), + is_wsl=_cast_bool(host_raw.get("is_wsl", False)), + has_docker=_cast_bool(host_raw.get("has_docker", False)), docker_version=str(host_raw.get("docker_version", "")), ctk=host_raw.get("ctk", "none"), + nvidia_gpu_name=str(host_raw.get("nvidia_gpu_name", "")), + nvidia_gpu_count=int(host_raw.get("nvidia_gpu_count", 0)), + nvidia_vram_gb=int(host_raw.get("nvidia_vram_gb", 0)), + nvidia_gpu_arch=str(host_raw.get("nvidia_gpu_arch", "")), + nvidia_gpu_list_csv=str(host_raw.get("nvidia_gpu_list_csv", "")), + amd_gpu_name=str(host_raw.get("amd_gpu_name", "")), + amd_gpu_count=int(host_raw.get("amd_gpu_count", 0)), + amd_vram_gb=int(host_raw.get("amd_vram_gb", 0)), + amd_gpu_arch=str(host_raw.get("amd_gpu_arch", "")), + amd_gpu_list_csv=str(host_raw.get("amd_gpu_list_csv", "")), + hybrid_runtime=_cast_bool(host_raw.get("hybrid_runtime", False)), ) # `[model]` is optional — legacy configs (pre-multi-model) carry no @@ -404,6 +453,7 @@ def _from_dict(raw: dict[str, Any]) -> Config: port=port, models_dir=models_dir, dflash=dflash, + placement=placement, host=host, model=model, ) @@ -491,13 +541,21 @@ def _write_runtime_doc( *, mode: str, source: str, + placement: PlacementRuntime | None = None, ) -> None: - """Replace ``[dflash]`` with one coherent resolved profile.""" + """Replace optimization/placement with one coherent resolved profile.""" from datetime import datetime doc.pop("dflash", None) for field, value in asdict(runtime).items(): _doc_set(doc, "dflash", field, _value_to_toml(value)) + # Placement belongs to the resolved runtime as one atomic unit. Clearing + # it even when an older caller supplies no replacement prevents a stale + # multi-GPU profile from surviving a heuristic/runtime reset. + doc.pop("placement", None) + if placement is not None: + for field, value in asdict(placement).items(): + _doc_set(doc, "placement", field, _value_to_toml(value)) _doc_set(doc, "autotune", "mode", mode) _doc_set(doc, "autotune", "source", source) _doc_set( @@ -511,6 +569,7 @@ def _write_runtime_doc( def write_optimization_runtime( runtime: DflashRuntime, *, + placement: PlacementRuntime | None = None, path: Path | None = None, mode: str = "automatic", source: str = "model+hardware", @@ -520,7 +579,52 @@ def write_optimization_runtime( raise ValueError(f"optimization mode must be automatic or custom, got {mode!r}") path = path or default_config_path() doc = load_doc(path) - _write_runtime_doc(doc, runtime, mode=mode, source=source) + _write_runtime_doc( + doc, + runtime, + mode=mode, + source=source, + placement=placement, + ) + path.parent.mkdir(parents=True, exist_ok=True) + _atomic_write_doc(path, doc) + + +def write_model_profile( + model: ModelMeta, + runtime: DflashRuntime, + placement: PlacementRuntime, + *, + path: Path | None = None, + mode: str = "automatic", + source: str = "model+hardware", +) -> None: + """Atomically activate a model and its validated execution profile.""" + if mode not in {"automatic", "custom"}: + raise ValueError(f"optimization mode must be automatic or custom, got {mode!r}") + preset = str(model.preset).strip() + if not preset: + raise ValueError("model preset must not be empty") + target_file = _cast_model_relative_path(model.target_file) + if not target_file: + raise ValueError("model target_file must not be empty") + draft_file = _cast_model_relative_path(model.draft_file) + + path = path or default_config_path() + doc = load_doc(path) + _doc_set(doc, "model", "preset", preset) + _doc_set(doc, "model", "target_file", target_file) + if draft_file: + _doc_set(doc, "model", "draft_file", draft_file) + else: + _doc_unset(doc, "model", "draft_file") + _write_runtime_doc( + doc, + runtime, + mode=mode, + source=source, + placement=placement, + ) path.parent.mkdir(parents=True, exist_ok=True) _atomic_write_doc(path, doc) @@ -551,7 +655,15 @@ def seed_optimization_from_config( if not force and mode == "custom": return False plan = autotune_mod.automatic_plan(cfg) - write_optimization_runtime(plan.runtime, path=path) + if not plan.placement.runnable: + raise ValueError( + f"automatic optimization has no runnable placement: {plan.placement.reason}" + ) + write_optimization_runtime( + plan.runtime, + placement=plan.placement.runtime, + path=path, + ) return True @@ -562,6 +674,8 @@ def _value_to_toml(value: Any) -> Any: """Make a Python value safe for tomli_w (no None, Path→str).""" if isinstance(value, Path): return str(value) + if isinstance(value, tuple): + return list(value) return value @@ -603,7 +717,8 @@ def config_set(key: str, value: Any, *, path: Path | None = None) -> None: try: _from_dict({"dflash": doc.get("dflash", {})}) except (TypeError, ValueError) as exc: - raise ValueError(f"invalid dflash profile after setting {key}: {exc}") from exc + raise ValueError(f"invalid runtime profile after setting {key}: {exc}") from exc + if section == "dflash": # A direct low-level edit transfers ownership from the automatic # planner to the user. Model changes will preserve it until they pick # Automatic again in ``lucebox optimize``. diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index 853128369..d8c29ae67 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -119,6 +119,33 @@ def _runtime_volumes(cfg: Config) -> tuple[BindMount, ...]: ) +def _placement_env(cfg: Config) -> list[tuple[str, str]]: + """Translate the portable placement profile to the entrypoint contract.""" + placement = cfg.placement + env: list[tuple[str, str]] = [("DFLASH_PLACEMENT_MODE", placement.mode)] + if placement.target_device: + env.append(("DFLASH_TARGET_DEVICE", placement.target_device)) + if placement.target_devices: + env.append(("DFLASH_TARGET_DEVICES", ",".join(placement.target_devices))) + env.append( + ( + "DFLASH_TARGET_LAYER_SPLIT", + ",".join(f"{weight:g}" for weight in placement.target_layer_split), + ) + ) + if placement.draft_device: + env.append(("DFLASH_DRAFT_DEVICE", placement.draft_device)) + if placement.remote_draft: + env.append(("DFLASH_REMOTE_DRAFT", "1")) + if placement.remote_target_shard: + env.append(("DFLASH_REMOTE_TARGET_SHARD", "1")) + if placement.peer_access: + env.append(("DFLASH_PEER_ACCESS", "1")) + if placement.remote_expert_device: + env.append(("DFLASH_REMOTE_EXPERT_DEVICE", placement.remote_expert_device)) + return env + + def _validate_model_relative_path(value: str, field: str) -> PurePosixPath: """Validate a config-provided path intended to live below models_dir.""" relative = PurePosixPath(value) @@ -261,6 +288,12 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: """Long-running OpenAI-compatible server. Foreground (systemd manages lifecycle), vendor-specific GPU devices, models bind-mounted, DFLASH_* propagated. """ + if cfg.placement.requires_hybrid_runtime: + raise ValueError( + "cross-vendor placement requires the paired native Lucebox runtime; " + "a single-backend Docker image cannot launch it" + ) + # LUCEBOX_HOST_* first so they ride out front in the rendered argv, # making it obvious in `print-run` output what host facts get forwarded. env: list[tuple[str, str]] = list(_host_facts_env()) @@ -274,6 +307,7 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: ("DFLASH_THINK_MAX", str(cfg.dflash.think_max)), ("DFLASH_PORT", "8080"), ] + env += _placement_env(cfg) # Keep the API model catalog aligned with the preset selected by the host # CLI. Client connectors use this id for explicit model selection and # clients such as Codex also discover it through ``/v1/models``. @@ -389,14 +423,15 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: # automatic tuning. This matters on an R9700 + Strix build where ROCm may # enumerate the integrated GPU before the larger discrete card. Explicit # CUDA/ROCR/HIP visibility supplied by an advanced user remains authoritative. - if gpu_vendor == "amd": + placement_is_explicit = bool(cfg.placement.target_device or cfg.placement.target_devices) + if gpu_vendor == "amd" and not placement_is_explicit: rocr_visible = os.environ.get("LUCEBOX_HOST_ROCR_VISIBLE_DEVICES", "").strip() hip_visible = os.environ.get("LUCEBOX_HOST_HIP_VISIBLE_DEVICES", "").strip() if rocr_visible: env.append(("ROCR_VISIBLE_DEVICES", rocr_visible)) elif hip_visible: env.append(("HIP_VISIBLE_DEVICES", hip_visible)) - elif gpu_vendor == "nvidia": + elif gpu_vendor == "nvidia" and not placement_is_explicit: visible = os.environ.get("LUCEBOX_HOST_CUDA_VISIBLE_DEVICES", "").strip() if visible: env.append(("CUDA_VISIBLE_DEVICES", visible)) diff --git a/lucebox/src/lucebox/download.py b/lucebox/src/lucebox/download.py index b39bd728a..5e68d7418 100644 --- a/lucebox/src/lucebox/download.py +++ b/lucebox/src/lucebox/download.py @@ -86,7 +86,9 @@ class ModelPreset: target_file: str draft_repo: str | None draft_file: str | None - approx_total_gb: int + approx_total_gb: float + approx_target_gb: float + approx_draft_gb: float architecture: str native_context: int description: str = "" @@ -109,6 +111,8 @@ def has_draft(self) -> bool: draft_repo="spiritbuun/Qwen3.6-27B-DFlash-GGUF", draft_file="dflash-draft-3.6-q4_k_m.gguf", approx_total_gb=17, + approx_target_gb=15.7, + approx_draft_gb=1.3, architecture="qwen35", native_context=262144, description="Qwen3.6 27B dense (Q4_K_M) + Qwen3.6 DFlash draft. Lucebox default.", @@ -120,6 +124,8 @@ def has_draft(self) -> bool: draft_repo="Lucebox/gemma-4-26B-A4B-it-DFlash-GGUF", draft_file="gemma-4-26B-A4B-it-DFlash-q8_0.gguf", approx_total_gb=18, + approx_target_gb=16.0, + approx_draft_gb=2.0, architecture="gemma4", native_context=262144, description="Gemma 4 26B-A4B IT MoE (Q4_K_M) + Lucebox DFlash q8_0 draft.", @@ -131,6 +137,8 @@ def has_draft(self) -> bool: draft_repo="Lucebox/gemma-4-31B-it-DFlash-GGUF", draft_file="gemma-4-31B-it-DFlash-q8_0.gguf", approx_total_gb=21, + approx_target_gb=19.0, + approx_draft_gb=2.0, architecture="gemma4", native_context=262144, description="Gemma 4 31B IT dense (Q4_K_M) + Lucebox DFlash q8_0 draft.", @@ -148,6 +156,8 @@ def has_draft(self) -> bool: draft_file=None, speculator_dir="laguna-xs2-speculator", approx_total_gb=20, + approx_target_gb=18.9, + approx_draft_gb=1.1, architecture="laguna", native_context=4096, description=( @@ -172,6 +182,8 @@ def has_draft(self) -> bool: draft_repo=None, draft_file=None, approx_total_gb=22, + approx_target_gb=22.0, + approx_draft_gb=0.0, architecture="qwen35moe", native_context=262144, description=( diff --git a/lucebox/src/lucebox/host_facts.py b/lucebox/src/lucebox/host_facts.py index c9556d246..f16db336b 100644 --- a/lucebox/src/lucebox/host_facts.py +++ b/lucebox/src/lucebox/host_facts.py @@ -60,4 +60,15 @@ def from_env() -> HostFacts: has_docker=_env_bool("LUCEBOX_HOST_HAS_DOCKER"), docker_version=os.environ.get("LUCEBOX_HOST_DOCKER_VERSION", ""), ctk=ctk, + nvidia_gpu_name=os.environ.get("LUCEBOX_HOST_NVIDIA_GPU_NAME", ""), + nvidia_gpu_count=_env_int("LUCEBOX_HOST_NVIDIA_GPU_COUNT"), + nvidia_vram_gb=_env_int("LUCEBOX_HOST_NVIDIA_VRAM_GB"), + nvidia_gpu_arch=os.environ.get("LUCEBOX_HOST_NVIDIA_GPU_ARCH", ""), + nvidia_gpu_list_csv=os.environ.get("LUCEBOX_HOST_NVIDIA_GPU_LIST_CSV", ""), + amd_gpu_name=os.environ.get("LUCEBOX_HOST_AMD_GPU_NAME", ""), + amd_gpu_count=_env_int("LUCEBOX_HOST_AMD_GPU_COUNT"), + amd_vram_gb=_env_int("LUCEBOX_HOST_AMD_VRAM_GB"), + amd_gpu_arch=os.environ.get("LUCEBOX_HOST_AMD_GPU_ARCH", ""), + amd_gpu_list_csv=os.environ.get("LUCEBOX_HOST_AMD_GPU_LIST_CSV", ""), + hybrid_runtime=_env_bool("LUCEBOX_HOST_HAS_HYBRID_RUNTIME"), ) diff --git a/lucebox/src/lucebox/placement.py b/lucebox/src/lucebox/placement.py new file mode 100644 index 000000000..86043ae64 --- /dev/null +++ b/lucebox/src/lucebox/placement.py @@ -0,0 +1,434 @@ +"""Capability-checked accelerator placement for Lucebox inference. + +Layer splitting is a capacity tool, not an automatic speed claim: the engine +runs contiguous layer groups sequentially and crosses a host/IPC boundary. +When a model fits the faster primary GPU, keeping the target monolithic is the +default. Secondary devices are used for an independently useful workload +(draft/scorer or Spark cold experts), or when splitting is required to make a +model runnable. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from pathlib import Path + +from lucebox.topology import GpuDevice, HardwareTopology, from_config +from lucebox.types import Config, DflashRuntime, PlacementRuntime + + +@dataclass(frozen=True, slots=True) +class ArchitecturePlacementCapabilities: + layer_split: bool + remote_draft: bool + draft_on_layer_split: bool + pflash_on_layer_split: bool + expert_offload: bool + + +# Mirrors server/src/common/model_capabilities.h for the product-level paths +# the packaged planner can actually launch. A parity test keeps this table +# explicit and reviewable when an engine architecture changes. +ARCHITECTURE_CAPABILITIES: dict[str, ArchitecturePlacementCapabilities] = { + "qwen35": ArchitecturePlacementCapabilities(True, True, True, True, False), + "qwen35moe": ArchitecturePlacementCapabilities(False, False, False, False, True), + "laguna": ArchitecturePlacementCapabilities(True, False, False, False, True), + "gemma4": ArchitecturePlacementCapabilities(True, False, False, False, False), + "deepseek4": ArchitecturePlacementCapabilities(True, False, False, False, False), + "qwen3": ArchitecturePlacementCapabilities(False, False, False, False, False), +} + + +@dataclass(frozen=True, slots=True) +class PlacementOption: + key: str + label: str + available: bool + recommended: bool + reason: str + + +@dataclass(frozen=True, slots=True) +class PlacementPlan: + runtime: PlacementRuntime + optimization_runtime: DflashRuntime + summary: str + reason: str + runnable: bool + topology: HardwareTopology + options: tuple[PlacementOption, ...] + + +def _artifact_size_gb(path: Path, fallback: float) -> float: + try: + size = path.stat().st_size / (1024**3) + except OSError: + return fallback + # Unit tests and interrupted downloads create tiny placeholders. They do + # not describe the memory footprint of a loaded model. + return size if size >= 0.25 else fallback + + +def _model_sizes(cfg: Config, preset: object, has_draft: bool) -> tuple[float, float]: + target_file = cfg.model.target_file or str(getattr(preset, "target_file", "")) + target_fallback = float(getattr(preset, "approx_target_gb", 0.0)) + target_gb = _artifact_size_gb(cfg.models_dir / target_file, target_fallback) + + draft_gb = 0.0 + if has_draft: + draft_file = cfg.model.draft_file or getattr(preset, "draft_file", None) + draft_fallback = float(getattr(preset, "approx_draft_gb", 0.0)) + if draft_file: + draft_gb = _artifact_size_gb(cfg.models_dir / "draft" / str(draft_file), draft_fallback) + else: + draft_gb = draft_fallback + return target_gb, draft_gb + + +def _primary_capacity(device: GpuDevice) -> float: + # Discrete GPUs need less host-side reserve; UMA devices share their pool + # with the OS and CPU engine buffers and therefore keep a wider margin. + reserve = 8.0 if device.unified_memory else 2.0 + return max(0.0, float(device.effective_memory_gb) - reserve) + + +def _secondary_score(device: GpuDevice, primary: GpuDevice) -> tuple[int, int, int, int]: + return ( + device.backend == primary.backend, + not device.unified_memory, + device.physical_vram_gb, + device.effective_memory_gb, + ) + + +def _best_secondary(topology: HardwareTopology) -> GpuDevice | None: + if topology.primary is None or not topology.companions: + return None + primary = topology.primary + return max( + topology.companions, + key=lambda device: _secondary_score(device, primary), + ) + + +def _cross_backend_ready( + cfg: Config, + primary: GpuDevice, + secondary: GpuDevice, +) -> bool: + """Whether the installed paired runtime supports this backend direction. + + The first packaged heterogeneous contract is deliberately directional: + a CUDA server drives a HIP companion daemon on RTX + Strix systems. A + future HIP-main/CUDA-daemon package can extend the host facts without the + planner accidentally claiming that today's binaries support it. + """ + return cfg.host.hybrid_runtime and primary.backend == "cuda" and secondary.backend == "hip" + + +def _split_devices( + cfg: Config, + topology: HardwareTopology, + target_gb: float, + primary_capacity: float, +) -> tuple[tuple[GpuDevice, ...], tuple[float, ...]]: + """Choose the smallest ordered device set that makes the target fit.""" + primary = topology.primary + if primary is None: + return (), () + + local = sorted( + (device for device in topology.companions if device.backend == primary.backend), + key=lambda device: _secondary_score(device, primary), + reverse=True, + ) + remote = sorted( + ( + device + for device in topology.companions + if device.backend != primary.backend and _cross_backend_ready(cfg, primary, device) + ), + key=lambda device: _secondary_score(device, primary), + reverse=True, + ) + selected = [primary] + capacities = [primary_capacity] + total = primary_capacity + for device in (*local, *remote): + if total >= target_gb: + break + capacity = _primary_capacity(device) + if capacity <= 0.0: + continue + selected.append(device) + capacities.append(capacity) + total += capacity + if len(selected) < 2 or total < target_gb: + return (), () + + # Assign only the portion needed from the final device, then normalize. + # A small floor prevents whole-layer rounding from creating an empty shard. + remaining = target_gb + allocations: list[float] = [] + for capacity in capacities: + allocation = min(capacity, remaining) + allocations.append(max(allocation, target_gb * 0.05)) + remaining -= allocation + total_allocation = sum(allocations) + weights = tuple(round(value / total_allocation, 4) for value in allocations) + # Make the serialized values sum exactly to one after rounding. + weights = (*weights[:-1], round(1.0 - sum(weights[:-1]), 4)) + return tuple(selected), weights + + +def _single_option(primary: GpuDevice, fits: bool) -> PlacementOption: + return PlacementOption( + key="single", + label="Primary GPU", + available=fits, + recommended=fits, + reason=( + f"the selected stack fits {primary.name}; avoiding a layer/IPC boundary is faster" + if fits + else f"the target does not fit the safe memory budget on {primary.name}" + ), + ) + + +def automatic_placement( + cfg: Config, + runtime: DflashRuntime, + preset: object, + *, + has_draft: bool, + optimizer_drafter_available: bool, +) -> PlacementPlan: + """Resolve a safe placement and every user-visible alternative.""" + topology = from_config(cfg) + primary = topology.primary + if primary is None: + return PlacementPlan( + runtime=PlacementRuntime(), + optimization_runtime=runtime, + summary="No accelerator selected", + reason="choose a CUDA or ROCm backend after a supported GPU is detected", + runnable=False, + topology=topology, + options=(), + ) + + architecture = str(getattr(preset, "architecture", "")) + capabilities = ARCHITECTURE_CAPABILITIES.get( + architecture, + ArchitecturePlacementCapabilities(False, False, False, False, False), + ) + secondary = _best_secondary(topology) + target_gb, draft_gb = _model_sizes(cfg, preset, has_draft) + scorer_gb = ( + 1.2 + if optimizer_drafter_available + and ( + runtime.prefill_mode != "off" + or (runtime.kvflash != "off" and runtime.kvflash_policy == "drafter") + ) + else 0.0 + ) + primary_capacity = _primary_capacity(primary) + target_fits = target_gb <= primary_capacity + stack_gb = target_gb + draft_gb + scorer_gb + # Spark's explicit purpose is to make an MoE target fit by leaving cold + # experts in host memory. Its load-time allocator remains the final source + # of truth for the exact byte budget. + stack_fits = stack_gb <= primary_capacity or (runtime.spark and capabilities.expert_offload) + + options: list[PlacementOption] = [_single_option(primary, stack_fits)] + base = PlacementRuntime(mode="single", target_device=primary.placement_name) + if secondary is None: + if stack_fits: + return PlacementPlan( + runtime=base, + optimization_runtime=runtime, + summary=primary.name, + reason="one accelerator is available and the selected stack fits its safe budget", + runnable=True, + topology=topology, + options=tuple(options), + ) + return PlacementPlan( + runtime=base, + optimization_runtime=runtime, + summary=primary.name, + reason="the selected target exceeds this GPU's safe memory budget", + runnable=False, + topology=topology, + options=tuple(options), + ) + + same_backend = primary.backend == secondary.backend + cross_backend_ready = _cross_backend_ready(cfg, primary, secondary) + draft_work_available = (has_draft and runtime.speculative_decode) or ( + optimizer_drafter_available and runtime.prefill_mode != "off" + ) + draft_offload_available = draft_work_available and ( + same_backend or (cross_backend_ready and capabilities.remote_draft) + ) + options.append( + PlacementOption( + key="draft-offload", + label="Secondary draft/scorer", + available=draft_offload_available, + recommended=draft_offload_available and not stack_fits and target_fits, + reason=( + "moves the independent draft/scorer workload off the primary GPU" + if draft_offload_available + else ( + "the selected model has no installed draft/scorer workload" + if not draft_work_available + else "the engine/runtime cannot execute this draft across the backend boundary" + ) + ), + ) + ) + + split_primary_capacity = primary_capacity + if capabilities.draft_on_layer_split and runtime.speculative_decode: + split_primary_capacity -= draft_gb + if capabilities.pflash_on_layer_split: + split_primary_capacity -= scorer_gb + split_devices, split_weights = _split_devices( + cfg, + topology, + target_gb, + max(0.0, split_primary_capacity), + ) + split_available = capabilities.layer_split and bool(split_devices) + options.append( + PlacementOption( + key="layer-split", + label="Target layer split", + available=split_available, + recommended=split_available and not target_fits, + reason=( + "capacity mode for a target that does not fit the primary GPU" + if split_available + else ( + f"{architecture or 'this model'} has no engine layer-split path" + if not capabilities.layer_split + else "a compatible runtime or enough combined safe memory is unavailable" + ) + ), + ) + ) + + remote_experts_available = ( + runtime.spark and capabilities.expert_offload and (same_backend or cross_backend_ready) + ) + options.append( + PlacementOption( + key="remote-experts", + label="Secondary Spark experts", + available=remote_experts_available, + recommended=remote_experts_available, + reason=( + "runs Spark cold experts on the secondary accelerator instead of CPU" + if remote_experts_available + else "available only when Spark is active and a compatible IPC runtime is installed" + ), + ) + ) + + if remote_experts_available: + placement = PlacementRuntime( + mode="heterogeneous", + target_device=primary.placement_name, + remote_expert_device=secondary.placement_name, + ) + return PlacementPlan( + runtime=placement, + optimization_runtime=runtime, + summary=f"{primary.name} target + {secondary.name} Spark experts", + reason="MoE memory pressure activated Spark; the secondary GPU handles cold experts", + runnable=True, + topology=topology, + options=tuple(options), + ) + + if stack_fits: + return PlacementPlan( + runtime=base, + optimization_runtime=runtime, + summary=primary.name, + reason=( + f"the full stack fits the faster primary; {secondary.name} remains available " + "because a sequential target split would reduce throughput" + ), + runnable=True, + topology=topology, + options=tuple(options), + ) + + if target_fits and draft_offload_available: + mixed = not same_backend + placement = PlacementRuntime( + mode="heterogeneous" if mixed else "draft-offload", + target_device=primary.placement_name, + draft_device=secondary.placement_name, + remote_draft=mixed, + ) + return PlacementPlan( + runtime=placement, + optimization_runtime=runtime, + summary=f"{primary.name} target + {secondary.name} draft/scorer", + reason="the target fits the primary, but its optional draft/scorer needs separate memory", + runnable=True, + topology=topology, + options=tuple(options), + ) + + if not target_fits and split_available: + mixed = len({device.backend for device in split_devices}) > 1 + adjusted = runtime + if not capabilities.draft_on_layer_split: + adjusted = replace(adjusted, speculative_decode=False) + if not capabilities.pflash_on_layer_split: + adjusted = replace( + adjusted, + prefill_mode="off", + prefill_drafter="", + ) + # Spark and target layer splitting are distinct placement systems. No + # architecture currently validates composing both in Automatic mode. + adjusted = replace(adjusted, spark=False, spark_vram_gb=0.0) + placement = PlacementRuntime( + mode="heterogeneous" if mixed else "layer-split", + target_devices=tuple(device.placement_name for device in split_devices), + target_layer_split=split_weights, + remote_target_shard=mixed, + # P2P is topology-specific and cannot be inferred from a shared + # backend name. The safe host-staged boundary is the automatic + # default; an explicit expert benchmark may still opt into P2P. + peer_access=False, + ) + device_names = " + ".join(device.name for device in split_devices) + return PlacementPlan( + runtime=placement, + optimization_runtime=adjusted, + summary=f"{device_names} target split", + reason="capacity mode: the target does not fit the primary GPU alone", + runnable=True, + topology=topology, + options=tuple(options), + ) + + return PlacementPlan( + runtime=base, + optimization_runtime=runtime, + summary=primary.name, + reason=( + "the target does not fit the primary and no validated secondary-device " + "placement is available" + ), + runnable=False, + topology=topology, + options=tuple(options), + ) diff --git a/lucebox/src/lucebox/topology.py b/lucebox/src/lucebox/topology.py new file mode 100644 index 000000000..f92f5593c --- /dev/null +++ b/lucebox/src/lucebox/topology.py @@ -0,0 +1,230 @@ +"""Normalized accelerator inventory for placement and tuning. + +The host wrapper is the only component that can reliably see every physical +GPU. It exports vendor-specific CSV snapshots; this module turns those opaque +probe strings into a small typed topology. No device is selected by summing +VRAM: the planner must name an explicit placement for every device it uses. +""" + +from __future__ import annotations + +import csv +import io +import re +from dataclasses import dataclass + +from lucebox.types import Config, GpuVendor, HostFacts + +_INTEGER_RE = re.compile(r"([0-9]+)") + + +@dataclass(frozen=True, slots=True) +class GpuDevice: + vendor: GpuVendor + index: int + name: str + architecture: str + physical_vram_gb: int + effective_memory_gb: int + unified_memory: bool = False + + @property + def backend(self) -> str: + return "cuda" if self.vendor == "nvidia" else "hip" + + @property + def placement_name(self) -> str: + return f"{self.backend}:{self.index}" + + @property + def label(self) -> str: + arch = f", {self.architecture}" if self.architecture else "" + memory_kind = " shared" if self.unified_memory else " VRAM" + return f"{self.name or self.vendor} ({self.effective_memory_gb} GB{memory_kind}{arch})" + + +@dataclass(frozen=True, slots=True) +class HardwareTopology: + devices: tuple[GpuDevice, ...] + primary: GpuDevice | None + + @property + def companions(self) -> tuple[GpuDevice, ...]: + if self.primary is None: + return self.devices + return tuple(device for device in self.devices if device != self.primary) + + @property + def heterogeneous(self) -> bool: + return len({device.vendor for device in self.devices}) > 1 or len(self.devices) > 1 + + +def _int_from_cell(value: str) -> int: + match = _INTEGER_RE.search(value.replace(",", "")) + return int(match.group(1)) if match else 0 + + +def _rows(raw: str) -> list[list[str]]: + return [ + [cell.strip() for cell in row] + for row in csv.reader(io.StringIO(raw)) + if row and any(cell.strip() for cell in row) + ] + + +def _nvidia_devices(host: HostFacts) -> list[GpuDevice]: + devices: list[GpuDevice] = [] + for row in _rows(host.nvidia_gpu_list_csv): + if len(row) < 6: + continue + try: + index = int(row[0]) + except ValueError: + continue + architecture = row[4].replace(".", "") + memory_gb = _int_from_cell(row[5]) // 1024 + devices.append( + GpuDevice( + vendor="nvidia", + index=index, + name=row[3], + architecture=architecture, + physical_vram_gb=memory_gb, + effective_memory_gb=memory_gb, + ) + ) + if not devices and (host.has_nvidia_gpu or host.gpu_vendor == "nvidia"): + name = host.nvidia_gpu_name + architecture = host.nvidia_gpu_arch + memory_gb = host.nvidia_vram_gb + if host.gpu_vendor == "nvidia": + name = name or host.gpu_name + architecture = architecture or host.gpu_sm + memory_gb = memory_gb or host.vram_gb + devices.append( + GpuDevice( + vendor="nvidia", + index=0, + name=name or "NVIDIA GPU", + architecture=architecture, + physical_vram_gb=memory_gb, + effective_memory_gb=memory_gb, + ) + ) + return devices + + +def _amd_devices(host: HostFacts) -> list[GpuDevice]: + devices: list[GpuDevice] = [] + for row in _rows(host.amd_gpu_list_csv): + # probe_host emits: index, blank, blank, name, gfx arch, MiB, blank + if len(row) < 6: + continue + try: + index = int(row[0]) + except ValueError: + continue + physical_gb = _int_from_cell(row[5]) // 1024 + architecture = row[4] + # gfx1151 is the Strix Halo integrated GPU. Firmware/driver versions + # disagree on whether SMI reports only the carve-out or a larger UMA + # aperture, so the architecture — not the reported VRAM size — is the + # stable signal that model memory comes from the shared system pool. + unified = architecture == "gfx1151" + # Keep 16 GB for the OS, file cache, and CPU-side engine buffers. The + # HIP allocator performs its own live free-memory check at load time; + # this number is only the planner's conservative capacity ceiling. + effective_gb = max(physical_gb, max(0, host.ram_gb - 16)) if unified else physical_gb + devices.append( + GpuDevice( + vendor="amd", + index=index, + name=row[3], + architecture=architecture, + physical_vram_gb=physical_gb, + effective_memory_gb=effective_gb, + unified_memory=unified, + ) + ) + if not devices and (host.has_amd_gpu or host.gpu_vendor == "amd"): + name = host.amd_gpu_name + architecture = host.amd_gpu_arch + memory_gb = host.amd_vram_gb + if host.gpu_vendor == "amd": + name = name or host.gpu_name + architecture = architecture or host.gpu_sm + memory_gb = memory_gb or host.vram_gb + unified = architecture == "gfx1151" + effective_gb = max(memory_gb, max(0, host.ram_gb - 16)) if unified else memory_gb + devices.append( + GpuDevice( + vendor="amd", + index=0, + name=name or "AMD GPU", + architecture=architecture, + physical_vram_gb=0 if unified else memory_gb, + effective_memory_gb=effective_gb, + unified_memory=unified, + ) + ) + return devices + + +def _selected_vendor(cfg: Config) -> GpuVendor: + variant = cfg.variant.lower() + if "rocm" in variant or "hip" in variant: + return "amd" + if "cuda" in variant: + return "nvidia" + return cfg.host.gpu_vendor + + +def _pick_primary(cfg: Config, devices: list[GpuDevice]) -> GpuDevice | None: + vendor = _selected_vendor(cfg) + candidates = [device for device in devices if device.vendor == vendor] + if not candidates: + return None + + # The wrapper's selected facts encode policy as well as capacity. In + # particular, R9700 remains primary over a larger-capacity Strix UMA GPU + # because the discrete card is substantially faster for a fitting model. + selected_name = cfg.host.gpu_name + selected_arch = cfg.host.gpu_sm + for device in candidates: + if selected_name and device.name == selected_name: + return device + for device in candidates: + if selected_arch and device.architecture == selected_arch: + return device + return max( + candidates, + key=lambda device: ( + not device.unified_memory, + device.physical_vram_gb, + device.effective_memory_gb, + -device.index, + ), + ) + + +def from_config(cfg: Config) -> HardwareTopology: + devices = _nvidia_devices(cfg.host) + _amd_devices(cfg.host) + # Older config snapshots and direct Python callers may only provide the + # historical generic GPU fields. The selected image variant still gives us + # enough information to form a single-device topology without pretending a + # second accelerator exists. + if not devices and cfg.host.vram_gb > 0: + vendor = _selected_vendor(cfg) + if vendor != "none": + devices.append( + GpuDevice( + vendor=vendor, + index=0, + name=cfg.host.gpu_name or f"{vendor.upper()} GPU", + architecture=cfg.host.gpu_sm, + physical_vram_gb=cfg.host.vram_gb, + effective_memory_gb=cfg.host.vram_gb, + ) + ) + devices.sort(key=lambda device: (device.vendor, device.index)) + return HardwareTopology(tuple(devices), _pick_primary(cfg, devices)) diff --git a/lucebox/src/lucebox/types.py b/lucebox/src/lucebox/types.py index 6164fc6c7..0e1946526 100644 --- a/lucebox/src/lucebox/types.py +++ b/lucebox/src/lucebox/types.py @@ -9,6 +9,7 @@ import math import os +import re from dataclasses import dataclass, field from pathlib import Path from typing import Literal @@ -31,6 +32,12 @@ def default_models_dir() -> Path: GpuVendor = Literal["nvidia", "amd", "none"] +PlacementMode = Literal[ + "single", + "draft-offload", + "layer-split", + "heterogeneous", +] @dataclass(frozen=True, slots=True) @@ -60,6 +67,25 @@ class HostFacts: has_docker: bool = False docker_version: str = "" ctk: CtkStatus = "none" + # Per-vendor inventory is kept in addition to the selected accelerator. + # On an RTX + Strix host the generic ``gpu_*`` fields describe the RTX, + # while these fields retain the AMD companion for placement planning. On + # an R9700 + Strix host the AMD CSV contains both devices. + nvidia_gpu_name: str = "" + nvidia_gpu_count: int = 0 + nvidia_vram_gb: int = 0 + nvidia_gpu_arch: str = "" + nvidia_gpu_list_csv: str = "" + amd_gpu_name: str = "" + amd_gpu_count: int = 0 + amd_vram_gb: int = 0 + amd_gpu_arch: str = "" + amd_gpu_list_csv: str = "" + # The first packaged cross-vendor runtime is a CUDA server paired with a + # HIP companion daemon. The host wrapper proves that contract before + # setting this fact; the planner never assumes a CUDA image can execute a + # HIP child process. + hybrid_runtime: bool = False def __post_init__(self) -> None: """Reject malformed persisted host snapshots. @@ -76,6 +102,124 @@ def __post_init__(self) -> None: ) +_PLACEMENT_DEVICE_RE = re.compile(r"^(cuda|hip):([0-9]+)$") + + +def _placement_backend(device: str) -> str: + match = _PLACEMENT_DEVICE_RE.fullmatch(device) + if not match: + raise ValueError(f"placement device must use cuda:N or hip:N syntax; got {device!r}") + return match.group(1) + + +@dataclass(frozen=True, slots=True) +class PlacementRuntime: + """Validated accelerator placement passed atomically to the engine. + + Paths to backend IPC binaries deliberately do not live in user config. + They are installation details resolved by the host wrapper after it has + verified the executable and its backend runtime. These fields describe + only the portable execution intent. + """ + + mode: PlacementMode = "single" + target_device: str = "" + target_devices: tuple[str, ...] = () + target_layer_split: tuple[float, ...] = () + draft_device: str = "" + remote_draft: bool = False + remote_target_shard: bool = False + peer_access: bool = False + remote_expert_device: str = "" + + def __post_init__(self) -> None: + if self.mode not in { + "single", + "draft-offload", + "layer-split", + "heterogeneous", + }: + raise ValueError(f"unknown placement mode {self.mode!r}") + if self.target_device and self.target_devices: + raise ValueError("target_device and target_devices are mutually exclusive") + + devices = self.target_devices or ((self.target_device,) if self.target_device else ()) + if not devices and self.mode != "single": + raise ValueError("non-single placement requires a target device") + if len(set(devices)) != len(devices): + raise ValueError("target devices must be unique") + backends = tuple(_placement_backend(device) for device in devices) + if self.draft_device: + draft_backend = _placement_backend(self.draft_device) + else: + draft_backend = backends[0] if backends else "" + if self.remote_expert_device: + _placement_backend(self.remote_expert_device) + + if self.target_layer_split: + if len(self.target_layer_split) != len(self.target_devices): + raise ValueError("target_layer_split must contain one weight per target device") + if len(self.target_layer_split) < 2: + raise ValueError("target layer split requires at least two devices") + if any( + not math.isfinite(weight) or weight <= 0.0 for weight in self.target_layer_split + ): + raise ValueError("target layer split weights must be finite and positive") + elif len(self.target_devices) > 1: + raise ValueError("multiple target devices require target_layer_split") + + mixed_target = len(set(backends)) > 1 + if mixed_target and not self.remote_target_shard: + raise ValueError("mixed-backend target placement requires remote_target_shard") + if self.remote_target_shard and not mixed_target: + raise ValueError("remote_target_shard requires mixed-backend target placement") + mixed_draft = bool(draft_backend and backends and draft_backend != backends[0]) + if mixed_draft and not self.remote_draft: + raise ValueError("mixed-backend draft placement requires remote_draft") + if self.remote_draft and not mixed_draft: + raise ValueError("remote_draft requires mixed-backend draft placement") + if self.draft_device and self.draft_device in devices: + raise ValueError("draft_device must differ from the target device") + if self.remote_expert_device and self.remote_expert_device in devices: + raise ValueError("remote_expert_device must differ from target devices") + if self.peer_access and (len(backends) < 2 or len(set(backends)) != 1): + raise ValueError("peer_access requires a same-backend target layer split") + if self.mode == "single" and ( + len(devices) > 1 + or self.draft_device + or self.remote_expert_device + or self.remote_draft + or self.remote_target_shard + ): + raise ValueError("single placement cannot contain secondary-device settings") + if self.mode == "draft-offload" and not self.draft_device: + raise ValueError("draft-offload placement requires draft_device") + if self.mode == "layer-split" and not self.target_layer_split: + raise ValueError("layer-split placement requires target_layer_split") + if self.mode == "heterogeneous" and not ( + self.remote_draft or self.remote_target_shard or self.remote_expert_device + ): + raise ValueError("heterogeneous placement requires a cross-device workload") + + @property + def uses_multiple_devices(self) -> bool: + return bool(len(self.target_devices) > 1 or self.draft_device or self.remote_expert_device) + + @property + def target_backend(self) -> str: + devices = self.target_devices or ((self.target_device,) if self.target_device else ()) + return _placement_backend(devices[0]) if devices else "" + + @property + def requires_hybrid_runtime(self) -> bool: + """Whether one process must launch a daemon from another backend.""" + if self.remote_draft or self.remote_target_shard: + return True + if not self.remote_expert_device or not self.target_backend: + return False + return _placement_backend(self.remote_expert_device) != self.target_backend + + @dataclass(frozen=True, slots=True) class DflashRuntime: """Typed inference and optimization settings persisted under ``[dflash]``. @@ -217,5 +361,6 @@ class Config: port: int = 8080 models_dir: Path = field(default_factory=default_models_dir) dflash: DflashRuntime = field(default_factory=DflashRuntime) + placement: PlacementRuntime = field(default_factory=PlacementRuntime) host: HostFacts = field(default_factory=HostFacts) model: ModelMeta = field(default_factory=ModelMeta) diff --git a/lucebox/tests/test_check.py b/lucebox/tests/test_check.py index ba48b4d7d..ddf34058a 100644 --- a/lucebox/tests/test_check.py +++ b/lucebox/tests/test_check.py @@ -33,6 +33,7 @@ def test_check_prints_host_facts_section(monkeypatch: pytest.MonkeyPatch) -> Non "LUCEBOX_HOST_GPU_LIST_CSV", "0, GPU-abc, 00000000:01:00.0, NVIDIA RTX 5090, 12.0, 24576 MiB, 175.00 W", ) + # Stub HostFacts so the pass/fail checks succeed at least minimally. # `cli.check` imports `from_env` into its module namespace, so patch # both names. @@ -53,6 +54,7 @@ def stub() -> HostFacts: docker_version="29.1.3", ctk="runtime", ) + monkeypatch.setattr("lucebox.host_facts.from_env", stub) monkeypatch.setattr("lucebox.cli.from_env", stub) result = CliRunner().invoke(app, ["check"]) @@ -93,6 +95,7 @@ def test_check_exit_code_independent_of_host_facts( facts block prints. """ monkeypatch.setenv("LUCEBOX_HOST_OS_PRETTY", "Bare Linux") + def stub() -> HostFacts: return HostFacts( nproc=8, @@ -110,6 +113,7 @@ def stub() -> HostFacts: docker_version="", ctk="none", # also fail ) + monkeypatch.setattr("lucebox.host_facts.from_env", stub) monkeypatch.setattr("lucebox.cli.from_env", stub) result = CliRunner().invoke(app, ["check"]) diff --git a/lucebox/tests/test_cli.py b/lucebox/tests/test_cli.py index 7fb146ffd..e7d44aa10 100644 --- a/lucebox/tests/test_cli.py +++ b/lucebox/tests/test_cli.py @@ -50,8 +50,7 @@ def test_deferred_verbs_are_not_registered(verb: str) -> None: def test_core_verbs_present_in_app() -> None: """The core launch/serve surface stays wired into the Typer command table.""" registered = { - c.name or (c.callback.__name__ if c.callback else "") - for c in app.registered_commands + c.name or (c.callback.__name__ if c.callback else "") for c in app.registered_commands } for verb in ("check", "pull", "optimize", "print-run", "print-serve-argv", "version"): assert verb in registered diff --git a/lucebox/tests/test_config.py b/lucebox/tests/test_config.py index a1ec51043..ac9a8bb1a 100644 --- a/lucebox/tests/test_config.py +++ b/lucebox/tests/test_config.py @@ -435,3 +435,100 @@ def test_model_switch_preserves_custom_profile(tmp_path: Path) -> None: loaded = config.load(path) assert loaded is not None assert loaded.dflash.max_ctx == 8192 + + +def test_optimization_and_placement_are_persisted_atomically(tmp_path: Path) -> None: + from lucebox.types import DflashRuntime, PlacementRuntime + + path = tmp_path / "config.toml" + placement = PlacementRuntime( + mode="layer-split", + target_devices=("hip:0", "hip:1"), + target_layer_split=(0.7, 0.3), + peer_access=True, + ) + + config.write_optimization_runtime( + DflashRuntime(max_ctx=32768), + placement=placement, + path=path, + ) + loaded = config.load(path) + + assert loaded is not None + assert loaded.dflash.max_ctx == 32768 + assert loaded.placement == placement + assert config.optimization_mode(path=path) == "automatic" + + +def test_runtime_reset_clears_stale_placement(tmp_path: Path) -> None: + from lucebox.types import DflashRuntime, PlacementRuntime + + path = tmp_path / "config.toml" + config.write_optimization_runtime( + DflashRuntime(), + placement=PlacementRuntime( + mode="layer-split", + target_devices=("hip:0", "hip:1"), + target_layer_split=(0.7, 0.3), + ), + path=path, + ) + + config.write_optimization_runtime(DflashRuntime(max_ctx=4096), path=path) + + assert "placement" not in config.load_doc(path) + + +def test_model_and_execution_profile_switch_atomically(tmp_path: Path) -> None: + from lucebox.types import DflashRuntime, ModelMeta, PlacementRuntime + + path = tmp_path / "config.toml" + path.write_text( + '[model]\npreset = "old"\ntarget_file = "old.gguf"\ndraft_file = "old-draft.gguf"\n' + ) + placement = PlacementRuntime(mode="single", target_device="cuda:0") + + config.write_model_profile( + ModelMeta(preset="new", target_file="new.gguf"), + DflashRuntime(max_ctx=32768), + placement, + path=path, + ) + + loaded = config.load(path) + assert loaded is not None + assert loaded.model == ModelMeta(preset="new", target_file="new.gguf") + assert loaded.dflash.max_ctx == 32768 + assert loaded.placement == placement + assert "draft_file" not in config.load_doc(path)["model"] + + +def test_automatic_seed_refuses_unrunnable_placement(tmp_path: Path) -> None: + from lucebox.types import Config, HostFacts, ModelMeta + + path = tmp_path / "config.toml" + cfg = Config( + models_dir=tmp_path / "models", + host=HostFacts(), + model=ModelMeta(preset="qwen3.6-27b"), + ) + + with pytest.raises(ValueError, match="no runnable placement"): + config.seed_optimization_from_config(cfg, path=path) + + assert not path.exists() + + +def test_config_rejects_unpaired_mixed_backend_placement(tmp_path: Path) -> None: + path = tmp_path / "config.toml" + path.write_text( + "[placement]\n" + 'mode = "heterogeneous"\n' + 'target_device = "cuda:0"\n' + 'draft_device = "hip:0"\n' + "remote_draft = false\n" + ) + + with pytest.raises(ValueError, match="requires remote_draft"): + config.load(path) diff --git a/lucebox/tests/test_config_cli.py b/lucebox/tests/test_config_cli.py index a9a727723..c028cb071 100644 --- a/lucebox/tests/test_config_cli.py +++ b/lucebox/tests/test_config_cli.py @@ -14,9 +14,7 @@ def _set_config_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: return tmp_path / "config.toml" -def test_config_set_then_get_round_trip( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_config_set_then_get_round_trip(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: cfg_path = _set_config_path(tmp_path, monkeypatch) set_result = CliRunner().invoke(app, ["config", "set", "dflash.budget=12"]) assert set_result.exit_code == 0 @@ -38,9 +36,7 @@ def test_config_get_with_no_key_lists_every_registered_key( assert key in result.stdout -def test_config_unset_drops_key( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_config_unset_drops_key(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: cfg_path = _set_config_path(tmp_path, monkeypatch) CliRunner().invoke(app, ["config", "set", "dflash.budget=9"]) assert "budget = 9" in cfg_path.read_text() @@ -50,17 +46,13 @@ def test_config_unset_drops_key( assert "budget" not in body -def test_config_set_unknown_key_errors( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_config_set_unknown_key_errors(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: _set_config_path(tmp_path, monkeypatch) result = CliRunner().invoke(app, ["config", "set", "totally.unknown=1"]) assert result.exit_code == 2 -def test_config_set_rejects_missing_equals( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_config_set_rejects_missing_equals(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: _set_config_path(tmp_path, monkeypatch) result = CliRunner().invoke(app, ["config", "set", "dflash.budget"]) assert result.exit_code == 2 @@ -108,8 +100,7 @@ def test_load_or_build_env_overrides_persisted_config( # Write a config.toml WITHOUT an image.registry line — the # bug-trigger shape on sindri. cfg_path.write_text( - '[image]\nvariant = "cuda12"\n[runtime]\nport = 9090\n' - '[dflash]\nbudget = 22\n' + '[image]\nvariant = "cuda12"\n[runtime]\nport = 9090\n[dflash]\nbudget = 22\n' ) # Env should override what config.toml says (and what dataclass # defaults fill in for missing keys). @@ -118,8 +109,8 @@ def test_load_or_build_env_overrides_persisted_config( monkeypatch.setenv("LUCEBOX_CONTAINER", "lucebox-test") cfg = _load_or_build() assert cfg.image == "ghcr.io/myfork/lucebox-hub" # env beats dataclass default - assert cfg.port == 7777 # env beats config.toml - assert cfg.container_name == "lucebox-test" # env applied + assert cfg.port == 7777 # env beats config.toml + assert cfg.container_name == "lucebox-test" # env applied # variant is in config.toml — config.toml value (no env override). assert cfg.variant == "cuda12" # dflash IS persisted in config.toml — env doesn't touch it (no DFLASH_* diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index 251a6028c..f8412e56a 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -12,7 +12,13 @@ import pytest from lucebox.download import PRESETS -from lucebox.types import Config, DflashRuntime, HostFacts, ModelMeta +from lucebox.types import ( + Config, + DflashRuntime, + HostFacts, + ModelMeta, + PlacementRuntime, +) from lucebox import docker_run @@ -44,9 +50,7 @@ def test_argv_flags_and_ordering() -> None: detach=True, remove=False, port_publish=(8080, 8080), - volumes=( - docker_run.BindMount("/host/models", "/opt/lucebox-hub/server/models"), - ), + volumes=(docker_run.BindMount("/host/models", "/opt/lucebox-hub/server/models"),), env=(("DFLASH_BUDGET", "22"),), entrypoint_args=("serve",), extra=("--shm-size", "1g"), @@ -57,9 +61,7 @@ def test_argv_flags_and_ordering() -> None: assert "--gpus" not in argv # gpus=False assert ["-p", "8080:8080"] == argv[argv.index("-p") : argv.index("-p") + 2] mount_arg = "type=bind,source=/host/models,target=/opt/lucebox-hub/server/models" - assert ["--mount", mount_arg] == argv[ - argv.index("--mount") : argv.index("--mount") + 2 - ] + assert ["--mount", mount_arg] == argv[argv.index("--mount") : argv.index("--mount") + 2] assert ["-e", "DFLASH_BUDGET=22"] == argv[argv.index("-e") : argv.index("-e") + 2] # extra flags precede the image; entrypoint_args follow it. assert argv[-1] == "serve" @@ -85,9 +87,7 @@ def test_argv_amd_uses_rocm_device_contract() -> None: spec = docker_run.DockerRunSpec(image="img:rocm", name="box", gpu_vendor="amd") argv = spec.argv() assert "--gpus" not in argv - assert ["--device", "/dev/kfd"] == argv[ - argv.index("--device") : argv.index("--device") + 2 - ] + assert ["--device", "/dev/kfd"] == argv[argv.index("--device") : argv.index("--device") + 2] assert "/dev/dri" in argv assert ["--group-add", "video"] == argv[ argv.index("--group-add") : argv.index("--group-add") + 2 @@ -126,12 +126,8 @@ def test_runtime_volumes_mounts_only_models_and_config( monkeypatch.delenv("LUCEBOX_HOME", raising=False) cfg = Config(models_dir=tmp_path / "models") vols = docker_run._runtime_volumes(cfg) - assert docker_run.BindMount( - str(tmp_path / "models"), "/opt/lucebox-hub/server/models" - ) in vols - assert docker_run.BindMount( - str(home / ".lucebox"), str(home / ".lucebox") - ) in vols + assert docker_run.BindMount(str(tmp_path / "models"), "/opt/lucebox-hub/server/models") in vols + assert docker_run.BindMount(str(home / ".lucebox"), str(home / ".lucebox")) in vols assert all(mount.source != str(home) for mount in vols) @@ -143,18 +139,11 @@ def test_runtime_volumes_dedupes_when_models_is_home(monkeypatch, tmp_path: Path # models_dir is mounted at the image's canonical path, while config keeps # its same-path mount. The parent home directory itself is never exposed. assert len(vols) == 2 - assert docker_run.BindMount( - str(tmp_path / ".lucebox"), str(tmp_path / ".lucebox") - ) in vols - assert all( - mount.source != str(tmp_path) or mount.target != str(tmp_path) - for mount in vols - ) + assert docker_run.BindMount(str(tmp_path / ".lucebox"), str(tmp_path / ".lucebox")) in vols + assert all(mount.source != str(tmp_path) or mount.target != str(tmp_path) for mount in vols) -def test_runtime_volumes_mounts_custom_config_home( - monkeypatch, tmp_path: Path -) -> None: +def test_runtime_volumes_mounts_custom_config_home(monkeypatch, tmp_path: Path) -> None: home = tmp_path / "home" config_home = tmp_path / "config-outside-home" monkeypatch.setattr(Path, "home", staticmethod(lambda: home)) @@ -242,9 +231,7 @@ def test_server_run_spec_top_level_shape(tmp_path: Path) -> None: assert spec.remove is True assert spec.detach is False assert spec.port_publish == (9000, 8080) - assert docker_run.BindMount( - str(tmp_path), "/opt/lucebox-hub/server/models" - ) in spec.volumes + assert docker_run.BindMount(str(tmp_path), "/opt/lucebox-hub/server/models") in spec.volumes def test_server_run_spec_mounts_selected_symlink_target_narrowly_read_only( @@ -264,11 +251,14 @@ def test_server_run_spec_mounts_selected_symlink_target_narrowly_read_only( Config(models_dir=models, model=ModelMeta(target_file="selected.gguf")) ) - assert docker_run.BindMount( - str(target), - "/opt/lucebox-resolved/target/target.gguf", - read_only=True, - ) in spec.volumes + assert ( + docker_run.BindMount( + str(target), + "/opt/lucebox-resolved/target/target.gguf", + read_only=True, + ) + in spec.volumes + ) assert _env(spec)["DFLASH_TARGET"] == "/opt/lucebox-resolved/target/target.gguf" assert all(mount.source != str(home) for mount in spec.volumes) argv = spec.argv() @@ -292,11 +282,14 @@ def test_server_run_spec_resolves_symlinked_model_parent(tmp_path: Path) -> None Config(models_dir=models, model=ModelMeta(target_file="selected/nested.gguf")) ) - assert docker_run.BindMount( - str(target), - "/opt/lucebox-resolved/target/nested.gguf", - read_only=True, - ) in spec.volumes + assert ( + docker_run.BindMount( + str(target), + "/opt/lucebox-resolved/target/nested.gguf", + read_only=True, + ) + in spec.volumes + ) assert _env(spec)["DFLASH_TARGET"] == "/opt/lucebox-resolved/target/nested.gguf" @@ -314,10 +307,7 @@ def test_server_run_spec_does_not_remount_file_for_symlinked_models_dir( ) assert _env(spec)["DFLASH_TARGET"] == "/opt/lucebox-hub/server/models/target.gguf" - assert all( - mount.target != "/opt/lucebox-resolved/target/target.gguf" - for mount in spec.volumes - ) + assert all(mount.target != "/opt/lucebox-resolved/target/target.gguf" for mount in spec.volumes) def test_server_run_spec_mounts_symlinked_speculator_directory_read_only( @@ -329,19 +319,20 @@ def test_server_run_spec_mounts_symlinked_speculator_directory_read_only( draft_root.mkdir(parents=True) external.mkdir() (external / "model.safetensors").write_bytes(b"speculator") - (draft_root / "laguna-xs2-speculator").symlink_to( - external, target_is_directory=True - ) + (draft_root / "laguna-xs2-speculator").symlink_to(external, target_is_directory=True) spec = docker_run.server_run_spec( Config(models_dir=models, model=ModelMeta(preset="laguna-xs.2")) ) - assert docker_run.BindMount( - str(external), - "/opt/lucebox-resolved/draft-dir", - read_only=True, - ) in spec.volumes + assert ( + docker_run.BindMount( + str(external), + "/opt/lucebox-resolved/draft-dir", + read_only=True, + ) + in spec.volumes + ) assert _env(spec)["DFLASH_DRAFT"] == "/opt/lucebox-resolved/draft-dir" @@ -472,15 +463,90 @@ def test_server_run_spec_forwards_primary_rocm_device( assert "HIP_VISIBLE_DEVICES" not in env +def test_server_run_spec_forwards_explicit_single_gpu_placement( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("LUCEBOX_HOST_ROCR_VISIBLE_DEVICES", "GPU-primary") + cfg = Config( + variant="rocm", + models_dir=tmp_path, + placement=PlacementRuntime(target_device="hip:1"), + ) + + env = _env(docker_run.server_run_spec(cfg)) + + assert env["DFLASH_TARGET_DEVICE"] == "hip:1" + # Explicit engine placement uses the physical inventory. A visibility pin + # would renumber the selected device to hip:0 and invalidate the plan. + assert "ROCR_VISIBLE_DEVICES" not in env + assert "HIP_VISIBLE_DEVICES" not in env + + +def test_server_run_spec_forwards_same_backend_layer_split(tmp_path: Path) -> None: + cfg = Config( + variant="rocm", + models_dir=tmp_path, + placement=PlacementRuntime( + mode="layer-split", + target_devices=("hip:0", "hip:1"), + target_layer_split=(0.75, 0.25), + peer_access=True, + ), + ) + + env = _env(docker_run.server_run_spec(cfg)) + + assert env["DFLASH_TARGET_DEVICES"] == "hip:0,hip:1" + assert env["DFLASH_TARGET_LAYER_SPLIT"] == "0.75,0.25" + assert env["DFLASH_PEER_ACCESS"] == "1" + + +def test_server_run_spec_forwards_same_backend_spark_companion( + tmp_path: Path, +) -> None: + cfg = Config( + variant="rocm", + models_dir=tmp_path, + dflash=DflashRuntime(spark=True), + placement=PlacementRuntime( + mode="heterogeneous", + target_device="hip:0", + remote_expert_device="hip:1", + ), + ) + + env = _env(docker_run.server_run_spec(cfg)) + + assert env["DFLASH_SPARK"] == "1" + assert env["DFLASH_REMOTE_EXPERT_DEVICE"] == "hip:1" + + +def test_server_run_spec_rejects_cross_backend_docker_placement( + tmp_path: Path, +) -> None: + cfg = Config( + variant="cuda12", + models_dir=tmp_path, + dflash=DflashRuntime(spark=True), + placement=PlacementRuntime( + mode="heterogeneous", + target_device="cuda:0", + remote_expert_device="hip:0", + ), + ) + + with pytest.raises(ValueError, match="paired native Lucebox runtime"): + docker_run.server_run_spec(cfg) + + def test_server_run_spec_resolves_target_and_draft_paths(tmp_path: Path) -> None: pres = PRESETS["qwen3.6-27b"] cfg = Config(models_dir=tmp_path, model=ModelMeta(preset="qwen3.6-27b")) env = _env(docker_run.server_run_spec(cfg)) assert env["DFLASH_TARGET"] == f"/opt/lucebox-hub/server/models/{pres.target_file}" if pres.draft_file: - assert env["DFLASH_DRAFT"] == ( - f"/opt/lucebox-hub/server/models/draft/{pres.draft_file}" - ) + assert env["DFLASH_DRAFT"] == (f"/opt/lucebox-hub/server/models/draft/{pres.draft_file}") def test_server_run_spec_can_disable_preset_dflash_draft(tmp_path: Path) -> None: diff --git a/lucebox/tests/test_placement.py b/lucebox/tests/test_placement.py new file mode 100644 index 000000000..784a0c0ba --- /dev/null +++ b/lucebox/tests/test_placement.py @@ -0,0 +1,292 @@ +import re +from pathlib import Path +from types import SimpleNamespace + +import pytest +from lucebox.autotune import automatic_plan +from lucebox.placement import ARCHITECTURE_CAPABILITIES, automatic_placement +from lucebox.topology import from_config +from lucebox.types import Config, DflashRuntime, HostFacts, ModelMeta + +from lucebox import download + + +def _install_optimizer_drafter(cfg: Config) -> None: + path = download.optimizer_drafter_path(cfg) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"test-scorer") + + +def _nvidia_csv(*memory_mib: int) -> str: + return "\n".join( + f"{index}, GPU-{index}, 0000:{index:02x}:00.0, RTX {index}, 8.6, {memory} MiB, 350 W" + for index, memory in enumerate(memory_mib) + ) + + +def _amd_csv(*rows: tuple[str, str, int]) -> str: + return "\n".join( + f"{index}, , , {name}, {architecture}, {memory_mib} MiB," + for index, (name, architecture, memory_mib) in enumerate(rows) + ) + + +def test_r9700_strix_keeps_fitting_qwen_target_on_r9700(tmp_path: Path) -> None: + host = HostFacts( + gpu_vendor="amd", + has_amd_gpu=True, + gpu_name="AMD Radeon AI PRO R9700", + gpu_count=2, + vram_gb=31, + gpu_sm="gfx1201", + ram_gb=125, + amd_gpu_name="AMD Radeon AI PRO R9700", + amd_gpu_count=2, + amd_vram_gb=31, + amd_gpu_arch="gfx1201", + amd_gpu_list_csv=_amd_csv( + ("AMD Radeon AI PRO R9700", "gfx1201", 32624), + ("AMD Radeon Graphics", "gfx1151", 512), + ), + ) + cfg = Config( + variant="rocm", + models_dir=tmp_path, + host=host, + model=ModelMeta(preset="qwen3.6-27b"), + ) + + topology = from_config(cfg) + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert topology.primary is not None + assert topology.primary.name == "AMD Radeon AI PRO R9700" + assert topology.companions[0].unified_memory is True + assert topology.companions[0].effective_memory_gb == 109 + assert plan.placement.runtime.target_device == "hip:0" + assert plan.placement.runtime.uses_multiple_devices is False + assert "full stack fits" in plan.placement.reason + + +def test_three_same_backend_gpus_are_used_only_when_capacity_requires_them( + tmp_path: Path, +) -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + gpu_name="RTX 0", + gpu_count=3, + vram_gb=24, + gpu_sm="86", + nvidia_gpu_name="RTX 0", + nvidia_gpu_count=3, + nvidia_vram_gb=24, + nvidia_gpu_arch="86", + nvidia_gpu_list_csv=_nvidia_csv(24576, 24576, 24576), + ) + cfg = Config(variant="cuda12", models_dir=tmp_path, host=host) + preset = SimpleNamespace( + architecture="qwen35", + target_file="large.gguf", + approx_target_gb=55.0, + approx_draft_gb=0.0, + ) + + plan = automatic_placement( + cfg, + DflashRuntime(speculative_decode=False), + preset, + has_draft=False, + optimizer_drafter_available=False, + ) + + assert plan.runnable is True + assert plan.runtime.mode == "layer-split" + assert plan.runtime.target_devices == ("cuda:0", "cuda:1", "cuda:2") + assert sum(plan.runtime.target_layer_split) == pytest.approx(1.0) + assert plan.runtime.peer_access is False + + +def test_rtx_strix_moe_uses_paired_runtime_for_remote_spark_experts( + tmp_path: Path, +) -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + has_amd_gpu=True, + gpu_name="RTX 3090", + gpu_count=1, + vram_gb=24, + gpu_sm="86", + ram_gb=125, + nvidia_gpu_name="RTX 3090", + nvidia_gpu_count=1, + nvidia_vram_gb=24, + nvidia_gpu_arch="86", + nvidia_gpu_list_csv=("0, GPU-0, 0000:01:00.0, RTX 3090, 8.6, 24576 MiB, 350 W"), + amd_gpu_name="AMD Radeon Graphics", + amd_gpu_count=1, + amd_vram_gb=125, + amd_gpu_arch="gfx1151", + amd_gpu_list_csv=_amd_csv(("AMD Radeon Graphics", "gfx1151", 512)), + hybrid_runtime=True, + ) + cfg = Config( + variant="cuda12", + models_dir=tmp_path, + host=host, + model=ModelMeta(preset="qwen3.6-moe"), + ) + _install_optimizer_drafter(cfg) + + plan = automatic_plan(cfg) + + assert plan.runtime.spark is True + assert plan.placement.runtime.target_device == "cuda:0" + assert plan.placement.runtime.remote_expert_device == "hip:0" + assert plan.placement.runtime.requires_hybrid_runtime is True + assert "Spark experts" in plan.placement.summary + + +def test_rtx_strix_without_paired_runtime_keeps_spark_on_cpu(tmp_path: Path) -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + has_amd_gpu=True, + gpu_name="RTX 3090", + vram_gb=24, + gpu_sm="86", + ram_gb=125, + nvidia_gpu_list_csv=("0, GPU-0, 0000:01:00.0, RTX 3090, 8.6, 24576 MiB, 350 W"), + amd_gpu_list_csv=_amd_csv(("AMD Radeon Graphics", "gfx1151", 512)), + ) + cfg = Config( + variant="cuda12", + models_dir=tmp_path, + host=host, + model=ModelMeta(preset="qwen3.6-moe"), + ) + _install_optimizer_drafter(cfg) + + plan = automatic_plan(cfg) + + assert plan.runtime.spark is True + assert plan.placement.runtime.target_device == "cuda:0" + assert plan.placement.runtime.remote_expert_device == "" + assert plan.placement.runtime.requires_hybrid_runtime is False + remote_option = next( + option for option in plan.placement.options if option.key == "remote-experts" + ) + assert remote_option.available is False + + +def test_paired_runtime_does_not_claim_unsupported_hip_to_cuda_direction( + tmp_path: Path, +) -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + has_amd_gpu=True, + gpu_name="RTX 3090", + vram_gb=24, + gpu_sm="86", + ram_gb=125, + nvidia_gpu_list_csv=("0, GPU-0, 0000:01:00.0, RTX 3090, 8.6, 24576 MiB, 350 W"), + amd_gpu_name="AMD Radeon Graphics", + amd_gpu_count=1, + amd_vram_gb=125, + amd_gpu_arch="gfx1151", + amd_gpu_list_csv=_amd_csv(("AMD Radeon Graphics", "gfx1151", 512)), + hybrid_runtime=True, + ) + cfg = Config( + variant="rocm", + models_dir=tmp_path, + host=host, + model=ModelMeta(preset="qwen3.6-moe"), + ) + _install_optimizer_drafter(cfg) + + plan = automatic_plan(cfg) + + assert plan.placement.runtime.target_device == "hip:0" + assert plan.placement.runtime.remote_expert_device == "" + assert plan.placement.runtime.requires_hybrid_runtime is False + + +def test_strix_is_uma_even_when_driver_reports_large_aperture(tmp_path: Path) -> None: + host = HostFacts( + gpu_vendor="amd", + has_amd_gpu=True, + gpu_name="AMD Radeon Graphics", + vram_gb=16, + gpu_sm="gfx1151", + ram_gb=128, + amd_gpu_list_csv=_amd_csv(("AMD Radeon Graphics", "gfx1151", 16384)), + ) + + topology = from_config(Config(variant="rocm", models_dir=tmp_path, host=host)) + + assert topology.primary is not None + assert topology.primary.unified_memory is True + assert topology.primary.physical_vram_gb == 16 + assert topology.primary.effective_memory_gb == 112 + + +def test_rtx_strix_uses_remote_target_shard_when_capacity_requires_it( + tmp_path: Path, +) -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + has_amd_gpu=True, + gpu_name="RTX 3090", + vram_gb=24, + gpu_sm="86", + ram_gb=128, + nvidia_gpu_list_csv=("0, GPU-0, 0000:01:00.0, RTX 3090, 8.6, 24576 MiB, 350 W"), + amd_gpu_list_csv=_amd_csv(("AMD Radeon Graphics", "gfx1151", 512)), + hybrid_runtime=True, + ) + cfg = Config(variant="cuda12", models_dir=tmp_path, host=host) + preset = SimpleNamespace( + architecture="qwen35", + target_file="large.gguf", + approx_target_gb=40.0, + approx_draft_gb=0.0, + ) + + plan = automatic_placement( + cfg, + DflashRuntime(speculative_decode=False), + preset, + has_draft=False, + optimizer_drafter_available=False, + ) + + assert plan.runnable is True + assert plan.runtime.mode == "heterogeneous" + assert plan.runtime.target_devices == ("cuda:0", "hip:0") + assert plan.runtime.remote_target_shard is True + assert plan.runtime.peer_access is False + + +def test_python_architecture_capabilities_match_engine_table() -> None: + header = (Path(__file__).parents[2] / "server/src/common/model_capabilities.h").read_text() + row_pattern = re.compile( + r'\{"(?P[^\"]+)",\s*' + r"(?Ptrue|false),\s*(?Ptrue|false),\s*" + r"(?Ptrue|false),\s*(?Ptrue|false),\s*" + r"(?PkNever|kMono|kBoth)," + ) + engine_rows = {match["arch"]: match.groupdict() for match in row_pattern.finditer(header)} + + assert set(ARCHITECTURE_CAPABILITIES) == set(engine_rows) + for architecture, capability in ARCHITECTURE_CAPABILITIES.items(): + row = engine_rows[architecture] + split = row["split"] == "true" + assert capability.layer_split is split + assert capability.remote_draft is (row["remote"] == "true") + assert capability.draft_on_layer_split is (row["draft"] == "kBoth") + assert capability.pflash_on_layer_split is (split and row["pflash"] == "true") + assert capability.expert_offload is (row["offload"] == "true") diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index 22eef4ea8..791f9609d 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -199,6 +199,15 @@ for dockerfile in Dockerfile Dockerfile.rocm; do report fail "$dockerfile includes CMake discovery helpers" \ "server/cmake/DiscoverCppUnitTests.cmake would be missing from the image build context" fi + if grep -qF -- "--target test_dflash dflash_server backend_ipc_daemon" \ + "$ROOT/$dockerfile" \ + && grep -qF -- "test -x /opt/lucebox-hub/server/build/backend_ipc_daemon" \ + "$ROOT/$dockerfile"; then + report ok "$dockerfile packages the backend IPC companion" + else + report fail "$dockerfile packages the backend IPC companion" \ + "multi-device Spark/draft placement would be saved but could not launch" + fi done # ── 3. Trivial subcommands (zero-exit expected) ─────────────────────────── @@ -272,7 +281,7 @@ if [ "$native_rc" -ne 0 ]; then report fail "native CUDA build dispatch" "exit=$native_rc output=$(head -3 <<<"$native_out")" elif ! grep -qF -- "-DDFLASH27B_GPU_BACKEND=cuda" "$native_log" \ || ! grep -qF -- "-DCMAKE_CUDA_ARCHITECTURES=86" "$native_log" \ - || ! grep -qF -- "--target dflash_server -j 2" "$native_log"; then + || ! grep -qF -- "--target dflash_server backend_ipc_daemon -j 2" "$native_log"; then report fail "native CUDA build dispatch" "unexpected cmake argv: $(tr '\n' ' ' < "$native_log")" else report ok "native CUDA build dispatch" @@ -303,6 +312,72 @@ else fi rm -rf "$native_tmp" +test_paired_runtime_packaging() { + local label="paired runtime packaging creates a relocatable factory layout" + local sandbox repo destination out rc package_backend + sandbox=$(mktemp -d "${TMPDIR:-/tmp}/lucebox-runtime-package.XXXXXX") + repo="$sandbox/repo" + destination="$sandbox/output/lucebox-runtime" + mkdir -p "$repo/harness" "$repo/server/scripts" \ + "$repo/server/share/model_cards" + : > "$repo/server/CMakeLists.txt" + printf '#!/usr/bin/env bash\nexit 0\n' > "$repo/server/scripts/entrypoint.sh" + printf '{}\n' > "$repo/server/share/model_cards/test.json" + chmod +x "$repo/server/scripts/entrypoint.sh" + + for package_backend in cuda hip; do + mkdir -p "$repo/server/build-$package_backend/deps/runtime" + printf '#!/usr/bin/env bash\nexit 0\n' \ + > "$repo/server/build-$package_backend/dflash_server" + printf '#!/usr/bin/env bash\nexit 0\n' \ + > "$repo/server/build-$package_backend/backend_ipc_daemon" + chmod +x "$repo/server/build-$package_backend/dflash_server" \ + "$repo/server/build-$package_backend/backend_ipc_daemon" + printf 'shared library\n' \ + > "$repo/server/build-$package_backend/deps/runtime/libtest.so.1" + printf 'delete me\n' \ + > "$repo/server/build-$package_backend/deps/runtime/build.txt" + done + + rc=0 + out=$(LUCEBOX_REPO="$repo" LUCEBOX_BUILD_DIR= NO_COLOR=1 \ + bash "$SCRIPT" package-runtime "$destination" 2>&1) || rc=$? + if [ "$rc" -ne 0 ]; then + report fail "$label" "exit=$rc output=$(head -3 <<<"$out")" + rm -rf "$sandbox" + return + fi + for expected in \ + cuda/dflash_server cuda/backend_ipc_daemon \ + hip/dflash_server hip/backend_ipc_daemon \ + server/scripts/entrypoint.sh server/share/model_cards/test.json \ + cuda/deps/runtime/libtest.so.1 hip/deps/runtime/libtest.so.1 \ + MANIFEST; do + if [ ! -f "$destination/$expected" ]; then + report fail "$label" "missing packaged file: $expected" + rm -rf "$sandbox" + return + fi + done + if [ -e "$destination/cuda/deps/runtime/build.txt" ] \ + || [ -e "$destination/hip/deps/runtime/build.txt" ]; then + report fail "$label" "non-runtime build artifacts were retained" + rm -rf "$sandbox" + return + fi + rc=0 + LUCEBOX_REPO="$repo" LUCEBOX_BUILD_DIR= NO_COLOR=1 \ + bash "$SCRIPT" package-runtime "$destination" >/dev/null 2>&1 || rc=$? + if [ "$rc" -eq 0 ]; then + report fail "$label" "existing destination was overwritten" + rm -rf "$sandbox" + return + fi + report ok "$label" + rm -rf "$sandbox" +} +test_paired_runtime_packaging + # Engine client connectors: use already-installed binaries, preserve every # normal client config, and point a Lucebox-owned profile at the local API. # Fake binaries make this an execution-path test without installing or @@ -736,10 +811,20 @@ _make_entrypoint_sandbox() { #!/usr/bin/env bash printf '[shim] dflash_server' for a in "$@"; do printf ' %q' "$a"; done +if [ -n "${DFLASH_MOE_EXPERT_COMPUTE_IPC_BIN:-}" ]; then + printf ' MOE_IPC=%q' "$DFLASH_MOE_EXPERT_COMPUTE_IPC_BIN" + printf ' MOE_GPU=%q' "${DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU:-}" + printf ' MOE_REQUIRED=%q' "${DFLASH_MOE_EXPERT_COMPUTE_IPC_REQUIRED:-}" +fi printf '\n' exit 0 STUB chmod +x "$bin_dir/dflash_server" + cat > "$bin_dir/backend_ipc_daemon" <<'STUB' +#!/usr/bin/env bash +exit 0 +STUB + chmod +x "$bin_dir/backend_ipc_daemon" # `nvidia-smi` shim — pretend we have a 24 GB GPU so the autotune # block runs but doesn't pick the under-12-GB warn tier. cat > "$shim_dir/nvidia-smi" <<'STUB' @@ -836,6 +921,97 @@ test_entrypoint_optimization_flags() { test_entrypoint_optimization_flags \ "entrypoint forwards PFlash, KVFlash, and Spark exactly once" +test_entrypoint_placement_flags() { + local label="$1" sandbox draft_dir models_dir bin_dir shim_dir out + _make_entrypoint_sandbox + touch "$models_dir/Qwen3.6-27B-Q4_K_M.gguf" + out=$(PATH="$shim_dir:$PATH" \ + DFLASH_DIR="$sandbox" \ + DFLASH_SERVER_BIN="$bin_dir/dflash_server" \ + DFLASH_BACKEND_IPC_BIN="$bin_dir/backend_ipc_daemon" \ + DFLASH_TARGET="$models_dir/Qwen3.6-27B-Q4_K_M.gguf" \ + DFLASH_DRAFT="$models_dir/.no-draft" \ + DFLASH_TARGET_DEVICES="cuda:0,hip:0" \ + DFLASH_TARGET_LAYER_SPLIT="0.8,0.2" \ + DFLASH_REMOTE_TARGET_SHARD=1 \ + timeout 10 bash "$ENTRYPOINT" serve 2>&1 || true) + rm -rf "$sandbox" + local required missing="" + for required in \ + "--target-devices cuda:0\\,hip:0" \ + "--target-layer-split 0.8\\,0.2" \ + "--target-shard-ipc-bin $bin_dir/backend_ipc_daemon"; do + [[ "$out" == *"$required"* ]] || missing+=" $required" + done + if [ -n "$missing" ]; then + report fail "$label" "missing argv/env:$missing; output: $(tail -3 <<<"$out")" + else + report ok "$label" + fi +} +test_entrypoint_placement_flags \ + "entrypoint forwards mixed target layer placement" + +test_entrypoint_remote_spark_flags() { + local label="$1" sandbox draft_dir models_dir bin_dir shim_dir out + _make_entrypoint_sandbox + touch "$models_dir/Qwen3.6-MoE-Q4_K_M.gguf" + out=$(PATH="$shim_dir:$PATH" \ + DFLASH_DIR="$sandbox" \ + DFLASH_SERVER_BIN="$bin_dir/dflash_server" \ + DFLASH_BACKEND_IPC_BIN="$bin_dir/backend_ipc_daemon" \ + DFLASH_TARGET="$models_dir/Qwen3.6-MoE-Q4_K_M.gguf" \ + DFLASH_DRAFT="$models_dir/.no-draft" \ + DFLASH_TARGET_DEVICE="cuda:0" \ + DFLASH_REMOTE_EXPERT_DEVICE="hip:0" \ + DFLASH_SPARK=1 \ + timeout 10 bash "$ENTRYPOINT" serve 2>&1 || true) + rm -rf "$sandbox" + local required missing="" + for required in \ + "--target-device cuda:0" \ + "--spark" \ + "MOE_IPC=$bin_dir/backend_ipc_daemon" \ + "MOE_GPU=0" \ + "MOE_REQUIRED=1"; do + [[ "$out" == *"$required"* ]] || missing+=" $required" + done + if [ -n "$missing" ]; then + report fail "$label" "missing argv/env:$missing; output: $(tail -3 <<<"$out")" + else + report ok "$label" + fi +} +test_entrypoint_remote_spark_flags \ + "entrypoint forwards remote Spark expert placement" + +test_entrypoint_remote_draft_flags() { + local label="$1" sandbox draft_dir models_dir bin_dir shim_dir out + _make_entrypoint_sandbox + touch "$models_dir/Qwen3.6-27B-Q4_K_M.gguf" + touch "$draft_dir/dflash-draft-3.6-test.gguf" + out=$(PATH="$shim_dir:$PATH" \ + DFLASH_DIR="$sandbox" \ + DFLASH_SERVER_BIN="$bin_dir/dflash_server" \ + DFLASH_BACKEND_IPC_BIN="$bin_dir/backend_ipc_daemon" \ + DFLASH_TARGET="$models_dir/Qwen3.6-27B-Q4_K_M.gguf" \ + DFLASH_DRAFT="$draft_dir/dflash-draft-3.6-test.gguf" \ + DFLASH_TARGET_DEVICE="cuda:0" \ + DFLASH_DRAFT_DEVICE="hip:0" \ + DFLASH_REMOTE_DRAFT=1 \ + timeout 10 bash "$ENTRYPOINT" serve 2>&1 || true) + rm -rf "$sandbox" + if [[ "$out" != *"--target-device cuda:0"* ]] \ + || [[ "$out" != *"--draft-device hip:0"* ]] \ + || [[ "$out" != *"--draft-ipc-bin $bin_dir/backend_ipc_daemon"* ]]; then + report fail "$label" "placement flags missing: $(tail -3 <<<"$out")" + else + report ok "$label" + fi +} +test_entrypoint_remote_draft_flags \ + "entrypoint forwards cross-backend draft placement" + test_harness_draft_directory_validation() { local label="$1" sandbox target server ambiguous empty out rc out_empty rc_empty sandbox=$(mktemp -d -t lucebox-harness-draft.XXXXXX) @@ -1251,6 +1427,16 @@ models = "/srv/models#fast" # an actual TOML comment [dflash] budget = 22 lazy = false + +[placement] +target_devices = [ + "hip:0", + "hip:1", # generated arrays may retain a trailing comma +] +target_layer_split = [ + 0.7, + 0.3, +] TOML # Exercise both helpers + the resolver via a subshell that sources @@ -1285,6 +1471,42 @@ TOML } test_config_toml_reader_and_resolve "config.toml reader + env > toml > default resolution (7 cases)" +test_multiline_placement_arrays_to_csv() { + local label="$1" tmp devices split + tmp=$(mktemp -d -t lucebox-placement-cfg.XXXXXX) + cat > "$tmp/config.toml" <<'TOML' +[placement] +target_devices = [ + "hip:0", + "hip:1", +] +target_layer_split = [ + 0.7, + 0.3, +] +TOML + devices=$(LUCEBOX_HOME="$tmp" bash -c ' + '"$(sed -n "/^_lucebox_config_path()/,/^}/p" "$SCRIPT")"' + '"$(sed -n "/^_lucebox_config_get()/,/^}/p" "$SCRIPT")"' + '"$(sed -n "/^_toml_array_to_csv()/,/^}/p" "$SCRIPT")"' + _toml_array_to_csv "$(_lucebox_config_get placement.target_devices)" + ') + split=$(LUCEBOX_HOME="$tmp" bash -c ' + '"$(sed -n "/^_lucebox_config_path()/,/^}/p" "$SCRIPT")"' + '"$(sed -n "/^_lucebox_config_get()/,/^}/p" "$SCRIPT")"' + '"$(sed -n "/^_toml_array_to_csv()/,/^}/p" "$SCRIPT")"' + _toml_array_to_csv "$(_lucebox_config_get placement.target_layer_split)" + ') + rm -rf "$tmp" + if [ "$devices" != "hip:0,hip:1" ] || [ "$split" != "0.7,0.3" ]; then + report fail "$label" "devices=$devices split=$split" + return + fi + report ok "$label" +} +test_multiline_placement_arrays_to_csv \ + "tomli_w multi-line placement arrays reach native runtime as CSV" + # ── cmd_serve under systemd: INVOCATION_ID short-circuits is-active ────── # When systemd invokes the wrapper as a unit's ExecStart, it sets # $INVOCATION_ID. The wrapper must NOT then refuse "already running under @@ -1842,6 +2064,48 @@ STUB } test_mixed_rtx_strix_probe_prefers_cuda "mixed RTX 3090 + Strix probe selects CUDA" +test_r9700_strix_probe_prefers_discrete_gpu() { + local label="$1" sandbox shim_dir out + sandbox=$(mktemp -d -t lucebox-amd-primary.XXXXXX) + shim_dir="$sandbox/bin" + mkdir -p "$shim_dir" + cat > "$shim_dir/nvidia-smi" <<'STUB' +#!/usr/bin/env bash +exit 1 +STUB + cat > "$shim_dir/amd-smi" <<'STUB' +#!/usr/bin/env bash +if [ "${1:-}" = "version" ]; then + echo "AMDSMI Tool: test | ROCm version: 7.2.4 | Platform: Linux Baremetal" + exit 0 +fi +cat <<'CSV' +gpu,market_name,vendor_id,vendor_name,subvendor_id,device_id,subsystem_id,rev_id,asic_serial,oam_id,num_compute_units,target_graphics_version,type,vendor,size,bit_width,max_bandwidth +0,AMD Radeon AI PRO R9700,0x1002,AMD,0xf111,0x7551,0x000a,0xc0,0xE099917E6553AFAA,N/A,64,gfx1201,GDDR6,SAMSUNG,32624,256,N/A +1,AMD Radeon Graphics,0x1002,AMD,0xf111,0x1586,0x000a,0xc1,0x0000000000000000,N/A,40,gfx1151,GDDR7,UNKNOWN,98304,256,N/A +CSV +STUB + cat > "$shim_dir/docker" <<'STUB' +#!/usr/bin/env bash +case "${1:-}" in + ps) exit 0 ;; + version) echo "29.1.3" ;; +esac +exit 0 +STUB + chmod +x "$shim_dir"/* + out=$(HOME="$sandbox" LUCEBOX_HOME="$sandbox/.lucebox" \ + PATH="$shim_dir:$PATH" NO_COLOR=1 bash "$SCRIPT" check 2>&1) + rm -rf "$sandbox" + if ! grep -qF "primary AMD Radeon AI PRO R9700 (gfx1201" <<<"$out"; then + report fail "$label" "large Strix aperture displaced the R9700: $(printf '%s' "$out" | head -20)" + return + fi + report ok "$label" +} +test_r9700_strix_probe_prefers_discrete_gpu \ + "R9700 remains primary when Strix reports a large UMA aperture" + test_cross_vendor_docker_args() { local label="$1" helpers cuda_args rocm_args pinned_rocm_args selected_rocm_args helpers=$(awk '/^_variant_is_rocm\(\) \{/,/^\}/' "$SCRIPT")$'\n' @@ -1903,6 +2167,84 @@ test_selected_backend_facts() { } test_selected_backend_facts "mixed-build autotune facts follow the selected ROCm backend" +test_native_placement_visibility() { + local label="$1" helpers explicit legacy + helpers=$(awk '/^_toml_array_to_csv\(\) \{/,/^\}/' "$SCRIPT")$'\n' + helpers+=$(awk '/^_export_native_config\(\) \{/,/^\}/' "$SCRIPT") + explicit=$(bash -c "$helpers"$'\n'' + _lucebox_config_get() { + [ "$1" = "placement.target_device" ] && printf "cuda:0" + } + CUDA_VISIBLE_DEVICES=7 + HIP_VISIBLE_DEVICES=6 + ROCR_VISIBLE_DEVICES=5 + LUCEBOX_HOST_CUDA_VISIBLE_DEVICES=7 + LUCEBOX_HOST_HIP_VISIBLE_DEVICES=6 + LUCEBOX_HOST_ROCR_VISIBLE_DEVICES=5 + _export_native_config + printf "%s|%s|%s" "${CUDA_VISIBLE_DEVICES-unset}" \ + "${HIP_VISIBLE_DEVICES-unset}" "${ROCR_VISIBLE_DEVICES-unset}" + ') + legacy=$(bash -c "$helpers"$'\n'' + _lucebox_config_get() { :; } + LUCEBOX_HOST_CUDA_VISIBLE_DEVICES=7 + LUCEBOX_HOST_ROCR_VISIBLE_DEVICES=GPU-primary + _export_native_config + printf "%s|%s|%s" "${CUDA_VISIBLE_DEVICES-unset}" \ + "${HIP_VISIBLE_DEVICES-unset}" "${ROCR_VISIBLE_DEVICES-unset}" + ') + if [ "$explicit" != "unset|unset|unset" ]; then + report fail "$label" "explicit physical placement retained a visibility mask: $explicit" + return + fi + if [ "$legacy" != "7|unset|GPU-primary" ]; then + report fail "$label" "legacy primary isolation changed: $legacy" + return + fi + report ok "$label" +} +test_native_placement_visibility "native placement owns the full physical GPU inventory" + +test_hybrid_profile_direction() { + local label="$1" helpers out rc + helpers=$(awk '/^_toml_array_to_csv\(\) \{/,/^\}/' "$SCRIPT")$'\n' + helpers+=$(awk '/^_validate_hybrid_profile\(\) \{/,/^\}/' "$SCRIPT") + out=$(bash -c "$helpers"$'\n'' + die() { printf "%s" "$1" >&2; exit 9; } + _lucebox_config_get() { + case "$1" in + placement.target_devices) printf '\''["cuda:0", "hip:0"]'\'' ;; + placement.remote_target_shard) printf true ;; + esac + } + _validate_hybrid_profile + ' 2>&1) || rc=$? + rc="${rc:-0}" + if [ "$rc" -ne 0 ]; then + report fail "$label" "valid CUDA→HIP target split failed: $out" + return + fi + + rc=0 + out=$(bash -c "$helpers"$'\n'' + die() { printf "%s" "$1" >&2; exit 9; } + _lucebox_config_get() { + case "$1" in + placement.target_device) printf hip:0 ;; + placement.draft_device) printf cuda:0 ;; + placement.remote_draft) printf true ;; + esac + } + _validate_hybrid_profile + ' 2>&1) || rc=$? + if [ "$rc" -ne 9 ] || ! grep -qF "requires a CUDA target server" <<<"$out"; then + report fail "$label" "unsupported HIP→CUDA direction was not rejected: rc=$rc out=$out" + return + fi + report ok "$label" +} +test_hybrid_profile_direction "hybrid runtime accepts CUDA→HIP and rejects HIP→CUDA" + # ── TTY flag selection. Regression guard for the process-substitution bug: # _set_tty_flags must run in the CALLER's scope so `[ -t 1 ]` inspects the # real terminal. If it is ever moved back behind `< <(...)` or `$(...)`, diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index 500f03f06..423112016 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -397,6 +397,18 @@ fi # the KV footprint. Only emitted to the server CLI when nonzero so # unset reproduces the server's own default unchanged. : "${DFLASH_FA_WINDOW:=0}" +# Accelerator placement is resolved by the host CLI. Empty values preserve the +# server's historical auto:0 defaults for direct Docker users. +: "${DFLASH_PLACEMENT_MODE:=single}" +: "${DFLASH_TARGET_DEVICE:=}" +: "${DFLASH_TARGET_DEVICES:=}" +: "${DFLASH_TARGET_LAYER_SPLIT:=}" +: "${DFLASH_DRAFT_DEVICE:=}" +: "${DFLASH_REMOTE_DRAFT:=0}" +: "${DFLASH_REMOTE_TARGET_SHARD:=0}" +: "${DFLASH_PEER_ACCESS:=0}" +: "${DFLASH_REMOTE_EXPERT_DEVICE:=}" +: "${DFLASH_BACKEND_IPC_BIN:=$DFLASH_DIR/build/backend_ipc_daemon}" # ── auto-detect target ───────────────────────────────────────────────────── # Target .gguf is typically 10-30 GB (Q4_K_M). Drafts are 1-2 GB (Q8_0 / Q4) @@ -576,7 +588,11 @@ elif [ -n "$DFLASH_DRAFT" ] && [ ! -f "$DFLASH_DRAFT" ]; then DRAFT_ARG="" fi -[ "$GPU_COUNT" -gt 1 ] && warn "${GPU_COUNT} GPUs detected — native server layer sharding is not auto-enabled" +if [ "$GPU_COUNT" -gt 1 ] \ + && [ -z "$DFLASH_TARGET_DEVICE" ] \ + && [ -z "$DFLASH_TARGET_DEVICES" ]; then + warn "${GPU_COUNT} GPUs detected but no placement profile was supplied; using the server default" +fi # ── build + exec native server ──────────────────────────────────────────── CMD=("$DFLASH_SERVER_BIN" "$DFLASH_TARGET" @@ -585,6 +601,45 @@ CMD=("$DFLASH_SERVER_BIN" "$DFLASH_TARGET" --max-ctx "$DFLASH_MAX_CTX" --think-max-tokens "$DFLASH_THINK_MAX") +if [ -n "$DFLASH_TARGET_DEVICES" ]; then + [ -z "$DFLASH_TARGET_DEVICE" ] \ + || die "DFLASH_TARGET_DEVICE conflicts with DFLASH_TARGET_DEVICES" + [ -n "$DFLASH_TARGET_LAYER_SPLIT" ] \ + || die "DFLASH_TARGET_DEVICES requires DFLASH_TARGET_LAYER_SPLIT" + CMD+=(--target-devices "$DFLASH_TARGET_DEVICES" + --target-layer-split "$DFLASH_TARGET_LAYER_SPLIT") +elif [ -n "$DFLASH_TARGET_DEVICE" ]; then + CMD+=(--target-device "$DFLASH_TARGET_DEVICE") +fi +[ -n "$DFLASH_DRAFT_DEVICE" ] && CMD+=(--draft-device "$DFLASH_DRAFT_DEVICE") + +if [ "$DFLASH_REMOTE_DRAFT" = "1" ] \ + || [ "$DFLASH_REMOTE_TARGET_SHARD" = "1" ] \ + || [ -n "$DFLASH_REMOTE_EXPERT_DEVICE" ]; then + [ -x "$DFLASH_BACKEND_IPC_BIN" ] \ + || die "backend IPC daemon missing or not executable at $DFLASH_BACKEND_IPC_BIN" +fi +[ "$DFLASH_REMOTE_DRAFT" = "1" ] \ + && CMD+=(--draft-ipc-bin "$DFLASH_BACKEND_IPC_BIN") +[ "$DFLASH_REMOTE_TARGET_SHARD" = "1" ] \ + && CMD+=(--target-shard-ipc-bin "$DFLASH_BACKEND_IPC_BIN") +[ "$DFLASH_PEER_ACCESS" = "1" ] && CMD+=(--peer-access) + +if [ -n "$DFLASH_REMOTE_EXPERT_DEVICE" ]; then + [ "$DFLASH_SPARK" = "1" ] \ + || die "DFLASH_REMOTE_EXPERT_DEVICE requires DFLASH_SPARK=1" + case "$DFLASH_REMOTE_EXPERT_DEVICE" in + cuda:[0-9]*|hip:[0-9]*) ;; + *) die "bad DFLASH_REMOTE_EXPERT_DEVICE (expected cuda:N or hip:N)" ;; + esac + remote_expert_gpu="${DFLASH_REMOTE_EXPERT_DEVICE##*:}" + [[ "$remote_expert_gpu" =~ ^[0-9]+$ ]] \ + || die "bad DFLASH_REMOTE_EXPERT_DEVICE GPU index" + export DFLASH_MOE_EXPERT_COMPUTE_IPC_BIN="$DFLASH_BACKEND_IPC_BIN" + export DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU="$remote_expert_gpu" + export DFLASH_MOE_EXPERT_COMPUTE_IPC_REQUIRED=1 +fi + # Keep cache defaults owned by dflash_server. In particular, omitting # DFLASH_PREFIX_CACHE_SLOTS preserves the native nonzero default instead of # silently disabling multi-turn prefix reuse in the container. Explicit diff --git a/server/src/common/dflash_layer_split_runtime.h b/server/src/common/dflash_layer_split_runtime.h index b5451b6a1..00577ab98 100644 --- a/server/src/common/dflash_layer_split_runtime.h +++ b/server/src/common/dflash_layer_split_runtime.h @@ -12,7 +12,9 @@ #include "ggml.h" #include "ggml-alloc.h" #include "ggml-backend.h" +#include "peer_access.h" +#include #include namespace dflash::common { @@ -47,6 +49,34 @@ struct ActivationBuffer { ggml_type type = GGML_TYPE_F32; }; +inline void copy_layer_split_tensor(const ggml_tensor * src, + int src_device, + ggml_tensor * dst, + int dst_device) { + GGML_ASSERT(src && dst); + GGML_ASSERT(src->type == dst->type); + GGML_ASSERT(ggml_are_same_shape(src, dst)); + GGML_ASSERT(ggml_are_same_stride(src, dst)); + if (src_device == dst_device) { + ggml_backend_tensor_copy(src, dst); + return; + } + + // Route every cross-device transfer through the shared transport policy. + // It uses reusable pinned host staging by default and P2P only after an + // explicit --peer-access opt-in plus a successful capability check. + if (copy_peer_async(dst->data, dst_device, src->data, src_device, + ggml_nbytes(src))) { + return; + } + // Preserve a backend-generic last resort if the runtime-level transfer + // fails (for example, pinned allocation pressure). + thread_local std::vector staging; + staging.resize(ggml_nbytes(src)); + ggml_backend_tensor_get(src, staging.data(), 0, staging.size()); + ggml_backend_tensor_set(dst, staging.data(), 0, staging.size()); +} + inline bool set_activation_tensor_from_f32(ggml_tensor * dst, const float * src, size_t offset, diff --git a/server/src/common/peer_access.cpp b/server/src/common/peer_access.cpp index 9658dce9d..37a9ee686 100644 --- a/server/src/common/peer_access.cpp +++ b/server/src/common/peer_access.cpp @@ -1,10 +1,53 @@ #include "peer_access.h" #include "internal.h" // dflash_cuda_copy_between_devices +#include #include +#include namespace dflash::common { +namespace { + +std::mutex g_peer_pair_cache_mutex; + +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) +std::mutex g_host_staging_mutex; +void * g_host_staging = nullptr; +size_t g_host_staging_capacity = 0; + +bool copy_between_hip_devices_via_host(void * dst, int dst_device, + const void * src, int src_device, + size_t bytes) { + std::lock_guard lock(g_host_staging_mutex); + if (g_host_staging_capacity < bytes) { + if (g_host_staging) { + if (cudaFreeHost(g_host_staging) != cudaSuccess) { + return false; + } + g_host_staging = nullptr; + g_host_staging_capacity = 0; + } + if (cudaMallocHost(&g_host_staging, bytes) != cudaSuccess) { + return false; + } + g_host_staging_capacity = bytes; + } + + if (cudaSetDevice(src_device) != cudaSuccess || + cudaMemcpy(g_host_staging, src, bytes, cudaMemcpyDeviceToHost) != cudaSuccess) { + return false; + } + if (cudaSetDevice(dst_device) != cudaSuccess || + cudaMemcpy(dst, g_host_staging, bytes, cudaMemcpyHostToDevice) != cudaSuccess) { + return false; + } + return true; +} +#endif + +} // namespace + // ── global state ──────────────────────────────────────────────── bool g_peer_access_opt_in = false; std::unordered_map g_peer_pair_ok_cache; @@ -19,7 +62,7 @@ bool enable_peer_access_one_way(int device, int peer) { if (err != cudaSuccess) return false; err = cudaDeviceEnablePeerAccess(peer, 0); if (err == cudaErrorPeerAccessAlreadyEnabled) { - cudaGetLastError(); + (void) cudaGetLastError(); return true; } return err == cudaSuccess; @@ -39,17 +82,18 @@ static std::uint64_t peer_pair_key(int a, int b) { } static void log_staged_cross_gpu_once() { - static bool logged = false; - if (logged) return; - logged = true; - std::fprintf(stderr, - "[dflash] Using safe (slower) cross-GPU copy via host staging " - "(--peer-access not set or P2P unavailable for this device pair).\n"); + static std::once_flag once; + std::call_once(once, [] { + std::fprintf(stderr, + "[dflash] Using safe (slower) cross-GPU copy via host staging " + "(--peer-access not set or P2P unavailable for this device pair).\n"); + }); } bool cross_device_peer_memcpy_ok(int src_device, int dst_device) { if (src_device == dst_device) return true; if (!g_peer_access_opt_in) return false; + std::lock_guard lock(g_peer_pair_cache_mutex); const std::uint64_t k = peer_pair_key(src_device, dst_device); const auto it = g_peer_pair_ok_cache.find(k); if (it != g_peer_pair_ok_cache.end()) return it->second; @@ -86,14 +130,13 @@ bool copy_peer_async(void * dst, int dst_device, } log_staged_cross_gpu_once(); #if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) - err = cudaSetDevice(dst_device); - if (err != cudaSuccess) return false; - err = cudaMemcpyPeerAsync(dst, dst_device, src, src_device, bytes, stream); - if (err != cudaSuccess) return false; - if (stream) { - return cudaStreamSynchronize(stream) == cudaSuccess; - } - return cudaDeviceSynchronize() == cudaSuccess; + // ROCm can report peer access for heterogeneous devices while returning + // corrupted data from hipMemcpyPeerAsync (observed on gfx1201 + gfx1151). + // Keep the default path correct and reasonably fast with a reusable pinned + // buffer. P2P remains available above when the operator explicitly opts in. + (void)stream; + return copy_between_hip_devices_via_host( + dst, dst_device, src, src_device, bytes); #else return dflash_cuda_copy_between_devices(src_device, src, dst_device, dst, bytes, nullptr, stream); diff --git a/server/src/common/peer_access.h b/server/src/common/peer_access.h index f6e2465c7..abbe28d8d 100644 --- a/server/src/common/peer_access.h +++ b/server/src/common/peer_access.h @@ -1,4 +1,4 @@ -// CUDA peer-access helpers for multi-GPU inference. +// Cross-device transport helpers for CUDA and HIP multi-GPU inference. // // Provides enable_peer_access_pair(), cross_device_peer_memcpy_ok(), and // copy_peer_async() — used by DraftFeatureMirror and the speculative-decode diff --git a/server/src/gemma4/gemma4_layer_split_adapter.cpp b/server/src/gemma4/gemma4_layer_split_adapter.cpp index d89fca3a8..bf3c9b6e6 100644 --- a/server/src/gemma4/gemma4_layer_split_adapter.cpp +++ b/server/src/gemma4/gemma4_layer_split_adapter.cpp @@ -600,8 +600,10 @@ bool Gemma4LayerSplitAdapter::run_forward( return false; } ggml_backend_synchronize(current_shard->backend); - ggml_backend_tensor_copy(act_in, next_acts.a); - ggml_backend_tensor_copy(orig.tensor, next_orig.tensor); + copy_layer_split_tensor( + act_in, current_shard->gpu, next_acts.a, shard->gpu); + copy_layer_split_tensor( + orig.tensor, current_shard->gpu, next_orig.tensor, shard->gpu); ggml_backend_synchronize(shard->backend); activation_pair_free(acts); activation_buffer_free(orig); @@ -852,8 +854,10 @@ bool Gemma4LayerSplitAdapter::run_mixed_forward( return false; } ggml_backend_synchronize(current_shard->backend); - ggml_backend_tensor_copy(act_in, next_acts.a); - ggml_backend_tensor_copy(orig.tensor, next_orig.tensor); + copy_layer_split_tensor( + act_in, current_shard->gpu, next_acts.a, shard->gpu); + copy_layer_split_tensor( + orig.tensor, current_shard->gpu, next_orig.tensor, shard->gpu); ggml_backend_synchronize(shard->backend); activation_pair_free(acts); activation_buffer_free(orig); @@ -1467,8 +1471,10 @@ int run_gemma4_target_shard_ipc_daemon(const char * target_path, break; } ggml_backend_synchronize(current_shard->backend); - ggml_backend_tensor_copy(act_in, next_acts.a); - ggml_backend_tensor_copy(orig.tensor, next_orig.tensor); + copy_layer_split_tensor( + act_in, current_shard->gpu, next_acts.a, shard->gpu); + copy_layer_split_tensor( + orig.tensor, current_shard->gpu, next_orig.tensor, shard->gpu); ggml_backend_synchronize(shard->backend); activation_pair_free(acts); activation_buffer_free(orig); diff --git a/server/src/ipc/backend_ipc_main.cpp b/server/src/ipc/backend_ipc_main.cpp index b2fae7791..91d39ab81 100644 --- a/server/src/ipc/backend_ipc_main.cpp +++ b/server/src/ipc/backend_ipc_main.cpp @@ -138,6 +138,7 @@ int main(int argc, char ** argv) { argv[0], argv[0], argv[0], + argv[0], argv[0]); return 2; } diff --git a/server/src/laguna/laguna_layer_split_adapter.cpp b/server/src/laguna/laguna_layer_split_adapter.cpp index 5c00f6666..81ed4b0af 100644 --- a/server/src/laguna/laguna_layer_split_adapter.cpp +++ b/server/src/laguna/laguna_layer_split_adapter.cpp @@ -496,7 +496,8 @@ bool LagunaLayerSplitAdapter::run_forward( return false; } ggml_backend_synchronize(current_shard->backend); - ggml_backend_tensor_copy(act_in, next_acts.a); + copy_layer_split_tensor( + act_in, current_shard->gpu, next_acts.a, shard->gpu); ggml_backend_synchronize(shard->backend); activation_pair_free(acts); acts = next_acts; @@ -698,7 +699,8 @@ bool LagunaLayerSplitAdapter::run_mixed_forward( return false; } ggml_backend_synchronize(current_shard->backend); - ggml_backend_tensor_copy(act_in, next_acts.a); + copy_layer_split_tensor( + act_in, current_shard->gpu, next_acts.a, shard->gpu); ggml_backend_synchronize(shard->backend); activation_pair_free(acts); acts = next_acts; @@ -1410,7 +1412,8 @@ int run_laguna_target_shard_ipc_daemon(const char * target_path, break; } ggml_backend_synchronize(current_shard->backend); - ggml_backend_tensor_copy(act_in, next_acts.a); + copy_layer_split_tensor( + act_in, current_shard->gpu, next_acts.a, shard->gpu); ggml_backend_synchronize(shard->backend); activation_pair_free(acts); acts = next_acts; diff --git a/server/src/qwen35/layer_split_forward.cpp b/server/src/qwen35/layer_split_forward.cpp index a66383cc9..22e953138 100644 --- a/server/src/qwen35/layer_split_forward.cpp +++ b/server/src/qwen35/layer_split_forward.cpp @@ -258,7 +258,8 @@ bool run_qwen35_layer_split_forward( return false; } ggml_backend_synchronize(current_shard->backend); - ggml_backend_tensor_copy(act_in, next_acts.a); + copy_layer_split_tensor( + act_in, current_shard->gpu, next_acts.a, shard->gpu); ggml_backend_synchronize(shard->backend); activation_pair_free(acts); acts = next_acts; @@ -455,7 +456,8 @@ bool run_qwen35_layer_split_layers_from_activation( return false; } ggml_backend_synchronize(current_shard->backend); - ggml_backend_tensor_copy(act_in, next_acts.a); + copy_layer_split_tensor( + act_in, current_shard->gpu, next_acts.a, shard->gpu); ggml_backend_synchronize(shard->backend); activation_pair_free(acts); acts = next_acts; From e7f2a0c2192687fa33676c46b1d1cf1da2800d96 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:50:06 +0200 Subject: [PATCH 14/25] fix(lucebox): keep automatic plans within safe memory --- lucebox/src/lucebox/autotune.py | 123 ++++++++++++++++++++- lucebox/src/lucebox/placement.py | 85 ++++++++++----- lucebox/tests/test_autotune.py | 124 +++++++++++++++++++++ lucebox/tests/test_placement.py | 181 +++++++++++++++++++++++++++++++ 4 files changed, 485 insertions(+), 28 deletions(-) diff --git a/lucebox/src/lucebox/autotune.py b/lucebox/src/lucebox/autotune.py index 18b43f293..feb64795f 100644 --- a/lucebox/src/lucebox/autotune.py +++ b/lucebox/src/lucebox/autotune.py @@ -151,6 +151,70 @@ def _cap_exact_context( return runtime +CapacityAdjustment = Literal["pflash", "kvflash-qk", "kvflash-off", "dflash"] + + +def _uses_optimizer_scorer(runtime: DflashRuntime) -> bool: + return runtime.prefill_mode != "off" or ( + runtime.kvflash != "off" and runtime.kvflash_policy == "drafter" + ) + + +def _without_unused_scorer(runtime: DflashRuntime) -> DflashRuntime: + """Clear the scorer path once no enabled feature consumes it.""" + return runtime if _uses_optimizer_scorer(runtime) else replace(runtime, prefill_drafter="") + + +def _capacity_fallbacks( + runtime: DflashRuntime, + *, + architecture: str, +) -> tuple[tuple[DflashRuntime, CapacityAdjustment], ...]: + """Return progressively smaller optional stacks, in performance order. + + Installed optional assets must never make a fitting target unrunnable. + Placement gets the full recommended stack first; only when that has no + valid device plan do we remove independent optional workloads. Qwen dense + can retain bounded KV residency without the 1.2 GB scorer by switching to + its validated QK policy. + """ + candidates: list[tuple[DflashRuntime, CapacityAdjustment]] = [] + current = runtime + + if current.prefill_mode != "off": + defaults = DflashRuntime() + current = _without_unused_scorer( + replace( + current, + prefill_mode="off", + prefill_keep_ratio=defaults.prefill_keep_ratio, + prefill_threshold=defaults.prefill_threshold, + prefill_drafter=( + current.prefill_drafter + if current.kvflash != "off" and current.kvflash_policy == "drafter" + else "" + ), + ) + ) + candidates.append((current, "pflash")) + + if current.kvflash != "off" and current.kvflash_policy == "drafter": + if architecture == "qwen35": + current = replace(current, kvflash_policy="qk") + action: CapacityAdjustment = "kvflash-qk" + else: + current = replace(current, kvflash="off") + action = "kvflash-off" + current = _without_unused_scorer(current) + candidates.append((current, action)) + + if current.speculative_decode: + current = replace(current, speculative_decode=False, lazy=False) + candidates.append((current, "dflash")) + + return tuple(candidates) + + def automatic_plan( cfg: Config, *, @@ -322,20 +386,67 @@ def automatic_plan( has_draft=has_draft, optimizer_drafter_available=optimizer_drafter_available, ) + capacity_adjustments: set[CapacityAdjustment] = set() + if not placement.runnable: + for fallback_runtime, adjustment in _capacity_fallbacks( + runtime, + architecture=preset.architecture, + ): + capacity_adjustments.add(adjustment) + placement = automatic_placement( + cfg, + fallback_runtime, + preset, + has_draft=has_draft and fallback_runtime.speculative_decode, + optimizer_drafter_available=optimizer_drafter_available, + ) + runtime = fallback_runtime + if placement.runnable: + break + adjusted = placement.optimization_runtime if dflash.enabled and not adjusted.speculative_decode: + reason = ( + "disabled because the draft would exceed the safe GPU memory budget" + if "dflash" in capacity_adjustments + else "disabled because the selected target placement has no compatible draft path" + ) dflash = OptimizationDecision( "DFlash", False, dflash.available, - "disabled because the selected target placement has no compatible draft path", + reason, ) if pflash.enabled and adjusted.prefill_mode == "off": + reason = ( + "disabled because the scorer would exceed the safe GPU memory budget" + if "pflash" in capacity_adjustments + else "disabled because the selected target placement has no compatible compression path" + ) pflash = OptimizationDecision( "PFlash", False, pflash.available, - "disabled because the selected target placement has no compatible compression path", + reason, + ) + if kvflash.enabled and adjusted.kvflash == "off": + reason = ( + "disabled because its scorer would exceed the safe GPU memory budget" + if "kvflash-off" in capacity_adjustments + else "disabled because the selected target placement has no compatible KV path" + ) + kvflash = OptimizationDecision( + "KVFlash", + False, + kvflash.available, + reason, + ) + elif "kvflash-qk" in capacity_adjustments: + kvflash = OptimizationDecision( + "KVFlash", + True, + kvflash.available, + "memory is tight; QK policy avoids a separate scorer allocation", ) if spark.enabled and not adjusted.spark: spark = OptimizationDecision( @@ -344,6 +455,14 @@ def automatic_plan( spark.available, "disabled because Automatic does not compose Spark with target layer splitting", ) + if needs_optimizer_drafter: + # Do not ask a first-time user to download 1.2 GB only to discard the + # scorer during capacity fallback. The scorer-present branch cannot + # recurse again because its availability makes ``needs_*`` false. + scorer_plan = automatic_plan(cfg, optimizer_drafter_available=True) + needs_optimizer_drafter = scorer_plan.placement.runnable and _uses_optimizer_scorer( + scorer_plan.runtime + ) return OptimizationPlan( runtime=adjusted, model_name=preset.name, diff --git a/lucebox/src/lucebox/placement.py b/lucebox/src/lucebox/placement.py index 86043ae64..29011a8da 100644 --- a/lucebox/src/lucebox/placement.py +++ b/lucebox/src/lucebox/placement.py @@ -134,7 +134,7 @@ def _split_devices( ) -> tuple[tuple[GpuDevice, ...], tuple[float, ...]]: """Choose the smallest ordered device set that makes the target fit.""" primary = topology.primary - if primary is None: + if primary is None or primary_capacity <= 0.0: return (), () local = sorted( @@ -172,7 +172,10 @@ def _split_devices( allocations: list[float] = [] for capacity in capacities: allocation = min(capacity, remaining) - allocations.append(max(allocation, target_gb * 0.05)) + # Keep a non-trivial shard when capacity allows it, without ever + # assigning a device more target weight than its safe byte budget. + allocation = min(capacity, max(allocation, target_gb * 0.05)) + allocations.append(allocation) remaining -= allocation total_allocation = sum(allocations) weights = tuple(round(value / total_allocation, 4) for value in allocations) @@ -181,17 +184,26 @@ def _split_devices( return tuple(selected), weights -def _single_option(primary: GpuDevice, fits: bool) -> PlacementOption: +def _single_option( + primary: GpuDevice, + *, + target_fits: bool, + stack_fits: bool, +) -> PlacementOption: + if stack_fits: + reason = f"the selected stack fits {primary.name}; avoiding a layer/IPC boundary is faster" + elif target_fits: + reason = ( + f"the target fits {primary.name}, but its optional workloads exceed the safe budget" + ) + else: + reason = f"the target does not fit the safe memory budget on {primary.name}" return PlacementOption( key="single", label="Primary GPU", - available=fits, - recommended=fits, - reason=( - f"the selected stack fits {primary.name}; avoiding a layer/IPC boundary is faster" - if fits - else f"the target does not fit the safe memory budget on {primary.name}" - ), + available=stack_fits, + recommended=stack_fits, + reason=reason, ) @@ -241,7 +253,9 @@ def automatic_placement( # of truth for the exact byte budget. stack_fits = stack_gb <= primary_capacity or (runtime.spark and capabilities.expert_offload) - options: list[PlacementOption] = [_single_option(primary, stack_fits)] + options: list[PlacementOption] = [ + _single_option(primary, target_fits=target_fits, stack_fits=stack_fits) + ] base = PlacementRuntime(mode="single", target_device=primary.placement_name) if secondary is None: if stack_fits: @@ -258,7 +272,12 @@ def automatic_placement( runtime=base, optimization_runtime=runtime, summary=primary.name, - reason="the selected target exceeds this GPU's safe memory budget", + reason=( + "the target fits, but its selected optional workloads exceed this GPU's " + "safe memory budget" + if target_fits + else "the selected target exceeds this GPU's safe memory budget" + ), runnable=False, topology=topology, options=tuple(options), @@ -266,27 +285,38 @@ def automatic_placement( same_backend = primary.backend == secondary.backend cross_backend_ready = _cross_backend_ready(cfg, primary, secondary) - draft_work_available = (has_draft and runtime.speculative_decode) or ( - optimizer_drafter_available and runtime.prefill_mode != "off" + decode_draft_work = has_draft and runtime.speculative_decode + pflash_scorer_work = optimizer_drafter_available and runtime.prefill_mode != "off" + draft_work_available = decode_draft_work or pflash_scorer_work + draft_work_gb = (draft_gb if decode_draft_work else 0.0) + ( + scorer_gb if pflash_scorer_work else 0.0 ) - draft_offload_available = draft_work_available and ( - same_backend or (cross_backend_ready and capabilities.remote_draft) + secondary_capacity = _primary_capacity(secondary) + draft_offload_fits = draft_work_gb <= secondary_capacity + draft_offload_compatible = same_backend or (cross_backend_ready and capabilities.remote_draft) + draft_offload_available = ( + draft_work_available and draft_offload_fits and draft_offload_compatible ) + if draft_offload_available: + draft_offload_reason = "moves the independent draft/scorer workload off the primary GPU" + elif not draft_work_available: + draft_offload_reason = "the selected model has no installed draft/scorer workload" + elif not draft_offload_fits: + draft_offload_reason = ( + f"the optional workload needs {draft_work_gb:g} GB, but " + f"{secondary.name} has {secondary_capacity:g} GB of safe capacity" + ) + else: + draft_offload_reason = ( + "the engine/runtime cannot execute this draft across the backend boundary" + ) options.append( PlacementOption( key="draft-offload", label="Secondary draft/scorer", available=draft_offload_available, recommended=draft_offload_available and not stack_fits and target_fits, - reason=( - "moves the independent draft/scorer workload off the primary GPU" - if draft_offload_available - else ( - "the selected model has no installed draft/scorer workload" - if not draft_work_available - else "the engine/runtime cannot execute this draft across the backend boundary" - ) - ), + reason=draft_offload_reason, ) ) @@ -425,7 +455,10 @@ def automatic_placement( optimization_runtime=runtime, summary=primary.name, reason=( - "the target does not fit the primary and no validated secondary-device " + "the target fits the primary, but its optional workloads do not and no validated " + "secondary-device placement is available" + if target_fits + else "the target does not fit the primary and no validated secondary-device " "placement is available" ), runnable=False, diff --git a/lucebox/tests/test_autotune.py b/lucebox/tests/test_autotune.py index 8b481c1a0..b74094d39 100644 --- a/lucebox/tests/test_autotune.py +++ b/lucebox/tests/test_autotune.py @@ -1,5 +1,6 @@ from pathlib import Path +import pytest from lucebox.autotune import automatic_plan, runtime_from_host from lucebox.types import Config, HostFacts, ModelMeta @@ -247,3 +248,126 @@ def test_known_gpu_specific_ddtree_budgets() -> None: assert runtime_from_host(HostFacts(vram_gb=24, gpu_sm="gfx1100")).budget == 8 assert runtime_from_host(HostFacts(vram_gb=32, gpu_sm="120")).budget == 40 assert runtime_from_host(HostFacts(vram_gb=31, gpu_sm="gfx1201")).budget == 22 + + +@pytest.mark.parametrize( + ("vendor", "variant", "architecture"), + [ + ("nvidia", "cuda12", "86"), + ("amd", "rocm", "gfx1201"), + ], +) +def test_20gb_automatic_drops_scorer_before_blocking_a_fitting_target( + tmp_path: Path, + vendor: str, + variant: str, + architecture: str, +) -> None: + cfg = Config( + variant=variant, + models_dir=tmp_path / vendor, + host=HostFacts(gpu_vendor=vendor, vram_gb=20, ram_gb=64, gpu_sm=architecture), + model=ModelMeta(preset="qwen3.6-27b"), + ) + _install_decode_draft(cfg) + _install_optimizer_drafter(cfg) + + plan = automatic_plan(cfg) + + assert plan.placement.runnable is True + assert plan.dflash.enabled is True + assert plan.pflash.enabled is False + assert plan.kvflash.enabled is True + assert plan.runtime.speculative_decode is True + assert plan.runtime.prefill_mode == "off" + assert plan.runtime.prefill_keep_ratio == 0.05 + assert plan.runtime.prefill_threshold == 32000 + assert plan.runtime.kvflash_policy == "qk" + assert plan.runtime.prefill_drafter == "" + assert "scorer" in plan.pflash.reason + + +def test_18gb_automatic_drops_draft_instead_of_blocking_qwen(tmp_path: Path) -> None: + cfg = _cfg( + tmp_path, + "qwen3.6-27b", + HostFacts(gpu_vendor="nvidia", vram_gb=18, ram_gb=64, gpu_sm="86"), + ) + _install_decode_draft(cfg) + _install_optimizer_drafter(cfg) + + plan = automatic_plan(cfg) + + assert plan.placement.runnable is True + assert plan.dflash.enabled is False + assert plan.pflash.enabled is False + assert plan.kvflash.enabled is True + assert plan.runtime.speculative_decode is False + assert plan.runtime.kvflash_policy == "qk" + assert "memory budget" in plan.dflash.reason + + +def test_18gb_automatic_does_not_offer_a_scorer_it_cannot_keep(tmp_path: Path) -> None: + cfg = _cfg( + tmp_path, + "qwen3.6-27b", + HostFacts(gpu_vendor="nvidia", vram_gb=18, ram_gb=64, gpu_sm="86"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.placement.runnable is True + assert plan.runtime.kvflash_policy == "qk" + assert plan.needs_optimizer_drafter is False + + +def test_18gb_automatic_drops_gemma_draft_when_target_alone_fits(tmp_path: Path) -> None: + cfg = Config( + variant="rocm", + models_dir=tmp_path / "models", + host=HostFacts(gpu_vendor="amd", vram_gb=18, ram_gb=64, gpu_sm="gfx1201"), + model=ModelMeta(preset="gemma-4-26b"), + ) + _install_decode_draft(cfg) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.placement.runnable is True + assert plan.dflash.enabled is False + assert plan.runtime.speculative_decode is False + assert "memory budget" in plan.dflash.reason + + +def test_24gb_moe_drops_scored_kvflash_when_host_cannot_run_spark( + tmp_path: Path, +) -> None: + cfg = _cfg( + tmp_path, + "qwen3.6-moe", + HostFacts(gpu_vendor="nvidia", vram_gb=24, ram_gb=16, gpu_sm="86"), + ) + _install_optimizer_drafter(cfg) + + plan = automatic_plan(cfg) + + assert plan.placement.runnable is True + assert plan.kvflash.enabled is False + assert plan.runtime.kvflash == "off" + assert plan.runtime.prefill_drafter == "" + assert "memory budget" in plan.kvflash.reason + + +def test_24gb_moe_does_not_offer_an_unusable_scorer_without_spark_ram( + tmp_path: Path, +) -> None: + cfg = _cfg( + tmp_path, + "qwen3.6-moe", + HostFacts(gpu_vendor="nvidia", vram_gb=24, ram_gb=16, gpu_sm="86"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.placement.runnable is True + assert plan.runtime.kvflash == "off" + assert plan.needs_optimizer_drafter is False diff --git a/lucebox/tests/test_placement.py b/lucebox/tests/test_placement.py index 784a0c0ba..da91c980c 100644 --- a/lucebox/tests/test_placement.py +++ b/lucebox/tests/test_placement.py @@ -1,9 +1,12 @@ +import os import re from pathlib import Path from types import SimpleNamespace import pytest from lucebox.autotune import automatic_plan +from lucebox.docker_run import server_run_spec +from lucebox.host_facts import from_env from lucebox.placement import ARCHITECTURE_CAPABILITIES, automatic_placement from lucebox.topology import from_config from lucebox.types import Config, DflashRuntime, HostFacts, ModelMeta @@ -31,6 +34,14 @@ def _amd_csv(*rows: tuple[str, str, int]) -> str: ) +def _set_host_env(monkeypatch: pytest.MonkeyPatch, values: dict[str, str | int]) -> None: + for key in tuple(os.environ): + if key.startswith("LUCEBOX_HOST_"): + monkeypatch.delenv(key, raising=False) + for key, value in values.items(): + monkeypatch.setenv(key, str(value)) + + def test_r9700_strix_keeps_fitting_qwen_target_on_r9700(tmp_path: Path) -> None: host = HostFacts( gpu_vendor="amd", @@ -107,6 +118,79 @@ def test_three_same_backend_gpus_are_used_only_when_capacity_requires_them( assert plan.runtime.peer_access is False +def test_undersized_secondary_does_not_claim_draft_offload(tmp_path: Path) -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + gpu_name="RTX 0", + gpu_count=2, + vram_gb=18, + gpu_sm="86", + nvidia_gpu_name="RTX 0", + nvidia_gpu_count=2, + nvidia_vram_gb=18, + nvidia_gpu_arch="86", + nvidia_gpu_list_csv=_nvidia_csv(18432, 1024), + ) + cfg = Config( + variant="cuda12", + models_dir=tmp_path, + host=host, + model=ModelMeta(preset="qwen3.6-27b"), + ) + preset = download.PRESETS["qwen3.6-27b"] + + placement = automatic_placement( + cfg, + DflashRuntime(speculative_decode=True), + preset, + has_draft=True, + optimizer_drafter_available=False, + ) + + draft_option = next(option for option in placement.options if option.key == "draft-offload") + assert placement.runnable is False + assert draft_option.available is False + assert "0 GB of safe capacity" in draft_option.reason + + +def test_layer_split_never_assigns_target_to_a_zero_capacity_primary( + tmp_path: Path, +) -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + gpu_name="RTX 0", + gpu_count=2, + vram_gb=4, + gpu_sm="86", + nvidia_gpu_name="RTX 0", + nvidia_gpu_count=2, + nvidia_vram_gb=4, + nvidia_gpu_arch="86", + nvidia_gpu_list_csv=_nvidia_csv(4096, 24576), + ) + cfg = Config(variant="cuda12", models_dir=tmp_path, host=host) + preset = SimpleNamespace( + architecture="qwen35", + target_file="large.gguf", + approx_target_gb=20.0, + approx_draft_gb=3.0, + ) + + placement = automatic_placement( + cfg, + DflashRuntime(speculative_decode=True), + preset, + has_draft=True, + optimizer_drafter_available=False, + ) + + split_option = next(option for option in placement.options if option.key == "layer-split") + assert placement.runnable is False + assert split_option.available is False + + def test_rtx_strix_moe_uses_paired_runtime_for_remote_spark_experts( tmp_path: Path, ) -> None: @@ -290,3 +374,100 @@ def test_python_architecture_capabilities_match_engine_table() -> None: assert capability.draft_on_layer_split is (row["draft"] == "kBoth") assert capability.pflash_on_layer_split is (split and row["pflash"] == "true") assert capability.expert_offload is (row["offload"] == "true") + + +def test_env_driven_dual_nvidia_plan_reaches_server_launch_contract( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _set_host_env( + monkeypatch, + { + "LUCEBOX_HOST_NPROC": 16, + "LUCEBOX_HOST_RAM_GB": 64, + "LUCEBOX_HOST_GPU_VENDOR": "nvidia", + "LUCEBOX_HOST_HAS_NVIDIA_GPU": 1, + "LUCEBOX_HOST_GPU_NAME": "RTX 0", + "LUCEBOX_HOST_GPU_COUNT": 2, + "LUCEBOX_HOST_VRAM_GB": 12, + "LUCEBOX_HOST_GPU_SM": "86", + "LUCEBOX_HOST_NVIDIA_GPU_LIST_CSV": _nvidia_csv(12288, 12288), + }, + ) + cfg = Config( + variant="cuda12", + models_dir=tmp_path, + host=from_env(), + model=ModelMeta(preset="qwen3.6-27b"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + launch = server_run_spec( + Config( + variant=cfg.variant, + models_dir=cfg.models_dir, + host=cfg.host, + model=cfg.model, + dflash=plan.runtime, + placement=plan.placement.runtime, + ) + ) + launch_env = dict(launch.env) + + assert plan.placement.runnable is True + assert plan.placement.runtime.mode == "layer-split" + assert plan.placement.runtime.target_devices == ("cuda:0", "cuda:1") + assert launch_env["DFLASH_TARGET_DEVICES"] == "cuda:0,cuda:1" + assert launch_env["DFLASH_TARGET_LAYER_SPLIT"] + assert launch_env["DFLASH_KVFLASH"] == "auto" + + +def test_env_driven_lucebox_plan_uses_r9700_without_unnecessary_split( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _set_host_env( + monkeypatch, + { + "LUCEBOX_HOST_NPROC": 32, + "LUCEBOX_HOST_RAM_GB": 125, + "LUCEBOX_HOST_GPU_VENDOR": "amd", + "LUCEBOX_HOST_HAS_AMD_GPU": 1, + "LUCEBOX_HOST_GPU_NAME": "AMD Radeon AI PRO R9700", + "LUCEBOX_HOST_GPU_COUNT": 2, + "LUCEBOX_HOST_VRAM_GB": 31, + "LUCEBOX_HOST_GPU_SM": "gfx1201", + "LUCEBOX_HOST_AMD_GPU_LIST_CSV": _amd_csv( + ("AMD Radeon AI PRO R9700", "gfx1201", 32624), + ("AMD Radeon Graphics", "gfx1151", 512), + ), + }, + ) + cfg = Config( + variant="rocm", + models_dir=tmp_path, + host=from_env(), + model=ModelMeta(preset="qwen3.6-27b"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + launch = server_run_spec( + Config( + variant=cfg.variant, + models_dir=cfg.models_dir, + host=cfg.host, + model=cfg.model, + dflash=plan.runtime, + placement=plan.placement.runtime, + ) + ) + launch_env = dict(launch.env) + + assert plan.placement.runnable is True + assert plan.placement.runtime.mode == "single" + assert plan.placement.runtime.target_device == "hip:0" + assert plan.placement.runtime.uses_multiple_devices is False + assert plan.placement.topology.companions[0].unified_memory is True + assert launch.gpu_vendor == "amd" + assert launch_env["DFLASH_TARGET_DEVICE"] == "hip:0" + assert "DFLASH_TARGET_DEVICES" not in launch_env From 0fb9a4c93624708b157d12beab63f24e9cddd27e Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 31 Jul 2026 10:06:14 +0200 Subject: [PATCH 15/25] feat(lucebox): focus the guided model library --- README.md | 6 +++ lucebox.sh | 17 ++++--- lucebox/README.md | 7 +++ lucebox/src/lucebox/cli.py | 77 +++++++++++++++++++++++++------- lucebox/src/lucebox/download.py | 71 ++++++++++++++++++++++++++--- lucebox/tests/test_download.py | 26 ++++++++++- lucebox/tests/test_models_cli.py | 28 +++++++++++- lucebox/tests/test_placement.py | 59 ++++++++++++++++++++++++ scripts/test_lucebox_sh.sh | 2 +- 9 files changed, 262 insertions(+), 31 deletions(-) diff --git a/README.md b/README.md index c8fe80116..f351df4df 100644 --- a/README.md +++ b/README.md @@ -173,6 +173,12 @@ lucebox start ./lucebox.sh package-runtime # factory-installable paired runtime ``` +The first-run picker focuses on the model families Lucebox qualifies most +heavily: Qwen3.6 27B, Qwen3.6 35B-A3B, Laguna XS.2, and DeepSeek V4 Flash. +Before any download, it checks whether the detected single-GPU, multi-GPU, or +Lucebox heterogeneous placement can actually run the selected model. Gemma +presets remain available from `lucebox models list` for existing users. + ## Quick Start On Harnesses [`harness/`](harness/) contains RTX 3090 client launchers and regression tests diff --git a/lucebox.sh b/lucebox.sh index 11addb8bb..b331ae1bb 100755 --- a/lucebox.sh +++ b/lucebox.sh @@ -284,12 +284,14 @@ die() { err "$*"; exit 1; } print_logo() { printf '%b' "$C_BRAND" cat <<'EOF' - · ╱ - ·──✦──· █ █ █ ▄▀▀ █▀▀ █▀▀▄ ▄▀▀▄ █ █ - ╱ · █ █ █ █ █▀▀ █▀▀▄ █ █ █ - · ▀▀ ▀▀ ▀▀ ▀▀▀ ▀▀▀ ▀▀ ▀ ▀ + ░██ ░██ + ░██ ░██ ░██ ░███████ ░███████ ░████████ ░███████ ░██ ░██ + ░██ ░██ ░██ ░██ ░██ ░██ ░██ ░██ ░██ ░██ ░██ ░██ ░██ + ░██ ░██ ░██ ░██ ░█████████ ░██ ░██ ░██ ░██ ░█████ + ░██ ░██ ░███ ░██ ░██ ░██ ░███ ░██ ░██ ░██ ░██ ░██ + ░██████████ ░█████░██ ░███████ ░███████ ░██░█████ ░███████ ░██ ░██ EOF - printf '%b local inference, made simple%b\n\n' "$C_DIM" "$C_RST" + printf '%b computers for agents%b\n\n' "$C_DIM" "$C_RST" } # Find a source checkout for contributor-only actions. An explicit path wins; @@ -3172,13 +3174,16 @@ _optimization_summary() { fi case "$decode" in false|0|no|off) ;; *) case "$model" in - qwen3.6-27b|gemma-4-26b|gemma-4-31b) + qwen3.6-27b|gemma-4-26b|gemma-4-31b|deepseek-v4-flash) draft_file=$(_lucebox_config_get model.draft_file) if [ -z "$draft_file" ]; then case "$model" in qwen3.6-27b) draft_file="dflash-draft-3.6-q4_k_m.gguf" ;; gemma-4-26b) draft_file="gemma-4-26B-A4B-it-DFlash-q8_0.gguf" ;; gemma-4-31b) draft_file="gemma-4-31B-it-DFlash-q8_0.gguf" ;; + deepseek-v4-flash) + draft_file="DeepSeek-V4-Flash-DSpark-draft-Q4RMFP4-denseF16.gguf" + ;; esac fi [ -s "$DEFAULT_MODELS_DIR/draft/$draft_file" ] && active="DFlash" diff --git a/lucebox/README.md b/lucebox/README.md index c37dc2953..8148218b2 100644 --- a/lucebox/README.md +++ b/lucebox/README.md @@ -20,6 +20,13 @@ download, optimization and placement settings, and construction of the final server command. Host facts are passed through `LUCEBOX_HOST_*` environment variables so the package never has to guess the host configuration. +The guided picker leads with Lucebox's four featured model paths: Qwen3.6 +27B, Qwen3.6 35B-A3B, Laguna XS.2, and DeepSeek V4 Flash. It checks the +resolved placement before downloading, so a machine without enough compatible +accelerator memory cannot accidentally begin DeepSeek's roughly 114 GB +target-plus-draft download. Older supported Gemma presets remain available by +name and through `lucebox models list`. + Automatic mode considers the selected model and the complete accelerator topology. It shows plain-language decisions for DFlash, PFlash, KVFlash, Spark, and execution placement. It keeps fitting targets on the faster primary GPU, diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index a5471d41a..426a9fa10 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -330,13 +330,40 @@ def _print_installed_presets() -> None: console.print(f"[dim]Total disk usage: {total:.1f} GB[/dim]") -def _activate_preset(cfg: Config, preset: download_mod.ModelPreset) -> None: - """Activate one preset together with a runnable execution profile.""" - selected_model = ModelMeta( +def _model_meta(preset: download_mod.ModelPreset) -> ModelMeta: + """Build the persisted model selection for one catalog preset.""" + return ModelMeta( preset=preset.name, target_file=preset.target_file, draft_file=preset.draft_file if preset.has_draft and preset.draft_file else "", ) + + +def _preset_placement(cfg: Config, preset: download_mod.ModelPreset) -> PlacementPlan: + """Plan the placement that activating ``preset`` would persist.""" + selected_cfg = replace(cfg, model=_model_meta(preset)) + if config_mod.optimization_mode() == "custom": + return autotune_mod.placement_for_runtime(selected_cfg, cfg.dflash) + return autotune_mod.automatic_plan(selected_cfg).placement + + +def _require_runnable_preset(cfg: Config, preset: download_mod.ModelPreset) -> None: + """Stop before a large download when this machine cannot run the model.""" + placement = _preset_placement(cfg, preset) + if placement.runnable: + return + error_console.print( + f"[red]{escape(preset.label)} cannot run on the detected hardware.[/red]\n" + f"[dim]{escape(placement.reason)}[/dim]\n" + "[dim]No model files were downloaded. Run `lucebox check` to review " + "the detected accelerators.[/dim]" + ) + raise typer.Exit(code=2) + + +def _activate_preset(cfg: Config, preset: download_mod.ModelPreset) -> None: + """Activate one preset together with a runnable execution profile.""" + selected_model = _model_meta(preset) selected_cfg = replace( cfg, model=selected_model, @@ -406,17 +433,23 @@ def models_list() -> None: cfg = _load_or_build() active = cfg.model.preset table = Table() + table.add_column("model") table.add_column("preset") table.add_column("status") table.add_column("size (GB)") table.add_column("description") - for name in sorted(download_mod.PRESETS): - pres = download_mod.PRESETS[name] - marker = "* " if name == active else " " + for pres in download_mod.catalog_presets(): + marker = "* " if pres.name == active else " " status = download_mod.installed_status(cfg, pres) size = download_mod.installed_size_gb(cfg, pres) size_text = f"{size:.1f}" if size > 0 else f"~{pres.approx_total_gb}*" - table.add_row(f"{marker}{name}", status, size_text, pres.description or "") + table.add_row( + f"{marker}{pres.label}", + pres.name, + status, + size_text, + pres.description or "", + ) console.print(table) @@ -434,7 +467,8 @@ def models_select( """Choose, download, and activate a model in one guided step.""" cfg = _load_or_build() if not preset: - names = sorted(download_mod.PRESETS) + candidates = download_mod.catalog_presets(featured_only=True) + names = [candidate.name for candidate in candidates] recommended = download_mod.recommend_preset(cfg.host) default_name = cfg.model.preset or recommended or names[0] default_index = names.index(default_name) + 1 if default_name in names else 1 @@ -444,22 +478,29 @@ def models_select( table.add_column("model") table.add_column("download") table.add_column("status") - table.add_column("notes") - for index, name in enumerate(names, start=1): - candidate = download_mod.PRESETS[name] + table.add_column("this machine") + for index, candidate in enumerate(candidates, start=1): labels: list[str] = [] - if name == cfg.model.preset: + if candidate.name == cfg.model.preset: labels.append("active") - if name == recommended: + if candidate.name == recommended: labels.append("recommended") + placement = _preset_placement(cfg, candidate) + fit = "ready" if placement.runnable else "not enough compatible memory" table.add_row( str(index), - name, + candidate.label, f"~{candidate.approx_total_gb} GB", download_mod.installed_status(cfg, candidate), - ", ".join(labels) or candidate.description, + ", ".join((*labels, fit)) if labels else fit, ) console.print(table) + extra = [item.name for item in download_mod.catalog_presets() if not item.featured] + if extra: + console.print( + "[dim]More supported presets: " + f"{escape(', '.join(extra))}. See `lucebox models list`.[/dim]" + ) try: answer = typer.prompt( "Model number or name", @@ -486,9 +527,10 @@ def models_select( # already present locally. _activate_preset(cfg, selected) return + _require_runnable_preset(cfg, selected) if state != "installed" and not yes: if not typer.confirm( - f"Download about {selected.approx_total_gb} GB and activate {selected.name}?", + f"Download about {selected.approx_total_gb} GB and activate {selected.label}?", default=True, ): console.print("[dim]No model changed.[/dim]") @@ -538,6 +580,9 @@ def models_download( console.print(f"[red]{escape(str(exc))}[/red]") raise typer.Exit(code=2) from exc + if activate: + _require_runnable_preset(cfg, pres) + current = download_mod.status(cfg, pres) console.print(f"Models dir: [bold]{cfg.models_dir}[/bold]") console.print(f"Preset: [bold]{pres.name}[/bold]") diff --git a/lucebox/src/lucebox/download.py b/lucebox/src/lucebox/download.py index 5e68d7418..2673a52e2 100644 --- a/lucebox/src/lucebox/download.py +++ b/lucebox/src/lucebox/download.py @@ -93,28 +93,42 @@ class ModelPreset: native_context: int description: str = "" speculator_dir: str | None = None + display_name: str = "" + featured_rank: int | None = None @property def has_draft(self) -> bool: return bool(self.draft_repo and self.draft_file) + @property + def label(self) -> str: + """Human-readable catalog label; ``name`` remains the stable CLI id.""" + return self.display_name or self.name + + @property + def featured(self) -> bool: + """Whether this preset belongs in the focused first-run picker.""" + return self.featured_rank is not None + # Registry of supported models. Keyed by preset name; the CLI surface # exposes these via `lucebox models download ` and the # `lucebox models list` table. The values come straight from the model -# cards under share/model_cards/ — keep them in sync. +# cards and published Lucebox artifacts — keep them in sync. PRESETS: dict[str, ModelPreset] = { "qwen3.6-27b": ModelPreset( name="qwen3.6-27b", target_repo="unsloth/Qwen3.6-27B-GGUF", target_file="Qwen3.6-27B-Q4_K_M.gguf", - draft_repo="spiritbuun/Qwen3.6-27B-DFlash-GGUF", + draft_repo="Lucebox/Qwen3.6-27B-DFlash-GGUF", draft_file="dflash-draft-3.6-q4_k_m.gguf", approx_total_gb=17, approx_target_gb=15.7, approx_draft_gb=1.3, architecture="qwen35", native_context=262144, + display_name="Qwen3.6 27B", + featured_rank=1, description="Qwen3.6 27B dense (Q4_K_M) + Qwen3.6 DFlash draft. Lucebox default.", ), "gemma-4-26b": ModelPreset( @@ -160,6 +174,8 @@ def has_draft(self) -> bool: approx_draft_gb=1.1, architecture="laguna", native_context=4096, + display_name="Laguna XS.2", + featured_rank=3, description=( "Laguna-XS.2 MoE code model (Q4_K_M). " "DFlash safetensors speculator in draft/laguna-xs2-speculator/ " @@ -186,16 +202,58 @@ def has_draft(self) -> bool: approx_draft_gb=0.0, architecture="qwen35moe", native_context=262144, + display_name="Qwen3.6 35B-A3B", + featured_rank=2, description=( "Qwen3.6 35B-A3B MoE (3B active per token), Q4_K_M unsloth " "dynamic quant. Target-only — no DFlash MoE draft published " "yet. Uses lucebox's qwen35moe arch backend." ), ), + "deepseek-v4-flash": ModelPreset( + name="deepseek-v4-flash", + target_repo="Lucebox/DeepSeek-V4-Flash-ROCMFPX", + target_file="DeepSeek-V4-Flash-ROCMFP2-STRIX.gguf", + draft_repo="Lucebox/DeepSeek-V4-Flash-DSpark-Drafter-GGUF", + draft_file="DeepSeek-V4-Flash-DSpark-draft-Q4RMFP4-denseF16.gguf", + approx_total_gb=114, + approx_target_gb=102.4, + approx_draft_gb=11.3, + architecture="deepseek4", + native_context=1048576, + display_name="DeepSeek V4 Flash", + featured_rank=4, + description=( + "DeepSeek V4 Flash ROCmFPX target + Lucebox DSpark draft. " + "Large-system profile: about 103 GB for the target and 114 GB " + "with speculative decoding." + ), + ), } DEFAULT_PRESET = PRESETS["qwen3.6-27b"] + +def catalog_presets(*, featured_only: bool = False) -> tuple[ModelPreset, ...]: + """Return presets in stable product order. + + The guided picker stays deliberately focused on the models Lucebox tunes + and qualifies most heavily. Older supported presets remain addressable by + their stable names and visible through ``lucebox models list``. + """ + candidates = (preset for preset in PRESETS.values() if preset.featured or not featured_only) + return tuple( + sorted( + candidates, + key=lambda preset: ( + preset.featured_rank is None, + preset.featured_rank if preset.featured_rank is not None else 0, + preset.name, + ), + ) + ) + + # Shared scorer used by PFlash and by KVFlash's drafter residency policy. # It is deliberately separate from each target preset: one ~1.2 GB file is # reused by every installed model and can be factory-preloaded on a Lucebox. @@ -214,10 +272,14 @@ def resolve_preset(name: str | None) -> ModelPreset: return DEFAULT_PRESET if name in PRESETS: return PRESETS[name] + normalized = name.strip().casefold() + for preset in PRESETS.values(): + if normalized == preset.label.casefold(): + return preset # Build a suggestion list — show every known preset; the user's # search space is small, so listing them all is # cheaper and clearer than a fuzzy-match heuristic. - known = ", ".join(sorted(PRESETS.keys())) + known = ", ".join(preset.name for preset in catalog_presets()) raise KeyError(f"unknown preset {name!r}. Known presets: {known}") @@ -553,8 +615,7 @@ def installed_presets(cfg: Config) -> list[ModelPreset]: default ``lucebox models`` view stays uncluttered. """ out: list[ModelPreset] = [] - for name in sorted(PRESETS): - pres = PRESETS[name] + for pres in catalog_presets(): if installed_status(cfg, pres) == "installed": out.append(pres) return out diff --git a/lucebox/tests/test_download.py b/lucebox/tests/test_download.py index bb5b1e4b9..bae056c29 100644 --- a/lucebox/tests/test_download.py +++ b/lucebox/tests/test_download.py @@ -16,6 +16,7 @@ from lucebox.download import ( DEFAULT_PRESET, PRESETS, + catalog_presets, recommend_preset, resolve_preset, status, @@ -42,7 +43,7 @@ def test_import_does_not_mutate_huggingface_environment() -> None: def test_default_preset_uses_quantized_gguf_draft(): - assert DEFAULT_PRESET.draft_repo == "spiritbuun/Qwen3.6-27B-DFlash-GGUF" + assert DEFAULT_PRESET.draft_repo == "Lucebox/Qwen3.6-27B-DFlash-GGUF" assert DEFAULT_PRESET.draft_file == "dflash-draft-3.6-q4_k_m.gguf" @@ -91,6 +92,29 @@ def test_resolve_preset_picks_qwen36_moe_target_only(): assert not pres.has_draft +def test_featured_catalog_has_stable_product_order(): + assert [preset.name for preset in catalog_presets(featured_only=True)] == [ + "qwen3.6-27b", + "qwen3.6-moe", + "laguna-xs.2", + "deepseek-v4-flash", + ] + + +def test_resolve_preset_accepts_display_name(): + assert resolve_preset("Qwen3.6 35B-A3B") is PRESETS["qwen3.6-moe"] + + +def test_resolve_preset_picks_deepseek_target_and_dspark_draft(): + pres = resolve_preset("deepseek-v4-flash") + assert pres.target_repo == "Lucebox/DeepSeek-V4-Flash-ROCMFPX" + assert pres.target_file == "DeepSeek-V4-Flash-ROCMFP2-STRIX.gguf" + assert pres.draft_repo == "Lucebox/DeepSeek-V4-Flash-DSpark-Drafter-GGUF" + assert pres.draft_file == "DeepSeek-V4-Flash-DSpark-draft-Q4RMFP4-denseF16.gguf" + assert pres.has_draft + assert pres.approx_target_gb == pytest.approx(102.4) + + def test_download_preset_target_only_qwen36_moe_skips_draft(tmp_path, monkeypatch): """qwen3.6-moe behaves identically to laguna-xs.2: target only, no draft fetch.""" cfg = SimpleNamespace(models_dir=tmp_path) diff --git a/lucebox/tests/test_models_cli.py b/lucebox/tests/test_models_cli.py index a9f86e81c..762446254 100644 --- a/lucebox/tests/test_models_cli.py +++ b/lucebox/tests/test_models_cli.py @@ -163,14 +163,38 @@ def test_models_select_numbered_picker_downloads_and_activates( "status", lambda cfg, pres: {"target_present": False, "draft_present": False}, ) - # qwen3.6-27b is the fourth entry in the stable alphabetical menu. - result = CliRunner().invoke(app, ["models", "select"], input="4\ny\n") + # Qwen3.6 27B leads the stable, product-focused menu. + result = CliRunner().invoke(app, ["models", "select"], input="1\ny\n") assert result.exit_code == 0 assert "Choose a model" in result.output entries = config_mod.config_get(path=cfg_path) assert entries["model.preset"] == ("qwen3.6-27b", "file") +def test_models_select_blocks_incompatible_model_before_download( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + attempted = False + + def _unexpected_download(cfg, preset): # noqa: ARG001 + nonlocal attempted + attempted = True + return 0 + + monkeypatch.setattr(download_mod, "download_preset", _unexpected_download) + result = CliRunner().invoke( + app, + ["models", "select", "deepseek-v4-flash", "--yes"], + ) + + assert result.exit_code == 2 + assert "cannot run on the detected hardware" in result.output + assert "No model files were downloaded" in result.output + assert attempted is False + + def test_optimize_resets_to_hardware_profile( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/lucebox/tests/test_placement.py b/lucebox/tests/test_placement.py index da91c980c..d9c8cbf63 100644 --- a/lucebox/tests/test_placement.py +++ b/lucebox/tests/test_placement.py @@ -355,6 +355,65 @@ def test_rtx_strix_uses_remote_target_shard_when_capacity_requires_it( assert plan.runtime.peer_access is False +def test_deepseek_uses_r9700_and_strix_capacity_automatically(tmp_path: Path) -> None: + host = HostFacts( + gpu_vendor="amd", + has_amd_gpu=True, + gpu_name="AMD Radeon AI PRO R9700", + gpu_count=2, + vram_gb=31, + gpu_sm="gfx1201", + ram_gb=125, + amd_gpu_list_csv=_amd_csv( + ("AMD Radeon AI PRO R9700", "gfx1201", 32624), + ("AMD Radeon Graphics", "gfx1151", 512), + ), + ) + cfg = Config( + variant="rocm", + models_dir=tmp_path, + host=host, + model=ModelMeta(preset="deepseek-v4-flash"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.placement.runnable is True + assert plan.placement.runtime.mode == "layer-split" + assert plan.placement.runtime.target_devices == ("hip:0", "hip:1") + assert sum(plan.placement.runtime.target_layer_split) == pytest.approx(1.0) + assert plan.runtime.speculative_decode is False + + +def test_deepseek_uses_rtx_strix_paired_runtime_automatically(tmp_path: Path) -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + has_amd_gpu=True, + gpu_name="RTX 3090", + vram_gb=24, + gpu_sm="86", + ram_gb=128, + nvidia_gpu_list_csv=("0, GPU-0, 0000:01:00.0, RTX 3090, 8.6, 24576 MiB, 350 W"), + amd_gpu_list_csv=_amd_csv(("AMD Radeon Graphics", "gfx1151", 512)), + hybrid_runtime=True, + ) + cfg = Config( + variant="cuda12", + models_dir=tmp_path, + host=host, + model=ModelMeta(preset="deepseek-v4-flash"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.placement.runnable is True + assert plan.placement.runtime.mode == "heterogeneous" + assert plan.placement.runtime.target_devices == ("cuda:0", "hip:0") + assert plan.placement.runtime.remote_target_shard is True + assert plan.placement.runtime.requires_hybrid_runtime is True + + def test_python_architecture_capabilities_match_engine_table() -> None: header = (Path(__file__).parents[2] / "server/src/common/model_capabilities.h").read_text() row_pattern = re.compile( diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index 791f9609d..7e51a3a09 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -218,7 +218,7 @@ assert_runs "no args (non-interactive help)" \ "bash '$SCRIPT' Date: Fri, 31 Jul 2026 17:37:31 +0200 Subject: [PATCH 16/25] feat(cli): centralize model optimization planning --- README.md | 16 +- lucebox.sh | 24 +- lucebox/README.md | 20 +- lucebox/src/lucebox/autotune.py | 349 +++++++++++++++++++++++----- lucebox/src/lucebox/capabilities.py | 316 +++++++++++++++++++++++++ lucebox/src/lucebox/cli.py | 99 ++++++-- lucebox/src/lucebox/config.py | 9 + lucebox/src/lucebox/docker_run.py | 2 + lucebox/src/lucebox/download.py | 5 +- lucebox/src/lucebox/placement.py | 104 +++++---- lucebox/src/lucebox/types.py | 9 + lucebox/tests/test_autotune.py | 48 +++- lucebox/tests/test_capabilities.py | 58 +++++ lucebox/tests/test_config.py | 3 + lucebox/tests/test_docker_run.py | 3 + lucebox/tests/test_placement.py | 99 +++++++- scripts/test_lucebox_sh.sh | 15 +- server/scripts/entrypoint.sh | 6 + variables.md | 1 + 19 files changed, 1032 insertions(+), 154 deletions(-) create mode 100644 lucebox/src/lucebox/capabilities.py create mode 100644 lucebox/tests/test_capabilities.py diff --git a/README.md b/README.md index f351df4df..198ea51b7 100644 --- a/README.md +++ b/README.md @@ -115,12 +115,16 @@ lucebox `lucebox` opens the branded menu. **Quick setup** detects NVIDIA CUDA or AMD ROCm, lets you choose and download a model, applies a model- and hardware-aware **Automatic** profile, and starts the inference service. Automatic explains its -DFlash, PFlash, KVFlash, and Spark choices and prefers exact, all-GPU execution -when it fits. It enables paging or expert offload only when the selected model -and primary GPU need them; Spark also requires enough host RAM for its cold -experts. **Review optimizations** offers an Advanced mode for users who want to -override those four choices. Wi-Fi and device provisioning -are intentionally outside this inference-only CLI. +DFlash, PFlash, KVFlash, and Spark choices, and also prints the resolved +prefill, decode, and KV-cache strategy. A typed capability contract prevents +unsupported model/backend combinations; preview paths are labeled and stay off +in Automatic. Exact, all-GPU execution remains preferred when it fits. Paging +or expert offload activates only when the selected model and primary GPU need +it; Spark also requires enough host RAM for its cold experts. **Review +optimizations** exposes only model-compatible overrides. DeepSeek sparse +prefill is available there as an approximate, monolithic-HIP preview; exact MLA +prefill remains the default. Wi-Fi and device provisioning are intentionally +outside this inference-only CLI. Automatic inventories every NVIDIA and AMD accelerator, then writes one validated execution plan together with the optimization profile. A model that diff --git a/lucebox.sh b/lucebox.sh index b331ae1bb..9ccae2a92 100755 --- a/lucebox.sh +++ b/lucebox.sh @@ -33,6 +33,17 @@ set -euo pipefail +# ROCm packages on minimal Ubuntu images install their tools under +# /opt/rocm/bin without necessarily adding that directory to a login shell's +# PATH. Keep probing and native builds self-contained; this changes PATH only +# for the wrapper process and its children. +if [ -d /opt/rocm/bin ]; then + case ":${PATH:-}:" in + *:/opt/rocm/bin:*) ;; + *) export PATH="/opt/rocm/bin:${PATH:-/usr/bin:/bin}" ;; + esac +fi + VERSION="0.2.0" SCRIPT_PATH="$(readlink -f "$0" 2>/dev/null || realpath "$0" 2>/dev/null || echo "$0")" SCRIPT_NAME="$(basename "$SCRIPT_PATH")" @@ -2313,6 +2324,7 @@ dflash.kvflash_policy|DFLASH_KVFLASH_POLICY dflash.kvflash_tau|DFLASH_KVFLASH_TAU dflash.spark|DFLASH_SPARK dflash.spark_vram_gb|DFLASH_SPARK_VRAM_GB +dflash.ds4_prefill|DFLASH_DS4_PREFILL dflash.think_max|DFLASH_THINK_MAX dflash.fa_window|DFLASH_FA_WINDOW dflash.think_soft_close_min_ratio|DFLASH_THINK_SOFT_CLOSE_MIN_RATIO @@ -2356,8 +2368,16 @@ cmd_native_build() { [ -f "$repo/server/deps/llama.cpp/ggml/CMakeLists.txt" ] \ || die "git submodules are missing — run: git -C '$repo' submodule update --init --recursive" - local configure=( - cmake -S "$repo/server" -B "$build_dir" + local configure=(cmake) + # Prefer Ninja for a fresh build when it is available. Minimal buyer + # images often ship Ninja with the ROCm SDK but omit GNU make; relying on + # CMake's Unix Makefiles default makes an otherwise complete toolchain + # fail before compiler detection. Preserve an existing build's generator. + if [ ! -f "$build_dir/CMakeCache.txt" ] && command -v ninja >/dev/null 2>&1; then + configure+=(-G Ninja) + fi + configure+=( + -S "$repo/server" -B "$build_dir" -DCMAKE_BUILD_TYPE=Release -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON ) diff --git a/lucebox/README.md b/lucebox/README.md index 8148218b2..e6a3bccb0 100644 --- a/lucebox/README.md +++ b/lucebox/README.md @@ -27,15 +27,17 @@ accelerator memory cannot accidentally begin DeepSeek's roughly 114 GB target-plus-draft download. Older supported Gemma presets remain available by name and through `lucebox models list`. -Automatic mode considers the selected model and the complete accelerator -topology. It shows plain-language decisions for DFlash, PFlash, KVFlash, Spark, -and execution placement. It keeps fitting targets on the faster primary GPU, -uses supported secondary-device workloads under memory pressure, and enables -memory-saving paths only when needed. Advanced mode lets experienced users -override the four optimization choices while placement invariants, context, -cache, and DDTree defaults remain guarded. A factory-preloaded Lucebox can ship -the models, shared optimizer scorer, and paired runtime, so a buyer does not -need to download or compile anything during first setup. +Automatic mode combines a typed model capability contract with the complete +accelerator topology. It prints the selected prefill, decode, and KV-cache +strategies, then explains each DFlash, PFlash, KVFlash, Spark, and placement +decision. A fitting target stays on the faster primary GPU; qualified +secondary-device and memory-saving paths activate only when needed. Advanced +mode exposes only policies that are legal for that model and backend while +placement, context, cache, and DDTree invariants remain guarded. Preview paths +are labeled and never enabled silently. In particular, DeepSeek sparse prefill +is an explicit approximate HIP preview and exact MLA prefill remains Automatic. +A factory-preloaded Lucebox can ship the models, shared optimizer scorer, and +paired runtime, so a buyer does not download or compile during first setup. Spark is enabled automatically only when an MoE model is under GPU-memory pressure and the machine reports at least 32 GB of host RAM; its cold experts diff --git a/lucebox/src/lucebox/autotune.py b/lucebox/src/lucebox/autotune.py index feb64795f..e6a24169f 100644 --- a/lucebox/src/lucebox/autotune.py +++ b/lucebox/src/lucebox/autotune.py @@ -11,7 +11,15 @@ from dataclasses import dataclass, replace from typing import Literal +from lucebox.capabilities import ( + FeatureCapability, + KvFlashCapability, + ModelOptimizationProfile, + Qualification, + model_profile, +) from lucebox.placement import PlacementPlan, automatic_placement +from lucebox.topology import from_config from lucebox.types import Config, DflashRuntime, HostFacts @@ -23,6 +31,7 @@ class OptimizationDecision: enabled: bool available: bool reason: str + qualification: Qualification = Qualification.QUALIFIED @dataclass(frozen=True, slots=True) @@ -36,16 +45,32 @@ class OptimizationPlan: pflash: OptimizationDecision kvflash: OptimizationDecision spark: OptimizationDecision + prefill_alternative: OptimizationDecision | None = None + prefill_strategy: str = "not selected" + decode_strategy: str = "not selected" + kv_strategy: str = "not selected" needs_optimizer_drafter: bool = False @property def decisions(self) -> tuple[OptimizationDecision, ...]: - return (self.dflash, self.pflash, self.kvflash, self.spark) + core = (self.dflash, self.pflash, self.kvflash, self.spark) + if self.prefill_alternative is None: + return core + return (*core, self.prefill_alternative) @property def active_names(self) -> tuple[str, ...]: return tuple(item.name for item in self.decisions if item.enabled) + @property + def phase_strategies(self) -> tuple[tuple[str, str], ...]: + """Resolved implementation for each performance-critical phase.""" + return ( + ("Prefill", self.prefill_strategy), + ("Decode", self.decode_strategy), + ("KV cache", self.kv_strategy), + ) + def _budget_for_host(host: HostFacts) -> int: """Use the documented per-architecture DDTree sweet spot when known.""" @@ -134,15 +159,17 @@ def _cap_exact_context( runtime: DflashRuntime, *, headroom_gb: float, - arch: str, + profile: ModelOptimizationProfile, + backend: str, ) -> DflashRuntime: """Keep exact-cache families inside conservative primary-GPU headroom. - Qwen families have an automatic bounded-residency path below; other - families currently keep their exact cache in Automatic mode. Those models - get a smaller context when their weights/draft leave little working room. + Models with a qualified automatic bounded-residency path are handled by + the planner below. Other models keep exact cache semantics in Automatic + mode and receive a smaller context when weights leave little working room. """ - if arch in {"qwen35", "qwen35moe"}: + kvflash = profile.kvflash + if kvflash is not None and kvflash.feature.automatic_on(backend): return runtime if headroom_gb < 4: return replace(runtime, max_ctx=min(runtime.max_ctx, 8192)) @@ -151,7 +178,12 @@ def _cap_exact_context( return runtime -CapacityAdjustment = Literal["pflash", "kvflash-qk", "kvflash-off", "dflash"] +CapacityAdjustment = Literal[ + "pflash", + "kvflash-scorerless", + "kvflash-off", + "dflash", +] def _uses_optimizer_scorer(runtime: DflashRuntime) -> bool: @@ -168,7 +200,7 @@ def _without_unused_scorer(runtime: DflashRuntime) -> DflashRuntime: def _capacity_fallbacks( runtime: DflashRuntime, *, - architecture: str, + kvflash: KvFlashCapability | None, ) -> tuple[tuple[DflashRuntime, CapacityAdjustment], ...]: """Return progressively smaller optional stacks, in performance order. @@ -199,9 +231,10 @@ def _capacity_fallbacks( candidates.append((current, "pflash")) if current.kvflash != "off" and current.kvflash_policy == "drafter": - if architecture == "qwen35": - current = replace(current, kvflash_policy="qk") - action: CapacityAdjustment = "kvflash-qk" + scorerless_policy = kvflash.scorerless_policy if kvflash is not None else None + if scorerless_policy is not None: + current = replace(current, kvflash_policy=scorerless_policy) + action: CapacityAdjustment = "kvflash-scorerless" else: current = replace(current, kvflash="off") action = "kvflash-off" @@ -215,6 +248,36 @@ def _capacity_fallbacks( return tuple(candidates) +def _selected_backend(cfg: Config) -> str: + """Return the backend of the runtime's primary device.""" + primary = from_config(cfg).primary + if primary is not None: + return primary.backend + # Config snapshots predating the inventory fields can still carry the + # selected generic vendor. Placement remains blocked when the configured + # image and inventory disagree; this fallback is only for feature support + # reporting and migration of those snapshots. + if cfg.host.gpu_vendor == "nvidia": + return "cuda" + if cfg.host.gpu_vendor == "amd": + return "hip" + return "" + + +def _qualification(feature: FeatureCapability | None, backend: str) -> Qualification: + if feature is None: + return Qualification.UNAVAILABLE + return feature.qualification_on(backend) + + +def _preview_reason(feature: FeatureCapability, backend: str) -> str: + backend_label = "CUDA" if backend == "cuda" else "HIP" + return ( + f"{feature.label} is available as a {backend_label} preview; " + "Automatic keeps the qualified baseline" + ) + + def automatic_plan( cfg: Config, *, @@ -222,11 +285,10 @@ def automatic_plan( ) -> OptimizationPlan: """Resolve the recommended profile for the selected model and primary GPU. - Automatic mode favors exact/full-cache execution whenever it fits. PFlash - is enabled only for its validated Qwen3.6 dense profile. KVFlash is enabled - only under real memory pressure and only on qwen-family paths with a tuned - scorer. Spark is enabled only for supported MoE models that would otherwise - leave too little VRAM headroom. Advanced mode can override these choices. + Automatic mode favors exact/full-cache execution whenever it fits. The + model contract limits each feature to legal, qualified backends; measured + memory pressure decides whether a qualified feature is useful. Device + placement is then resolved independently from the detected topology. """ from lucebox import download as download_mod @@ -241,6 +303,7 @@ def automatic_plan( False, False, "choose a model first; it activates when that model has a matching draft", + Qualification.UNAVAILABLE, ) unavailable = "choose a model first" placement = automatic_placement( @@ -255,11 +318,24 @@ def automatic_plan( model_name=cfg.model.preset or "not selected", placement=placement, dflash=dflash, - pflash=OptimizationDecision("PFlash", False, False, unavailable), - kvflash=OptimizationDecision("KVFlash", False, False, unavailable), - spark=OptimizationDecision("Spark", False, False, unavailable), + pflash=OptimizationDecision( + "PFlash", False, False, unavailable, Qualification.UNAVAILABLE + ), + kvflash=OptimizationDecision( + "KVFlash", False, False, unavailable, Qualification.UNAVAILABLE + ), + spark=OptimizationDecision( + "Spark", False, False, unavailable, Qualification.UNAVAILABLE + ), ) + profile = model_profile(preset.name) + if profile.architecture != preset.architecture: + raise ValueError( + f"optimization profile for {preset.name!r} targets {profile.architecture!r}, " + f"but the model catalog declares {preset.architecture!r}" + ) + backend = _selected_backend(cfg) runtime = replace(runtime, max_ctx=min(runtime.max_ctx, preset.native_context)) vram_known = cfg.host.vram_gb > 0 headroom = cfg.host.vram_gb - preset.approx_total_gb if vram_known else 0 @@ -267,43 +343,100 @@ def automatic_plan( runtime = _cap_exact_context( runtime, headroom_gb=headroom, - arch=preset.architecture, + profile=profile, + backend=backend, ) has_draft = draft_available(cfg, preset) - runtime = replace(runtime, speculative_decode=has_draft) - if has_draft: + decode_feature = profile.speculative_decode + decode_qualification = _qualification(decode_feature, backend) + decode_supported = decode_feature is not None and decode_feature.available_on(backend) + dflash_available = has_draft and decode_supported + dflash_enabled = ( + dflash_available + and decode_feature is not None + and decode_feature.automatic_on(backend) + ) + runtime = replace(runtime, speculative_decode=dflash_enabled) + if dflash_available and not dflash_enabled and decode_feature is not None: + dflash_reason = _preview_reason(decode_feature, backend) + elif dflash_enabled: dflash_reason = "matching speculative draft is available for this model" + elif has_draft and not decode_supported: + dflash_reason = "the installed draft is not supported by the selected backend" elif preset.speculator_dir: dflash_reason = "optional model speculator is not installed" elif preset.has_draft: dflash_reason = "matching speculative draft is not installed" else: dflash_reason = "no compatible speculative draft is published for this model" - dflash = OptimizationDecision("DFlash", has_draft, has_draft, dflash_reason) + dflash = OptimizationDecision( + decode_feature.label if decode_feature is not None else "DFlash", + dflash_enabled, + dflash_available, + dflash_reason, + decode_qualification, + ) + + alternative_capability = profile.deepseek_prefill + if alternative_capability is None: + prefill_alternative = None + else: + alternative_feature = alternative_capability.feature + alternative_qualification = alternative_feature.qualification_on(backend) + alternative_available = alternative_feature.available_on(backend) + if alternative_available: + alternative_reason = _preview_reason(alternative_feature, backend) + else: + alternative_reason = ( + "this approximate prefill path is unavailable on the selected backend" + ) + prefill_alternative = OptimizationDecision( + alternative_feature.label, + False, + alternative_available, + alternative_reason, + alternative_qualification, + ) long_context = runtime.max_ctx >= 32768 - # PFlash automatic mode is deliberately narrow: this is the profile with - # end-to-end long-context quality validation. Other architectures remain - # unavailable until the engine exposes a compatible compression path. - pflash_supported = preset.architecture == "qwen35" - pflash_available = pflash_supported and long_context - pflash_profile = pflash_available and vram_known + pflash_capability = profile.pflash + pflash_feature = pflash_capability.feature if pflash_capability is not None else None + pflash_qualification = _qualification(pflash_feature, backend) + pflash_supported = ( + pflash_feature is not None and pflash_feature.available_on(backend) + ) + pflash_available = ( + pflash_supported + and pflash_capability is not None + and runtime.max_ctx >= pflash_capability.minimum_context + ) + pflash_profile = ( + pflash_available + and vram_known + and pflash_feature is not None + and pflash_feature.automatic_on(backend) + ) pflash_enabled = pflash_profile and optimizer_drafter_available if pflash_enabled: pflash_reason = "long prompts use the installed scorer above 32K tokens" + assert pflash_capability is not None runtime = replace( runtime, prefill_mode="auto", - prefill_keep_ratio=0.10, - prefill_threshold=32768, + prefill_keep_ratio=pflash_capability.keep_ratio, + prefill_threshold=pflash_capability.minimum_context, prefill_drafter=download_mod.optimizer_drafter_container_path(), ) elif pflash_profile: pflash_reason = "shared 1.2 GB scorer is not installed" + elif pflash_available and pflash_feature is not None: + pflash_reason = _preview_reason(pflash_feature, backend) elif not pflash_supported: - pflash_reason = "this model architecture has no PFlash compression path" + pflash_reason = ( + f"this model has no PFlash production path; {profile.prefill_baseline} remains active" + ) elif not long_context: pflash_reason = "the safe context for this GPU is below PFlash's long-prompt threshold" else: @@ -313,60 +446,106 @@ def automatic_plan( pflash_enabled, pflash_available, pflash_reason, + pflash_qualification, ) - # Bounded KV residency changes long-context retrieval semantics, so it is - # not a blanket speed toggle. Use it only when the selected stack leaves - # less than 5 GB for KV/compute and a tuned qwen-family policy is available. + # Bounded residency affects long-context retrieval semantics, so it is + # activated only by a model contract and only under real memory pressure. + kvflash_capability = profile.kvflash + kvflash_feature = ( + kvflash_capability.feature if kvflash_capability is not None else None + ) + kvflash_qualification = _qualification(kvflash_feature, backend) + kvflash_supported = ( + kvflash_feature is not None and kvflash_feature.available_on(backend) + ) + kvflash_available = ( + kvflash_supported + and kvflash_capability is not None + and runtime.max_ctx >= kvflash_capability.minimum_context + ) kvflash_profile = ( vram_known - and long_context - and headroom < 5 - and preset.architecture in {"qwen35", "qwen35moe"} - ) - kvflash_enabled = kvflash_profile and ( - optimizer_drafter_available or preset.architecture == "qwen35" + and kvflash_available + and kvflash_capability is not None + and kvflash_feature is not None + and kvflash_feature.automatic_on(backend) + and headroom < kvflash_capability.pressure_headroom_gb ) + kvflash_policy = None + if kvflash_profile and kvflash_capability is not None: + kvflash_policy = ( + kvflash_capability.preferred_policy + if optimizer_drafter_available + else kvflash_capability.scorerless_policy + ) + kvflash_enabled = kvflash_policy is not None if kvflash_enabled: - policy: Literal["drafter", "qk"] = "drafter" if optimizer_drafter_available else "qk" - runtime = replace(runtime, kvflash="auto", kvflash_policy=policy) - if optimizer_drafter_available and not runtime.prefill_drafter: + assert kvflash_policy is not None + runtime = replace(runtime, kvflash="auto", kvflash_policy=kvflash_policy) + if kvflash_policy == "drafter" and not runtime.prefill_drafter: runtime = replace( runtime, prefill_drafter=download_mod.optimizer_drafter_container_path(), ) - kvflash_reason = f"only {headroom} GB VRAM headroom; {policy} policy bounds KV residency" + kvflash_reason = ( + f"only {headroom} GB VRAM headroom; {kvflash_policy} policy bounds KV residency" + ) elif kvflash_profile: kvflash_reason = "memory pressure detected, but no quality-safe scorer is installed" - elif not long_context: + elif kvflash_available and kvflash_feature is not None: + if kvflash_feature.qualification_on(backend) is Qualification.PREVIEW: + kvflash_reason = _preview_reason(kvflash_feature, backend) + else: + kvflash_reason = "full KV cache fits; exact full-cache execution is preferred" + elif kvflash_capability is not None and runtime.max_ctx < kvflash_capability.minimum_context: kvflash_reason = "the selected context does not need bounded KV residency" - elif preset.architecture not in {"qwen35", "qwen35moe"}: - kvflash_reason = "cross-tokenizer residency remains an Advanced choice" else: - kvflash_reason = "full KV cache fits; exact full-cache execution is preferred" + kvflash_reason = ( + f"this model uses {profile.kv_baseline} instead of generic KVFlash" + ) kvflash = OptimizationDecision( "KVFlash", kvflash_enabled, - long_context, + kvflash_available, kvflash_reason, + kvflash_qualification, ) - spark_available = preset.architecture in {"laguna", "qwen35moe"} - spark_pressure = spark_available and vram_known and headroom < 6 + spark_capability = profile.spark + spark_feature = spark_capability.feature if spark_capability is not None else None + spark_qualification = _qualification(spark_feature, backend) + spark_available = spark_feature is not None and spark_feature.available_on(backend) + spark_pressure = ( + spark_available + and vram_known + and spark_capability is not None + and headroom < spark_capability.pressure_headroom_gb + ) # Cold experts live in host RAM. A 20–22 GB preset plus the OS, runtime, # KV spill, and working buffers is not a safe automatic fit below 32 GB. # Advanced mode can still let an informed operator opt in. - spark_host_ready = cfg.host.ram_gb >= 32 - spark_enabled = spark_pressure and spark_host_ready + minimum_host_ram = ( + spark_capability.minimum_host_ram_gb if spark_capability is not None else 0 + ) + spark_host_ready = cfg.host.ram_gb >= minimum_host_ram + spark_enabled = ( + spark_pressure + and spark_host_ready + and spark_feature is not None + and spark_feature.automatic_on(backend) + ) runtime = replace(runtime, spark=spark_enabled) if spark_enabled: spark_reason = f"MoE weights leave {headroom} GB headroom; expert residency self-tunes" + elif spark_pressure and spark_feature is not None and not spark_feature.automatic_on(backend): + spark_reason = _preview_reason(spark_feature, backend) elif spark_pressure and cfg.host.ram_gb <= 0: spark_reason = "GPU memory is tight, but host RAM is unknown; automatic offload stays off" elif spark_pressure: spark_reason = ( f"GPU memory is tight, but {cfg.host.ram_gb} GB host RAM is below " - "the 32 GB automatic offload floor" + f"the {minimum_host_ram} GB automatic offload floor" ) elif spark_available and vram_known: spark_reason = "the model fits the primary GPU; all-GPU execution is faster and simpler" @@ -374,30 +553,41 @@ def automatic_plan( spark_reason = "GPU memory is unknown; automatic mode will not assume offload" else: spark_reason = "this is not a Spark-compatible MoE architecture" - spark = OptimizationDecision("Spark", spark_enabled, spark_available, spark_reason) + spark = OptimizationDecision( + "Spark", + spark_enabled, + spark_available, + spark_reason, + spark_qualification, + ) needs_optimizer_drafter = not optimizer_drafter_available and ( - pflash_profile or (kvflash_profile and preset.architecture == "qwen35moe") + pflash_profile + or ( + kvflash_profile + and kvflash_capability is not None + and kvflash_capability.scorerless_policy is None + ) ) placement = automatic_placement( cfg, runtime, preset, - has_draft=has_draft, + has_draft=dflash_enabled, optimizer_drafter_available=optimizer_drafter_available, ) capacity_adjustments: set[CapacityAdjustment] = set() if not placement.runnable: for fallback_runtime, adjustment in _capacity_fallbacks( runtime, - architecture=preset.architecture, + kvflash=kvflash_capability, ): capacity_adjustments.add(adjustment) placement = automatic_placement( cfg, fallback_runtime, preset, - has_draft=has_draft and fallback_runtime.speculative_decode, + has_draft=dflash_enabled and fallback_runtime.speculative_decode, optimizer_drafter_available=optimizer_drafter_available, ) runtime = fallback_runtime @@ -412,10 +602,11 @@ def automatic_plan( else "disabled because the selected target placement has no compatible draft path" ) dflash = OptimizationDecision( - "DFlash", + dflash.name, False, dflash.available, reason, + dflash.qualification, ) if pflash.enabled and adjusted.prefill_mode == "off": reason = ( @@ -428,6 +619,7 @@ def automatic_plan( False, pflash.available, reason, + pflash.qualification, ) if kvflash.enabled and adjusted.kvflash == "off": reason = ( @@ -440,13 +632,15 @@ def automatic_plan( False, kvflash.available, reason, + kvflash.qualification, ) - elif "kvflash-qk" in capacity_adjustments: + elif "kvflash-scorerless" in capacity_adjustments: kvflash = OptimizationDecision( "KVFlash", True, kvflash.available, - "memory is tight; QK policy avoids a separate scorer allocation", + f"memory is tight; {adjusted.kvflash_policy} policy avoids a separate scorer allocation", + kvflash.qualification, ) if spark.enabled and not adjusted.spark: spark = OptimizationDecision( @@ -454,7 +648,28 @@ def automatic_plan( False, spark.available, "disabled because Automatic does not compose Spark with target layer splitting", + spark.qualification, ) + if ( + prefill_alternative is not None + and prefill_alternative.available + and alternative_capability is not None + ): + alternative_placement = automatic_placement( + cfg, + replace(adjusted, ds4_prefill=alternative_capability.mode), + preset, + has_draft=adjusted.speculative_decode, + optimizer_drafter_available=optimizer_drafter_available, + ) + if not alternative_placement.runnable: + prefill_alternative = OptimizationDecision( + prefill_alternative.name, + False, + False, + alternative_placement.reason, + prefill_alternative.qualification, + ) if needs_optimizer_drafter: # Do not ask a first-time user to download 1.2 GB only to discard the # scorer during capacity fallback. The scorer-present branch cannot @@ -471,5 +686,13 @@ def automatic_plan( pflash=pflash, kvflash=kvflash, spark=spark, + prefill_alternative=prefill_alternative, + prefill_strategy=(pflash.name if pflash.enabled else profile.prefill_baseline), + decode_strategy=(dflash.name if dflash.enabled else profile.decode_baseline), + kv_strategy=( + f"{kvflash.name} ({adjusted.kvflash_policy})" + if kvflash.enabled + else profile.kv_baseline + ), needs_optimizer_drafter=needs_optimizer_drafter, ) diff --git a/lucebox/src/lucebox/capabilities.py b/lucebox/src/lucebox/capabilities.py new file mode 100644 index 000000000..1cd79d803 --- /dev/null +++ b/lucebox/src/lucebox/capabilities.py @@ -0,0 +1,316 @@ +"""Product capability contracts for supported models and engine architectures. + +This module deliberately contains *static* facts only: which optimization an +engine/model pair can run, which backends have been qualified, and which +policies are legal. Hardware inventory, memory pressure, and device placement +remain runtime decisions owned by :mod:`lucebox.autotune` and +:mod:`lucebox.placement`. + +Keeping those concerns separate prevents a common class of planner bugs: a +model name must never imply a particular GPU layout, and detecting a large GPU +must never make an unsupported model feature legal. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from enum import StrEnum +from types import MappingProxyType +from typing import Literal + +Backend = Literal["cuda", "hip"] +KvFlashPolicy = Literal["drafter", "lru", "qk"] + + +class Qualification(StrEnum): + """Product support level for one model/backend implementation.""" + + UNAVAILABLE = "unavailable" + PREVIEW = "preview" + QUALIFIED = "qualified" + + +@dataclass(frozen=True, slots=True) +class BackendSupport: + """Qualification of the same implementation on CUDA and HIP.""" + + cuda: Qualification + hip: Qualification + + def for_backend(self, backend: str) -> Qualification: + if backend == "cuda": + return self.cuda + if backend == "hip": + return self.hip + return Qualification.UNAVAILABLE + + def available_on(self, backend: str) -> bool: + return self.for_backend(backend) is not Qualification.UNAVAILABLE + + def qualified_on(self, backend: str) -> bool: + return self.for_backend(backend) is Qualification.QUALIFIED + + +QUALIFIED_BOTH = BackendSupport(Qualification.QUALIFIED, Qualification.QUALIFIED) +PREVIEW_BOTH = BackendSupport(Qualification.PREVIEW, Qualification.PREVIEW) +HIP_PREVIEW = BackendSupport(Qualification.UNAVAILABLE, Qualification.PREVIEW) + + +@dataclass(frozen=True, slots=True) +class FeatureCapability: + """A launchable optimization and its automatic-mode policy.""" + + label: str + support: BackendSupport + automatic: bool = True + + def __post_init__(self) -> None: + if not self.label.strip(): + raise ValueError("feature label must not be empty") + + def available_on(self, backend: str) -> bool: + return self.support.available_on(backend) + + def automatic_on(self, backend: str) -> bool: + return self.automatic and self.support.qualified_on(backend) + + def qualification_on(self, backend: str) -> Qualification: + return self.support.for_backend(backend) + + +@dataclass(frozen=True, slots=True) +class PFlashCapability: + feature: FeatureCapability + minimum_context: int = 32_768 + keep_ratio: float = 0.10 + + def __post_init__(self) -> None: + if self.minimum_context <= 0: + raise ValueError("PFlash minimum context must be positive") + if not 0.0 < self.keep_ratio <= 1.0: + raise ValueError("PFlash keep ratio must be in (0.0, 1.0]") + + +@dataclass(frozen=True, slots=True) +class KvFlashCapability: + feature: FeatureCapability + policies: tuple[KvFlashPolicy, ...] + preferred_policy: KvFlashPolicy + scorerless_policy: KvFlashPolicy | None + minimum_context: int = 32_768 + pressure_headroom_gb: float = 5.0 + + def __post_init__(self) -> None: + if not self.policies: + raise ValueError("KVFlash must expose at least one legal policy") + if self.preferred_policy not in self.policies: + raise ValueError("preferred KVFlash policy must be legal for the model") + if self.scorerless_policy is not None and self.scorerless_policy not in self.policies: + raise ValueError("scorerless KVFlash policy must be legal for the model") + if self.minimum_context <= 0: + raise ValueError("KVFlash minimum context must be positive") + if self.pressure_headroom_gb < 0.0: + raise ValueError("KVFlash pressure headroom must not be negative") + + +@dataclass(frozen=True, slots=True) +class SparkCapability: + feature: FeatureCapability + pressure_headroom_gb: float = 6.0 + minimum_host_ram_gb: int = 32 + + def __post_init__(self) -> None: + if self.pressure_headroom_gb < 0.0: + raise ValueError("Spark pressure headroom must not be negative") + if self.minimum_host_ram_gb <= 0: + raise ValueError("Spark minimum host RAM must be positive") + + +@dataclass(frozen=True, slots=True) +class DeepSeekPrefillCapability: + """Approximate DeepSeek4 prefill mode exposed by the native server.""" + + feature: FeatureCapability + mode: Literal["dense", "sparse"] + + +@dataclass(frozen=True, slots=True) +class ModelOptimizationProfile: + """The optimization contract for one downloadable model preset. + + The three baseline labels make phase coverage explicit even when a model + does not use a branded Lucebox optimization. For example, DeepSeek uses + native MLA-compressed KV state rather than generic KVFlash. + """ + + preset: str + architecture: str + prefill_baseline: str + decode_baseline: str + kv_baseline: str + speculative_decode: FeatureCapability | None = None + pflash: PFlashCapability | None = None + kvflash: KvFlashCapability | None = None + spark: SparkCapability | None = None + deepseek_prefill: DeepSeekPrefillCapability | None = None + + def __post_init__(self) -> None: + if not self.preset.strip() or not self.architecture.strip(): + raise ValueError("model profile preset and architecture must not be empty") + if not all(label.strip() for label in self.phase_baselines): + raise ValueError("every model profile must name all three baseline strategies") + + @property + def phase_baselines(self) -> tuple[str, str, str]: + return (self.prefill_baseline, self.decode_baseline, self.kv_baseline) + + +@dataclass(frozen=True, slots=True) +class ArchitecturePlacementCapabilities: + """Engine placement operations that are legal for one architecture.""" + + layer_split: bool + remote_draft: bool + draft_on_layer_split: bool + pflash_on_layer_split: bool + expert_offload: bool + + +# Mirrors server/src/common/model_capabilities.h. Placement imports this table +# rather than maintaining a second architecture registry. +ARCHITECTURE_CAPABILITIES: Mapping[str, ArchitecturePlacementCapabilities] = ( + MappingProxyType( + { + "qwen35": ArchitecturePlacementCapabilities(True, True, True, True, False), + "qwen35moe": ArchitecturePlacementCapabilities(False, False, False, False, True), + "laguna": ArchitecturePlacementCapabilities(True, False, False, False, True), + "gemma4": ArchitecturePlacementCapabilities(True, False, False, False, False), + "deepseek4": ArchitecturePlacementCapabilities(True, False, False, False, False), + "qwen3": ArchitecturePlacementCapabilities(False, False, False, False, False), + } + ) +) + +_DFLASH = FeatureCapability("DFlash", QUALIFIED_BOTH) +_DSPARK = FeatureCapability("DSpark", QUALIFIED_BOTH) +_PFLASH = FeatureCapability("PFlash", QUALIFIED_BOTH) +_KVFLASH = FeatureCapability("KVFlash", QUALIFIED_BOTH) +_KVFLASH_PREVIEW = FeatureCapability("KVFlash", PREVIEW_BOTH, automatic=False) +_SPARK = FeatureCapability("Spark", QUALIFIED_BOTH) + + +MODEL_OPTIMIZATION_PROFILES: Mapping[str, ModelOptimizationProfile] = MappingProxyType( + { + "qwen3.6-27b": ModelOptimizationProfile( + preset="qwen3.6-27b", + architecture="qwen35", + prefill_baseline="Exact prefill", + decode_baseline="Autoregressive decode", + kv_baseline="Full KV cache", + speculative_decode=_DFLASH, + pflash=PFlashCapability(_PFLASH), + kvflash=KvFlashCapability( + _KVFLASH, + policies=("drafter", "qk", "lru"), + preferred_policy="drafter", + scorerless_policy="qk", + ), + ), + "gemma-4-26b": ModelOptimizationProfile( + preset="gemma-4-26b", + architecture="gemma4", + prefill_baseline="Exact prefill", + decode_baseline="Autoregressive decode", + kv_baseline="Full KV cache", + speculative_decode=_DFLASH, + kvflash=KvFlashCapability( + _KVFLASH_PREVIEW, + policies=("drafter", "lru"), + preferred_policy="drafter", + scorerless_policy="lru", + ), + ), + "gemma-4-31b": ModelOptimizationProfile( + preset="gemma-4-31b", + architecture="gemma4", + prefill_baseline="Exact prefill", + decode_baseline="Autoregressive decode", + kv_baseline="Full KV cache", + speculative_decode=_DFLASH, + kvflash=KvFlashCapability( + _KVFLASH_PREVIEW, + policies=("drafter", "lru"), + preferred_policy="drafter", + scorerless_policy="lru", + ), + ), + "laguna-xs.2": ModelOptimizationProfile( + preset="laguna-xs.2", + architecture="laguna", + prefill_baseline="Exact sparse-attention prefill", + decode_baseline="Autoregressive decode", + kv_baseline="Full hybrid-attention KV cache", + speculative_decode=_DFLASH, + # The cross-tokenizer path launches, but remains opt-in until its + # retrieval/quality qualification is complete on both backends. + kvflash=KvFlashCapability( + _KVFLASH_PREVIEW, + policies=("drafter", "lru"), + preferred_policy="drafter", + scorerless_policy="lru", + ), + spark=SparkCapability(_SPARK), + ), + "qwen3.6-moe": ModelOptimizationProfile( + preset="qwen3.6-moe", + architecture="qwen35moe", + prefill_baseline="Exact prefill", + decode_baseline="Autoregressive decode", + kv_baseline="Full KV cache", + kvflash=KvFlashCapability( + _KVFLASH, + policies=("drafter", "lru"), + preferred_policy="drafter", + scorerless_policy=None, + ), + spark=SparkCapability(_SPARK), + ), + "deepseek-v4-flash": ModelOptimizationProfile( + preset="deepseek-v4-flash", + architecture="deepseek4", + prefill_baseline="Exact MLA prefill", + decode_baseline="Autoregressive decode", + kv_baseline="Native MLA-compressed KV cache", + speculative_decode=_DSPARK, + # Sparse DeepSeek prefill is engine-backed but approximate and + # currently restricted to the monolithic HIP path. It is recorded + # here so the contract is complete, but Automatic does not enable + # it until product qualification is promoted. + deepseek_prefill=DeepSeekPrefillCapability( + feature=FeatureCapability( + "DeepSeek sparse prefill", + HIP_PREVIEW, + automatic=False, + ), + mode="sparse", + ), + ), + } +) + + +def model_profile(preset: str) -> ModelOptimizationProfile: + """Return a preset contract, failing loudly if catalog metadata drifted.""" + try: + return MODEL_OPTIMIZATION_PROFILES[preset] + except KeyError as exc: + raise KeyError(f"no optimization profile is registered for preset {preset!r}") from exc + + +def architecture_capabilities(architecture: str) -> ArchitecturePlacementCapabilities: + """Return legal placement operations; unknown architectures are inert.""" + return ARCHITECTURE_CAPABILITIES.get( + architecture, + ArchitecturePlacementCapabilities(False, False, False, False, False), + ) diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index 426a9fa10..864bcfff9 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -26,6 +26,7 @@ from rich.table import Table import lucebox.autotune as autotune_mod +import lucebox.capabilities as capabilities_mod import lucebox.config as config_mod import lucebox.docker_run as docker_run import lucebox.download as download_mod @@ -621,6 +622,11 @@ def _render_optimization_plan( placement_state = "[green]ready[/green]" if plan.placement.runnable else "[red]blocked[/red]" console.print(f"Execution: [bold]{escape(plan.placement.summary)}[/bold] ({placement_state})") console.print(f"[dim]{escape(plan.placement.reason)}[/dim]") + strategies = " · ".join( + f"{phase}: [bold]{escape(strategy)}[/bold]" + for phase, strategy in plan.phase_strategies + ) + console.print(strategies) table = Table(show_header=True, box=None, pad_edge=False) table.add_column("Optimization", style="bold") table.add_column("State") @@ -628,6 +634,11 @@ def _render_optimization_plan( for decision in plan.decisions: if decision.enabled: state = "[green]ON[/green]" + elif ( + decision.available + and decision.qualification is capabilities_mod.Qualification.PREVIEW + ): + state = "[yellow]preview[/yellow]" elif decision.available: state = "[dim]off[/dim]" else: @@ -651,12 +662,14 @@ def _render_custom_runtime( table = Table(show_header=True, box=None, pad_edge=False) table.add_column("Optimization", style="bold") table.add_column("Selected") - selections = ( + selections = [ ("DFlash", runtime.speculative_decode), ("PFlash", runtime.prefill_mode != "off"), ("KVFlash", runtime.kvflash != "off"), ("Spark", runtime.spark), - ) + ] + if runtime.ds4_prefill != "exact": + selections.append((f"DeepSeek {runtime.ds4_prefill} prefill", True)) for name, enabled in selections: table.add_row(name, "[green]ON[/green]" if enabled else "[dim]off[/dim]") console.print(table) @@ -705,6 +718,8 @@ def _customize_runtime( ) -> DflashRuntime: """Prompt only for product-level choices; keep low-level knobs automatic.""" preset = download_mod.PRESETS.get(cfg.model.preset) + profile = capabilities_mod.model_profile(preset.name) if preset is not None else None + kvflash_capability = profile.kvflash if profile is not None else None runtime = plan.runtime console.print("\n[bold]Customize optimizations[/bold]") console.print("[dim]Context, cache format, and budgets remain hardware-tuned.[/dim]\n") @@ -728,8 +743,21 @@ def _customize_runtime( else: console.print(f"PFlash: [dim]unavailable — {escape(plan.pflash.reason)}[/dim]") + ds4_prefill: Literal["exact", "dense", "sparse"] = "exact" + if plan.prefill_alternative is not None and plan.prefill_alternative.available: + console.print( + "[yellow]DeepSeek sparse prefill is approximate and currently a HIP preview.[/yellow]" + ) + if typer.confirm("Enable DeepSeek sparse prefill?", default=False): + assert profile is not None and profile.deepseek_prefill is not None + ds4_prefill = profile.deepseek_prefill.mode + kvflash = False - kvflash_policy: Literal["drafter", "lru", "qk"] = "drafter" + kvflash_policy: Literal["drafter", "lru", "qk"] = ( + kvflash_capability.preferred_policy + if kvflash_capability is not None + else "drafter" + ) if preset is not None and plan.kvflash.available: kvflash = typer.confirm( "Enable KVFlash bounded long-context memory?", default=plan.kvflash.enabled @@ -738,14 +766,14 @@ def _customize_runtime( console.print(f"KVFlash: [dim]unavailable — {escape(plan.kvflash.reason)}[/dim]") if kvflash: has_scorer = download_mod.optimizer_drafter_installed(cfg) - if has_scorer: - policy_choices = "drafter" - if preset is not None and preset.architecture == "qwen35": - policy_choices += "/qk/lru" - else: - policy_choices += "/lru" + policies = kvflash_capability.policies if kvflash_capability is not None else () + if has_scorer and policies and kvflash_capability is not None: + policy_choices = "/".join(policies) answer = ( - typer.prompt(f"KVFlash policy ({policy_choices})", default="drafter") + typer.prompt( + f"KVFlash policy ({policy_choices})", + default=kvflash_capability.preferred_policy, + ) .strip() .lower() ) @@ -759,21 +787,35 @@ def _customize_runtime( kvflash_policy = "lru" else: kvflash_policy = "drafter" - elif preset is not None and preset.architecture == "qwen35": - kvflash_policy = "qk" - console.print("KVFlash policy: [bold]qk[/bold] (no extra scorer required)") else: - console.print( - "[yellow]KVFlash would use recency-only LRU without the shared scorer; " - "older context can be evicted.[/yellow]" + scorerless_policy = ( + kvflash_capability.scorerless_policy + if kvflash_capability is not None + else None ) - if typer.confirm("Install the scorer instead?", default=True): - if _ensure_optimizer_drafter(cfg): - kvflash_policy = "drafter" - else: - kvflash_policy = "lru" + if scorerless_policy == "qk": + kvflash_policy = "qk" + console.print("KVFlash policy: [bold]qk[/bold] (no extra scorer required)") else: - kvflash_policy = "lru" + if scorerless_policy == "lru": + console.print( + "[yellow]Without the shared scorer, KVFlash uses recency-only LRU; " + "older context can be evicted.[/yellow]" + ) + if typer.confirm("Install the scorer instead?", default=True): + if _ensure_optimizer_drafter(cfg): + kvflash_policy = "drafter" + elif scorerless_policy is not None: + kvflash_policy = scorerless_policy + else: + kvflash = False + elif scorerless_policy is not None: + kvflash_policy = scorerless_policy + else: + kvflash = False + console.print( + "[yellow]KVFlash left off because this model requires its scorer.[/yellow]" + ) spark = False if plan.spark.available: @@ -793,13 +835,22 @@ def _customize_runtime( speculative_decode=speculative_decode, lazy=False if not speculative_decode else runtime.lazy, prefill_mode="auto" if pflash else "off", - prefill_keep_ratio=0.10 if pflash else runtime.prefill_keep_ratio, - prefill_threshold=32768, + prefill_keep_ratio=( + profile.pflash.keep_ratio + if pflash and profile is not None and profile.pflash is not None + else runtime.prefill_keep_ratio + ), + prefill_threshold=( + profile.pflash.minimum_context + if profile is not None and profile.pflash is not None + else runtime.prefill_threshold + ), prefill_drafter=prefill_drafter, kvflash="auto" if kvflash else "off", kvflash_policy=kvflash_policy, spark=spark, spark_vram_gb=0.0, + ds4_prefill=ds4_prefill, fa_window=0 if kvflash else runtime.fa_window, ) diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py index aaedcc3c3..ae4d2170c 100644 --- a/lucebox/src/lucebox/config.py +++ b/lucebox/src/lucebox/config.py @@ -93,6 +93,13 @@ def _cast_kvflash_policy(v: Any) -> Literal["drafter", "lru", "qk"]: return cast(Literal["drafter", "lru", "qk"], value) +def _cast_ds4_prefill(v: Any) -> Literal["exact", "dense", "sparse"]: + value = str(v).strip().lower() + if value not in {"exact", "dense", "sparse"}: + raise ValueError(f"ds4_prefill must be exact/dense/sparse, got {value!r}") + return cast(Literal["exact", "dense", "sparse"], value) + + def _cast_positive_int(v: Any) -> int: value = int(v) if value <= 0: @@ -221,6 +228,7 @@ def _cast_float_tuple(v: Any) -> tuple[float, ...]: "dflash.kvflash_tau": (("dflash", "kvflash_tau"), _cast_positive_int), "dflash.spark": (("dflash", "spark"), _cast_bool), "dflash.spark_vram_gb": (("dflash", "spark_vram_gb"), _cast_nonnegative_float), + "dflash.ds4_prefill": (("dflash", "ds4_prefill"), _cast_ds4_prefill), "dflash.think_max": (("dflash", "think_max"), int), "dflash.fa_window": (("dflash", "fa_window"), int), "dflash.think_soft_close_min_ratio": ( @@ -372,6 +380,7 @@ def _from_dict(raw: dict[str, Any]) -> Config: kvflash_tau=_cast_positive_int(df.get("kvflash_tau", 64)), spark=_cast_bool(df.get("spark", False)), spark_vram_gb=_cast_nonnegative_float(df.get("spark_vram_gb", 0.0)), + ds4_prefill=_cast_ds4_prefill(df.get("ds4_prefill", "exact")), think_max=int(df.get("think_max", 15488)), fa_window=int(df.get("fa_window", 0)), think_soft_close_min_ratio=_cast_think_soft_close_min_ratio( diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index d8c29ae67..e973535a9 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -388,6 +388,8 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: env.append(("DFLASH_SPARK", "1")) if cfg.dflash.spark_vram_gb > 0.0: env.append(("DFLASH_SPARK_VRAM_GB", f"{cfg.dflash.spark_vram_gb:g}")) + if cfg.dflash.ds4_prefill != "exact": + env.append(("DFLASH_DS4_PREFILL", cfg.dflash.ds4_prefill)) # fa_window=0 is the server's own default (full attention); only emit # the env when the operator has selected a sparse decode window. The # entrypoint mirrors this guard so an unset env reproduces the diff --git a/lucebox/src/lucebox/download.py b/lucebox/src/lucebox/download.py index 2673a52e2..c50f8fec3 100644 --- a/lucebox/src/lucebox/download.py +++ b/lucebox/src/lucebox/download.py @@ -173,7 +173,10 @@ def featured(self) -> bool: approx_target_gb=18.9, approx_draft_gb=1.1, architecture="laguna", - native_context=4096, + # Poolside's model contract is 262K. Hardware-aware planning may pick + # a smaller exact-cache context, but the catalog must not erase the + # model's native capability (the previous 4K value did exactly that). + native_context=262144, display_name="Laguna XS.2", featured_rank=3, description=( diff --git a/lucebox/src/lucebox/placement.py b/lucebox/src/lucebox/placement.py index 29011a8da..b4ae3d533 100644 --- a/lucebox/src/lucebox/placement.py +++ b/lucebox/src/lucebox/placement.py @@ -13,32 +13,11 @@ from dataclasses import dataclass, replace from pathlib import Path +from lucebox.capabilities import architecture_capabilities from lucebox.topology import GpuDevice, HardwareTopology, from_config from lucebox.types import Config, DflashRuntime, PlacementRuntime -@dataclass(frozen=True, slots=True) -class ArchitecturePlacementCapabilities: - layer_split: bool - remote_draft: bool - draft_on_layer_split: bool - pflash_on_layer_split: bool - expert_offload: bool - - -# Mirrors server/src/common/model_capabilities.h for the product-level paths -# the packaged planner can actually launch. A parity test keeps this table -# explicit and reviewable when an engine architecture changes. -ARCHITECTURE_CAPABILITIES: dict[str, ArchitecturePlacementCapabilities] = { - "qwen35": ArchitecturePlacementCapabilities(True, True, True, True, False), - "qwen35moe": ArchitecturePlacementCapabilities(False, False, False, False, True), - "laguna": ArchitecturePlacementCapabilities(True, False, False, False, True), - "gemma4": ArchitecturePlacementCapabilities(True, False, False, False, False), - "deepseek4": ArchitecturePlacementCapabilities(True, False, False, False, False), - "qwen3": ArchitecturePlacementCapabilities(False, False, False, False, False), -} - - @dataclass(frozen=True, slots=True) class PlacementOption: key: str @@ -86,9 +65,12 @@ def _model_sizes(cfg: Config, preset: object, has_draft: bool) -> tuple[float, f def _primary_capacity(device: GpuDevice) -> float: - # Discrete GPUs need less host-side reserve; UMA devices share their pool - # with the OS and CPU engine buffers and therefore keep a wider margin. - reserve = 8.0 if device.unified_memory else 2.0 + # ``topology`` already removes 16 GB from a UMA device's shared system + # memory for the OS and CPU-side engine buffers. Subtracting another UMA + # reserve here would double-count that margin and incorrectly reject the + # Strix-specific 102.4 GB DeepSeek target on a 128 GB machine. Discrete + # devices still retain a small allocator/working-buffer reserve. + reserve = 0.0 if device.unified_memory else 2.0 return max(0.0, float(device.effective_memory_gb) - reserve) @@ -230,10 +212,28 @@ def automatic_placement( ) architecture = str(getattr(preset, "architecture", "")) - capabilities = ARCHITECTURE_CAPABILITIES.get( - architecture, - ArchitecturePlacementCapabilities(False, False, False, False, False), - ) + capabilities = architecture_capabilities(architecture) + base = PlacementRuntime(mode="single", target_device=primary.placement_name) + if runtime.ds4_prefill != "exact": + if architecture != "deepseek4": + reason = "DeepSeek prefill modes can only be applied to a DeepSeek4 target" + elif primary.backend != "hip": + reason = ( + f"DeepSeek {runtime.ds4_prefill} prefill is a HIP-only preview; " + f"the selected primary uses {primary.backend}" + ) + else: + reason = "" + if reason: + return PlacementPlan( + runtime=base, + optimization_runtime=runtime, + summary=primary.name, + reason=reason, + runnable=False, + topology=topology, + options=(), + ) secondary = _best_secondary(topology) target_gb, draft_gb = _model_sizes(cfg, preset, has_draft) scorer_gb = ( @@ -256,7 +256,6 @@ def automatic_placement( options: list[PlacementOption] = [ _single_option(primary, target_fits=target_fits, stack_fits=stack_fits) ] - base = PlacementRuntime(mode="single", target_device=primary.placement_name) if secondary is None: if stack_fits: return PlacementPlan( @@ -273,10 +272,15 @@ def automatic_placement( optimization_runtime=runtime, summary=primary.name, reason=( - "the target fits, but its selected optional workloads exceed this GPU's " - "safe memory budget" - if target_fits - else "the selected target exceeds this GPU's safe memory budget" + "approximate DeepSeek prefill requires the target to fit one HIP device; " + "choose exact prefill to enable layer splitting" + if runtime.ds4_prefill != "exact" and not target_fits + else ( + "the target fits, but its selected optional workloads exceed this GPU's " + "safe memory budget" + if target_fits + else "the selected target exceeds this GPU's safe memory budget" + ) ), runnable=False, topology=topology, @@ -331,7 +335,11 @@ def automatic_placement( target_gb, max(0.0, split_primary_capacity), ) - split_available = capabilities.layer_split and bool(split_devices) + split_available = ( + capabilities.layer_split + and bool(split_devices) + and runtime.ds4_prefill == "exact" + ) options.append( PlacementOption( key="layer-split", @@ -342,9 +350,14 @@ def automatic_placement( "capacity mode for a target that does not fit the primary GPU" if split_available else ( - f"{architecture or 'this model'} has no engine layer-split path" - if not capabilities.layer_split - else "a compatible runtime or enough combined safe memory is unavailable" + "approximate DeepSeek prefill is monolithic-only; choose exact prefill " + "to use a target layer split" + if runtime.ds4_prefill != "exact" and bool(split_devices) + else ( + f"{architecture or 'this model'} has no engine layer-split path" + if not capabilities.layer_split + else "a compatible runtime or enough combined safe memory is unavailable" + ) ) ), ) @@ -455,11 +468,16 @@ def automatic_placement( optimization_runtime=runtime, summary=primary.name, reason=( - "the target fits the primary, but its optional workloads do not and no validated " - "secondary-device placement is available" - if target_fits - else "the target does not fit the primary and no validated secondary-device " - "placement is available" + "approximate DeepSeek prefill requires the target to fit one HIP device; " + "choose exact prefill to enable layer splitting" + if runtime.ds4_prefill != "exact" and not target_fits + else ( + "the target fits the primary, but its optional workloads do not and no validated " + "secondary-device placement is available" + if target_fits + else "the target does not fit the primary and no validated secondary-device " + "placement is available" + ) ), runnable=False, topology=topology, diff --git a/lucebox/src/lucebox/types.py b/lucebox/src/lucebox/types.py index 0e1946526..b0982e725 100644 --- a/lucebox/src/lucebox/types.py +++ b/lucebox/src/lucebox/types.py @@ -248,6 +248,10 @@ class DflashRuntime: kvflash_tau: int = 64 spark: bool = False spark_vram_gb: float = 0.0 + # DeepSeek4's architecture-specific prefill implementation. ``exact`` is + # the quality-safe default. ``dense`` and ``sparse`` are approximate, + # monolithic-HIP preview paths and are never selected silently. + ds4_prefill: Literal["exact", "dense", "sparse"] = "exact" # Phase-1 (thinking) cap when a request opts into thinking. Default mirrors # antirez/ds4 ds4_eval.c: think_max_tokens = max_tokens - hard_limit_reply # budget = 16000 - 512 = 15488. The server's own hardcoded default is 10000. @@ -320,6 +324,11 @@ def __post_init__(self) -> None: "spark_vram_gb must be finite and zero (automatic) or positive; " f"got {self.spark_vram_gb!r}" ) + if self.ds4_prefill not in {"exact", "dense", "sparse"}: + raise ValueError( + "ds4_prefill must be exact, dense, or sparse; " + f"got {self.ds4_prefill!r}" + ) if self.prefill_threshold <= 0: raise ValueError(f"prefill_threshold must be positive; got {self.prefill_threshold!r}") if self.think_max < 0: diff --git a/lucebox/tests/test_autotune.py b/lucebox/tests/test_autotune.py index b74094d39..b284abf81 100644 --- a/lucebox/tests/test_autotune.py +++ b/lucebox/tests/test_autotune.py @@ -103,6 +103,11 @@ def test_qwen_24gb_enables_dflash_and_pflash_but_prefers_full_kv( assert plan.runtime.prefill_keep_ratio == 0.10 assert plan.runtime.max_ctx == 98304 assert plan.active_names == ("DFlash", "PFlash") + assert dict(plan.phase_strategies) == { + "Prefill": "PFlash", + "Decode": "DFlash", + "KV cache": "Full KV cache", + } def test_published_draft_is_not_claimed_until_it_is_installed(tmp_path: Path) -> None: @@ -181,11 +186,52 @@ def test_laguna_uses_native_context_and_spark_under_24gb_pressure( plan = automatic_plan(cfg, optimizer_drafter_available=False) - assert plan.runtime.max_ctx == 4096 + # Laguna advertises 262K native context. The planner keeps an exact-cache + # 32K cap on this tight 24 GB profile rather than the old, incorrect 4K + # catalog limit. + assert plan.runtime.max_ctx == 32768 assert plan.dflash.enabled is False assert plan.pflash.available is False assert plan.kvflash.enabled is False + assert plan.kvflash.available is True + assert plan.kvflash.qualification.value == "preview" assert plan.spark.enabled is True + assert dict(plan.phase_strategies) == { + "Prefill": "Exact sparse-attention prefill", + "Decode": "Autoregressive decode", + "KV cache": "Full hybrid-attention KV cache", + } + + +def test_deepseek_reports_native_mla_instead_of_generic_flash_features( + tmp_path: Path, +) -> None: + cfg = Config( + variant="rocm", + models_dir=tmp_path / "models", + host=HostFacts( + gpu_vendor="amd", + vram_gb=128, + ram_gb=256, + gpu_sm="gfx1201", + ), + model=ModelMeta(preset="deepseek-v4-flash"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.pflash.available is False + assert plan.kvflash.available is False + assert plan.prefill_alternative is not None + assert plan.prefill_alternative.available is True + assert plan.prefill_alternative.enabled is False + assert plan.runtime.ds4_prefill == "exact" + assert "Native MLA-compressed KV cache" in plan.kvflash.reason + assert dict(plan.phase_strategies) == { + "Prefill": "Exact MLA prefill", + "Decode": "Autoregressive decode", + "KV cache": "Native MLA-compressed KV cache", + } def test_spark_stays_off_when_host_ram_cannot_hold_cold_experts( diff --git a/lucebox/tests/test_capabilities.py b/lucebox/tests/test_capabilities.py new file mode 100644 index 000000000..0b77feb46 --- /dev/null +++ b/lucebox/tests/test_capabilities.py @@ -0,0 +1,58 @@ +"""Integrity tests for the product optimization contract.""" + +import pytest +from lucebox.capabilities import ( + ARCHITECTURE_CAPABILITIES, + MODEL_OPTIMIZATION_PROFILES, + Qualification, + model_profile, +) +from lucebox.download import PRESETS + + +def test_every_catalog_model_has_exactly_one_optimization_profile() -> None: + assert set(MODEL_OPTIMIZATION_PROFILES) == set(PRESETS) + + +@pytest.mark.parametrize("preset_name", PRESETS) +def test_profile_matches_catalog_architecture_and_covers_every_phase( + preset_name: str, +) -> None: + preset = PRESETS[preset_name] + profile = model_profile(preset_name) + + assert profile.preset == preset_name + assert profile.architecture == preset.architecture + assert profile.architecture in ARCHITECTURE_CAPABILITIES + assert all(label.strip() for label in profile.phase_baselines) + + +def test_deepseek_uses_native_mla_instead_of_claiming_generic_flash_paths() -> None: + profile = model_profile("deepseek-v4-flash") + + assert profile.pflash is None + assert profile.kvflash is None + assert "MLA" in profile.prefill_baseline + assert "MLA" in profile.kv_baseline + assert profile.deepseek_prefill is not None + assert ( + profile.deepseek_prefill.feature.qualification_on("hip") + is Qualification.PREVIEW + ) + assert profile.deepseek_prefill.feature.automatic_on("hip") is False + assert profile.deepseek_prefill.mode == "sparse" + + +def test_laguna_kvflash_is_explicitly_preview_not_an_automatic_claim() -> None: + profile = model_profile("laguna-xs.2") + + assert profile.kvflash is not None + assert profile.kvflash.policies == ("drafter", "lru") + assert profile.kvflash.feature.available_on("cuda") is True + assert profile.kvflash.feature.available_on("hip") is True + assert profile.kvflash.feature.automatic_on("cuda") is False + assert profile.kvflash.feature.automatic_on("hip") is False + + +def test_laguna_catalog_uses_the_published_native_context() -> None: + assert PRESETS["laguna-xs.2"].native_context == 262_144 diff --git a/lucebox/tests/test_config.py b/lucebox/tests/test_config.py index ac9a8bb1a..398340099 100644 --- a/lucebox/tests/test_config.py +++ b/lucebox/tests/test_config.py @@ -334,6 +334,7 @@ def test_optimization_fields_round_trip_and_validate(tmp_path: Path) -> None: "kvflash_tau = 96\n" "spark = true\n" "spark_vram_gb = 14.5\n" + 'ds4_prefill = "sparse"\n' ) loaded = config.load(path) @@ -344,6 +345,7 @@ def test_optimization_fields_round_trip_and_validate(tmp_path: Path) -> None: assert loaded.dflash.kvflash_tau == 96 assert loaded.dflash.spark is True assert loaded.dflash.spark_vram_gb == 14.5 + assert loaded.dflash.ds4_prefill == "sparse" @pytest.mark.parametrize( @@ -356,6 +358,7 @@ def test_optimization_fields_round_trip_and_validate(tmp_path: Path) -> None: ("dflash.spark_vram_gb", "-1"), ("dflash.spark_vram_gb", "nan"), ("dflash.spark_vram_gb", "inf"), + ("dflash.ds4_prefill", "fast"), ], ) def test_config_rejects_invalid_optimization_values(tmp_path: Path, key: str, value: str) -> None: diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index f8412e56a..8c43ab73a 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -397,6 +397,7 @@ def test_server_run_spec_optional_env_off_by_default(tmp_path: Path) -> None: "DFLASH_KVFLASH_TAU", "DFLASH_SPARK", "DFLASH_SPARK_VRAM_GB", + "DFLASH_DS4_PREFILL", "DFLASH_FA_WINDOW", "DFLASH_THINK_SOFT_CLOSE_MIN_RATIO", "DFLASH_DEBUG_THINKING_LOGITS", @@ -423,6 +424,7 @@ def test_server_run_spec_optional_env_emitted_when_set(tmp_path: Path) -> None: kvflash_tau=96, spark=True, spark_vram_gb=14.5, + ds4_prefill="sparse", think_soft_close_min_ratio=0.5, debug_thinking_logits=True, ), @@ -440,6 +442,7 @@ def test_server_run_spec_optional_env_emitted_when_set(tmp_path: Path) -> None: assert env["DFLASH_KVFLASH_TAU"] == "96" assert env["DFLASH_SPARK"] == "1" assert env["DFLASH_SPARK_VRAM_GB"] == "14.5" + assert env["DFLASH_DS4_PREFILL"] == "sparse" assert env["DFLASH_THINK_SOFT_CLOSE_MIN_RATIO"] == "0.5" assert env["DFLASH_DEBUG_THINKING_LOGITS"] == "1" diff --git a/lucebox/tests/test_placement.py b/lucebox/tests/test_placement.py index d9c8cbf63..1371068e8 100644 --- a/lucebox/tests/test_placement.py +++ b/lucebox/tests/test_placement.py @@ -5,9 +5,10 @@ import pytest from lucebox.autotune import automatic_plan +from lucebox.capabilities import ARCHITECTURE_CAPABILITIES from lucebox.docker_run import server_run_spec from lucebox.host_facts import from_env -from lucebox.placement import ARCHITECTURE_CAPABILITIES, automatic_placement +from lucebox.placement import automatic_placement from lucebox.topology import from_config from lucebox.types import Config, DflashRuntime, HostFacts, ModelMeta @@ -79,6 +80,39 @@ def test_r9700_strix_keeps_fitting_qwen_target_on_r9700(tmp_path: Path) -> None: assert "full stack fits" in plan.placement.reason +def test_strix_128gb_does_not_double_reserve_memory_for_deepseek( + tmp_path: Path, +) -> None: + host = HostFacts( + gpu_vendor="amd", + has_amd_gpu=True, + gpu_name="Radeon 8060S", + gpu_count=1, + vram_gb=125, + gpu_sm="gfx1151", + ram_gb=125, + amd_gpu_name="Radeon 8060S", + amd_gpu_count=1, + amd_vram_gb=125, + amd_gpu_arch="gfx1151", + amd_gpu_list_csv=_amd_csv(("Radeon 8060S", "gfx1151", 0)), + ) + cfg = Config( + variant="rocm", + models_dir=tmp_path, + host=host, + model=ModelMeta(preset="deepseek-v4-flash"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.placement.runnable is True + assert plan.placement.runtime.mode == "single" + assert plan.placement.runtime.target_device == "hip:0" + assert plan.prefill_alternative is not None + assert plan.prefill_alternative.available is True + + def test_three_same_backend_gpus_are_used_only_when_capacity_requires_them( tmp_path: Path, ) -> None: @@ -383,6 +417,69 @@ def test_deepseek_uses_r9700_and_strix_capacity_automatically(tmp_path: Path) -> assert plan.placement.runtime.target_devices == ("hip:0", "hip:1") assert sum(plan.placement.runtime.target_layer_split) == pytest.approx(1.0) assert plan.runtime.speculative_decode is False + assert plan.prefill_alternative is not None + assert plan.prefill_alternative.available is False + assert "requires the target to fit one HIP device" in plan.prefill_alternative.reason + + +def test_deepseek_sparse_prefill_runs_only_when_target_fits_one_hip_gpu( + tmp_path: Path, +) -> None: + host = HostFacts( + gpu_vendor="amd", + has_amd_gpu=True, + gpu_name="Large HIP GPU", + vram_gb=128, + gpu_sm="gfx1201", + ram_gb=256, + amd_gpu_list_csv=_amd_csv(("Large HIP GPU", "gfx1201", 131072)), + ) + cfg = Config( + variant="rocm", + models_dir=tmp_path, + host=host, + model=ModelMeta(preset="deepseek-v4-flash"), + ) + + placement = automatic_placement( + cfg, + DflashRuntime(speculative_decode=False, ds4_prefill="sparse"), + download.PRESETS["deepseek-v4-flash"], + has_draft=False, + optimizer_drafter_available=False, + ) + + assert placement.runnable is True + assert placement.runtime.mode == "single" + assert placement.runtime.target_device == "hip:0" + + +def test_deepseek_sparse_prefill_rejects_required_layer_split(tmp_path: Path) -> None: + host = HostFacts( + gpu_vendor="amd", + has_amd_gpu=True, + gpu_name="AMD Radeon AI PRO R9700", + gpu_count=2, + vram_gb=31, + gpu_sm="gfx1201", + ram_gb=125, + amd_gpu_list_csv=_amd_csv( + ("AMD Radeon AI PRO R9700", "gfx1201", 32624), + ("AMD Radeon Graphics", "gfx1151", 512), + ), + ) + cfg = Config(variant="rocm", models_dir=tmp_path, host=host) + + placement = automatic_placement( + cfg, + DflashRuntime(speculative_decode=False, ds4_prefill="sparse"), + download.PRESETS["deepseek-v4-flash"], + has_draft=False, + optimizer_drafter_available=False, + ) + + assert placement.runnable is False + assert "requires the target to fit one HIP device" in placement.reason def test_deepseek_uses_rtx_strix_paired_runtime_automatically(tmp_path: Path) -> None: diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index 7e51a3a09..3f1468d09 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -261,7 +261,11 @@ cat > "$native_tmp/bin/cmake" <<'STUB' #!/usr/bin/env bash printf '%s\n' "$*" >> "${LUCEBOX_TEST_CMAKE_LOG:?}" STUB -chmod +x "$native_tmp/bin/cmake" +cat > "$native_tmp/bin/ninja" <<'STUB' +#!/usr/bin/env bash +exit 0 +STUB +chmod +x "$native_tmp/bin/cmake" "$native_tmp/bin/ninja" native_log="$native_tmp/cmake.log" native_out="" native_rc=0 @@ -281,6 +285,7 @@ if [ "$native_rc" -ne 0 ]; then report fail "native CUDA build dispatch" "exit=$native_rc output=$(head -3 <<<"$native_out")" elif ! grep -qF -- "-DDFLASH27B_GPU_BACKEND=cuda" "$native_log" \ || ! grep -qF -- "-DCMAKE_CUDA_ARCHITECTURES=86" "$native_log" \ + || ! grep -qF -- "-G Ninja" "$native_log" \ || ! grep -qF -- "--target dflash_server backend_ipc_daemon -j 2" "$native_log"; then report fail "native CUDA build dispatch" "unexpected cmake argv: $(tr '\n' ' ' < "$native_log")" else @@ -305,7 +310,8 @@ native_out=$( if [ "$native_rc" -ne 0 ]; then report fail "native ROCm build dispatch" "exit=$native_rc output=$(head -3 <<<"$native_out")" elif ! grep -qF -- "-DDFLASH27B_GPU_BACKEND=hip" "$native_log" \ - || ! grep -qF -- "-DDFLASH27B_HIP_ARCHITECTURES=gfx1201" "$native_log"; then + || ! grep -qF -- "-DDFLASH27B_HIP_ARCHITECTURES=gfx1201" "$native_log" \ + || ! grep -qF -- "-G Ninja" "$native_log"; then report fail "native ROCm build dispatch" "unexpected cmake argv: $(tr '\n' ' ' < "$native_log")" else report ok "native ROCm build dispatch" @@ -901,13 +907,14 @@ test_entrypoint_optimization_flags() { DFLASH_KVFLASH_TAU=96 \ DFLASH_SPARK=1 \ DFLASH_SPARK_VRAM_GB=14 \ + DFLASH_DS4_PREFILL=sparse \ timeout 10 bash "$ENTRYPOINT" serve 2>&1 || true) rm -rf "$sandbox" local required missing="" for required in \ "--prefill-drafter" "--prefill-compression auto" \ "--kvflash auto" "--kvflash-policy qk" "--kvflash-tau 96" \ - "--spark" "--spark-vram 14"; do + "--spark" "--spark-vram 14" "--ds4-prefill sparse"; do [[ "$out" == *"$required"* ]] || missing+=" $required" done if [ -n "$missing" ]; then @@ -919,7 +926,7 @@ test_entrypoint_optimization_flags() { fi } test_entrypoint_optimization_flags \ - "entrypoint forwards PFlash, KVFlash, and Spark exactly once" + "entrypoint forwards optimization modes exactly once" test_entrypoint_placement_flags() { local label="$1" sandbox draft_dir models_dir bin_dir shim_dir out diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index 423112016..383f29ac7 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -367,6 +367,7 @@ fi : "${DFLASH_KVFLASH_TAU:=64}" : "${DFLASH_SPARK:=0}" : "${DFLASH_SPARK_VRAM_GB:=0}" +: "${DFLASH_DS4_PREFILL:=exact}" # Optional server default for requests that omit max_tokens. When unset, # the C++ server uses the model-card default. : "${DFLASH_DEFAULT_MAX_TOKENS:=}" @@ -680,6 +681,11 @@ if [ "$DFLASH_SPARK" = "1" ]; then *) CMD+=(--spark-vram "$DFLASH_SPARK_VRAM_GB") ;; esac fi +case "$DFLASH_DS4_PREFILL" in + exact) ;; + dense|sparse) CMD+=(--ds4-prefill "$DFLASH_DS4_PREFILL") ;; + *) die "DFLASH_DS4_PREFILL must be exact, dense, or sparse" ;; +esac # Soft-close ratio: emit only when nonzero. The default-string compare # guards against the floating-point quirks of `[` numeric tests for # values like 0.0/0/0.00 — anything non-"0.0" passes through to the diff --git a/variables.md b/variables.md index fb06395b6..98a418b05 100644 --- a/variables.md +++ b/variables.md @@ -137,6 +137,7 @@ Untagged variables are operational tuning knobs. | Variable | Purpose | |---|---| +| `DFLASH_DS4_PREFILL` | DeepSeek4 prefill mode: `exact` (default), `dense`, or `sparse`. Approximate modes require a monolithic HIP placement. | | `DFLASH_FP_USE_BSA` | Use block-sparse attention in flash prefill. | | `DFLASH_FP_ALPHA` | FlashPrefill alpha parameter. | | `DFLASH_FP_CHUNK_S` | FlashPrefill chunk size. | From 8215391cd9972e89d761d9f42bda945f18df8efd Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sat, 1 Aug 2026 13:01:43 +0200 Subject: [PATCH 17/25] feat(lucebox): finalize model-aware automatic runtime --- harness/clients/common.sh | 16 +- lucebox.sh | 325 +++++++++++++++++++++------- lucebox/README.md | 8 +- lucebox/src/lucebox/autotune.py | 10 + lucebox/src/lucebox/capabilities.py | 11 + lucebox/src/lucebox/cli.py | 102 ++++++++- lucebox/src/lucebox/config.py | 15 +- lucebox/src/lucebox/docker_run.py | 43 +++- lucebox/src/lucebox/download.py | 110 +++++++--- lucebox/src/lucebox/host_facts.py | 62 ++++++ lucebox/src/lucebox/topology.py | 20 +- lucebox/src/lucebox/types.py | 24 +- lucebox/tests/test_autotune.py | 257 +++++++++++++++++++++- lucebox/tests/test_config.py | 35 +++ lucebox/tests/test_docker_run.py | 53 ++++- lucebox/tests/test_download.py | 37 +++- lucebox/tests/test_models_cli.py | 89 ++++++++ lucebox/tests/test_placement.py | 102 ++++++++- scripts/test_lucebox_sh.sh | 155 +++++++++++-- 19 files changed, 1298 insertions(+), 176 deletions(-) diff --git a/harness/clients/common.sh b/harness/clients/common.sh index d8dd28d8a..2a21d9065 100755 --- a/harness/clients/common.sh +++ b/harness/clients/common.sh @@ -33,6 +33,9 @@ elif [[ "$TARGET_WAS_EXPLICIT" == "1" ]]; then else DRAFT="$DEFAULT_DRAFT" fi +# ``external`` is the production CLI contract: the canonical Lucebox launch +# path owns the engine and this harness only exercises the client protocol. +# ``lucebox`` and ``llamacpp`` remain explicit standalone comparison modes. MODEL_SERVER="${MODEL_SERVER:-lucebox}" DFLASH_SERVER_BIN="${DFLASH_SERVER_BIN:-$REPO_DIR/server/build/dflash_server}" LLAMA_BUILD_DIR="${LLAMA_BUILD_DIR:-$CLIENT_WORK_DIR/llama-cpp-server-build}" @@ -77,7 +80,7 @@ fi STAMP="${STAMP:-$(date +%Y%m%d-%H%M%S)}" BASE_URL="http://$HOST:$PORT" LOG_DIR="$RUN_DIR/$STAMP" -SERVER_LOG="$LOG_DIR/server.log" +SERVER_LOG="${SERVER_LOG:-$LOG_DIR/server.log}" mkdir -p "$LOG_DIR" @@ -113,12 +116,15 @@ draft_enabled() { } start_lucebox_server() { + if [[ "$MODEL_SERVER" == "external" ]]; then + return 0 + fi if [[ "$MODEL_SERVER" == "llamacpp" ]]; then start_llamacpp_server return fi if [[ "$MODEL_SERVER" != "lucebox" ]]; then - echo "unknown MODEL_SERVER=$MODEL_SERVER; expected lucebox or llamacpp" >&2 + echo "unknown MODEL_SERVER=$MODEL_SERVER; expected external, lucebox, or llamacpp" >&2 return 1 fi start_dflash_native_server @@ -298,7 +304,8 @@ wait_lucebox_server() { return 0 fi sleep 1 - if ! kill -0 "$SERVER_PID" 2>/dev/null; then + if [[ "$MODEL_SERVER" != "external" ]] \ + && ! kill -0 "$SERVER_PID" 2>/dev/null; then echo "server exited early; log: $SERVER_LOG" >&2 tail -n 160 "$SERVER_LOG" >&2 || true return 1 @@ -310,6 +317,9 @@ wait_lucebox_server() { } stop_lucebox_server() { + if [[ "$MODEL_SERVER" == "external" ]]; then + return 0 + fi if [[ -n "${SERVER_PID:-}" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then kill "$SERVER_PID" 2>/dev/null || true wait "$SERVER_PID" 2>/dev/null || true diff --git a/lucebox.sh b/lucebox.sh index 9ccae2a92..92559d84a 100755 --- a/lucebox.sh +++ b/lucebox.sh @@ -33,16 +33,19 @@ set -euo pipefail -# ROCm packages on minimal Ubuntu images install their tools under -# /opt/rocm/bin without necessarily adding that directory to a login shell's -# PATH. Keep probing and native builds self-contained; this changes PATH only -# for the wrapper process and its children. -if [ -d /opt/rocm/bin ]; then +# GPU SDK packages on minimal Ubuntu images do not always update a login +# shell's PATH. Keep probing and native builds self-contained; this changes +# PATH only for the wrapper process and its children. +for lucebox_toolchain_bin in /opt/rocm/bin /usr/local/cuda/bin; do + [ -d "$lucebox_toolchain_bin" ] || continue case ":${PATH:-}:" in - *:/opt/rocm/bin:*) ;; - *) export PATH="/opt/rocm/bin:${PATH:-/usr/bin:/bin}" ;; + *:"$lucebox_toolchain_bin":*) ;; + # Respect explicit caller paths (including test shims and operator + # toolchain pins); SDK defaults are fallbacks, not overrides. + *) export PATH="${PATH:-/usr/bin:/bin}:$lucebox_toolchain_bin" ;; esac -fi +done +unset lucebox_toolchain_bin VERSION="0.2.0" SCRIPT_PATH="$(readlink -f "$0" 2>/dev/null || realpath "$0" 2>/dev/null || echo "$0")" @@ -63,6 +66,8 @@ UNIT_PATH="${XDG_CONFIG_HOME:-$HOME/.config}/systemd/user/$UNIT_NAME" # CUDA driver floor for the prebuilt CUDA 12 image. # shellcheck disable=SC2034 MIN_DRIVER_CUDA12=525 +MIN_DRIVER_CUDA128=570 +MIN_DRIVER_CUDA13=580 # Canonical source of `lucebox.sh`. The bootstrap installer (`install.sh`) # rewrites this line at install time to record which URL the user actually @@ -241,6 +246,7 @@ CONNECTOR_SELECTION_FILE="$CONNECTOR_STATE_DIR/selected" : "${LUCEBOX_HOST_NVIDIA_VRAM_GB:=0}" : "${LUCEBOX_HOST_NVIDIA_GPU_ARCH:=}" : "${LUCEBOX_HOST_NVIDIA_GPU_LIST_CSV:=}" +: "${LUCEBOX_HOST_NVIDIA_UNIFIED_MEMORY:=0}" : "${LUCEBOX_HOST_ROCM_VERSION:=}" : "${LUCEBOX_HOST_HAS_KFD:=0}" : "${LUCEBOX_HOST_HAS_DRI:=0}" @@ -463,6 +469,7 @@ probe_host() { LUCEBOX_HOST_NVIDIA_VRAM_GB=0 LUCEBOX_HOST_NVIDIA_GPU_ARCH="" LUCEBOX_HOST_NVIDIA_GPU_LIST_CSV="" + LUCEBOX_HOST_NVIDIA_UNIFIED_MEMORY=0 LUCEBOX_HOST_ROCM_VERSION="" LUCEBOX_HOST_HAS_KFD=0 LUCEBOX_HOST_HAS_DRI=0 @@ -484,14 +491,28 @@ probe_host() { LUCEBOX_HOST_GPU_VENDOR="nvidia" LUCEBOX_HOST_HAS_NVIDIA_GPU=1 LUCEBOX_HOST_GPU_NAME=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $1}') - local mem_mib - mem_mib=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $2}') - LUCEBOX_HOST_VRAM_GB=$((mem_mib / 1024)) LUCEBOX_HOST_DRIVER_VERSION=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $3}') LUCEBOX_HOST_DRIVER_MAJOR=${LUCEBOX_HOST_DRIVER_VERSION%%.*} - local cc + local cc mem_mib cc=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $4}') LUCEBOX_HOST_GPU_SM="${cc//./}" + mem_mib=$(printf '%s\n' "$q" | head -1 | awk -F', ' '{print $2}') + if [[ "$mem_mib" =~ ^[0-9]+$ ]]; then + LUCEBOX_HOST_VRAM_GB=$((mem_mib / 1024)) + elif [ "$LUCEBOX_HOST_GPU_SM" = "121" ] \ + && [[ "$LUCEBOX_HOST_GPU_NAME" == *GB10* ]] \ + && [ "$LUCEBOX_HOST_RAM_GB" -gt 16 ]; then + # GB10 exposes one coherent CPU/GPU memory pool. NVML reports + # memory.total as [N/A], while the CUDA runtime sees nearly all + # system RAM. Reserve 16 GB for the OS and CPU-side buffers, + # matching the Strix UMA planner policy. + LUCEBOX_HOST_VRAM_GB=$((LUCEBOX_HOST_RAM_GB - 16)) + LUCEBOX_HOST_NVIDIA_UNIFIED_MEMORY=1 + else + # Unknown/non-numeric NVML memory must never enter arithmetic + # or be guessed as system RAM on a discrete card. + LUCEBOX_HOST_VRAM_GB=0 + fi LUCEBOX_HOST_GPU_COUNT=$(printf '%s\n' "$q" | wc -l) LUCEBOX_HOST_NVIDIA_GPU_NAME="$LUCEBOX_HOST_GPU_NAME" LUCEBOX_HOST_NVIDIA_GPU_COUNT="$LUCEBOX_HOST_GPU_COUNT" @@ -760,6 +781,7 @@ probe_host() { export LUCEBOX_HOST_NVIDIA_GPU_NAME LUCEBOX_HOST_NVIDIA_GPU_COUNT export LUCEBOX_HOST_NVIDIA_VRAM_GB LUCEBOX_HOST_NVIDIA_GPU_ARCH export LUCEBOX_HOST_NVIDIA_GPU_LIST_CSV + export LUCEBOX_HOST_NVIDIA_UNIFIED_MEMORY export LUCEBOX_HOST_HAS_HYBRID_RUNTIME export LUCEBOX_HOST_HYBRID_SERVER_BIN LUCEBOX_HOST_HYBRID_IPC_BIN export LUCEBOX_HOST_HYBRID_DFLASH_DIR LUCEBOX_HOST_HYBRID_ENTRYPOINT @@ -784,12 +806,27 @@ pick_variant() { fi configured=$(_lucebox_config_get image.variant) if [ -n "$configured" ]; then + # cuda12 used to contain sm_120. New releases keep that newer toolkit + # in cuda128 so RTX 20/30/40 users retain the r525 driver floor. Migrate + # the old moving tag in memory; explicit LUCEBOX_VARIANT still wins. + if [ "$configured" = "cuda12" ]; then + ensure_probed + if [ "$LUCEBOX_HOST_GPU_SM" = "120" ]; then + printf 'cuda128' + return + fi + if [ "$LUCEBOX_HOST_GPU_SM" = "121" ]; then + case "$(uname -m 2>/dev/null || echo unknown)" in + aarch64|arm64) printf 'cuda13'; return ;; + esac + fi + fi printf '%s' "$configured" return fi ensure_probed if [ "$LUCEBOX_HOST_HAS_NVIDIA_GPU" = "1" ]; then - printf 'cuda12' + _default_cuda_variant elif [ "$LUCEBOX_HOST_HAS_AMD_GPU" = "1" ]; then printf 'rocm' else @@ -799,6 +836,24 @@ pick_variant() { fi } +_default_cuda_variant() { + # Each Blackwell target needs a newer compiler/runtime than the broad + # CUDA-12 image. Keep that newer driver floor isolated from Turing through + # Hopper hosts, which remain on cuda12. + local machine + machine=$(uname -m 2>/dev/null || echo unknown) + if [ "$LUCEBOX_HOST_GPU_SM" = "121" ]; then + case "$machine" in + aarch64|arm64) printf 'cuda13'; return ;; + esac + fi + if [ "$LUCEBOX_HOST_GPU_SM" = "120" ]; then + printf 'cuda128' + return + fi + printf 'cuda12' +} + _variant_is_rocm() { # Variant names are not limited to the moving `rocm` tag. Releases and # CI produce tags such as `0.3.0-rocm` and `pr-335-rocm`; treat any tag @@ -811,6 +866,20 @@ _variant_is_rocm() { esac } +_variant_is_cuda13() { + case "$1" in + *[Cc][Uu][Dd][Aa]13*) return 0 ;; + *) return 1 ;; + esac +} + +_variant_is_cuda128() { + case "$1" in + *[Cc][Uu][Dd][Aa]128*) return 0 ;; + *) return 1 ;; + esac +} + # ── prereq checks (host-only) ───────────────────────────────────────────── # Print-and-exit on anything that needs root to install. The Python CLI does # the richer reporting; this is the bare minimum to make `docker run` viable. @@ -1164,11 +1233,7 @@ cmd_hybrid_serve() { export DFLASH_SERVER_BIN="$LUCEBOX_HOST_HYBRID_SERVER_BIN" export DFLASH_BACKEND_IPC_BIN="$LUCEBOX_HOST_HYBRID_IPC_BIN" export DFLASH_TARGET="$target" - if [ "$draft" = "none" ]; then - export DFLASH_DRAFT="$DEFAULT_MODELS_DIR/.lucebox-no-draft" - else - export DFLASH_DRAFT="$draft" - fi + _export_selected_decode_companion "$model_id" "$draft" export DFLASH_HOST="${LUCEBOX_NATIVE_HOST:-127.0.0.1}" export DFLASH_PORT="$DEFAULT_PORT" export DFLASH_MODEL_NAME="$model_id" @@ -1583,7 +1648,7 @@ cmd_pull() { _native_backend() { local requested="${1:-}" variant case "$requested" in - cuda|cuda12|nvidia) printf 'cuda'; return ;; + cuda|cuda12|cuda128|cuda13|nvidia) printf 'cuda'; return ;; rocm|hip|amd) printf 'rocm'; return ;; "") variant=$(pick_variant) @@ -1673,6 +1738,7 @@ _selected_model_paths() { gemma-4-31b) target_file="google_gemma-4-31B-it-Q4_K_M.gguf" ;; laguna-xs.2) target_file="laguna-xs2-Q4_K_M.gguf" ;; qwen3.6-moe) target_file="Qwen3.6-35B-A3B-UD-Q4_K_M.gguf" ;; + deepseek-v4-flash) target_file="DeepSeek-V4-Flash-ROCMFP2-STRIX.gguf" ;; esac fi if [ -z "$draft_file" ]; then @@ -1680,6 +1746,7 @@ _selected_model_paths() { qwen3.6-27b) draft_file="dflash-draft-3.6-q4_k_m.gguf" ;; gemma-4-26b) draft_file="gemma-4-26B-A4B-it-DFlash-q8_0.gguf" ;; gemma-4-31b) draft_file="gemma-4-31B-it-DFlash-q8_0.gguf" ;; + deepseek-v4-flash) draft_file="DeepSeek-V4-Flash-DSpark-draft-Q4RMFP4-denseF16.gguf" ;; esac fi speculative_decode=$(_lucebox_config_get dflash.speculative_decode) @@ -1698,12 +1765,33 @@ _selected_model_paths() { && [ "$speculative_decode" != "0" ] \ && [ "$speculative_decode" != "no" ] \ && [ "$speculative_decode" != "off" ] \ - && [ -d "$DEFAULT_MODELS_DIR/draft/laguna-xs2-speculator" ]; then + && [ -s "$DEFAULT_MODELS_DIR/draft/laguna-xs2-speculator/model.safetensors" ] \ + && [ -s "$DEFAULT_MODELS_DIR/draft/laguna-xs2-speculator/config.json" ]; then draft_path="$DEFAULT_MODELS_DIR/draft/laguna-xs2-speculator" fi printf '%s\n%s\n%s\n' "$target_path" "$draft_path" "${preset:-lucebox}" } +_export_selected_decode_companion() { + # Generic DFlash consumes --draft. DeepSeek's architecture-specific + # DSpark backend consumes a separate switch/path and deliberately ignores + # --draft, so native and packaged-hybrid launches must mirror the Python + # container planner instead of relying on generic draft discovery. + local model_id="$1" draft="$2" + unset DFLASH_DS4_SPEC DFLASH_DS4_DRAFT + if [ "$model_id" = "deepseek-v4-flash" ]; then + export DFLASH_DRAFT="$DEFAULT_MODELS_DIR/.lucebox-no-draft" + if [ "$draft" != "none" ]; then + export DFLASH_DS4_SPEC=1 + export DFLASH_DS4_DRAFT="$draft" + fi + elif [ "$draft" = "none" ]; then + export DFLASH_DRAFT="$DEFAULT_MODELS_DIR/.lucebox-no-draft" + else + export DFLASH_DRAFT="$draft" + fi +} + # ── engine client connectors ───────────────────────────────────────────── # # A connector uses a client the user already installed and points only that @@ -1839,22 +1927,25 @@ _connector_api_root() { } _connector_api_ready() { - local api_root="$1" response + local api_root="$1" expected_model="${2:-}" response compact command -v curl >/dev/null 2>&1 || return 1 response=$(curl -fsS --connect-timeout 2 --max-time 4 \ "$api_root/v1/models" 2>/dev/null) || return 1 # A generic 200 response from another process on the configured port is # not sufficient. Require an OpenAI-compatible model catalog shape. - [[ "$response" == *'"data"'* || "$response" == *'"models"'* ]] + [[ "$response" == *'"data"'* || "$response" == *'"models"'* ]] || return 1 + [ -n "$expected_model" ] || return 0 + compact=$(printf '%s' "$response" | tr -d '[:space:]') + [[ "$compact" == *"\"id\":\"$expected_model\""* ]] } _connector_ensure_api() { - local api_root="$1" i state + local api_root="$1" expected_model="${2:-}" i state if ! command -v curl >/dev/null 2>&1; then err "curl is required to verify the local Lucebox API" return 1 fi - if _connector_api_ready "$api_root"; then + if _connector_api_ready "$api_root" "$expected_model"; then return 0 fi ensure_probed @@ -1866,7 +1957,7 @@ _connector_ensure_api() { bash "$SCRIPT_PATH" start || return $? state=running for i in 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15; do - _connector_api_ready "$api_root" && return 0 + _connector_api_ready "$api_root" "$expected_model" && return 0 sleep 1 done fi @@ -2248,7 +2339,7 @@ EOF return 0 fi - _connector_ensure_api "$api_root" || return $? + _connector_ensure_api "$api_root" "$model" || return $? info "Opening $label with the Lucebox model" exec "${CONNECTOR_LAUNCH_ARGV[@]}" } @@ -2501,11 +2592,7 @@ cmd_native_serve() { export DFLASH_SERVER_BIN="$binary" export DFLASH_BACKEND_IPC_BIN="$build_dir/backend_ipc_daemon" export DFLASH_TARGET="$target" - if [ "$draft" = "none" ]; then - export DFLASH_DRAFT="$DEFAULT_MODELS_DIR/.lucebox-no-draft" - else - export DFLASH_DRAFT="$draft" - fi + _export_selected_decode_companion "$model_id" "$draft" export DFLASH_HOST="${LUCEBOX_NATIVE_HOST:-127.0.0.1}" export DFLASH_PORT="$DEFAULT_PORT" export DFLASH_MODEL_NAME="$model_id" @@ -2514,8 +2601,39 @@ cmd_native_serve() { exec "$repo/server/scripts/entrypoint.sh" serve } +HARNESS_ENGINE_PID="" +HARNESS_ENGINE_LOG="" + +_stop_harness_engine() { + if [ -n "$HARNESS_ENGINE_PID" ] \ + && kill -0 "$HARNESS_ENGINE_PID" 2>/dev/null; then + kill "$HARNESS_ENGINE_PID" 2>/dev/null || true + wait "$HARNESS_ENGINE_PID" 2>/dev/null || true + fi + HARNESS_ENGINE_PID="" +} + +_wait_for_harness_engine() { + local api_root="$1" model="$2" attempts=0 + while [ "$attempts" -lt 300 ]; do + attempts=$((attempts + 1)) + _connector_api_ready "$api_root" "$model" && return 0 + if [ -n "$HARNESS_ENGINE_PID" ] \ + && ! kill -0 "$HARNESS_ENGINE_PID" 2>/dev/null; then + err "the inference engine exited while loading" + [ ! -f "$HARNESS_ENGINE_LOG" ] \ + || tail -n 120 "$HARNESS_ENGINE_LOG" >&2 + return 1 + fi + sleep 1 + done + err "the inference engine did not become ready within 5 minutes" + [ ! -f "$HARNESS_ENGINE_LOG" ] || tail -n 120 "$HARNESS_ENGINE_LOG" >&2 + return 1 +} + cmd_harness() { - local repo name="${1:-}" backend build_dir binary script selected=() target draft + local repo name="${1:-}" backend script api_root model max_ctx log_dir rc=0 repo=$(_find_repo_root) \ || die "harness launchers are contributor tools; run this command inside the lucebox repository" if [ -z "$name" ]; then @@ -2545,40 +2663,45 @@ EOF script="$repo/harness/clients/run_${name}.sh" [ -x "$script" ] || die "harness launcher is missing: $script" ensure_probed - if _config_requires_hybrid_runtime; then - die "this profile uses the paired CUDA + HIP runtime; start it with '$SCRIPT_NAME native', then connect the harness with '$SCRIPT_NAME connect $name'" - fi backend=$(_native_backend "${LUCEBOX_HARNESS_BACKEND:-}") - build_dir=$(_native_build_dir "$repo" "$backend") - binary="$build_dir/dflash_server" - [ -x "$binary" ] \ - || die "native engine is not built — run '$SCRIPT_NAME build $backend' first" - mapfile -t selected < <(_selected_model_paths) - target="${selected[0]:-}" - draft="${selected[1]:-none}" - _model_artifact_ready "$target" \ - || die "selected target is not installed: $target — run '$SCRIPT_NAME models select'" + api_root=$(_connector_api_root) + model=$(_connector_model_id) + max_ctx=$(_connector_context_size) + log_dir="$CONFIG_HOME/logs" + _connector_private_dir "$log_dir" + HARNESS_ENGINE_LOG="$log_dir/harness-engine.log" + + if _connector_api_ready "$api_root"; then + _connector_api_ready "$api_root" "$model" \ + || die "the API at $api_root is running a different model; stop it or select its advertised model" + # We do not own the reused process and therefore cannot promise a log + # path to the harness. Avoid exposing a stale log from an earlier + # CLI-owned launch if the client later reports an API error. + HARNESS_ENGINE_LOG=/dev/null + info "Reusing the running Lucebox API at $api_root" + else + # The canonical native path owns model resolution, placement, and all + # optimization flags. Harness scripts are protocol/client adapters; + # they must not reconstruct a second, divergent engine command line. + info "Starting the optimized Lucebox engine for $name" + bash "$SCRIPT_PATH" native "$backend" >"$HARNESS_ENGINE_LOG" 2>&1 & + HARNESS_ENGINE_PID=$! + trap _stop_harness_engine EXIT + _wait_for_harness_engine "$api_root" "$model" || return $? + fi - _export_native_config export REPO_DIR="$repo" - export DFLASH_SERVER_BIN="$binary" - export DFLASH_TARGET="$target" - export TARGET="$target" - export DFLASH_DRAFT="$draft" - export DRAFT="$draft" - if [ "$draft" != "none" ]; then - # Match `lucebox native` and the container entrypoint: a selected - # DFlash draft uses DDTree by default. An explicit contributor - # VERIFY_MODE remains authoritative for comparison experiments. - export VERIFY_MODE="${VERIFY_MODE:-ddtree}" - fi - local max_ctx budget - max_ctx=$(_lucebox_config_get dflash.max_ctx) - budget=$(_lucebox_config_get dflash.budget) - [ -z "$max_ctx" ] || export MAX_CTX="$max_ctx" - [ -z "$budget" ] || export BUDGET="$budget" - info "Launching $name with the selected Lucebox model" - exec "$script" + export MODEL_SERVER=external + export HOST=127.0.0.1 + export PORT="$DEFAULT_PORT" + export MODEL_ID="$model" + export MAX_CTX="$max_ctx" + export SERVER_LOG="$HARNESS_ENGINE_LOG" + info "Launching $name against the selected Lucebox model" + "$script" || rc=$? + _stop_harness_engine + trap - EXIT + return "$rc" } cmd_update() { @@ -2788,18 +2911,46 @@ cmd_check() { none|*) _row 0 "nvidia ctk" "not installed — https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html" ;; esac - if [ "$LUCEBOX_HOST_DRIVER_MAJOR" -ge "$MIN_DRIVER_CUDA12" ]; then - _row 1 "nvidia driver" "$LUCEBOX_HOST_DRIVER_VERSION (≥ $MIN_DRIVER_CUDA12 required for cuda12)" + local required_driver cuda_label memory_label + if _variant_is_cuda13 "$variant"; then + required_driver=$MIN_DRIVER_CUDA13 + cuda_label=cuda13 + elif _variant_is_cuda128 "$variant"; then + required_driver=$MIN_DRIVER_CUDA128 + cuda_label=cuda128 else - _row 0 "nvidia driver" "$LUCEBOX_HOST_DRIVER_VERSION (< $MIN_DRIVER_CUDA12 — cuda12 image will fail)" + required_driver=$MIN_DRIVER_CUDA12 + cuda_label=cuda12 + fi + if [ "$LUCEBOX_HOST_DRIVER_MAJOR" -ge "$required_driver" ]; then + _row 1 "nvidia driver" "$LUCEBOX_HOST_DRIVER_VERSION (≥ $required_driver required for $cuda_label)" + else + _row 0 "nvidia driver" "$LUCEBOX_HOST_DRIVER_VERSION (< $required_driver — $cuda_label image will fail)" + fi + memory_label="${LUCEBOX_HOST_VRAM_GB} GB VRAM" + if [ "$LUCEBOX_HOST_NVIDIA_UNIFIED_MEMORY" = "1" ]; then + memory_label="${LUCEBOX_HOST_VRAM_GB} GB usable shared memory" + fi + _row 1 "nvidia gpu" "$LUCEBOX_HOST_GPU_NAME × $LUCEBOX_HOST_GPU_COUNT (sm_$LUCEBOX_HOST_GPU_SM, $memory_label)" + if _variant_is_cuda13 "$variant"; then + case "$LUCEBOX_HOST_GPU_SM" in + 121) _row 1 "cuda13 arch" "sm_121 covered by the arm64 image" ;; + "") _row warn "cuda13 arch" "compute capability not detected" ;; + *) _row warn "cuda13 arch" "sm_$LUCEBOX_HOST_GPU_SM is not covered by the arm64 image (121)" ;; + esac + elif _variant_is_cuda128 "$variant"; then + case "$LUCEBOX_HOST_GPU_SM" in + 120) _row 1 "cuda128 arch" "sm_120 covered by the CUDA 12.8 image" ;; + "") _row warn "cuda128 arch" "compute capability not detected" ;; + *) _row warn "cuda128 arch" "sm_$LUCEBOX_HOST_GPU_SM is not covered by the CUDA 12.8 image (120)" ;; + esac + else + case "$LUCEBOX_HOST_GPU_SM" in + 75|80|86|89|90) _row 1 "cuda12 arch" "sm_$LUCEBOX_HOST_GPU_SM covered by image" ;; + "") _row warn "cuda12 arch" "compute capability not detected" ;; + *) _row warn "cuda12 arch" "sm_$LUCEBOX_HOST_GPU_SM not in image arch list (75;80;86;89;90)" ;; + esac fi - _row 1 "nvidia gpu" "$LUCEBOX_HOST_GPU_NAME × $LUCEBOX_HOST_GPU_COUNT (sm_$LUCEBOX_HOST_GPU_SM, ${LUCEBOX_HOST_VRAM_GB} GB VRAM)" - # cuda12 image arch coverage: sm_75;80;86;89;90;120 (see docker-bake.hcl) - case "$LUCEBOX_HOST_GPU_SM" in - 75|80|86|89|90|120) _row 1 "cuda12 arch" "sm_$LUCEBOX_HOST_GPU_SM covered by image" ;; - "") _row warn "cuda12 arch" "compute capability not detected" ;; - *) _row warn "cuda12 arch" "sm_$LUCEBOX_HOST_GPU_SM not in image arch list (75;80;86;89;90;120)" ;; - esac elif command -v nvidia-smi &>/dev/null; then _row 0 "nvidia driver" "nvidia-smi present but NVML calls fail — driver/library mismatch, try reboot" fi @@ -3063,6 +3214,24 @@ _menu_restart_if_running() { fi } +_ensure_configured_image() { + # Model activation may switch backends on a mixed machine when the current + # backend cannot place the selected model. Keep guided setup/menu one-step: + # install the newly selected runtime before optimize/start tries to use it. + local selected + selected=$(_lucebox_config_get image.variant) + [ -n "$selected" ] || selected=$(pick_variant) + if docker image inspect "${IMAGE_BASE}:${selected}" >/dev/null 2>&1; then + return 0 + fi + if ! _confirm "Download the ${selected} inference image required by this model?" 1; then + warn "The model is configured, but ${IMAGE_BASE}:${selected} is not installed." + hint "Run '$SCRIPT_NAME pull' before starting the engine." + return 1 + fi + LUCEBOX_VARIANT="$selected" bash "$SCRIPT_PATH" pull +} + cmd_setup() { [ -t 0 ] && [ -t 1 ] \ || die "guided setup needs a terminal — use '$SCRIPT_NAME --help' for non-interactive commands" @@ -3085,7 +3254,7 @@ cmd_setup() { IFS= read -r backend_choice || return 1 case "$backend_choice" in 2) variant=rocm ;; - *) variant=cuda12 ;; + *) variant=$(_default_cuda_variant) ;; esac fi info "Selected backend: $variant" @@ -3109,6 +3278,7 @@ cmd_setup() { warn "No model was selected; setup stopped without starting the engine." return 1 fi + _ensure_configured_image || return $? printf '\n' if _confirm "Enable recommended optimizations (may add a ~1.2 GB shared scorer)?" 1; then @@ -3316,7 +3486,10 @@ cmd_menu() { case "$choice" in 1) cmd_setup; _menu_pause ;; 2) - if _menu_run models select; then _menu_restart_if_running; fi + if _menu_run models select \ + && _ensure_configured_image; then + _menu_restart_if_running + fi _menu_pause ;; 3) @@ -3389,7 +3562,7 @@ Misc: Environment overrides: LUCEBOX_IMAGE image name without tag (default: ghcr.io/luce-org/lucebox-hub) - LUCEBOX_VARIANT image tag override (default: cuda12 on NVIDIA, rocm on AMD) + LUCEBOX_VARIANT image tag override (default: cuda13 on GB10, cuda128 on RTX 5090, cuda12 on other NVIDIA, rocm on AMD) LUCEBOX_PORT host port for the server (default: 8080) LUCEBOX_CONTAINER server container name (default: lucebox) LUCEBOX_MODELS host model directory (default: \$XDG_DATA_HOME/lucebox/models) diff --git a/lucebox/README.md b/lucebox/README.md index e6a3bccb0..c8ac35d60 100644 --- a/lucebox/README.md +++ b/lucebox/README.md @@ -13,8 +13,9 @@ which invokes the package in the appropriate container: lucebox optimize --advanced # override individual optimizations lucebox start -The wrapper inventories the host and selects CUDA for NVIDIA builds (including -RTX 3090 + Strix) or ROCm for AMD builds (including R9700 + Strix). The package +The wrapper inventories the host and selects the compatible CUDA image for +NVIDIA builds (`cuda12` through Hopper, `cuda128` for RTX 5090, `cuda13` for +GB10), or ROCm for AMD builds (including R9700 + Strix). The package then handles readiness checks, TOML configuration, model selection and download, optimization and placement settings, and construction of the final server command. Host facts are passed through `LUCEBOX_HOST_*` environment @@ -38,6 +39,9 @@ are labeled and never enabled silently. In particular, DeepSeek sparse prefill is an explicit approximate HIP preview and exact MLA prefill remains Automatic. A factory-preloaded Lucebox can ship the models, shared optimizer scorer, and paired runtime, so a buyer does not download or compile during first setup. +If model placement must switch backend on a mixed machine, activation persists +that backend atomically with its optimization plan and guided setup installs +the corresponding image before continuing. Spark is enabled automatically only when an MoE model is under GPU-memory pressure and the machine reports at least 32 GB of host RAM; its cold experts diff --git a/lucebox/src/lucebox/autotune.py b/lucebox/src/lucebox/autotune.py index e6a24169f..991c1fc5b 100644 --- a/lucebox/src/lucebox/autotune.py +++ b/lucebox/src/lucebox/autotune.py @@ -123,6 +123,12 @@ def draft_available(cfg: Config, preset: object) -> bool: root = cfg.models_dir / "draft" / str(speculator_dir) if not root.is_dir(): return False + required_files = tuple(getattr(preset, "speculator_files", ())) + if required_files: + return all( + local_artifact_present(root / str(filename)) + for filename in required_files + ) try: return any( local_artifact_present(candidate) @@ -168,6 +174,8 @@ def _cap_exact_context( the planner below. Other models keep exact cache semantics in Automatic mode and receive a smaller context when weights leave little working room. """ + if profile.compact_native_kv: + return runtime kvflash = profile.kvflash if kvflash is not None and kvflash.feature.automatic_on(backend): return runtime @@ -221,6 +229,7 @@ def _capacity_fallbacks( prefill_mode="off", prefill_keep_ratio=defaults.prefill_keep_ratio, prefill_threshold=defaults.prefill_threshold, + prefill_cache_slots=defaults.prefill_cache_slots, prefill_drafter=( current.prefill_drafter if current.kvflash != "off" and current.kvflash_policy == "drafter" @@ -427,6 +436,7 @@ def automatic_plan( prefill_mode="auto", prefill_keep_ratio=pflash_capability.keep_ratio, prefill_threshold=pflash_capability.minimum_context, + prefill_cache_slots=pflash_capability.exact_cache_slots, prefill_drafter=download_mod.optimizer_drafter_container_path(), ) elif pflash_profile: diff --git a/lucebox/src/lucebox/capabilities.py b/lucebox/src/lucebox/capabilities.py index 1cd79d803..ade2c9f99 100644 --- a/lucebox/src/lucebox/capabilities.py +++ b/lucebox/src/lucebox/capabilities.py @@ -84,12 +84,18 @@ class PFlashCapability: feature: FeatureCapability minimum_context: int = 32_768 keep_ratio: float = 0.10 + # Exact raw-prompt matches can skip both the scorer and target prefill. + # Keep this pool deliberately small: unlike turn-boundary snapshots, a + # full PFlash snapshot is useful only when the complete request repeats. + exact_cache_slots: int = 4 def __post_init__(self) -> None: if self.minimum_context <= 0: raise ValueError("PFlash minimum context must be positive") if not 0.0 < self.keep_ratio <= 1.0: raise ValueError("PFlash keep ratio must be in (0.0, 1.0]") + if self.exact_cache_slots <= 0: + raise ValueError("PFlash exact-cache slots must be positive") @dataclass(frozen=True, slots=True) @@ -149,6 +155,10 @@ class ModelOptimizationProfile: prefill_baseline: str decode_baseline: str kv_baseline: str + # The architecture stores an exact but substantially compressed native KV + # representation. It still grows with context, but does not need the blunt + # full-KV headroom cap used for ordinary attention caches. + compact_native_kv: bool = False speculative_decode: FeatureCapability | None = None pflash: PFlashCapability | None = None kvflash: KvFlashCapability | None = None @@ -282,6 +292,7 @@ class ArchitecturePlacementCapabilities: prefill_baseline="Exact MLA prefill", decode_baseline="Autoregressive decode", kv_baseline="Native MLA-compressed KV cache", + compact_native_kv=True, speculative_decode=_DSPARK, # Sparse DeepSeek prefill is engine-backed but approximate and # currently restricted to the monolithic HIP path. It is recorded diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index 864bcfff9..03044c2da 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -33,7 +33,7 @@ import lucebox.host_check as host_check from lucebox import __version__ from lucebox.config import config_get, config_set, config_unset, live_config -from lucebox.host_facts import from_env +from lucebox.host_facts import compatible_variant, for_variant, from_env, nvidia_variant from lucebox.placement import PlacementPlan from lucebox.types import Config, DflashRuntime, ModelMeta @@ -96,13 +96,21 @@ def _load_or_build() -> Config: try: cfg = config_mod.load() if cfg is None: - return live_config() + cfg = live_config() + variant = compatible_variant(cfg.host, cfg.variant) + return replace(cfg, variant=variant, host=for_variant(cfg.host, variant)) # Host facts exported by the wrapper take precedence over an absent or # stale persisted snapshot. A zero-filled environment means the CLI was # invoked directly, so retain any snapshot already in the config. live_host = from_env() host = live_host if live_host.vram_gb > 0 or live_host.nproc > 0 else cfg.host - return config_mod.overlay_env(replace(cfg, host=host)) + overlaid = config_mod.overlay_env(replace(cfg, host=host)) + variant = compatible_variant(overlaid.host, overlaid.variant) + return replace( + overlaid, + variant=variant, + host=for_variant(overlaid.host, variant), + ) except (OSError, ValueError) as exc: error_console.print(f"[red]Invalid configuration:[/red] {escape(str(exc))}") raise typer.Exit(code=2) from exc @@ -340,12 +348,47 @@ def _model_meta(preset: download_mod.ModelPreset) -> ModelMeta: ) +def _alternative_variants(cfg: Config) -> tuple[str, ...]: + """Return detected backend variants other than the currently selected one.""" + current = cfg.variant.casefold() + variants: list[str] = [] + if cfg.host.has_nvidia_gpu and "cuda" not in current: + variants.append(nvidia_variant(cfg.host)) + if cfg.host.has_amd_gpu and "rocm" not in current and "hip" not in current: + variants.append("rocm") + return tuple(variants) + + +def _automatic_preset_plan( + cfg: Config, + preset: download_mod.ModelPreset, +) -> tuple[Config, autotune_mod.OptimizationPlan]: + """Keep the active backend when possible, otherwise try detected peers.""" + model = _model_meta(preset) + selected = replace(cfg, model=model) + plan = autotune_mod.automatic_plan(selected) + if plan.placement.runnable: + return selected, plan + + for variant in _alternative_variants(cfg): + candidate = replace( + cfg, + variant=variant, + host=for_variant(cfg.host, variant), + model=model, + ) + candidate_plan = autotune_mod.automatic_plan(candidate) + if candidate_plan.placement.runnable: + return candidate, candidate_plan + return selected, plan + + def _preset_placement(cfg: Config, preset: download_mod.ModelPreset) -> PlacementPlan: """Plan the placement that activating ``preset`` would persist.""" selected_cfg = replace(cfg, model=_model_meta(preset)) if config_mod.optimization_mode() == "custom": return autotune_mod.placement_for_runtime(selected_cfg, cfg.dflash) - return autotune_mod.automatic_plan(selected_cfg).placement + return _automatic_preset_plan(cfg, preset)[1].placement def _require_runnable_preset(cfg: Config, preset: download_mod.ModelPreset) -> None: @@ -401,7 +444,7 @@ def _activate_preset(cfg: Config, preset: download_mod.ModelPreset) -> None: ) return - plan = autotune_mod.automatic_plan(selected_cfg) + selected_cfg, plan = _automatic_preset_plan(cfg, preset) if not plan.placement.runnable: error_console.print( "[red]Model not activated: Automatic found no runnable placement " @@ -412,7 +455,17 @@ def _activate_preset(cfg: Config, preset: download_mod.ModelPreset) -> None: selected_model, plan.runtime, plan.placement.runtime, + variant=selected_cfg.variant if selected_cfg.variant != cfg.variant else None, ) + if selected_cfg.variant != cfg.variant: + console.print( + f"[green]Backend:[/green] switched to {selected_cfg.variant} because " + f"{cfg.variant} cannot place this model on the detected hardware" + ) + console.print( + "[dim]If this runtime image is not installed yet, run `lucebox pull` " + "before starting the engine.[/dim]" + ) console.print(f"[green]Activated:[/green] model.preset = {preset.name}") active = ", ".join(plan.active_names) or "standard engine" console.print( @@ -487,7 +540,15 @@ def models_select( if candidate.name == recommended: labels.append("recommended") placement = _preset_placement(cfg, candidate) - fit = "ready" if placement.runnable else "not enough compatible memory" + primary = placement.topology.primary + backend = ( + "CUDA" if primary is not None and primary.backend == "cuda" else "ROCm" + ) + fit = ( + f"ready on {backend}" + if placement.runnable + else "not enough compatible memory" + ) table.add_row( str(index), candidate.label, @@ -545,8 +606,15 @@ def models_download( activate: Annotated[ bool, typer.Option("--activate", help="Also set as active preset (model.preset).") ] = False, + force: Annotated[ + bool, + typer.Option( + "--force", + help="Download for another machine even when this host cannot run the preset.", + ), + ] = False, ) -> None: - """Fetch a preset's GGUFs into the models dir. + """Fetch a preset's target and decode companion into the models dir. With no argument and no preset configured, recommends one for this host's VRAM tier and auto-activates it (the first-install path). @@ -581,7 +649,10 @@ def models_download( console.print(f"[red]{escape(str(exc))}[/red]") raise typer.Exit(code=2) from exc - if activate: + # A 100+ GB model is an expensive mistake on an incompatible host. Keep + # staging for another machine possible, but make that intent explicit. + # --force never bypasses the activation gate. + if activate or not force: _require_runnable_preset(cfg, pres) current = download_mod.status(cfg, pres) @@ -596,6 +667,12 @@ def models_download( f" draft ({pres.draft_repo}/{pres.draft_file}):" f" {'present' if current['draft_present'] else 'will download'}" ) + elif pres.has_speculator: + files = ", ".join(pres.speculator_files) + console.print( + f" draft ({pres.speculator_repo}: {files}):" + f" {'present' if current['draft_present'] else 'will download'}" + ) else: console.print(" draft [dim](none — target-only preset)[/dim]") @@ -651,6 +728,10 @@ def _render_optimization_plan( f"KV format: [bold]{cache}[/bold] · " f"DFlash budget: [bold]{plan.runtime.budget}[/bold]" ) + console.print( + f"Cache: [bold]{plan.runtime.prefix_cache_slots}[/bold] turn-prefix slots · " + f"[bold]{plan.runtime.prefill_cache_slots}[/bold] exact-prompt slots" + ) def _render_custom_runtime( @@ -845,6 +926,11 @@ def _customize_runtime( if profile is not None and profile.pflash is not None else runtime.prefill_threshold ), + prefill_cache_slots=( + profile.pflash.exact_cache_slots + if pflash and profile is not None and profile.pflash is not None + else 0 + ), prefill_drafter=prefill_drafter, kvflash="auto" if kvflash else "off", kvflash_policy=kvflash_policy, diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py index ae4d2170c..dd104e2c6 100644 --- a/lucebox/src/lucebox/config.py +++ b/lucebox/src/lucebox/config.py @@ -367,7 +367,7 @@ def _from_dict(raw: dict[str, Any]) -> Config: budget=int(df.get("budget", 22)), max_ctx=int(df.get("max_ctx", 16384)), lazy=_cast_bool(df.get("lazy", False)), - prefix_cache_slots=int(df.get("prefix_cache_slots", 0)), + prefix_cache_slots=int(df.get("prefix_cache_slots", 32)), prefill_cache_slots=int(df.get("prefill_cache_slots", 0)), cache_type_k=str(df.get("cache_type_k", "")), cache_type_v=str(df.get("cache_type_v", "")), @@ -381,7 +381,7 @@ def _from_dict(raw: dict[str, Any]) -> Config: spark=_cast_bool(df.get("spark", False)), spark_vram_gb=_cast_nonnegative_float(df.get("spark_vram_gb", 0.0)), ds4_prefill=_cast_ds4_prefill(df.get("ds4_prefill", "exact")), - think_max=int(df.get("think_max", 15488)), + think_max=int(df["think_max"]) if "think_max" in df else None, fa_window=int(df.get("fa_window", 0)), think_soft_close_min_ratio=_cast_think_soft_close_min_ratio( df.get("think_soft_close_min_ratio", 0.0) @@ -428,6 +428,7 @@ def _from_dict(raw: dict[str, Any]) -> Config: nvidia_vram_gb=int(host_raw.get("nvidia_vram_gb", 0)), nvidia_gpu_arch=str(host_raw.get("nvidia_gpu_arch", "")), nvidia_gpu_list_csv=str(host_raw.get("nvidia_gpu_list_csv", "")), + nvidia_unified_memory=_cast_bool(host_raw.get("nvidia_unified_memory", False)), amd_gpu_name=str(host_raw.get("amd_gpu_name", "")), amd_gpu_count=int(host_raw.get("amd_gpu_count", 0)), amd_vram_gb=int(host_raw.get("amd_vram_gb", 0)), @@ -557,6 +558,10 @@ def _write_runtime_doc( doc.pop("dflash", None) for field, value in asdict(runtime).items(): + # Optional fields represent "let the model/server decide" and TOML + # has no null value. Omitting the key preserves that ownership. + if value is None: + continue _doc_set(doc, "dflash", field, _value_to_toml(value)) # Placement belongs to the resolved runtime as one atomic unit. Clearing # it even when an older caller supplies no replacement prevents a stale @@ -604,6 +609,7 @@ def write_model_profile( runtime: DflashRuntime, placement: PlacementRuntime, *, + variant: str | None = None, path: Path | None = None, mode: str = "automatic", source: str = "model+hardware", @@ -621,6 +627,11 @@ def write_model_profile( path = path or default_config_path() doc = load_doc(path) + if variant is not None: + normalized_variant = variant.strip() + if not normalized_variant: + raise ValueError("image variant must not be empty") + _doc_set(doc, "image", "variant", normalized_variant) _doc_set(doc, "model", "preset", preset) _doc_set(doc, "model", "target_file", target_file) if draft_file: diff --git a/lucebox/src/lucebox/docker_run.py b/lucebox/src/lucebox/docker_run.py index e973535a9..1f1f8feda 100644 --- a/lucebox/src/lucebox/docker_run.py +++ b/lucebox/src/lucebox/docker_run.py @@ -91,7 +91,7 @@ def _resolve_model_files(cfg: Config) -> tuple[str, str, str]: draft = cfg.model.draft_file draft_dir = "" if (not target or not draft) and cfg.model.preset: - from lucebox.download import PRESETS + from lucebox.download import PRESETS, local_artifact_present pres = PRESETS.get(cfg.model.preset) if pres is not None: @@ -104,11 +104,24 @@ def _resolve_model_files(cfg: Config) -> tuple[str, str, str]: # is_dir() follows valid directory symlinks. It also rejects # dangling links and links to files, which cannot satisfy the # speculator-directory contract. - if spec_path.is_dir(): + required = tuple(spec_path / filename for filename in pres.speculator_files) + if spec_path.is_dir() and required and all( + local_artifact_present(path) for path in required + ): draft_dir = pres.speculator_dir return target, draft, draft_dir +def _selected_model_architecture(cfg: Config) -> str: + """Return the engine architecture declared by the active preset.""" + if not cfg.model.preset: + return "" + from lucebox.download import PRESETS + + preset = PRESETS.get(cfg.model.preset) + return preset.architecture if preset is not None else "" + + def _runtime_volumes(cfg: Config) -> tuple[BindMount, ...]: """Mount only the writable application data needed by the server.""" models = str(cfg.models_dir.absolute()) @@ -304,9 +317,10 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: ("DFLASH_MAX_CTX", str(cfg.dflash.max_ctx)), ("DFLASH_PREFIX_CACHE_SLOTS", str(cfg.dflash.prefix_cache_slots)), ("DFLASH_PREFILL_CACHE_SLOTS", str(cfg.dflash.prefill_cache_slots)), - ("DFLASH_THINK_MAX", str(cfg.dflash.think_max)), ("DFLASH_PORT", "8080"), ] + if cfg.dflash.think_max is not None: + env.append(("DFLASH_THINK_MAX", str(cfg.dflash.think_max))) env += _placement_env(cfg) # Keep the API model catalog aligned with the preset selected by the host # CLI. Client connectors use this id for explicit model selection and @@ -322,6 +336,7 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: # draft_dir is a subdirectory of models/draft/ holding a safetensors speculator; # it takes effect only when draft_file is empty and the directory exists on disk. target_file, draft_file, draft_dir = _resolve_model_files(cfg) + model_architecture = _selected_model_architecture(cfg) volumes = list(_runtime_volumes(cfg)) if target_file: target_path = _selected_model_path( @@ -334,7 +349,27 @@ def server_run_spec(cfg: Config) -> DockerRunSpec: mounts=volumes, ) env.append(("DFLASH_TARGET", target_path)) - if not cfg.dflash.speculative_decode: + if model_architecture == "deepseek4": + # DeepSeek's DSpark path is architecture-specific. Passing its GGUF as + # generic --draft is explicitly inert for deepseek4; the backend reads + # these two variables instead. Keep generic draft discovery disabled + # so a stale Qwen/Gemma draft cannot be attached by the entrypoint. + env.append(("DFLASH_DRAFT", "/opt/lucebox-hub/server/models/.lucebox-no-draft")) + if cfg.dflash.speculative_decode and draft_file: + draft_path = _selected_model_path( + cfg, + draft_file, + field="model.draft_file", + role="ds4-draft", + under_draft=True, + directory=False, + mounts=volumes, + ) + env += [ + ("DFLASH_DS4_SPEC", "1"), + ("DFLASH_DS4_DRAFT", draft_path), + ] + elif not cfg.dflash.speculative_decode: env.append(("DFLASH_DRAFT", "/opt/lucebox-hub/server/models/.lucebox-no-draft")) elif draft_file: draft_path = _selected_model_path( diff --git a/lucebox/src/lucebox/download.py b/lucebox/src/lucebox/download.py index c50f8fec3..fe86203ec 100644 --- a/lucebox/src/lucebox/download.py +++ b/lucebox/src/lucebox/download.py @@ -73,12 +73,10 @@ class ModelPreset: speculator). In that case the entrypoint runs target-only — DFlash speculative decoding is disabled but the server still works. - ``speculator_dir`` names a directory under ``models/draft/`` that holds - a safetensors-format speculator (e.g. ``model.safetensors``). When - present on disk the server launch sets ``DFLASH_DRAFT`` to that - directory; absent, the server runs target-only. Unlike ``draft_file`` - (which marks the preset as incomplete when missing), ``speculator_dir`` - is optional supplementary hardware and doesn't affect installed_status. + ``speculator_repo`` / ``speculator_files`` describe a safetensors-format + decode companion stored under ``models/draft/``. Some + loaders also require metadata next to ``model.safetensors``; listing every + required file here makes download, status, and activation one contract. """ name: str @@ -93,13 +91,34 @@ class ModelPreset: native_context: int description: str = "" speculator_dir: str | None = None + speculator_repo: str | None = None + speculator_files: tuple[str, ...] = () display_name: str = "" featured_rank: int | None = None + def __post_init__(self) -> None: + speculator_fields = ( + self.speculator_dir is not None, + self.speculator_repo is not None, + bool(self.speculator_files), + ) + if any(speculator_fields) and not all(speculator_fields): + raise ValueError( + "speculator_dir, speculator_repo, and speculator_files must be set together" + ) + @property def has_draft(self) -> bool: return bool(self.draft_repo and self.draft_file) + @property + def has_speculator(self) -> bool: + return bool(self.speculator_dir and self.speculator_repo and self.speculator_files) + + @property + def has_decode_companion(self) -> bool: + return self.has_draft or self.has_speculator + @property def label(self) -> str: """Human-readable catalog label; ``name`` remains the stable CLI id.""" @@ -161,14 +180,14 @@ def featured(self) -> bool: name="laguna-xs.2", target_repo="Lucebox/Laguna-XS.2-GGUF", target_file="laguna-xs2-Q4_K_M.gguf", - # Laguna's DFlash speculator is safetensors-format - # (poolside/Laguna-XS.2-speculator.dflash), downloaded manually - # into models/draft/laguna-xs2-speculator/. The download command - # doesn't fetch it automatically — it's opt-in. When present, - # speculator_dir wires it into DFLASH_DRAFT at server launch. + # Laguna's DFlash companion is safetensors-format. config.json is a + # runtime requirement: the native loader reads rope_theta beside the + # weights and rejects a lone model.safetensors. draft_repo=None, draft_file=None, speculator_dir="laguna-xs2-speculator", + speculator_repo="poolside/Laguna-XS.2-speculator.dflash", + speculator_files=("model.safetensors", "config.json"), approx_total_gb=20, approx_target_gb=18.9, approx_draft_gb=1.1, @@ -180,9 +199,8 @@ def featured(self) -> bool: display_name="Laguna XS.2", featured_rank=3, description=( - "Laguna-XS.2 MoE code model (Q4_K_M). " - "DFlash safetensors speculator in draft/laguna-xs2-speculator/ " - "is used automatically when present." + "Laguna-XS.2 MoE code model (Q4_K_M) + published DFlash " + "safetensors speculator." ), ), "qwen3.6-moe": ModelPreset( @@ -486,15 +504,15 @@ def _fetch( def download_preset(cfg: Config, preset: ModelPreset | None = None) -> int: - """Fetch the target GGUF + (optional) DFlash draft into cfg.models_dir. + """Fetch the target and its published decode companion into cfg.models_dir. Returns 0 on success, non-zero on failure. Verifies each file's size and (LFS) sha256 against the repo metadata before downloading, so a repeat run with the files already on disk is a no-op + sha256 walk. ``preset=None`` resolves to :data:`DEFAULT_PRESET` for back-compat; - presets with ``has_draft=False`` (e.g. Laguna) skip the draft fetch - entirely and let the server run target-only. + GGUF drafts live directly under ``models/draft``; safetensors speculators + retain their required sidecar metadata in a dedicated subdirectory. """ preset = preset or DEFAULT_PRESET console = Console() @@ -511,9 +529,16 @@ def download_preset(cfg: Config, preset: ModelPreset | None = None) -> int: # exactly the predicate that proves these aren't None. assert preset.draft_repo is not None and preset.draft_file is not None _fetch(api, preset.draft_repo, preset.draft_file, draft, console) + elif preset.has_speculator: + assert preset.speculator_repo is not None + assert preset.speculator_dir is not None + speculator_dir = draft / preset.speculator_dir + for filename in preset.speculator_files: + _fetch(api, preset.speculator_repo, filename, speculator_dir, console) else: console.print( - f" [dim]no DFlash draft published for {preset.name} — running target-only[/dim]" + f" [dim]no decode companion published for {preset.name} — " + "running target-only[/dim]" ) except Exception as exc: console.print(f"[red]download failed:[/red] {exc}") @@ -566,28 +591,30 @@ def _local_target_path(cfg: Config, preset: ModelPreset) -> Path: return cfg.models_dir / preset.target_file -def _local_draft_path(cfg: Config, preset: ModelPreset) -> Path | None: - if not (preset.has_draft and preset.draft_file): - return None - return cfg.models_dir / "draft" / preset.draft_file +def _local_decode_companion_paths(cfg: Config, preset: ModelPreset) -> tuple[Path, ...]: + if preset.has_draft and preset.draft_file: + return (cfg.models_dir / "draft" / preset.draft_file,) + if preset.has_speculator and preset.speculator_dir: + root = cfg.models_dir / "draft" / preset.speculator_dir + return tuple(root / filename for filename in preset.speculator_files) + return () def installed_status(cfg: Config, preset: ModelPreset) -> str: """Return ``"installed"`` / ``"partial"`` / ``"absent"`` for a preset. Presence-only — doesn't query the network or hash. ``"installed"`` - requires a non-empty target (and draft when one is published); + requires a non-empty target (and every decode-companion file when one is + published); ``"partial"`` means at least one of the two is present but the set is incomplete. """ target_exists = local_artifact_present(_local_target_path(cfg, preset)) - draft_path = _local_draft_path(cfg, preset) - if draft_path is None: - return "installed" if target_exists else "absent" - draft_exists = local_artifact_present(draft_path) - if target_exists and draft_exists: + companion_paths = _local_decode_companion_paths(cfg, preset) + companion_states = tuple(local_artifact_present(path) for path in companion_paths) + if target_exists and all(companion_states): return "installed" - if target_exists or draft_exists: + if target_exists or any(companion_states): return "partial" return "absent" @@ -601,10 +628,9 @@ def installed_size_gb(cfg: Config, preset: ModelPreset) -> float: total += target.stat().st_size except OSError: pass - draft = _local_draft_path(cfg, preset) - if draft is not None and draft.exists(): + for companion in _local_decode_companion_paths(cfg, preset): try: - total += draft.stat().st_size + total += companion.stat().st_size except OSError: pass return total / 1e9 @@ -627,10 +653,9 @@ def installed_presets(cfg: Config) -> list[ModelPreset]: def status(cfg: Config, preset: ModelPreset | None = None) -> dict[str, bool]: """Quick presence check — what's already on disk? Size-only, no sha256. - For presets without a published DFlash draft, ``draft_present`` is - reported as ``True`` (nothing to fetch → nothing missing). That - keeps the "all present, nothing to do" UX path uniform whether or - not a draft exists. + ``draft_present`` covers either a GGUF draft or every required file in a + safetensors speculator. For target-only presets it remains ``True`` because + there is no companion to fetch. """ preset = preset or DEFAULT_PRESET api = _new_hf_api() @@ -650,6 +675,19 @@ def status(cfg: Config, preset: ModelPreset | None = None) -> dict[str, bool]: out["draft_present"] = local.exists() and local.stat().st_size == size except Exception: out["draft_present"] = False + elif preset.has_speculator: + assert preset.speculator_repo is not None + assert preset.speculator_dir is not None + root = cfg.models_dir / "draft" / preset.speculator_dir + present = True + for filename in preset.speculator_files: + try: + size, _ = _file_meta(api, preset.speculator_repo, filename) + local = root / filename + present = present and local.exists() and local.stat().st_size == size + except Exception: + present = False + out["draft_present"] = present else: out["draft_present"] = True return out diff --git a/lucebox/src/lucebox/host_facts.py b/lucebox/src/lucebox/host_facts.py index f16db336b..a5271c5c1 100644 --- a/lucebox/src/lucebox/host_facts.py +++ b/lucebox/src/lucebox/host_facts.py @@ -10,11 +10,72 @@ from __future__ import annotations import os +from dataclasses import replace from typing import cast from lucebox.types import CtkStatus, GpuVendor, HostFacts +def nvidia_variant(host: HostFacts) -> str: + """Select the published CUDA image matching the detected architecture.""" + architecture = host.nvidia_gpu_arch or ( + host.gpu_sm if host.gpu_vendor == "nvidia" else "" + ) + name = host.nvidia_gpu_name or ( + host.gpu_name if host.gpu_vendor == "nvidia" else "" + ) + if architecture == "121" and "GB10" in name.upper(): + return "cuda13" + if architecture == "120": + return "cuda128" + return "cuda12" + + +def compatible_variant(host: HostFacts, variant: str) -> str: + """Migrate the former broad ``cuda12`` tag on newer architectures.""" + if variant.casefold() == "cuda12": + return nvidia_variant(host) + return variant + + +def for_variant(host: HostFacts, variant: str) -> HostFacts: + """Project a full host inventory onto the backend selected by an image. + + The wrapper normally performs this projection before entering a container. + Keeping the same operation in Python is necessary while switching models: + a CUDA container on an RTX + Strix machine must be able to evaluate whether + the already-detected AMD device can run a model before persisting ``rocm``. + """ + normalized = variant.casefold() + if "rocm" in normalized or "hip" in normalized: + if not (host.has_amd_gpu or host.gpu_vendor == "amd"): + return host + generic_is_amd = host.gpu_vendor == "amd" + return replace( + host, + gpu_vendor="amd", + gpu_name=host.amd_gpu_name or (host.gpu_name if generic_is_amd else ""), + gpu_count=host.amd_gpu_count or (host.gpu_count if generic_is_amd else 0), + vram_gb=host.amd_vram_gb or (host.vram_gb if generic_is_amd else 0), + gpu_sm=host.amd_gpu_arch or (host.gpu_sm if generic_is_amd else ""), + ) + if "cuda" in normalized: + if not (host.has_nvidia_gpu or host.gpu_vendor == "nvidia"): + return host + generic_is_nvidia = host.gpu_vendor == "nvidia" + return replace( + host, + gpu_vendor="nvidia", + gpu_name=host.nvidia_gpu_name or (host.gpu_name if generic_is_nvidia else ""), + gpu_count=( + host.nvidia_gpu_count or (host.gpu_count if generic_is_nvidia else 0) + ), + vram_gb=host.nvidia_vram_gb or (host.vram_gb if generic_is_nvidia else 0), + gpu_sm=host.nvidia_gpu_arch or (host.gpu_sm if generic_is_nvidia else ""), + ) + return host + + def _env_int(key: str, default: int = 0) -> int: raw = os.environ.get(key, "").strip() if not raw: @@ -65,6 +126,7 @@ def from_env() -> HostFacts: nvidia_vram_gb=_env_int("LUCEBOX_HOST_NVIDIA_VRAM_GB"), nvidia_gpu_arch=os.environ.get("LUCEBOX_HOST_NVIDIA_GPU_ARCH", ""), nvidia_gpu_list_csv=os.environ.get("LUCEBOX_HOST_NVIDIA_GPU_LIST_CSV", ""), + nvidia_unified_memory=_env_bool("LUCEBOX_HOST_NVIDIA_UNIFIED_MEMORY"), amd_gpu_name=os.environ.get("LUCEBOX_HOST_AMD_GPU_NAME", ""), amd_gpu_count=_env_int("LUCEBOX_HOST_AMD_GPU_COUNT"), amd_vram_gb=_env_int("LUCEBOX_HOST_AMD_VRAM_GB"), diff --git a/lucebox/src/lucebox/topology.py b/lucebox/src/lucebox/topology.py index f92f5593c..11bc98f58 100644 --- a/lucebox/src/lucebox/topology.py +++ b/lucebox/src/lucebox/topology.py @@ -83,14 +83,21 @@ def _nvidia_devices(host: HostFacts) -> list[GpuDevice]: continue architecture = row[4].replace(".", "") memory_gb = _int_from_cell(row[5]) // 1024 + unified = ( + memory_gb == 0 + and architecture == "121" + and "GB10" in row[3].upper() + ) or (index == 0 and host.nvidia_unified_memory) + effective_gb = max(0, host.ram_gb - 16) if unified else memory_gb devices.append( GpuDevice( vendor="nvidia", index=index, name=row[3], architecture=architecture, - physical_vram_gb=memory_gb, - effective_memory_gb=memory_gb, + physical_vram_gb=0 if unified else memory_gb, + effective_memory_gb=effective_gb, + unified_memory=unified, ) ) if not devices and (host.has_nvidia_gpu or host.gpu_vendor == "nvidia"): @@ -101,14 +108,19 @@ def _nvidia_devices(host: HostFacts) -> list[GpuDevice]: name = name or host.gpu_name architecture = architecture or host.gpu_sm memory_gb = memory_gb or host.vram_gb + unified = host.nvidia_unified_memory or ( + architecture == "121" and "GB10" in (name or "").upper() + ) + effective_gb = max(0, host.ram_gb - 16) if unified else memory_gb devices.append( GpuDevice( vendor="nvidia", index=0, name=name or "NVIDIA GPU", architecture=architecture, - physical_vram_gb=memory_gb, - effective_memory_gb=memory_gb, + physical_vram_gb=0 if unified else memory_gb, + effective_memory_gb=effective_gb, + unified_memory=unified, ) ) return devices diff --git a/lucebox/src/lucebox/types.py b/lucebox/src/lucebox/types.py index b0982e725..39e9c1a17 100644 --- a/lucebox/src/lucebox/types.py +++ b/lucebox/src/lucebox/types.py @@ -76,6 +76,7 @@ class HostFacts: nvidia_vram_gb: int = 0 nvidia_gpu_arch: str = "" nvidia_gpu_list_csv: str = "" + nvidia_unified_memory: bool = False amd_gpu_name: str = "" amd_gpu_count: int = 0 amd_vram_gb: int = 0 @@ -235,7 +236,11 @@ class DflashRuntime: budget: int = 22 max_ctx: int = 16384 lazy: bool = False - prefix_cache_slots: int = 0 + # Exact turn-boundary snapshots. This is safe for every model/backend and + # is the server's production default; 0 remains an explicit opt-out. + prefix_cache_slots: int = 32 + # Exact full-prompt snapshots, keyed by the raw prompt. PFlash can use + # these to skip both rescoring and target prefill when a request repeats. prefill_cache_slots: int = 0 cache_type_k: str = "" cache_type_v: str = "" @@ -252,10 +257,10 @@ class DflashRuntime: # the quality-safe default. ``dense`` and ``sparse`` are approximate, # monolithic-HIP preview paths and are never selected silently. ds4_prefill: Literal["exact", "dense", "sparse"] = "exact" - # Phase-1 (thinking) cap when a request opts into thinking. Default mirrors - # antirez/ds4 ds4_eval.c: think_max_tokens = max_tokens - hard_limit_reply - # budget = 16000 - 512 = 15488. The server's own hardcoded default is 10000. - think_max: int = 15488 + # Optional operator override for the phase-1 reasoning cap. ``None`` lets + # the selected model card choose its own safe value; one global DeepSeek + # default would silently truncate Qwen/Gemma reasoning workloads. + think_max: int | None = None # Flash-attention sliding-window on full-attention layers. 0 = full # attention (server default). On gemma4's hybrid iSWA the full-attn # layers grow KV linearly with max_ctx; a sparse fa_window keeps @@ -294,6 +299,13 @@ def __post_init__(self) -> None: raise ValueError( f"prefill_cache_slots must be zero or positive; got {self.prefill_cache_slots!r}" ) + if self.prefill_cache_slots > 0 and ( + self.prefix_cache_slots + self.prefill_cache_slots > 63 + ): + raise ValueError( + "prefix_cache_slots + prefill_cache_slots must not exceed 63 " + "when the exact prefill cache is enabled" + ) if not 0.0 < self.prefill_keep_ratio <= 1.0: raise ValueError( "prefill_keep_ratio must be in the interval (0.0, 1.0]; " @@ -331,7 +343,7 @@ def __post_init__(self) -> None: ) if self.prefill_threshold <= 0: raise ValueError(f"prefill_threshold must be positive; got {self.prefill_threshold!r}") - if self.think_max < 0: + if self.think_max is not None and self.think_max < 0: raise ValueError(f"think_max must be zero or positive; got {self.think_max!r}") if self.fa_window < 0: raise ValueError(f"fa_window must be zero or positive; got {self.fa_window!r}") diff --git a/lucebox/tests/test_autotune.py b/lucebox/tests/test_autotune.py index b284abf81..cb4dea8b5 100644 --- a/lucebox/tests/test_autotune.py +++ b/lucebox/tests/test_autotune.py @@ -1,8 +1,10 @@ +from dataclasses import replace from pathlib import Path import pytest from lucebox.autotune import automatic_plan, runtime_from_host -from lucebox.types import Config, HostFacts, ModelMeta +from lucebox.docker_run import server_run_spec +from lucebox.types import Config, GpuVendor, HostFacts, ModelMeta from lucebox import download @@ -20,7 +22,7 @@ def test_wsl_24gb_defaults_leave_cuda_headroom() -> None: # to False makes the host config match runtime behaviour. See the # `entrypoint.sh` warning emitted when the two are out-of-sync. assert runtime.lazy is False - assert runtime.prefix_cache_slots == 0 + assert runtime.prefix_cache_slots == 32 def test_native_24gb_caps_context_below_vmm_failure_boundary() -> None: @@ -29,7 +31,7 @@ def test_native_24gb_caps_context_below_vmm_failure_boundary() -> None: assert runtime.budget == 22 assert runtime.max_ctx == 98304 assert runtime.lazy is False # see WSL test above - assert runtime.prefix_cache_slots == 0 + assert runtime.prefix_cache_slots == 32 def test_no_heuristic_tier_sets_lazy_without_prefill_drafter() -> None: @@ -69,6 +71,15 @@ def _install_decode_draft(cfg: Config) -> None: path.write_bytes(b"test-draft") +def _server_env_for_plan(cfg: Config, plan) -> dict[str, str]: + launch_cfg = replace( + cfg, + dflash=plan.runtime, + placement=plan.placement.runtime, + ) + return dict(server_run_spec(launch_cfg).env) + + def test_unconfigured_plan_does_not_claim_an_active_optimization(tmp_path: Path) -> None: cfg = Config( models_dir=tmp_path / "models", @@ -101,6 +112,7 @@ def test_qwen_24gb_enables_dflash_and_pflash_but_prefers_full_kv( assert plan.spark.enabled is False assert plan.runtime.prefill_mode == "auto" assert plan.runtime.prefill_keep_ratio == 0.10 + assert plan.runtime.prefill_cache_slots == 4 assert plan.runtime.max_ctx == 98304 assert plan.active_names == ("DFlash", "PFlash") assert dict(plan.phase_strategies) == { @@ -108,6 +120,11 @@ def test_qwen_24gb_enables_dflash_and_pflash_but_prefers_full_kv( "Decode": "DFlash", "KV cache": "Full KV cache", } + env = _server_env_for_plan(cfg, plan) + assert env["DFLASH_PREFILL_MODE"] == "auto" + assert env["DFLASH_PREFILL_CACHE_SLOTS"] == "4" + assert "DFLASH_KVFLASH" not in env + assert "DFLASH_SPARK" not in env def test_published_draft_is_not_claimed_until_it_is_installed(tmp_path: Path) -> None: @@ -159,6 +176,11 @@ def test_constrained_qwen_moe_enables_spark_and_scored_kvflash( assert plan.kvflash.enabled is True assert plan.runtime.kvflash_policy == "drafter" assert plan.spark.enabled is True + env = _server_env_for_plan(cfg, plan) + assert env["DFLASH_KVFLASH"] == "auto" + assert env["DFLASH_KVFLASH_POLICY"] == "drafter" + assert env["DFLASH_SPARK"] == "1" + assert env["DFLASH_PREFILL_DRAFTER"].endswith("Qwen3-0.6B-BF16.gguf") def test_r9700_qwen_moe_keeps_exact_all_gpu_path(tmp_path: Path) -> None: @@ -183,6 +205,7 @@ def test_laguna_uses_native_context_and_spark_under_24gb_pressure( "laguna-xs.2", HostFacts(gpu_vendor="nvidia", vram_gb=24, ram_gb=64, gpu_sm="86"), ) + _install_any_decode_draft(cfg) plan = automatic_plan(cfg, optimizer_drafter_available=False) @@ -190,7 +213,7 @@ def test_laguna_uses_native_context_and_spark_under_24gb_pressure( # 32K cap on this tight 24 GB profile rather than the old, incorrect 4K # catalog limit. assert plan.runtime.max_ctx == 32768 - assert plan.dflash.enabled is False + assert plan.dflash.enabled is True assert plan.pflash.available is False assert plan.kvflash.enabled is False assert plan.kvflash.available is True @@ -198,9 +221,34 @@ def test_laguna_uses_native_context_and_spark_under_24gb_pressure( assert plan.spark.enabled is True assert dict(plan.phase_strategies) == { "Prefill": "Exact sparse-attention prefill", - "Decode": "Autoregressive decode", + "Decode": "DFlash", "KV cache": "Full hybrid-attention KV cache", } + env = _server_env_for_plan(cfg, plan) + assert env["DFLASH_DRAFT"].endswith("/draft/laguna-xs2-speculator") + assert env["DFLASH_SPARK"] == "1" + assert "DFLASH_KVFLASH" not in env + + +def test_laguna_dflash_requires_the_complete_speculator(tmp_path: Path) -> None: + cfg = _cfg( + tmp_path, + "laguna-xs.2", + HostFacts(gpu_vendor="nvidia", vram_gb=24, ram_gb=64, gpu_sm="86"), + ) + preset = download.PRESETS["laguna-xs.2"] + assert preset.speculator_dir is not None + speculator = cfg.models_dir / "draft" / preset.speculator_dir + speculator.mkdir(parents=True) + (speculator / "model.safetensors").write_bytes(b"partial-download") + + partial = automatic_plan(cfg, optimizer_drafter_available=False) + assert partial.dflash.enabled is False + assert "not installed" in partial.dflash.reason + + (speculator / "config.json").write_text("{}") + complete = automatic_plan(cfg, optimizer_drafter_available=False) + assert complete.dflash.enabled is True def test_deepseek_reports_native_mla_instead_of_generic_flash_features( @@ -226,6 +274,7 @@ def test_deepseek_reports_native_mla_instead_of_generic_flash_features( assert plan.prefill_alternative.available is True assert plan.prefill_alternative.enabled is False assert plan.runtime.ds4_prefill == "exact" + assert plan.runtime.max_ctx == 131072 assert "Native MLA-compressed KV cache" in plan.kvflash.reason assert dict(plan.phase_strategies) == { "Prefill": "Exact MLA prefill", @@ -234,6 +283,36 @@ def test_deepseek_reports_native_mla_instead_of_generic_flash_features( } +def test_deepseek_compact_mla_cache_avoids_the_generic_full_kv_context_cap( + tmp_path: Path, +) -> None: + cfg = Config( + variant="cuda13", + models_dir=tmp_path / "models", + host=HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + gpu_name="NVIDIA GB10", + gpu_count=1, + vram_gb=106, + ram_gb=122, + gpu_sm="121", + nvidia_gpu_name="NVIDIA GB10", + nvidia_gpu_count=1, + nvidia_vram_gb=106, + nvidia_gpu_arch="121", + nvidia_unified_memory=True, + ), + model=ModelMeta(preset="deepseek-v4-flash"), + ) + + plan = automatic_plan(cfg, optimizer_drafter_available=False) + + assert plan.placement.runnable is True + assert plan.runtime.max_ctx == 131072 + assert plan.runtime.ds4_prefill == "exact" + + def test_spark_stays_off_when_host_ram_cannot_hold_cold_experts( tmp_path: Path, ) -> None: @@ -264,6 +343,7 @@ def test_gemma_does_not_offer_unsupported_pflash(tmp_path: Path) -> None: assert plan.pflash.available is False assert plan.pflash.enabled is False assert plan.runtime.prefill_mode == "off" + assert plan.runtime.prefill_cache_slots == 0 assert "no PFlash" in plan.pflash.reason @@ -305,7 +385,7 @@ def test_known_gpu_specific_ddtree_budgets() -> None: ) def test_20gb_automatic_drops_scorer_before_blocking_a_fitting_target( tmp_path: Path, - vendor: str, + vendor: GpuVendor, variant: str, architecture: str, ) -> None: @@ -331,6 +411,12 @@ def test_20gb_automatic_drops_scorer_before_blocking_a_fitting_target( assert plan.runtime.kvflash_policy == "qk" assert plan.runtime.prefill_drafter == "" assert "scorer" in plan.pflash.reason + env = _server_env_for_plan(cfg, plan) + assert env["DFLASH_KVFLASH"] == "auto" + assert env["DFLASH_KVFLASH_POLICY"] == "qk" + assert "DFLASH_PREFILL_MODE" not in env + assert "DFLASH_PREFILL_DRAFTER" not in env + assert "DFLASH_SPARK" not in env def test_18gb_automatic_drops_draft_instead_of_blocking_qwen(tmp_path: Path) -> None: @@ -417,3 +503,162 @@ def test_24gb_moe_does_not_offer_an_unusable_scorer_without_spark_ram( assert plan.placement.runnable is True assert plan.runtime.kvflash == "off" assert plan.needs_optimizer_drafter is False + + +def _install_any_decode_draft(cfg: Config) -> None: + preset = download.PRESETS[cfg.model.preset] + if preset.draft_file: + _install_decode_draft(cfg) + return + assert preset.speculator_dir is not None + root = cfg.models_dir / "draft" / preset.speculator_dir + root.mkdir(parents=True, exist_ok=True) + for filename in preset.speculator_files: + (root / filename).write_bytes(b"test-speculator") + + +@pytest.mark.parametrize( + ("preset_name", "decode_enabled", "pflash_enabled"), + [ + ("qwen3.6-27b", True, True), + ("qwen3.6-moe", False, False), + ("laguna-xs.2", True, False), + # The 103 GB target fits Strix's safe 109 GB shared-memory budget; + # its 11.3 GB DSpark companion does not, so Automatic keeps exact AR. + ("deepseek-v4-flash", False, False), + ], +) +def test_featured_model_strix_only_contract_reaches_exact_server_environment( + tmp_path: Path, + preset_name: str, + decode_enabled: bool, + pflash_enabled: bool, +) -> None: + """A 128 GB Strix-only buyer gets a runnable, single-HIP launch contract.""" + host = HostFacts( + gpu_vendor="amd", + has_amd_gpu=True, + gpu_name="Radeon 8060S", + gpu_count=1, + vram_gb=125, + gpu_sm="gfx1151", + ram_gb=125, + amd_gpu_name="Radeon 8060S", + amd_gpu_count=1, + amd_vram_gb=125, + amd_gpu_arch="gfx1151", + amd_gpu_list_csv="0, , , Radeon 8060S, gfx1151, 0 MiB,", + ) + cfg = Config( + variant="rocm", + models_dir=tmp_path / preset_name, + host=host, + model=ModelMeta(preset=preset_name), + ) + preset = download.PRESETS[preset_name] + if preset.has_decode_companion: + _install_any_decode_draft(cfg) + _install_optimizer_drafter(cfg) + + plan = automatic_plan(cfg) + env = _server_env_for_plan(cfg, plan) + + assert plan.placement.runnable is True + assert plan.placement.runtime.mode == "single" + assert plan.placement.runtime.target_device == "hip:0" + assert plan.placement.runtime.uses_multiple_devices is False + assert plan.runtime.max_ctx == 131072 + assert plan.runtime.prefix_cache_slots == 32 + assert plan.runtime.prefill_cache_slots == (4 if pflash_enabled else 0) + assert plan.dflash.enabled is decode_enabled + assert plan.pflash.enabled is pflash_enabled + assert plan.kvflash.enabled is False + assert plan.spark.enabled is False + assert env["DFLASH_TARGET_DEVICE"] == "hip:0" + assert env["DFLASH_PREFIX_CACHE_SLOTS"] == "32" + assert env["DFLASH_PREFILL_CACHE_SLOTS"] == ("4" if pflash_enabled else "0") + assert "DFLASH_TARGET_DEVICES" not in env + assert ("DFLASH_PREFILL_MODE" in env) is pflash_enabled + assert "DFLASH_KVFLASH" not in env + assert "DFLASH_SPARK" not in env + + if preset_name == "deepseek-v4-flash": + assert "DFLASH_DS4_SPEC" not in env + assert "DFLASH_DS4_DRAFT" not in env + assert env["DFLASH_DRAFT"].endswith("/.lucebox-no-draft") + + +@pytest.mark.parametrize( + ("vendor", "variant", "architecture"), + [ + ("nvidia", "cuda12", "90"), + ("amd", "rocm", "gfx1201"), + ], +) +@pytest.mark.parametrize( + ("preset_name", "decode_name", "pflash_enabled"), + [ + ("qwen3.6-27b", "DFlash", True), + ("qwen3.6-moe", "", False), + ("laguna-xs.2", "DFlash", False), + ("deepseek-v4-flash", "DSpark", False), + ], +) +def test_featured_model_roomy_gpu_contract_reaches_exact_server_environment( + tmp_path: Path, + vendor: GpuVendor, + variant: str, + architecture: str, + preset_name: str, + decode_name: str, + pflash_enabled: bool, +) -> None: + """Every featured preset reaches its qualified CUDA/HIP launch contract.""" + cfg = Config( + variant=variant, + models_dir=tmp_path / f"{vendor}-{preset_name}", + host=HostFacts( + gpu_vendor=vendor, + vram_gb=128, + ram_gb=256, + gpu_sm=architecture, + ), + model=ModelMeta(preset=preset_name), + ) + if decode_name: + _install_any_decode_draft(cfg) + _install_optimizer_drafter(cfg) + + plan = automatic_plan(cfg) + launch_cfg = replace( + cfg, + dflash=plan.runtime, + placement=plan.placement.runtime, + ) + env = dict(server_run_spec(launch_cfg).env) + + assert plan.placement.runnable is True + assert plan.runtime.prefix_cache_slots == 32 + assert plan.runtime.prefill_cache_slots == (4 if pflash_enabled else 0) + assert plan.dflash.enabled is bool(decode_name) + assert plan.dflash.name == (decode_name or "DFlash") + assert plan.pflash.enabled is pflash_enabled + assert plan.kvflash.enabled is False + assert plan.spark.enabled is False + assert env["DFLASH_MODEL_NAME"] == preset_name + assert env["DFLASH_TARGET_DEVICE"] == ("cuda:0" if vendor == "nvidia" else "hip:0") + assert env["DFLASH_PREFIX_CACHE_SLOTS"] == "32" + assert env["DFLASH_PREFILL_CACHE_SLOTS"] == ("4" if pflash_enabled else "0") + assert ("DFLASH_PREFILL_MODE" in env) is pflash_enabled + assert "DFLASH_KVFLASH" not in env + assert "DFLASH_SPARK" not in env + + if preset_name == "deepseek-v4-flash": + assert env["DFLASH_DS4_SPEC"] == "1" + assert env["DFLASH_DS4_DRAFT"].endswith( + "DeepSeek-V4-Flash-DSpark-draft-Q4RMFP4-denseF16.gguf" + ) + assert env["DFLASH_DRAFT"].endswith("/.lucebox-no-draft") + else: + assert "DFLASH_DS4_SPEC" not in env + assert "DFLASH_DS4_DRAFT" not in env diff --git a/lucebox/tests/test_config.py b/lucebox/tests/test_config.py index 398340099..92f3aea06 100644 --- a/lucebox/tests/test_config.py +++ b/lucebox/tests/test_config.py @@ -380,6 +380,14 @@ def test_cache_slot_validation_names_the_invalid_field(field: str, message: str) DflashRuntime(**{field: -1}) +def test_exact_prefill_cache_reserves_the_internal_staging_slot() -> None: + from lucebox.types import DflashRuntime + + assert DflashRuntime(prefix_cache_slots=32, prefill_cache_slots=4) + with pytest.raises(ValueError, match="must not exceed 63"): + DflashRuntime(prefix_cache_slots=60, prefill_cache_slots=4) + + def test_direct_optimization_edit_marks_profile_custom(tmp_path: Path) -> None: path = tmp_path / "config.toml" @@ -464,6 +472,31 @@ def test_optimization_and_placement_are_persisted_atomically(tmp_path: Path) -> assert config.optimization_mode(path=path) == "automatic" +def test_automatic_runtime_preserves_model_card_thinking_budget(tmp_path: Path) -> None: + """An absent override must survive TOML persistence as ``None``.""" + from lucebox.types import DflashRuntime + + path = tmp_path / "config.toml" + config.write_optimization_runtime(DflashRuntime(), path=path) + + assert "think_max" not in config.load_doc(path)["dflash"] + loaded = config.load(path) + assert loaded is not None + assert loaded.dflash.think_max is None + assert loaded.dflash.prefix_cache_slots == 32 + + +def test_explicit_thinking_budget_round_trips(tmp_path: Path) -> None: + from lucebox.types import DflashRuntime + + path = tmp_path / "config.toml" + config.write_optimization_runtime(DflashRuntime(think_max=15488), path=path) + + loaded = config.load(path) + assert loaded is not None + assert loaded.dflash.think_max == 15488 + + def test_runtime_reset_clears_stale_placement(tmp_path: Path) -> None: from lucebox.types import DflashRuntime, PlacementRuntime @@ -496,12 +529,14 @@ def test_model_and_execution_profile_switch_atomically(tmp_path: Path) -> None: ModelMeta(preset="new", target_file="new.gguf"), DflashRuntime(max_ctx=32768), placement, + variant="rocm", path=path, ) loaded = config.load(path) assert loaded is not None assert loaded.model == ModelMeta(preset="new", target_file="new.gguf") + assert loaded.variant == "rocm" assert loaded.dflash.max_ctx == 32768 assert loaded.placement == placement assert "draft_file" not in config.load_doc(path)["model"] diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index 8c43ab73a..e09672e3b 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -213,6 +213,18 @@ def test_resolve_model_files_ignores_invalid_speculator_symlink( assert draft_dir == "" +def test_resolve_model_files_rejects_incomplete_speculator(tmp_path: Path) -> None: + speculator = tmp_path / "draft" / "laguna-xs2-speculator" + speculator.mkdir(parents=True) + (speculator / "model.safetensors").write_bytes(b"partial") + cfg = Config(models_dir=tmp_path, model=ModelMeta(preset="laguna-xs.2")) + + assert docker_run._resolve_model_files(cfg)[2] == "" + + (speculator / "config.json").write_text("{}") + assert docker_run._resolve_model_files(cfg)[2] == "laguna-xs2-speculator" + + # ── server_run_spec ────────────────────────────────────────────────────────── @@ -319,6 +331,7 @@ def test_server_run_spec_mounts_symlinked_speculator_directory_read_only( draft_root.mkdir(parents=True) external.mkdir() (external / "model.safetensors").write_bytes(b"speculator") + (external / "config.json").write_text("{}") (draft_root / "laguna-xs2-speculator").symlink_to(external, target_is_directory=True) spec = docker_run.server_run_spec( @@ -364,9 +377,9 @@ def test_server_run_spec_always_emits_core_dflash_env(tmp_path: Path) -> None: env = _env(docker_run.server_run_spec(cfg)) assert env["DFLASH_BUDGET"] == "22" assert env["DFLASH_MAX_CTX"] == "32768" - assert env["DFLASH_PREFIX_CACHE_SLOTS"] == "0" + assert env["DFLASH_PREFIX_CACHE_SLOTS"] == "32" assert env["DFLASH_PREFILL_CACHE_SLOTS"] == "0" - assert env["DFLASH_THINK_MAX"] == "15488" + assert "DFLASH_THINK_MAX" not in env assert env["DFLASH_PORT"] == "8080" assert env["LUCEBOX_HOME"] @@ -423,6 +436,7 @@ def test_server_run_spec_optional_env_emitted_when_set(tmp_path: Path) -> None: kvflash_policy="qk", kvflash_tau=96, spark=True, + think_max=15488, spark_vram_gb=14.5, ds4_prefill="sparse", think_soft_close_min_ratio=0.5, @@ -443,6 +457,7 @@ def test_server_run_spec_optional_env_emitted_when_set(tmp_path: Path) -> None: assert env["DFLASH_SPARK"] == "1" assert env["DFLASH_SPARK_VRAM_GB"] == "14.5" assert env["DFLASH_DS4_PREFILL"] == "sparse" + assert env["DFLASH_THINK_MAX"] == "15488" assert env["DFLASH_THINK_SOFT_CLOSE_MIN_RATIO"] == "0.5" assert env["DFLASH_DEBUG_THINKING_LOGITS"] == "1" @@ -564,6 +579,40 @@ def test_server_run_spec_can_disable_preset_dflash_draft(tmp_path: Path) -> None assert env["DFLASH_DRAFT"].endswith("/.lucebox-no-draft") +def test_server_run_spec_uses_deepseek_dspark_contract(tmp_path: Path) -> None: + preset = PRESETS["deepseek-v4-flash"] + assert preset.draft_file is not None + cfg = Config( + models_dir=tmp_path, + model=ModelMeta(preset=preset.name), + dflash=DflashRuntime(speculative_decode=True), + ) + + env = _env(docker_run.server_run_spec(cfg)) + + assert env["DFLASH_DS4_SPEC"] == "1" + assert env["DFLASH_DS4_DRAFT"] == ( + f"/opt/lucebox-hub/server/models/draft/{preset.draft_file}" + ) + # --draft is not DeepSeek's DSpark switch and the server warns that it is + # inert. Pin generic discovery off instead. + assert env["DFLASH_DRAFT"].endswith("/.lucebox-no-draft") + + +def test_server_run_spec_disables_deepseek_dspark_cleanly(tmp_path: Path) -> None: + cfg = Config( + models_dir=tmp_path, + model=ModelMeta(preset="deepseek-v4-flash"), + dflash=DflashRuntime(speculative_decode=False), + ) + + env = _env(docker_run.server_run_spec(cfg)) + + assert "DFLASH_DS4_SPEC" not in env + assert "DFLASH_DS4_DRAFT" not in env + assert env["DFLASH_DRAFT"].endswith("/.lucebox-no-draft") + + def test_server_run_spec_forwards_host_env(monkeypatch, tmp_path: Path) -> None: monkeypatch.setenv("LUCEBOX_HOST_OS_PRETTY", "Ubuntu 22.04") monkeypatch.setenv("LUCEBOX_HOST_GPU_NAME", "RTX 5090") diff --git a/lucebox/tests/test_download.py b/lucebox/tests/test_download.py index bae056c29..e23d50c88 100644 --- a/lucebox/tests/test_download.py +++ b/lucebox/tests/test_download.py @@ -67,11 +67,15 @@ def test_resolve_preset_picks_gemma_target_and_draft(): assert pres.has_draft -def test_resolve_preset_supports_target_only_laguna(): +def test_resolve_preset_includes_laguna_speculator_contract(): pres = resolve_preset("laguna-xs.2") assert pres.target_repo == "Lucebox/Laguna-XS.2-GGUF" assert pres.draft_repo is None assert not pres.has_draft + assert pres.has_speculator + assert pres.has_decode_companion + assert pres.speculator_repo == "poolside/Laguna-XS.2-speculator.dflash" + assert pres.speculator_files == ("model.safetensors", "config.json") def test_resolve_preset_picks_qwen36_moe_target_only(): @@ -116,7 +120,7 @@ def test_resolve_preset_picks_deepseek_target_and_dspark_draft(): def test_download_preset_target_only_qwen36_moe_skips_draft(tmp_path, monkeypatch): - """qwen3.6-moe behaves identically to laguna-xs.2: target only, no draft fetch.""" + """Qwen3.6 MoE has no published decode companion.""" cfg = SimpleNamespace(models_dir=tmp_path) pres = resolve_preset("qwen3.6-moe") assert not pres.has_draft @@ -313,17 +317,20 @@ def _stub_fetch(api, repo_id, filename, local_dir, console): # noqa: ARG001 assert (pres.draft_repo, pres.draft_file, str(tmp_path / "draft")) in fetches -def test_download_preset_target_only_skips_draft_fetch(tmp_path, monkeypatch): +def test_download_preset_fetches_laguna_speculator_and_required_config( + tmp_path, monkeypatch +): cfg = SimpleNamespace(models_dir=tmp_path) pres = resolve_preset("laguna-xs.2") assert not pres.has_draft - fetches: list[tuple[str, str]] = [] + assert pres.has_speculator + fetches: list[tuple[str, str, str]] = [] def _meta(_api, repo_id: str, filename: str) -> tuple[int, None]: return 10, None def _stub_fetch(api, repo_id, filename, local_dir, console): # noqa: ARG001 - fetches.append((repo_id, filename)) + fetches.append((repo_id, filename, str(local_dir))) out = local_dir / filename out.parent.mkdir(parents=True, exist_ok=True) with out.open("wb") as f: @@ -334,11 +341,15 @@ def _stub_fetch(api, repo_id, filename, local_dir, console): # noqa: ARG001 monkeypatch.setattr(download, "_fetch", _stub_fetch) assert download.download_preset(cfg, pres) == 0 - # Target fetched, no draft fetch attempted at all. - assert fetches == [(pres.target_repo, pres.target_file)] + speculator_dir = str(tmp_path / "draft" / "laguna-xs2-speculator") + assert fetches == [ + (pres.target_repo, pres.target_file, str(tmp_path)), + (pres.speculator_repo, "model.safetensors", speculator_dir), + (pres.speculator_repo, "config.json", speculator_dir), + ] -def test_status_target_only_preset_reports_draft_as_present(tmp_path, monkeypatch): +def test_status_requires_every_laguna_speculator_file(tmp_path, monkeypatch): cfg = SimpleNamespace(models_dir=tmp_path) pres = resolve_preset("laguna-xs.2") @@ -346,8 +357,14 @@ def _meta(_api, repo_id: str, filename: str) -> tuple[int, None]: return 1024, None monkeypatch.setattr(download, "_file_meta", _meta) - # Target absent → target_present False, draft_present True (nothing to download). - assert status(cfg, pres) == {"target_present": False, "draft_present": True} + assert status(cfg, pres) == {"target_present": False, "draft_present": False} + + root = tmp_path / "draft" / "laguna-xs2-speculator" + root.mkdir(parents=True) + (root / "model.safetensors").write_bytes(b"x" * 1024) + assert status(cfg, pres)["draft_present"] is False + (root / "config.json").write_bytes(b"x" * 1024) + assert status(cfg, pres)["draft_present"] is True def test_recommend_preset_tiers() -> None: diff --git a/lucebox/tests/test_models_cli.py b/lucebox/tests/test_models_cli.py index 762446254..99a576fb0 100644 --- a/lucebox/tests/test_models_cli.py +++ b/lucebox/tests/test_models_cli.py @@ -26,6 +26,33 @@ def _stub_host(monkeypatch: pytest.MonkeyPatch, vram_gb: int) -> None: monkeypatch.setattr("lucebox.cli.from_env", lambda: host) +def _stub_rtx_strix_host(monkeypatch: pytest.MonkeyPatch) -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + has_amd_gpu=True, + gpu_name="NVIDIA GeForce RTX 3090", + gpu_count=1, + vram_gb=24, + gpu_sm="86", + ram_gb=125, + nvidia_gpu_name="NVIDIA GeForce RTX 3090", + nvidia_gpu_count=1, + nvidia_vram_gb=24, + nvidia_gpu_arch="86", + nvidia_gpu_list_csv=( + "0, GPU-test, 0000:01:00.0, NVIDIA GeForce RTX 3090, 8.6, 24576 MiB," + ), + amd_gpu_name="AMD Radeon Graphics", + amd_gpu_count=1, + amd_vram_gb=125, + amd_gpu_arch="gfx1151", + amd_gpu_list_csv="0, , , AMD Radeon Graphics, gfx1151, 512 MiB,", + ) + monkeypatch.setattr("lucebox.host_facts.from_env", lambda: host) + monkeypatch.setattr("lucebox.cli.from_env", lambda: host) + + def test_models_list_shows_every_registered_preset( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -195,6 +222,62 @@ def _unexpected_download(cfg, preset): # noqa: ARG001 assert attempted is False +def test_models_select_switches_to_detected_backend_only_when_required( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + cfg_path = _set_config_path(tmp_path, monkeypatch) + _stub_rtx_strix_host(monkeypatch) + monkeypatch.setattr(download_mod, "download_preset", lambda cfg, pres: 0) + monkeypatch.setattr( + download_mod, + "status", + lambda cfg, pres: {"target_present": False, "draft_present": False}, + ) + + result = CliRunner().invoke( + app, + ["models", "select", "deepseek-v4-flash", "--yes"], + ) + + assert result.exit_code == 0, result.output + assert "switched to rocm" in result.output + entries = config_mod.config_get(path=cfg_path) + assert entries["variant"] == ("rocm", "file") + assert entries["model.preset"] == ("deepseek-v4-flash", "file") + assert entries["dflash.max_ctx"] == (131072, "file") + + +def test_models_download_blocks_incompatible_model_unless_staging_is_explicit( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _set_config_path(tmp_path, monkeypatch) + _stub_host(monkeypatch, vram_gb=24) + attempted = False + + def _download(cfg, preset): # noqa: ARG001 + nonlocal attempted + attempted = True + return 0 + + monkeypatch.setattr(download_mod, "download_preset", _download) + monkeypatch.setattr( + download_mod, + "status", + lambda cfg, pres: {"target_present": False, "draft_present": False}, + ) + + blocked = CliRunner().invoke(app, ["models", "download", "deepseek-v4-flash"]) + assert blocked.exit_code == 2 + assert attempted is False + + staged = CliRunner().invoke( + app, + ["models", "download", "deepseek-v4-flash", "--force"], + ) + assert staged.exit_code == 0 + assert attempted is True + + def test_optimize_resets_to_hardware_profile( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -324,5 +407,11 @@ def test_installed_helpers_track_presence(tmp_path: Path, monkeypatch: pytest.Mo # stat path without allocating and writing a 5 GB Python byte string. with target.open("wb") as sparse: sparse.truncate(5 * 10**9) + assert download_mod.installed_status(cfg, laguna) == "partial" + assert laguna.speculator_dir is not None + speculator_root = cfg.models_dir / "draft" / laguna.speculator_dir + speculator_root.mkdir(parents=True) + for filename in laguna.speculator_files: + (speculator_root / filename).write_bytes(b"installed") assert download_mod.installed_status(cfg, laguna) == "installed" assert download_mod.installed_size_gb(cfg, laguna) == pytest.approx(5.0, rel=0.01) diff --git a/lucebox/tests/test_placement.py b/lucebox/tests/test_placement.py index 1371068e8..43563bb89 100644 --- a/lucebox/tests/test_placement.py +++ b/lucebox/tests/test_placement.py @@ -7,7 +7,7 @@ from lucebox.autotune import automatic_plan from lucebox.capabilities import ARCHITECTURE_CAPABILITIES from lucebox.docker_run import server_run_spec -from lucebox.host_facts import from_env +from lucebox.host_facts import compatible_variant, for_variant, from_env, nvidia_variant from lucebox.placement import automatic_placement from lucebox.topology import from_config from lucebox.types import Config, DflashRuntime, HostFacts, ModelMeta @@ -351,6 +351,106 @@ def test_strix_is_uma_even_when_driver_reports_large_aperture(tmp_path: Path) -> assert topology.primary.effective_memory_gb == 112 +def test_gb10_nvml_na_memory_is_normalized_as_shared_memory(tmp_path: Path) -> None: + """DGX Spark reports ``[N/A]`` for NVML memory despite CUDA UMA.""" + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + gpu_name="NVIDIA GB10", + gpu_count=1, + vram_gb=105, + gpu_sm="121", + ram_gb=121, + nvidia_gpu_name="NVIDIA GB10", + nvidia_gpu_count=1, + nvidia_vram_gb=105, + nvidia_gpu_arch="121", + nvidia_gpu_list_csv=( + "0, GPU-test, 00000000:01:00.0, NVIDIA GB10, 12.1, [N/A], [N/A]" + ), + nvidia_unified_memory=True, + ) + + topology = from_config(Config(variant="cuda13", models_dir=tmp_path, host=host)) + + assert topology.primary is not None + assert topology.primary.backend == "cuda" + assert topology.primary.architecture == "121" + assert topology.primary.unified_memory is True + assert topology.primary.physical_vram_gb == 0 + assert topology.primary.effective_memory_gb == 105 + assert "105 GB shared" in topology.primary.label + + +def test_gb10_unified_memory_flag_is_read_from_host_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _set_host_env( + monkeypatch, + { + "LUCEBOX_HOST_GPU_VENDOR": "nvidia", + "LUCEBOX_HOST_NVIDIA_UNIFIED_MEMORY": 1, + }, + ) + + assert from_env().nvidia_unified_memory is True + + +def test_backend_projection_uses_the_selected_vendor_inventory() -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + has_amd_gpu=True, + gpu_name="RTX 3090", + gpu_count=1, + vram_gb=24, + gpu_sm="86", + nvidia_gpu_name="RTX 3090", + nvidia_gpu_count=1, + nvidia_vram_gb=24, + nvidia_gpu_arch="86", + amd_gpu_name="AMD Radeon Graphics", + amd_gpu_count=1, + amd_vram_gb=125, + amd_gpu_arch="gfx1151", + ) + + selected = for_variant(host, "rocm") + + assert selected.gpu_vendor == "amd" + assert selected.gpu_name == "AMD Radeon Graphics" + assert selected.gpu_count == 1 + assert selected.vram_gb == 125 + assert selected.gpu_sm == "gfx1151" + + +@pytest.mark.parametrize( + ("architecture", "name", "expected"), + [ + ("86", "NVIDIA GeForce RTX 3090", "cuda12"), + ("90", "NVIDIA H100", "cuda12"), + ("120", "NVIDIA GeForce RTX 5090", "cuda128"), + ("121", "NVIDIA GB10", "cuda13"), + ], +) +def test_nvidia_image_variant_tracks_the_toolkit_required_by_the_architecture( + architecture: str, + name: str, + expected: str, +) -> None: + host = HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + gpu_name=name, + gpu_sm=architecture, + nvidia_gpu_name=name, + nvidia_gpu_arch=architecture, + ) + + assert nvidia_variant(host) == expected + assert compatible_variant(host, "cuda12") == expected + + def test_rtx_strix_uses_remote_target_shard_when_capacity_requires_it( tmp_path: Path, ) -> None: diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index 3f1468d09..06679182d 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -1060,6 +1060,76 @@ test_harness_draft_directory_validation() { test_harness_draft_directory_validation \ "native harness rejects empty or ambiguous draft directories" +test_harness_external_mode_reuses_canonical_engine() { + local label="$1" sandbox server_log out rc=0 + sandbox=$(mktemp -d -t lucebox-harness-external.XXXXXX) + server_log="$sandbox/canonical-engine.log" + out=$(REPO_DIR="$ROOT" CLIENT_WORK_DIR="$sandbox/work" \ + MODEL_SERVER=external SERVER_LOG="$server_log" \ + bash -c ' + curl() { return 0; } + source "$1" + start_lucebox_server + wait_lucebox_server + stop_lucebox_server + printf "server_log=%s\n" "$SERVER_LOG" + ' _ "$HARNESS_COMMON" 2>&1) || rc=$? + rm -rf "$sandbox" + if [ "$rc" -ne 0 ] || [[ "$out" != *"server_log=$server_log"* ]]; then + report fail "$label" "external mode tried to own the server: $out" + else + report ok "$label" + fi +} +test_harness_external_mode_reuses_canonical_engine \ + "CLI harness mode reuses the canonical optimized engine" + +test_connector_api_checks_selected_model() { + local label="$1" fn runner rc=0 + fn=$(awk '/^_connector_api_ready\(\) \{/,/^\}/' "$SCRIPT") + runner=$'\ncurl(){ printf "%s" "$API_RESPONSE"; }\n_connector_api_ready http://127.0.0.1:8080 qwen3.6-27b' + API_RESPONSE='{"object":"list","data":[{"id":"qwen3.6-27b"}]}' \ + bash -c "$fn$runner" \ + || rc=$? + if [ "$rc" -ne 0 ]; then + report fail "$label" "selected model was rejected" + return + fi + rc=0 + API_RESPONSE='{"object":"list","data":[{"id":"laguna-xs.2"}]}' \ + bash -c "$fn$runner" \ + || rc=$? + if [ "$rc" -eq 0 ]; then + report fail "$label" "a different model was accepted" + else + report ok "$label" + fi +} +test_connector_api_checks_selected_model \ + "harness connector rejects an API serving the wrong model" + +test_native_decode_companion_contract() { + local label="$1" fn out runner + fn=$(awk '/^_export_selected_decode_companion\(\) \{/,/^\}/' "$SCRIPT") + + runner=$'\n_export_selected_decode_companion deepseek-v4-flash /models/draft/dspark.gguf\nprintf "%s|%s|%s" "$DFLASH_DRAFT" "$DFLASH_DS4_SPEC" "$DFLASH_DS4_DRAFT"' + out=$(DEFAULT_MODELS_DIR=/models bash -c "$fn$runner" 2>&1) + if [ "$out" != "/models/.lucebox-no-draft|1|/models/draft/dspark.gguf" ]; then + report fail "$label" "DeepSeek native contract was $out" + return + fi + + runner=$'\n_export_selected_decode_companion qwen3.6-27b /models/draft/qwen.gguf\nprintf "%s|%s|%s" "$DFLASH_DRAFT" "${DFLASH_DS4_SPEC:-}" "${DFLASH_DS4_DRAFT:-}"' + out=$(DEFAULT_MODELS_DIR=/models bash -c "$fn$runner" 2>&1) + if [ "$out" != "/models/draft/qwen.gguf||" ]; then + report fail "$label" "generic native contract was $out" + return + fi + report ok "$label" +} +test_native_decode_companion_contract \ + "native CLI maps DeepSeek to DSpark and other drafts to generic DFlash" + test_entrypoint_keeps_family_kv_default() { local label="$1" sandbox draft_dir models_dir bin_dir shim_dir out _make_entrypoint_sandbox @@ -1194,19 +1264,15 @@ test_entrypoint_draft_is_file "entrypoint serve: DFLASH_DRAFT is a file (no DRAF # test 10/11 already; this test pins the JSON shape independently. test_entrypoint_host_info_json() { local label="$1" - # Source the helper functions from the real entrypoint.sh. - # shellcheck disable=SC1090 - source <(awk '/^_json_escape\(\) \{/,/^\}/' "$ENTRYPOINT") - # shellcheck disable=SC1090 - source <(awk '/^_json_str_or_null\(\) \{/,/^\}/' "$ENTRYPOINT") - # shellcheck disable=SC1090 - source <(awk '/^_json_int_or_null\(\) \{/,/^\}/' "$ENTRYPOINT") - # shellcheck disable=SC1090 - source <(awk '/^_trim\(\) \{/,/^\}/' "$ENTRYPOINT") - # shellcheck disable=SC1090 - source <(awk '/^_emit_gpu_array\(\) \{/,/^\}/' "$ENTRYPOINT") - # shellcheck disable=SC1090 - source <(awk '/^_build_host_info_json\(\) \{/,/^\}/' "$ENTRYPOINT") + # Load the helper functions from the real entrypoint. Bash 3.2 on macOS + # can report success for `source <(...)` while reading an empty /dev/fd; + # command substitution keeps this contributor test portable. + eval "$(awk '/^_json_escape/,/^}/' "$ENTRYPOINT")" + eval "$(awk '/^_json_str_or_null/,/^}/' "$ENTRYPOINT")" + eval "$(awk '/^_json_int_or_null/,/^}/' "$ENTRYPOINT")" + eval "$(awk '/^_trim/,/^}/' "$ENTRYPOINT")" + eval "$(awk '/^_emit_gpu_array/,/^}/' "$ENTRYPOINT")" + eval "$(awk '/^_build_host_info_json/,/^}/' "$ENTRYPOINT")" local out LUCEBOX_HOST_OS_PRETTY="Ubuntu 22.04.3 LTS" \ @@ -2003,12 +2069,22 @@ test_amd_smi_parser "amd-smi parser recognizes R9700 + Strix" test_variant_autoselection() { local label="$1" fn common got - fn=$(awk '/^pick_variant\(\) \{/,/^\}/' "$SCRIPT") + fn=$(awk '/^pick_variant\(\) \{/,/^\}/' "$SCRIPT")$'\n' + fn+=$(awk '/^_default_cuda_variant\(\) \{/,/^\}/' "$SCRIPT") common=$'_lucebox_config_get(){ :; }\nensure_probed(){ :; }\nLUCEBOX_VARIANT=""\n' - got=$(bash -c "$fn"$'\n'"$common"$'LUCEBOX_HOST_HAS_NVIDIA_GPU=1\nLUCEBOX_HOST_HAS_AMD_GPU=1\npick_variant') + got=$(bash -c "$fn"$'\n'"$common"$'LUCEBOX_HOST_HAS_NVIDIA_GPU=1\nLUCEBOX_HOST_HAS_AMD_GPU=1\nLUCEBOX_HOST_GPU_SM=86\npick_variant') [ "$got" = "cuda12" ] || { report fail "$label" "mixed RTX + Strix chose $got"; return; } + got=$(bash -c "$fn"$'\n'"$common"$'uname(){ echo aarch64; }\nLUCEBOX_HOST_HAS_NVIDIA_GPU=1\nLUCEBOX_HOST_HAS_AMD_GPU=0\nLUCEBOX_HOST_GPU_SM=121\npick_variant') + [ "$got" = "cuda13" ] || { report fail "$label" "GB10 chose $got"; return; } + + got=$(bash -c "$fn"$'\n'"$common"$'LUCEBOX_HOST_HAS_NVIDIA_GPU=1\nLUCEBOX_HOST_HAS_AMD_GPU=0\nLUCEBOX_HOST_GPU_SM=120\npick_variant') + [ "$got" = "cuda128" ] || { report fail "$label" "RTX 5090 chose $got"; return; } + + got=$(bash -c "$fn"$'\n_lucebox_config_get(){ echo cuda12; }\nensure_probed(){ :; }\nLUCEBOX_VARIANT=""\nLUCEBOX_HOST_GPU_SM=120\npick_variant') + [ "$got" = "cuda128" ] || { report fail "$label" "configured RTX 5090 migration chose $got"; return; } + got=$(bash -c "$fn"$'\n'"$common"$'LUCEBOX_HOST_HAS_NVIDIA_GPU=0\nLUCEBOX_HOST_HAS_AMD_GPU=1\npick_variant') [ "$got" = "rocm" ] || { report fail "$label" "R9700 + Strix chose $got"; return; } @@ -2016,7 +2092,54 @@ test_variant_autoselection() { [ "$got" = "rocm" ] || { report fail "$label" "explicit override chose $got"; return; } report ok "$label" } -test_variant_autoselection "variant selection: RTX+Strix=cuda12, R9700+Strix=rocm" +test_variant_autoselection "variant selection: RTX=cuda12, RTX5090=cuda128, GB10=cuda13, R9700=rocm" + +test_gb10_na_memory_probe() { + local label="$1" sandbox shim_dir out + sandbox=$(mktemp -d -t lucebox-gb10.XXXXXX) + shim_dir="$sandbox/bin" + mkdir -p "$shim_dir" + cat > "$shim_dir/nvidia-smi" <<'STUB' +#!/usr/bin/env bash +case "$*" in + *"name,memory.total,driver_version,compute_cap"*) + echo "NVIDIA GB10, [N/A], 580.159.03, 12.1" ;; + *"index,uuid,pci.bus_id,name,compute_cap,memory.total,power.limit"*) + echo "0, GPU-test, 00000000:01:00.0, NVIDIA GB10, 12.1, [N/A], [N/A]" ;; + *"--query-gpu=name"*) echo "NVIDIA GB10" ;; + "-L") echo "GPU 0: NVIDIA GB10" ;; +esac +exit 0 +STUB + cat > "$shim_dir/docker" <<'STUB' +#!/usr/bin/env bash +case "${1:-}" in + ps) exit 0 ;; + version) echo "29.1.3" ;; +esac +exit 0 +STUB + cat > "$shim_dir/uname" <<'STUB' +#!/usr/bin/env bash +case "${1:-}" in + -m) echo "aarch64" ;; + *) /usr/bin/uname "$@" ;; +esac +STUB + printf '#!/usr/bin/env bash\nexit 0\n' > "$shim_dir/nvidia-container-runtime" + chmod +x "$shim_dir"/* + out=$(HOME="$sandbox" LUCEBOX_HOME="$sandbox/.lucebox" \ + PATH="$shim_dir:$PATH" NO_COLOR=1 bash "$SCRIPT" check 2>&1) + rm -rf "$sandbox" + if ! grep -qF "NVIDIA GB10" <<<"$out" \ + || ! grep -qF ":cuda13 (NVIDIA selected)" <<<"$out" \ + || ! grep -qF "sm_121 covered" <<<"$out"; then + report fail "$label" "GB10 probe output: $(printf '%s' "$out" | head -20)" + return + fi + report ok "$label" +} +test_gb10_na_memory_probe "GB10 [N/A] NVML memory selects CUDA 13 without crashing" test_mixed_rtx_strix_probe_prefers_cuda() { local label="$1" sandbox shim_dir out From 5229eba69d6b3dd4dc1f4c83ac234d9e85078569 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sat, 1 Aug 2026 13:02:06 +0200 Subject: [PATCH 18/25] fix(server): preserve multi-turn cache and tool replay --- README.md | 12 +- optimizations/pflash/README.md | 2 +- server/scripts/entrypoint.sh | 28 +- server/src/device_runtime.h | 1 + server/src/flashprefill.cpp | 33 +-- server/src/flashprefill_kernels.cu | 59 ++-- server/src/flashprefill_kernels.hip.cu | 13 + server/src/flashprefill_launchers.h | 31 ++ server/src/server/http_server.cpp | 279 ++++++++++++++---- server/src/server/http_server.h | 19 ++ server/test/test_entrypoint_cache_defaults.sh | 6 +- server/test/test_flashprefill_kernels.cpp | 50 ++-- server/test/test_server_unit.cpp | 175 +++++++++++ 13 files changed, 576 insertions(+), 132 deletions(-) create mode 100644 server/src/flashprefill_launchers.h diff --git a/README.md b/README.md index 198ea51b7..697f6e827 100644 --- a/README.md +++ b/README.md @@ -253,7 +253,9 @@ Prebuilt images on GHCR track `main`. No CUDA toolkit or build needed. Pull the | GPU | Image tag | |-----|-----------| -| NVIDIA (CUDA 12+) | `:cuda12` | +| NVIDIA Turing–Hopper (sm_75–sm_90) | `:cuda12` | +| NVIDIA RTX 5090 (sm_120) | `:cuda128` | +| NVIDIA GB10 / DGX Spark (sm_121) | `:cuda13` | | AMD (ROCm 6+) | `:rocm` | Drop a GGUF model target into `server/models/` first, then @@ -274,6 +276,8 @@ Drop a GGUF model target into `server/models/` first, then ```bash # 1. Pull the image for your GPU docker pull ghcr.io/luce-org/lucebox-hub:cuda12 # NVIDIA +docker pull ghcr.io/luce-org/lucebox-hub:cuda128 # RTX 5090 +docker pull ghcr.io/luce-org/lucebox-hub:cuda13 # GB10 / DGX Spark docker pull ghcr.io/luce-org/lucebox-hub:rocm # AMD # 2. Download a target model into server/models/ and the DFlash draft @@ -418,7 +422,7 @@ End-to-end repro: `DFLASH_SAMP=0.8,1.0,0,1.1,42 python server/scripts/bench_llm. | `DFLASH_FP_ALPHA=0.85` | `0.12` | Block-selection threshold; higher = stricter = fewer K-blocks | | `DFLASH_FP_PROFILE=1` | `0` | Per-stage timing log | -When compression is on, the request path picks one of three modes automatically, so they never stack: the first turn is sent verbatim (the system prompt stays as a stable cache anchor), multi-turn continuations use **FlowKV** (only the aged history is compressed, recent turns kept verbatim, so the disk prefix cache from `--prefix-cache-slots` keeps hitting), and a single oversized prompt with no prior turns uses whole-prompt PFlash. With `--prefill-compression off` the request path is identical to a build without compression. +When compression is on, the request policy keeps reusable state reusable. Structured system+user chats and tool requests preserve the target prefix so the live prefix cache can hit on later turns. A continuation also preserves that prefix unless the request explicitly enables FlowKV; in that mode only aged messages are compressed and recent turns remain verbatim. A true long one-shot prompt uses whole-prompt PFlash. Exact repeats can restore a completed PFlash/target snapshot from `--prefill-cache-slots`, skipping both scorer and target prefill. These paths are mutually exclusive; with compression off the request path is unchanged. **KV cache** @@ -429,13 +433,13 @@ When compression is on, the request path picks one of three modes automatically, | `DFLASH27B_KV_Q4=1` | off | Q4_0 K+V (4.5 bpv, legacy, ~128K ceiling) | | `--prefix-cache-slots N` | — | Live prefix-cache slot count | | `DFLASH_PREFIX_CACHE_SLOTS=N` | `32` | Container-entrypoint equivalent of `--prefix-cache-slots`; the native binary itself uses the CLI flag. | -| `DFLASH_PREFILL_CACHE_SLOTS=N` | `0` | Container-entrypoint equivalent of `--prefill-cache-slots`; the native binary itself uses the CLI flag. | +| `--prefill-cache-slots N` / `DFLASH_PREFILL_CACHE_SLOTS=N` | `0` | Exact full-prompt snapshot slots. Automatic assigns four for qualified PFlash models; exact repeats skip scorer and target prefill. | | `--kv-cache-dir ` | — | Persist prefix cache to disk | | `--kv-cache-budget N` | — | On-disk cache size cap | **Bounded KV residency (KVFlash)** -Pages the attention KV cache through a fixed pool of GPU slots; cold 64-token chunks live in host RAM, bit-exact and recallable. Decode speed stops depending on context length and resident KV stays pool-sized at any context. Off by default; works on every model family. Drafter-scored residency is the default on every family: the server finds the Qwen3-0.6B drafter next to the model (or via `--prefill-drafter`) and lazy-loads it as the relevance scorer that decides which chunks stay resident — non-qwen targets (laguna, gemma4) bridge the tokenizer gap by re-tokenizing the context text for the drafter. LRU is the fallback when no drafter is present, or the explicit choice via `--kvflash-policy lru`. Per-model numbers in [Luce KVFlash →](optimizations/kvflash/README.md). +Pages a conventional attention KV cache through a fixed pool of GPU slots; cold 64-token chunks live in host RAM, bit-exact and recallable. Decode speed stops depending on context length and resident KV stays pool-sized at any context. Automatic enables only qualified model/backend combinations when full KV no longer fits; Laguna's CUDA path remains preview, and DeepSeek uses its native MLA-compressed cache instead of generic KVFlash. Drafter-scored residency uses the shared Qwen3-0.6B scorer; LRU is the fallback when no scorer is present, while qwen35 can use target-native QK scoring. Per-model numbers in [Luce KVFlash →](optimizations/kvflash/README.md). | Flag / env | Default | Effect | |---|---|---| diff --git a/optimizations/pflash/README.md b/optimizations/pflash/README.md index 30d0c7371..1f0f3f1e0 100644 --- a/optimizations/pflash/README.md +++ b/optimizations/pflash/README.md @@ -114,7 +114,7 @@ When `--prefill-compression != off`, the server auto-sets `DFLASH27B_LM_HEAD_FIX --prefill-drafter server/models/Qwen3-0.6B-BF16.gguf ``` -Below the threshold the server runs the standard target generate (no compression). Above it, the server transparently runs `compress` on the daemon, swaps the prompt for the compressed text, and continues the normal `/v1/chat/completions` flow. Tool-calling requests (`req.tools` non-empty) skip compression so JSON tool definitions stay intact. +Below the threshold the server runs standard target prefill. Above it, a true long one-shot request transparently runs `compress`, swaps in the selected token stream, and continues the normal `/v1/chat/completions` flow. Structured system/tool chats stay verbatim so their target prefix remains reusable. Continuations also preserve that prefix unless the request explicitly enables FlowKV, which compresses only aged messages. Exact repeated prompts can restore a completed target snapshot and skip both scorer and target prefill. Validated end-to-end at 64K and 128K source on RTX 3090 (Qwen3.6-27B Q4_K_M target + Qwen3.5-DFlash draft + Qwen3-0.6B BF16 drafter). diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index 383f29ac7..fd1e61295 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -375,11 +375,10 @@ fi # share/model_cards/.json). When unset, the C++ server uses its default # ("dflash"). Lets an operator surface the real model id without a wrapper. : "${DFLASH_MODEL_NAME:=}" -# Phase-1 (thinking) cap when a request opts into thinking. Default mirrors -# antirez/ds4 ds4_eval.c: think_max_tokens = max_tokens(16000) - hard_limit -# reply budget(512) = 15488. The server's own hardcoded default is 10000; -# overriding here aligns ds4-eval and similar reasoning benches with upstream. -: "${DFLASH_THINK_MAX:=15488}" +# Optional phase-1 (thinking) cap. When unset, the native server resolves it +# from the selected model card, then its per-family and hard fallbacks. A +# global DeepSeek-derived cap is unsafe for models with larger card budgets. +: "${DFLASH_THINK_MAX:=}" # Soft-close thinking termination dial (PR #326). Lets the AR loop force # early when the close-token logit comes within this probability # ratio of the chosen-token logit. Range [0.0, 1.0]; 0.0 = disabled (server @@ -427,12 +426,18 @@ if [ -z "$DFLASH_TARGET" ] && [ -d "$DFLASH_DIR/models" ]; then # anything under models/draft/. Sort alphabetically for determinism. TARGET_CANDIDATES=() while IFS= read -r candidate; do - [ -n "$candidate" ] && TARGET_CANDIDATES+=("$candidate") + [ -n "$candidate" ] || continue + # Avoid GNU-only find predicates (`-printf`, `-size +5G`). The native + # entrypoint is also exercised from macOS contributor checkouts, and + # `wc -c` obtains a regular file's logical size without reading a + # sparse multi-gigabyte fixture into memory. + candidate_bytes=$(wc -c <"$candidate" 2>/dev/null || echo 0) + [ "$candidate_bytes" -gt 5368709120 ] 2>/dev/null \ + && TARGET_CANDIDATES+=("$candidate") done < <( find -L "$DFLASH_DIR/models" -maxdepth 4 -type f -name '*.gguf' \ - -size +5G \ - -not -path '*/draft/*' \ - -printf '%p\n' 2>/dev/null \ + ! -path '*/draft/*' \ + -print 2>/dev/null \ | sort ) case "${#TARGET_CANDIDATES[@]}" in @@ -599,8 +604,9 @@ fi CMD=("$DFLASH_SERVER_BIN" "$DFLASH_TARGET" --host "$DFLASH_HOST" --port "$DFLASH_PORT" - --max-ctx "$DFLASH_MAX_CTX" - --think-max-tokens "$DFLASH_THINK_MAX") + --max-ctx "$DFLASH_MAX_CTX") + +[ -n "$DFLASH_THINK_MAX" ] && CMD+=(--think-max-tokens "$DFLASH_THINK_MAX") if [ -n "$DFLASH_TARGET_DEVICES" ]; then [ -z "$DFLASH_TARGET_DEVICE" ] \ diff --git a/server/src/device_runtime.h b/server/src/device_runtime.h index 8226d7900..0d866a777 100644 --- a/server/src/device_runtime.h +++ b/server/src/device_runtime.h @@ -50,6 +50,7 @@ using __nv_bfloat16 = __hip_bfloat16; #define cudaEventDestroy hipEventDestroy #define cudaStreamSynchronize hipStreamSynchronize #define cudaGetLastError hipGetLastError +#define cudaPeekAtLastError hipPeekAtLastError #define cudaGetErrorString hipGetErrorString #define cudaDeviceSynchronize hipDeviceSynchronize #define cudaSetDevice hipSetDevice diff --git a/server/src/flashprefill.cpp b/server/src/flashprefill.cpp index 0745c66bc..8502f0059 100644 --- a/server/src/flashprefill.cpp +++ b/server/src/flashprefill.cpp @@ -5,6 +5,7 @@ // 4. sparse_flash_forward_bf16 (kernel) #include "flashprefill.h" +#include "flashprefill_launchers.h" #include #include @@ -23,24 +24,6 @@ namespace flashprefill { #if defined(DFLASH27B_HAVE_FLASHPREFILL) || defined(DFLASH27B_HAVE_SM80_FLASHPREFILL) extern "C" { -int launch_compute_mean_vector_bf16( - const void * K, void * mean_K, - int batch, int seq_len, int n_kv_heads, int head_dim, int block_size, - int s_K_b, int s_K_n, int s_K_h, int s_K_d, - int s_mK_b, int s_mK_m, int s_mK_h, int s_mK_d, - cudaStream_t stream); - -int launch_compute_block_score_bf16( - const void * Q, const void * mean_K, float sm_scale, - void * score, void * score_max, - int batch, int n_q_heads, int n_k_heads, - int seq_len, int head_dim, int block_size, - int s_Q_b, int s_Q_n, int s_Q_h, int s_Q_d, - int s_mK_b, int s_mK_m, int s_mK_h, int s_mK_d, - int s_S_b, int s_S_m, int s_S_n, int s_S_h, - int s_M_b, int s_M_m, int s_M_n, int s_M_h, - cudaStream_t stream); - #ifdef DFLASH27B_BACKEND_HIP // Phase 4 (HIP): mean_Q + tiled rocWMMA GEMM replaces the O(M²) scalar // block-score kernel. ~5-10× faster on the score step at 8K-32K context. @@ -270,6 +253,7 @@ int flash_prefill_forward_bf16( float * dS = nullptr, * dM = nullptr; int32_t * dIdx = nullptr, * dCnt = nullptr; cudaError_t e; + const char * failure_context = "scratch allocation"; #ifdef DFLASH27B_BACKEND_HIP if ((e = cudaMalloc(&dmK, (size_t)B * M_gemm * Hk * D * 2)) != cudaSuccess) goto err; // bf16, padded #else @@ -289,6 +273,7 @@ int flash_prefill_forward_bf16( if (prof) for (int i=0;i<5;i++) cudaEventCreate(&pE[i]); if (prof) cudaEventRecord(pE[0]); // 1. mean_K + failure_context = "mean-vector launch"; if (launch_compute_mean_vector_bf16( K, dmK, B, S, Hk, D, BLOCK, s_K_b, s_K_n, s_K_h, s_K_d, @@ -298,10 +283,12 @@ int flash_prefill_forward_bf16( // 2. block scores #ifdef DFLASH27B_BACKEND_HIP // Phase 4: mean_Q + rocWMMA GEMM replaces the O(M²) scalar kernel. + failure_context = "mean-query launch"; if (launch_compute_mean_vector_bf16( Q, dmQ, B, S, H, D, BLOCK, s_Q_b, s_Q_n, s_Q_h, s_Q_d, s_mQ_b, s_mQ_m, s_mQ_h, 1, 0) != 0) goto err; + failure_context = "block-score GEMM launch"; if (launch_compute_block_score_gemm_bf16( dmQ, dmK, scale, dS, B, H, Hk, M, D, @@ -309,6 +296,7 @@ int flash_prefill_forward_bf16( s_mK_b, s_mK_m, s_mK_h, s_S_b, s_S_m, s_S_n, s_S_h, 0) != 0) goto err; #else + failure_context = "block-score launch"; if (launch_compute_block_score_bf16( Q, dmK, scale, dS, dM, B, H, Hk, S, D, BLOCK, @@ -387,7 +375,14 @@ int flash_prefill_forward_bf16( if (dM) cudaFree(dM); if (dIdx) cudaFree(dIdx); if (dCnt) cudaFree(dCnt); - std::fprintf(stderr, "[flashprefill] cudaMalloc failed: %s\n", cudaGetErrorString(e)); + if (e != cudaSuccess) { + std::fprintf(stderr, "[flashprefill] %s failed: %s\n", + failure_context, cudaGetErrorString(e)); + } else { + // Shape/dispatch failures are reported by the launcher itself and do + // not necessarily set a CUDA/HIP runtime error. + std::fprintf(stderr, "[flashprefill] %s failed\n", failure_context); + } return -1; } diff --git a/server/src/flashprefill_kernels.cu b/server/src/flashprefill_kernels.cu index 22dd72bd6..e98552aba 100644 --- a/server/src/flashprefill_kernels.cu +++ b/server/src/flashprefill_kernels.cu @@ -33,8 +33,10 @@ #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 #include +#include #include #include "device_runtime.h" +#include "flashprefill_launchers.h" #if !defined(DFLASH27B_BACKEND_HIP) #include #endif @@ -101,24 +103,34 @@ __global__ void compute_mean_vector_kernel_bf16( } // Public launcher (called from C++). -extern "C" void launch_compute_mean_vector_bf16( +extern "C" int launch_compute_mean_vector_bf16( const void * K, void * mean_K, int batch, int seq_len, int n_kv_heads, int head_dim, int block_size, int s_K_b, int s_K_n, int s_K_h, int s_K_d, int s_mK_b, int s_mK_m, int s_mK_h, int s_mK_d, cudaStream_t stream) { + if (head_dim != 128 || block_size != 128) { + std::fprintf(stderr, + "[flashprefill] mean-vector unsupported shape: head_dim=%d block_size=%d\n", + head_dim, block_size); + return -1; + } const int n_k_blocks = (seq_len + block_size - 1) / block_size; dim3 grid(n_k_blocks, batch * n_kv_heads, 1); dim3 block(head_dim, 1, 1); - if (head_dim == 128 && block_size == 128) { - compute_mean_vector_kernel_bf16<128, 128><<>>( - (const __nv_bfloat16 *)K, (__nv_bfloat16 *)mean_K, - batch, seq_len, n_kv_heads, - s_K_b, s_K_n, s_K_h, s_K_d, - s_mK_b, s_mK_m, s_mK_h, s_mK_d); + compute_mean_vector_kernel_bf16<128, 128><<>>( + (const __nv_bfloat16 *)K, (__nv_bfloat16 *)mean_K, + batch, seq_len, n_kv_heads, + s_K_b, s_K_n, s_K_h, s_K_d, + s_mK_b, s_mK_m, s_mK_h, s_mK_d); + const cudaError_t error = cudaPeekAtLastError(); + if (error != cudaSuccess) { + std::fprintf(stderr, "[flashprefill] mean-vector launch failed: %s\n", + cudaGetErrorString(error)); + return -1; } - // Only D_HEAD=128 BLOCK=128 dispatched here. Add other combos when new heads/blocks needed. + return 0; } // ---- Kernel 2: compute_block_score ---- @@ -229,7 +241,7 @@ __global__ void compute_block_score_kernel_bf16( } } -extern "C" void launch_compute_block_score_bf16( +extern "C" int launch_compute_block_score_bf16( const void * Q, const void * mean_K, float sm_scale, void * score, void * score_max, int batch, int n_q_heads, int n_k_heads, @@ -240,20 +252,31 @@ extern "C" void launch_compute_block_score_bf16( int s_M_b, int s_M_m, int s_M_n, int s_M_h, cudaStream_t stream) { + if (head_dim != 128 || block_size != 128) { + std::fprintf(stderr, + "[flashprefill] block-score unsupported shape: head_dim=%d block_size=%d\n", + head_dim, block_size); + return -1; + } const int M = (seq_len + block_size - 1) / block_size; dim3 grid(M, batch * n_q_heads, 1); dim3 block(block_size, 1, 1); size_t smem = block_size * sizeof(float); - if (head_dim == 128 && block_size == 128) { - compute_block_score_kernel_bf16<128, 128, 1><<>>( - (const __nv_bfloat16 *)Q, (const __nv_bfloat16 *)mean_K, sm_scale, - (float *)score, (float *)score_max, - batch, n_q_heads, n_k_heads, M, M, - s_Q_b, s_Q_n, s_Q_h, s_Q_d, - s_mK_b, s_mK_m, s_mK_h, s_mK_d, - s_S_b, s_S_m, s_S_n, s_S_h, - s_M_b, s_M_m, s_M_n, s_M_h); + compute_block_score_kernel_bf16<128, 128, 1><<>>( + (const __nv_bfloat16 *)Q, (const __nv_bfloat16 *)mean_K, sm_scale, + (float *)score, (float *)score_max, + batch, n_q_heads, n_k_heads, M, M, + s_Q_b, s_Q_n, s_Q_h, s_Q_d, + s_mK_b, s_mK_m, s_mK_h, s_mK_d, + s_S_b, s_S_m, s_S_n, s_S_h, + s_M_b, s_M_m, s_M_n, s_M_h); + const cudaError_t error = cudaPeekAtLastError(); + if (error != cudaSuccess) { + std::fprintf(stderr, "[flashprefill] block-score launch failed: %s\n", + cudaGetErrorString(error)); + return -1; } + return 0; } // ---- Kernel 4: sparse_flash_forward ---- diff --git a/server/src/flashprefill_kernels.hip.cu b/server/src/flashprefill_kernels.hip.cu index eb3ce1cd2..ff2d46d34 100644 --- a/server/src/flashprefill_kernels.hip.cu +++ b/server/src/flashprefill_kernels.hip.cu @@ -22,6 +22,7 @@ #include #include #include +#include "flashprefill_launchers.h" // These kernels are WAVE32-ONLY BY DESIGN, not merely wave32-tuned. Kernel 4's // accumulator handling assumes the RDNA3 v_wmma_f32_16x16x16 fragment layout @@ -96,6 +97,12 @@ extern "C" int launch_compute_mean_vector_bf16( batch, seq_len, n_kv_heads, s_K_b, s_K_n, s_K_h, s_K_d, s_mK_b, s_mK_m, s_mK_h, s_mK_d); + const hipError_t error = hipPeekAtLastError(); + if (error != hipSuccess) { + fprintf(stderr, "[dflash] mean-vector launch failed: %s\n", + hipGetErrorString(error)); + return -1; + } return 0; } @@ -211,6 +218,12 @@ extern "C" int launch_compute_block_score_bf16( s_mK_b, s_mK_m, s_mK_h, s_mK_d, s_S_b, s_S_m, s_S_n, s_S_h, s_M_b, s_M_m, s_M_n, s_M_h); + const hipError_t error = hipPeekAtLastError(); + if (error != hipSuccess) { + fprintf(stderr, "[dflash] block-score launch failed: %s\n", + hipGetErrorString(error)); + return -1; + } return 0; } diff --git a/server/src/flashprefill_launchers.h b/server/src/flashprefill_launchers.h new file mode 100644 index 000000000..4be9b1e41 --- /dev/null +++ b/server/src/flashprefill_launchers.h @@ -0,0 +1,31 @@ +#pragma once + +// Internal ABI shared by the FlashPrefill orchestration layer and the +// architecture-specific CUDA/HIP launchers. Keep these declarations in one +// header: a return-type mismatch across separate `extern "C"` declarations is +// not diagnosed by the linker and previously made successful CUDA launches +// look like failures on aarch64/GB10. + +#include "device_runtime.h" + +namespace dflash::common::flashprefill { + +extern "C" int launch_compute_mean_vector_bf16( + const void * K, void * mean_K, + int batch, int seq_len, int n_kv_heads, int head_dim, int block_size, + int s_K_b, int s_K_n, int s_K_h, int s_K_d, + int s_mK_b, int s_mK_m, int s_mK_h, int s_mK_d, + cudaStream_t stream); + +extern "C" int launch_compute_block_score_bf16( + const void * Q, const void * mean_K, float sm_scale, + void * score, void * score_max, + int batch, int n_q_heads, int n_k_heads, + int seq_len, int head_dim, int block_size, + int s_Q_b, int s_Q_n, int s_Q_h, int s_Q_d, + int s_mK_b, int s_mK_m, int s_mK_h, int s_mK_d, + int s_S_b, int s_S_m, int s_S_n, int s_S_h, + int s_M_b, int s_M_m, int s_M_n, int s_M_h, + cudaStream_t stream); + +} // namespace dflash::common::flashprefill diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 46081fa23..69b34281a 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -790,30 +790,176 @@ std::string render_tool_call_xml(const std::string & name, const json & argument return out; } +struct ReplayToolCall { + std::string id; + std::string name; + json arguments = json::object(); +}; + +struct ToolReplay { + std::string text; + bool exact = false; +}; + +std::string extract_text_part(const json & part) { + if (!part.is_object()) return ""; + const std::string type = part.value("type", ""); + if (type == "text" || type == "input_text" || type == "output_text") { + return part.value("text", ""); + } + return ""; +} + +std::string extract_text_content(const json & content) { + if (content.is_string()) return content.get(); + if (!content.is_array()) return ""; + + std::string text; + for (const auto & part : content) { + text += extract_text_part(part); + } + return text; +} + +ToolReplay replay_or_render_tool_calls( + const std::vector & calls, + ToolMemory & tool_memory) { + std::vector ids; + ids.reserve(calls.size()); + for (const auto & call : calls) { + if (call.id.empty()) { + ids.clear(); + break; + } + ids.push_back(call.id); + } + if (!ids.empty()) { + std::string replay = tool_memory.lookup(ids); + if (!replay.empty()) return {std::move(replay), true}; + } + + std::string rendered; + for (const auto & call : calls) { + rendered += render_tool_call_xml(call.name, call.arguments); + } + return {std::move(rendered), false}; +} + +std::vector openai_tool_calls(const json & value) { + std::vector calls; + if (!value.is_array()) return calls; + for (const auto & item : value) { + if (!item.is_object()) continue; + const json function = item.value("function", json::object()); + calls.push_back({ + item.value("id", ""), + function.value("name", ""), + parse_responses_arguments(function), + }); + } + return calls; +} + +bool append_anthropic_content( + const json & message, + ToolMemory & tool_memory, + std::vector & chat_messages) { + if (!message.contains("content") || !message["content"].is_array()) { + return false; + } + const auto & content = message["content"]; + const std::string role = message.value("role", "user"); + + if (role == "assistant") { + std::vector calls; + for (const auto & part : content) { + if (!part.is_object() || part.value("type", "") != "tool_use") continue; + calls.push_back({ + part.value("id", ""), + part.value("name", ""), + part.value("input", json::object()), + }); + } + if (!calls.empty()) { + ToolReplay replay = replay_or_render_tool_calls(calls, tool_memory); + if (!replay.exact) { + replay.text = extract_text_content(content) + replay.text; + } + chat_messages.push_back({"assistant", std::move(replay.text)}); + return true; + } + } + + if (role != "user") return false; + bool has_tool_result = false; + for (const auto & part : content) { + if (part.is_object() && part.value("type", "") == "tool_result") { + has_tool_result = true; + break; + } + } + if (!has_tool_result) return false; + + std::string user_text; + const auto flush_user_text = [&chat_messages, &user_text]() { + if (user_text.empty()) return; + chat_messages.push_back({"user", std::move(user_text)}); + user_text.clear(); + }; + for (const auto & part : content) { + if (!part.is_object()) continue; + if (part.value("type", "") != "tool_result") { + user_text += extract_text_part(part); + continue; + } + flush_user_text(); + std::string output; + if (part.contains("content")) { + output = extract_text_content(part["content"]); + if (output.empty() && !part["content"].is_null() && + !part["content"].is_string()) { + output = part["content"].dump(); + } + } + chat_messages.push_back({ + "tool", + std::move(output), + part.value("tool_use_id", part.value("id", "")), + }); + } + flush_user_text(); + return true; +} + std::vector normalize_chat_messages( const json & messages, ApiFormat format, ToolMemory & tool_memory) { std::vector chat_msgs; std::vector system_parts; + std::vector pending_responses_calls; + const auto flush_responses_calls = [&]() { + if (pending_responses_calls.empty()) return; + chat_msgs.push_back({ + "assistant", + replay_or_render_tool_calls(pending_responses_calls, tool_memory).text, + }); + pending_responses_calls.clear(); + }; if (messages.is_array()) { for (const auto & m : messages) { if (format == ApiFormat::RESPONSES && m.is_object()) { std::string item_type = m.value("type", "message"); if (item_type == "function_call") { - std::string call_id = m.value("call_id", m.value("id", "")); - std::string raw; - if (!call_id.empty()) { - raw = tool_memory.lookup({call_id}); - } - if (raw.empty()) { - raw = render_tool_call_xml(m.value("name", ""), - parse_responses_arguments(m)); - } - chat_msgs.push_back({"assistant", raw}); + pending_responses_calls.push_back({ + m.value("call_id", m.value("id", "")), + m.value("name", ""), + parse_responses_arguments(m), + }); continue; } + flush_responses_calls(); if (item_type == "function_call_output") { std::string output; if (m.contains("output") && m["output"].is_string()) { @@ -827,36 +973,28 @@ std::vector normalize_chat_messages( } } + if (format == ApiFormat::ANTHROPIC && m.is_object() && + append_anthropic_content(m, tool_memory, chat_msgs)) { + continue; + } + ChatMessage cm; cm.role = m.value("role", "user"); bool replayed = false; if (cm.role == "assistant" && m.contains("tool_calls") && m["tool_calls"].is_array() && !m["tool_calls"].empty()) { - std::vector call_ids; - for (const auto & tc : m["tool_calls"]) { - std::string id = tc.value("id", ""); - if (!id.empty()) call_ids.push_back(id); - } - std::string raw = tool_memory.lookup(call_ids); - if (!raw.empty()) { - cm.content = raw; - replayed = true; + const auto calls = openai_tool_calls(m["tool_calls"]); + ToolReplay replay = replay_or_render_tool_calls(calls, tool_memory); + cm.content = std::move(replay.text); + if (!replay.exact && m.contains("content")) { + cm.content = extract_text_content(m["content"]) + cm.content; } + replayed = !cm.content.empty(); } if (!replayed) { - if (m.contains("content") && m["content"].is_string()) { - cm.content = m["content"].get(); - } else if (m.contains("content") && m["content"].is_array()) { - for (const auto & part : m["content"]) { - std::string ptype = part.value("type", ""); - if (ptype == "text" || ptype == "input_text" || - ptype == "output_text") { - cm.content += part.value("text", ""); - } - } - } + if (m.contains("content")) cm.content = extract_text_content(m["content"]); } if (format == ApiFormat::RESPONSES && @@ -866,6 +1004,7 @@ std::vector normalize_chat_messages( chat_msgs.push_back(std::move(cm)); } } + flush_responses_calls(); } else if (messages.is_string()) { chat_msgs.push_back({"user", messages.get()}); } @@ -2224,6 +2363,31 @@ bool is_continuation_request(const json & messages) { } // namespace +PflashRequestStrategy select_pflash_request_strategy( + ServerConfig::PflashMode mode, + int prompt_tokens, + int threshold, + bool continuation, + bool has_tools, + bool has_reusable_prefix, + bool flowkv_enabled) { + const bool eligible = + mode == ServerConfig::PflashMode::ALWAYS || + (mode == ServerConfig::PflashMode::AUTO && prompt_tokens >= threshold); + if (!eligible) return PflashRequestStrategy::Off; + + // FlowKV rewrites only aged message bodies and deliberately preserves the + // system/tool prefix. It is therefore the only compression mode that may + // run on a continuation without invalidating turn-boundary snapshots. + if (continuation && flowkv_enabled) { + return PflashRequestStrategy::FlowKv; + } + if (continuation || has_tools || has_reusable_prefix || flowkv_enabled) { + return PflashRequestStrategy::PreservePrefix; + } + return PflashRequestStrategy::WholePrompt; +} + void HttpServer::apply_flowkv_compression( const ParsedRequest & req, PreparedPrompt & prepared) { int hot_window = 2; @@ -2548,33 +2712,46 @@ HttpServer::PreparedPrompt HttpServer::prepare_prompt( if (config_.pflash_mode != ServerConfig::PflashMode::OFF && drafter_tokenizer_ != nullptr) { const int prompt_tokens = (int) req.prompt_tokens.size(); - bool should_compress = - config_.pflash_mode == ServerConfig::PflashMode::ALWAYS || - (config_.pflash_mode == ServerConfig::PflashMode::AUTO && - prompt_tokens >= config_.pflash_threshold); - const bool continuation = should_compress && - is_continuation_request(req.messages); - - if (should_compress && continuation && - req.disk_cache_policy.compress && req.messages.is_array()) { - // FlowKV owns continuation compression; falling back to whole- - // prompt compression would destroy the reusable prefix anchor. + const bool continuation = is_continuation_request(req.messages); + const bool has_tools = req.tools.is_array() && !req.tools.empty(); + // A system/developer message followed by a user message is a stable + // prefix even before the first assistant turn exists. Preserve it so + // turn two can restore target KV instead of inheriting a compressed, + // non-matching token stream. + const bool has_reusable_prefix = + has_tools || (req.messages.is_array() && req.messages.size() > 1); + const bool flowkv_enabled = + req.disk_cache_policy.compress && req.messages.is_array(); + const PflashRequestStrategy strategy = + select_pflash_request_strategy( + config_.pflash_mode, + prompt_tokens, + config_.pflash_threshold, + continuation, + has_tools, + has_reusable_prefix, + flowkv_enabled); + + if (strategy == PflashRequestStrategy::FlowKv) { apply_flowkv_compression(req, prepared); - should_compress = false; - } else if (should_compress && continuation) { - should_compress = false; + } else if (strategy == PflashRequestStrategy::PreservePrefix && + continuation) { std::fprintf(stderr, "[pflash] skip-compress (continuation: prior assistant/tool history)\n"); - } - - if (should_compress && req.disk_cache_policy.compress) { + } else if (strategy == PflashRequestStrategy::PreservePrefix && has_tools) { + // Whole-prompt compression would rewrite the tool schema and make + // the target snapshot unusable on the next agent turn. + std::fprintf(stderr, + "[pflash] skip-compress (tools: preserving reusable system/tool prefix)\n"); + } else if (strategy == PflashRequestStrategy::PreservePrefix && + has_reusable_prefix) { + std::fprintf(stderr, + "[pflash] skip-compress (chat: preserving reusable system prefix)\n"); + } else if (strategy == PflashRequestStrategy::PreservePrefix) { // Turn one stays verbatim so the next turn can reuse its KV prefix. - should_compress = false; std::fprintf(stderr, "[flowkv] turn-1 verbatim (system kept as cache anchor)\n"); - } - - if (should_compress) { + } else if (strategy == PflashRequestStrategy::WholePrompt) { prepared.error = apply_pflash_compression(req, prepared); if (!prepared.error.empty()) { prepared.error_status = 500; diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 56f555369..aa3c3aa6a 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -189,6 +189,25 @@ struct ServerConfig { std::string collect_routing_path; }; +// Request-level composition of long-prompt compression and reusable KV +// prefixes. Exposed as a pure policy function so the cache-sensitive cases +// stay model-free and unit-testable. +enum class PflashRequestStrategy { + Off, + WholePrompt, + FlowKv, + PreservePrefix, +}; + +PflashRequestStrategy select_pflash_request_strategy( + ServerConfig::PflashMode mode, + int prompt_tokens, + int threshold, + bool continuation, + bool has_tools, + bool has_reusable_prefix, + bool flowkv_enabled); + // ─── Parsed request ───────────────────────────────────────────────────── struct ParsedRequest { diff --git a/server/test/test_entrypoint_cache_defaults.sh b/server/test/test_entrypoint_cache_defaults.sh index cb797e981..6cd302e11 100755 --- a/server/test/test_entrypoint_cache_defaults.sh +++ b/server/test/test_entrypoint_cache_defaults.sh @@ -45,7 +45,7 @@ default_output="$( unset DFLASH_PREFIX_CACHE_SLOTS DFLASH_PREFILL_CACHE_SLOTS run_entrypoint )" -for flag in --prefix-cache-slots --prefill-cache-slots; do +for flag in --prefix-cache-slots --prefill-cache-slots --think-max-tokens; do if grep -Fq "SERVER_ARG=$flag" <<<"$default_output"; then echo "entrypoint overrides the native cache default with $flag" >&2 exit 1 @@ -58,9 +58,11 @@ assert_arg_pair "$disabled_output" --prefix-cache-slots 0 configured_output="$( run_entrypoint \ DFLASH_PREFIX_CACHE_SLOTS=4 \ - DFLASH_PREFILL_CACHE_SLOTS=2 + DFLASH_PREFILL_CACHE_SLOTS=2 \ + DFLASH_THINK_MAX=15488 )" assert_arg_pair "$configured_output" --prefix-cache-slots 4 assert_arg_pair "$configured_output" --prefill-cache-slots 2 +assert_arg_pair "$configured_output" --think-max-tokens 15488 echo "entrypoint cache defaults: PASS" diff --git a/server/test/test_flashprefill_kernels.cpp b/server/test/test_flashprefill_kernels.cpp index 7b3927269..f6ac88ea5 100644 --- a/server/test/test_flashprefill_kernels.cpp +++ b/server/test/test_flashprefill_kernels.cpp @@ -20,26 +20,12 @@ #include #include "../src/flashprefill.h" +#include "../src/flashprefill_launchers.h" -extern "C" { -void launch_compute_mean_vector_bf16( - const void * K, void * mean_K, - int batch, int seq_len, int n_kv_heads, int head_dim, int block_size, - int s_K_b, int s_K_n, int s_K_h, int s_K_d, - int s_mK_b, int s_mK_m, int s_mK_h, int s_mK_d, - cudaStream_t stream); - -void launch_compute_block_score_bf16( - const void * Q, const void * mean_K, float sm_scale, - void * score, void * score_max, - int batch, int n_q_heads, int n_k_heads, - int seq_len, int head_dim, int block_size, - int s_Q_b, int s_Q_n, int s_Q_h, int s_Q_d, - int s_mK_b, int s_mK_m, int s_mK_h, int s_mK_d, - int s_S_b, int s_S_m, int s_S_n, int s_S_h, - int s_M_b, int s_M_m, int s_M_n, int s_M_h, - cudaStream_t stream); +using dflash::common::flashprefill::launch_compute_block_score_bf16; +using dflash::common::flashprefill::launch_compute_mean_vector_bf16; +extern "C" { void launch_sparse_flash_forward_bf16( const void * Q, const void * K, const void * V, void * O, const int32_t * block_index, const int32_t * counts, @@ -110,10 +96,13 @@ int main() { int s_cnt_b = M * H, s_cnt_m = H, s_cnt_h = 1; // ── Kernel 1: compute_mean_vector ── - launch_compute_mean_vector_bf16( + if (launch_compute_mean_vector_bf16( dK, dmK, B, S, Hk, D, BLOCK, s_K_b, s_K_n, s_K_h, s_K_d, - s_mK_b, s_mK_m, s_mK_h, s_mK_d, 0); + s_mK_b, s_mK_m, s_mK_h, s_mK_d, 0) != 0) { + std::fprintf(stderr, "[fp-test] kernel 1 launcher failed\n"); + return 1; + } CK(cudaDeviceSynchronize()); // Verify host vs GPU mean for one block. @@ -140,13 +129,16 @@ int main() { // ── Kernel 2: compute_block_score ── float scale = 1.0f / std::sqrt((float)D); - launch_compute_block_score_bf16( + if (launch_compute_block_score_bf16( dQ, dmK, scale, dS, dM, B, H, Hk, S, D, BLOCK, s_Q_b, s_Q_n, s_Q_h, s_Q_d, s_mK_b, s_mK_m, s_mK_h, s_mK_d, s_S_b, s_S_m, s_S_n, s_S_h, - s_S_b, s_S_m, s_S_n, s_S_h, 0); + s_S_b, s_S_m, s_S_n, s_S_h, 0) != 0) { + std::fprintf(stderr, "[fp-test] kernel 2 launcher failed\n"); + return 1; + } CK(cudaDeviceSynchronize()); std::printf("[fp-test] kernel 2 (compute_block_score): launch ok\n"); @@ -256,9 +248,12 @@ int main() { cfg.alpha = 0.12f; // Warm-up - dflash::common::flashprefill::flash_prefill_forward_bf16( + if (dflash::common::flashprefill::flash_prefill_forward_bf16( bdQ, bdK, bdV, bdO, BB, BS, BH, BHk, BD, - 1.0f / std::sqrt((float)BD), cfg); + 1.0f / std::sqrt((float)BD), cfg) != 0) { + std::fprintf(stderr, "[fp-test] e2e warm-up failed\n"); + return 1; + } CK(cudaDeviceSynchronize()); cudaEvent_t e_a, e_b; @@ -266,9 +261,12 @@ int main() { cudaEventCreate(&e_b); cudaEventRecord(e_a); for (int it = 0; it < 5; ++it) { - dflash::common::flashprefill::flash_prefill_forward_bf16( + if (dflash::common::flashprefill::flash_prefill_forward_bf16( bdQ, bdK, bdV, bdO, BB, BS, BH, BHk, BD, - 1.0f / std::sqrt((float)BD), cfg); + 1.0f / std::sqrt((float)BD), cfg) != 0) { + std::fprintf(stderr, "[fp-test] e2e timed iteration failed\n"); + return 1; + } } cudaEventRecord(e_b); cudaEventSynchronize(e_b); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 06e3521d4..c211e77d8 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -2365,6 +2365,152 @@ TEST_CASE(ServerUnitFixture, test_normalize_responses_tool_followup_messages) { } } +TEST_CASE(ServerUnitFixture, test_normalize_anthropic_tool_followup_messages) { + ToolMemory tool_memory; + const std::string first_id = "toolu_read_001"; + const std::string second_id = "toolu_search_002"; + const std::string raw_tool_turn = + "\n\nREADME.md\n\n" + "\n\n\nDFlash\n\n" + "\n"; + tool_memory.remember({first_id, second_id}, raw_tool_turn); + + json messages = json::array({ + { + {"role", "assistant"}, + {"content", json::array({ + {{"type", "text"}, {"text", "I will inspect both files."}}, + { + {"type", "tool_use"}, + {"id", first_id}, + {"name", "read_file"}, + {"input", {{"path", "README.md"}}}, + }, + { + {"type", "tool_use"}, + {"id", second_id}, + {"name", "search"}, + {"input", {{"query", "DFlash"}}}, + }, + })}, + }, + { + {"role", "user"}, + {"content", json::array({ + { + {"type", "tool_result"}, + {"tool_use_id", first_id}, + {"content", "# Lucebox"}, + }, + { + {"type", "tool_result"}, + {"tool_use_id", second_id}, + {"content", json::array({{ + {"type", "text"}, + {"text", "Found 3 matches"}, + }})}, + }, + {{"type", "text"}, {"text", "Now summarize."}}, + })}, + }, + }); + + auto chat_msgs = normalize_chat_messages(messages, ApiFormat::ANTHROPIC, tool_memory); + TEST_ASSERT(chat_msgs.size() == 4); + if (chat_msgs.size() == 4) { + TEST_ASSERT(chat_msgs[0].role == "assistant"); + TEST_ASSERT(chat_msgs[0].content == raw_tool_turn); + TEST_ASSERT(chat_msgs[1].role == "tool"); + TEST_ASSERT(chat_msgs[1].tool_call_id == first_id); + TEST_ASSERT(chat_msgs[1].content == "# Lucebox"); + TEST_ASSERT(chat_msgs[2].role == "tool"); + TEST_ASSERT(chat_msgs[2].tool_call_id == second_id); + TEST_ASSERT(chat_msgs[2].content == "Found 3 matches"); + TEST_ASSERT(chat_msgs[3].role == "user"); + TEST_ASSERT(chat_msgs[3].content == "Now summarize."); + } +} + +TEST_CASE(ServerUnitFixture, test_normalize_responses_parallel_calls_replay_once) { + ToolMemory tool_memory; + const std::string first_id = "call_read_001"; + const std::string second_id = "call_search_002"; + const std::string raw_tool_turn = + "\n\nREADME.md\n\n" + "\n\n\nDFlash\n\n" + "\n"; + tool_memory.remember({first_id, second_id}, raw_tool_turn); + + json messages = json::array({ + { + {"type", "message"}, + {"role", "user"}, + {"content", json::array({{ + {"type", "input_text"}, + {"text", "Inspect the repository"}, + }})}, + }, + { + {"type", "function_call"}, + {"call_id", first_id}, + {"name", "read_file"}, + {"arguments", R"({"path":"README.md"})"}, + }, + { + {"type", "function_call"}, + {"call_id", second_id}, + {"name", "search"}, + {"arguments", R"({"query":"DFlash"})"}, + }, + { + {"type", "function_call_output"}, + {"call_id", first_id}, + {"output", "# Lucebox"}, + }, + { + {"type", "function_call_output"}, + {"call_id", second_id}, + {"output", "Found 3 matches"}, + }, + }); + + auto chat_msgs = normalize_chat_messages(messages, ApiFormat::RESPONSES, tool_memory); + TEST_ASSERT(chat_msgs.size() == 4); + if (chat_msgs.size() == 4) { + TEST_ASSERT(chat_msgs[0].role == "user"); + TEST_ASSERT(chat_msgs[1].role == "assistant"); + TEST_ASSERT(chat_msgs[1].content == raw_tool_turn); + TEST_ASSERT(chat_msgs[2].role == "tool"); + TEST_ASSERT(chat_msgs[2].tool_call_id == first_id); + TEST_ASSERT(chat_msgs[3].role == "tool"); + TEST_ASSERT(chat_msgs[3].tool_call_id == second_id); + } +} + +TEST_CASE(ServerUnitFixture, test_normalize_openai_tool_fallback_survives_restart) { + ToolMemory empty_memory; + json messages = json::array({{ + {"role", "assistant"}, + {"content", "I will inspect it."}, + {"tool_calls", json::array({{ + {"id", "call_missing_after_restart"}, + {"type", "function"}, + {"function", { + {"name", "read_file"}, + {"arguments", R"({"path":"README.md"})"}, + }}, + }})}, + }}); + + auto chat_msgs = normalize_chat_messages(messages, ApiFormat::OPENAI_CHAT, empty_memory); + TEST_ASSERT(chat_msgs.size() == 1); + if (chat_msgs.size() == 1) { + TEST_ASSERT(chat_msgs[0].content.find("I will inspect it.") != std::string::npos); + TEST_ASSERT(chat_msgs[0].content.find("") != std::string::npos); + TEST_ASSERT(chat_msgs[0].content.find("README.md") != std::string::npos); + } +} + // ═══════════════════════════════════════════════════════════════════════ // Placement config tests // ═══════════════════════════════════════════════════════════════════════ @@ -4576,6 +4722,35 @@ TEST_CASE(ServerUnitFixture, test_prefix_key_stable_across_header_change) { // FlowKV + disk-cache compose tests (T1–T7) +TEST_CASE(ServerUnitFixture, test_pflash_request_policy_preserves_agent_prefixes) { + using Mode = ServerConfig::PflashMode; + using Strategy = PflashRequestStrategy; + + TEST_ASSERT(select_pflash_request_strategy( + Mode::AUTO, 16000, 32000, false, false, false, false) == Strategy::Off); + TEST_ASSERT(select_pflash_request_strategy( + Mode::AUTO, 64000, 32000, false, false, false, false) == + Strategy::WholePrompt); + TEST_ASSERT(select_pflash_request_strategy( + Mode::AUTO, 64000, 32000, false, true, true, false) == + Strategy::PreservePrefix); + TEST_ASSERT(select_pflash_request_strategy( + Mode::AUTO, 64000, 32000, false, false, true, false) == + Strategy::PreservePrefix); + TEST_ASSERT(select_pflash_request_strategy( + Mode::AUTO, 64000, 32000, true, false, true, false) == + Strategy::PreservePrefix); + TEST_ASSERT(select_pflash_request_strategy( + Mode::AUTO, 64000, 32000, true, true, true, true) == + Strategy::FlowKv); + TEST_ASSERT(select_pflash_request_strategy( + Mode::AUTO, 64000, 32000, false, false, false, true) == + Strategy::PreservePrefix); + TEST_ASSERT(select_pflash_request_strategy( + Mode::ALWAYS, 1, 32000, false, false, false, false) == + Strategy::WholePrompt); +} + // T4 (compress=false): policy name has no "+compress" suffix. TEST_CASE(ServerUnitFixture, test_flowkv_T4_compress_false_policy_name_no_suffix) { DiskPrefixCachePolicy p; From 6e085c57e8138810cd04d59115e26fff25c11a43 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sat, 1 Aug 2026 13:02:22 +0200 Subject: [PATCH 19/25] ci(images): split CUDA targets by architecture --- .github/workflows/docker.yml | 21 +++++-- .github/workflows/speed-profile.yml | 67 ++++++++++++++++++--- Dockerfile | 19 +++--- docker-bake.hcl | 93 +++++++++++++++++++++++------ 4 files changed, 159 insertions(+), 41 deletions(-) diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index 462b1b938..c67f3bd13 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -1,7 +1,7 @@ name: Docker prebuilds -# Builds the cuda12 lucebox-hub Docker image defined in docker-bake.hcl -# and pushes it to GHCR. The bake file is the source of +# Builds the cuda12, cuda128/RTX5090, cuda13/GB10, and ROCm images defined in +# docker-bake.hcl and pushes them to GHCR. The bake file is the source of # truth for arch matrices and CUDA pinning; this workflow only handles # fetching submodules, freeing runner disk, signing in to the registry, and # wiring the cache. @@ -12,7 +12,7 @@ on: # event=tag` + `type=semver` rules below. release: types: [published] - # Build + push the rolling `:cuda12` / `:rocm` tags when an image-affecting + # Build + push the rolling CUDA / ROCm tags when an image-affecting # file changes on main, so the public images track main without a ~2h # rebuild on every unrelated commit (docs, harness, server tweaks that # don't reach the image). Same paths as the PR guard below. The @@ -90,14 +90,23 @@ jobs: # halve wall time. # • A self-hosted runner with the host's nvcc avoids the # containerised CUDA toolkit pull entirely. - runs-on: ubuntu-latest + runs-on: ${{ matrix.runner }} permissions: contents: read packages: write strategy: fail-fast: false matrix: - variant: [cuda12, rocm] + include: + - variant: cuda12 + runner: ubuntu-latest + - variant: cuda128 + runner: ubuntu-latest + - variant: cuda13 + # Native arm64 compilation avoids hours of QEMU-emulated nvcc. + runner: ubuntu-24.04-arm + - variant: rocm + runner: ubuntu-latest steps: - name: Free runner disk space # The default ubuntu-latest image keeps ~25 GB of preinstalled @@ -190,7 +199,7 @@ jobs: # under the concurrency group's pre-emption window. Release / main / # dispatch builds keep the full consumer-GPU list so the published # image runs on every supported card. - DFLASH_CUDA_ARCHES: ${{ github.event_name == 'pull_request' && '86' || '75;80;86;89;90;120' }} + DFLASH_CUDA_ARCHES: ${{ github.event_name == 'pull_request' && '86' || '75;80;86;89;90' }} # Same split for HIP: PR builds compile gfx1151 only (the verified # reference, Strix Halo); main / release builds widen to consumer # RDNA so the published :rocm runs on RX 7900 (gfx1100) and both RDNA4 diff --git a/.github/workflows/speed-profile.yml b/.github/workflows/speed-profile.yml index c6e5991e1..4f43c4a46 100644 --- a/.github/workflows/speed-profile.yml +++ b/.github/workflows/speed-profile.yml @@ -2,8 +2,8 @@ name: Speed Profile # Report-only speed profile for the inference engine. Runs on the self-hosted # RTX 3090 (lucebox3) on PRs that touch the engine or the optimizations, and on -# manual dispatch. It NEVER blocks a PR (continue-on-error: true) — it publishes a -# report to the run summary + uploads the JSON / markdown / nsys trace as artifacts. +# manual dispatch. Benchmark/runtime failures are warnings, not a red PR check; +# the job publishes a summary plus JSON / markdown / nsys artifacts for triage. # # Why report-only: perf has run-to-run variance (thermals, clocks, scheduling). # Gating a merge on a noisy absolute number produces false failures. We surface the @@ -30,7 +30,9 @@ jobs: name: Speed profile (self-hosted RTX 3090, sm_86) runs-on: [self-hosted, gpu, sm86] timeout-minutes: 30 - continue-on-error: true # report-only: a slow/failed profile must not block the PR + # Individual benchmark failures are converted into warnings below. Keeping + # the job itself green avoids a misleading red PR check for a report-only + # measurement while still preserving logs and artifacts for triage. # Model paths live on the runner, not in the repo (multi-GB weights). They are # overridable via repo variables so the runner owner can point at whatever is @@ -51,7 +53,7 @@ jobs: - name: GPU info (and pin clocks to cut variance, if permitted) run: | - nvidia-smi --query-gpu=name,driver_version,memory.total,power.limit --format=csv + nvidia-smi --query-gpu=name,driver_version,memory.total,memory.free,power.limit --format=csv # Locking clocks makes the numbers comparable run-to-run. Safe to skip if the # runner user can't run nvidia-smi -lgc; the profiler still records the power cap. sudo nvidia-smi -lgc 1395 2>/dev/null || echo "clock lock not permitted; continuing" @@ -94,8 +96,43 @@ jobs: echo "::warning title=Speed profile skipped::Model weights not found under $MODELS — see the run summary." fi - - name: Build engine binaries (sm_86, Release) + - name: Wait for exclusive GPU capacity + id: capacity if: steps.models.outputs.present == 'true' + env: + MIN_FREE_MIB: ${{ vars.LUCEBOX_SPEED_MIN_FREE_MIB || '22000' }} + run: | + # The 27B target + draft fit on a clean 24 GB 3090, but fail with a + # cryptic cudaMalloc error when another process temporarily owns VRAM. + # Never kill a user's process: wait briefly, then skip this report. + ready=false + free_mib=0 + for attempt in {1..30}; do + free_mib=$(nvidia-smi --query-gpu=memory.free \ + --format=csv,noheader,nounits | sed -n '1p' | tr -d '[:space:]') + if [[ "$free_mib" =~ ^[0-9]+$ ]] && [ "$free_mib" -ge "$MIN_FREE_MIB" ]; then + ready=true + break + fi + if [ "$attempt" -eq 1 ]; then + echo "Only ${free_mib:-unknown} MiB is free; waiting for ${MIN_FREE_MIB} MiB." + fi + sleep 10 + done + echo "ready=$ready" >> "$GITHUB_OUTPUT" + echo "free_mib=${free_mib:-0}" >> "$GITHUB_OUTPUT" + if [ "$ready" != true ]; then + echo "::warning title=Speed profile skipped::GPU remained busy (${free_mib:-unknown} MiB free; ${MIN_FREE_MIB} MiB required)." + { + echo "## 🏎️ Speed profile — skipped (GPU busy)" + echo "" + echo "The RTX 3090 had only \`${free_mib:-unknown} MiB\` free after a five-minute wait;" + echo "the model-backed profile requires \`${MIN_FREE_MIB} MiB\`. No process was killed." + } >> "$GITHUB_STEP_SUMMARY" + fi + + - name: Build engine binaries (sm_86, Release) + if: steps.models.outputs.present == 'true' && steps.capacity.outputs.ready == 'true' run: | cd server cmake -B build \ @@ -106,7 +143,7 @@ jobs: cmake --build build --target test_dflash test_generate -j"$(nproc)" - name: Install profiler Python deps (isolated, pinned venv) - if: steps.models.outputs.present == 'true' + if: steps.models.outputs.present == 'true' && steps.capacity.outputs.ready == 'true' run: | cd server # The profiler only needs a tokenizer, so we use a tiny isolated venv @@ -123,7 +160,8 @@ jobs: protobuf==6.31.1 - name: Run speed profiler - if: steps.models.outputs.present == 'true' + id: profile + if: steps.models.outputs.present == 'true' && steps.capacity.outputs.ready == 'true' run: | cd server # Use a committed baseline if one is staged so the report can flag a @@ -143,6 +181,7 @@ jobs: # The nsys pass adds a separate short profiled run; tok/s is measured on clean passes. Run 5 # timing reps by default so the report can distinguish real deltas from # thermal/clock jitter; repo variables can trim this for temporary smoke runs. + set +e .profiler-venv/bin/python scripts/profile.py \ --target "$MODELS/$TARGET_MODEL" \ --draft "$MODELS/$DRAFT_MODEL" \ @@ -153,6 +192,18 @@ jobs: --nsys --check-lossless \ "${baseline_arg[@]}" \ --out-json profile.json --out-md profile.md + profile_rc=$? + set -e + echo "exit_code=$profile_rc" >> "$GITHUB_OUTPUT" + if [ "$profile_rc" -ne 0 ]; then + echo "::warning title=Speed profiler failed::Profiler exited $profile_rc; see its logs and uploaded artifacts." + { + echo "## 🏎️ Speed profile — profiler failed" + echo "" + echo "The report-only profiler exited with code \`$profile_rc\`. See the job log for the root cause." + } >> "$GITHUB_STEP_SUMMARY" + fi + exit 0 env: LUCEBOX_SPEED_BASELINE: ${{ vars.LUCEBOX_SPEED_BASELINE || '' }} LUCEBOX_SPEED_REGRESS_PCT: ${{ vars.LUCEBOX_SPEED_REGRESS_PCT || '' }} @@ -161,7 +212,7 @@ jobs: LUCEBOX_SPEED_NOISE_RSD_PCT: ${{ vars.LUCEBOX_SPEED_NOISE_RSD_PCT || '' }} - name: Publish report to the run summary - if: always() && steps.models.outputs.present == 'true' + if: always() && steps.models.outputs.present == 'true' && steps.capacity.outputs.ready == 'true' run: | if [ -f server/profile.md ]; then { echo "## 🏎️ Speed profile"; echo ""; cat server/profile.md; } >> "$GITHUB_STEP_SUMMARY" diff --git a/Dockerfile b/Dockerfile index cfeb6f29d..86c0125e7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -2,30 +2,31 @@ # ─── Stage 1: builder ─────────────────────────────────────────────────────── # CUDA_VERSION / UBUNTU_VERSION / DFLASH_CUDA_ARCHES are build args so the -# same Dockerfile can be repinned later. The prebuilt image is the -# CUDA 12.8 path: -# • lucebox-hub:cuda12 — CUDA 12.8.1, sm_75;80;86;89;90;120 +# same Dockerfile serves all published CUDA variants: +# • lucebox-hub:cuda12 — CUDA 12.0.1, x86_64, sm_75;80;86;89;90 +# • lucebox-hub:cuda128 — CUDA 12.8.1, x86_64, sm_120 (RTX 5090) +# • lucebox-hub:cuda13 — CUDA 13.0.1, arm64, sm_121 (GB10/DGX Spark) # See docker-bake.hcl for the canonical invocation. -ARG CUDA_VERSION=12.8.1 +ARG CUDA_VERSION=12.0.1 ARG UBUNTU_VERSION=22.04 FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION} AS builder ARG DEBIAN_FRONTEND=noninteractive -# Fat-binary CUDA arch list, semicolon-separated. Defaults cover the CUDA 12.8 -# image. dflash-supported arches in this image: +# Fat-binary CUDA arch list, semicolon-separated. Defaults cover the broadly +# compatible CUDA 12 image. dflash-supported arches in this image: # 75 Turing RTX 2080 Ti # 80 Ampere A100 # 86 Ampere RTX 3090, A40, A10 # 89 Ada RTX 4090, L40 # 90 Hopper H100 -# 120 Blackwell RTX 5090, RTX 5090 Laptop -# Thor and GB10 prebuilt-image coverage is intentionally omitted. +# RTX 5090 (sm_120) and GB10 (sm_121) are built by docker-bake.hcl's +# CUDA 12.8 and CUDA 13 targets, respectively. # Pre-Turing arches (sm_60/61/70/72) are intentionally excluded — dflash's # BF16/WMMA paths have no fallback below sm_75. Each arch adds ~50-200 MB # of fat-binary kernel code and ~3-5 min of nvcc time per .cu translation # unit. -ARG DFLASH_CUDA_ARCHES="75;80;86;89;90;120" +ARG DFLASH_CUDA_ARCHES="75;80;86;89;90" RUN apt-get update && apt-get install -y --no-install-recommends \ build-essential \ diff --git a/docker-bake.hcl b/docker-bake.hcl index ab9348a25..85989de01 100644 --- a/docker-bake.hcl +++ b/docker-bake.hcl @@ -1,12 +1,13 @@ # docker-bake.hcl — Lucebox hub prebuild matrix. # -# Single CUDA 12 image from one Dockerfile. Additional CUDA stacks are -# intentionally omitted. +# CUDA 12 covers Turing through Hopper on the broad CUDA-12 driver floor. +# CUDA 12.8 covers x86 Blackwell sm_120, while CUDA 13 provides the native +# arm64/sm_121 image required by GB10/DGX Spark. # # scripts/build_image.sh # version-derived local build (preferred) # docker buildx bake cuda12-local # raw local build; tagged lucebox-hub:cuda12 # docker buildx bake cuda12 # CI target; tags come from metadata-action -# # Arches: sm_75;80;86;89;90;120 +# # Arches: sm_75;80;86;89;90 # # Pre-Turing arches (Pascal sm_60/61, Volta sm_70) are intentionally # excluded — dflash's kernels assume sm_75+ with no fallback below @@ -32,17 +33,24 @@ variable "VERSION" { default = "" } # but new callers should use `VERSION`. variable "TAG" { default = "" } -# Fat-binary CUDA arch list. Defaults to all supported arches so the -# released image runs on every consumer/datacenter GPU we target. Local +# Fat-binary CUDA arch list for the broad-driver image. Local # dev builds can narrow this to the host's compute capability to skip the # 5-6× CUDA template recompile cost: # -# DFLASH_CUDA_ARCHES=120 docker buildx bake cuda12-local --load +# DFLASH_CUDA_ARCHES=86 docker buildx bake cuda12-local --load # -# (RTX 5090 / 5090 Laptop = 120, RTX 4090 = 89, RTX 3090 = 86, H100 = 90, -# A100 = 80, RTX 2080 Ti = 75.) Use a semicolon-separated list to include -# multiple arches. -variable "DFLASH_CUDA_ARCHES" { default = "75;80;86;89;90;120" } +# (RTX 4090 = 89, RTX 3090 = 86, H100 = 90, A100 = 80, RTX 2080 Ti = 75.) +# Use a semicolon-separated list to include multiple arches. +variable "DFLASH_CUDA_ARCHES" { default = "75;80;86;89;90" } + +# RTX 5090 is the only currently qualified x86 sm_120 target. Keeping it in a +# narrow CUDA 12.8 image avoids raising the driver floor for every older GPU. +variable "DFLASH_CUDA128_ARCHES" { default = "120" } + +# GB10 is an arm64 Blackwell part with compute capability 12.1. nvcc gained +# sm_121 support after CUDA 12.8, so it intentionally lives in a separate +# CUDA 13 image instead of raising the driver floor for existing x86 users. +variable "DFLASH_CUDA13_ARCHES" { default = "121" } # Fat-binary HIP/gfx arch list for the rocm variant (semicolon-separated). # Default is gfx1151 (Strix Halo, the lucebox appliance iGPU) only, to keep the @@ -90,10 +98,10 @@ group "default" { targets = ["cuda12-local"] } -# Build every published variant locally (cuda + rocm). CI builds these as a -# matrix; this group is the local equivalent for a full two-image build. +# Build every published variant locally (CUDA 12, CUDA 12.8, CUDA 13, and +# ROCm). CI builds these as a matrix; this group is the local equivalent. group "all" { - targets = ["cuda12-local", "rocm-local"] + targets = ["cuda12-local", "cuda128-local", "cuda13-local", "rocm-local"] } # CI integration. docker/metadata-action in .github/workflows/docker.yml @@ -103,15 +111,15 @@ group "all" { # file, so this empty target keeps inheritance valid. target "docker-metadata-action" {} -# ── CUDA 12.8 ─────────────────────────────────────────────────────────────── -# CUDA 12.8 matches the uv-managed PyTorch cu128 stack and carries current-gen -# consumer Blackwell sm_120 coverage. Thor/GB10 variants stay out of this -# build matrix. +# ── CUDA 12 (broad driver compatibility) ─────────────────────────────────── +# CUDA 12.0 is the oldest toolkit that builds every qualified pre-Blackwell +# architecture and runs on the CUDA-12 minor-compatibility driver floor +# (r525+). This is the default for RTX 20/30/40, A-series, A100, and H100. target "_cuda12-base" { context = "." dockerfile = "Dockerfile" args = { - CUDA_VERSION = "12.8.1" + CUDA_VERSION = "12.0.1" UBUNTU_VERSION = "22.04" DFLASH_CUDA_ARCHES = DFLASH_CUDA_ARCHES # /props.build identity. CI passes these as env vars from the @@ -132,6 +140,55 @@ target "cuda12-local" { tags = image_tags("cuda12") } +# ── CUDA 12.8 / x86 Blackwell ────────────────────────────────────────────── +# sm_120 first appeared in CUDA 12.8. It has a newer driver floor, so it must +# not share the image selected for older NVIDIA cards. +target "_cuda128-base" { + context = "." + dockerfile = "Dockerfile" + args = { + CUDA_VERSION = "12.8.1" + UBUNTU_VERSION = "22.04" + DFLASH_CUDA_ARCHES = DFLASH_CUDA128_ARCHES + GIT_SHA = GIT_SHA + IMAGE_TAG = IMAGE_TAG + BUILD_TIME = BUILD_TIME + } +} + +target "cuda128" { + inherits = ["_cuda128-base", "docker-metadata-action"] +} + +target "cuda128-local" { + inherits = ["_cuda128-base"] + tags = image_tags("cuda128") +} + +# ── CUDA 13 / GB10 arm64 ────────────────────────────────────────────────── +target "_cuda13-base" { + context = "." + dockerfile = "Dockerfile" + platforms = ["linux/arm64"] + args = { + CUDA_VERSION = "13.0.1" + UBUNTU_VERSION = "24.04" + DFLASH_CUDA_ARCHES = DFLASH_CUDA13_ARCHES + GIT_SHA = GIT_SHA + IMAGE_TAG = IMAGE_TAG + BUILD_TIME = BUILD_TIME + } +} + +target "cuda13" { + inherits = ["_cuda13-base", "docker-metadata-action"] +} + +target "cuda13-local" { + inherits = ["_cuda13-base"] + tags = image_tags("cuda13") +} + # ── ROCm / HIP ─────────────────────────────────────────────────────────────── # AMD GPU build from Dockerfile.rocm: gfx1151 (Strix Halo) by default, widen via # DFLASH_HIP_ARCHES for a broadly-runnable image. Block-Sparse-Attention is From 92d31d1ff05b52cc48f613fc78721e5618d3246a Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sat, 1 Aug 2026 13:37:46 +0200 Subject: [PATCH 20/25] test(lucebox): assert Strix UMA reserve semantics --- lucebox/tests/test_placement.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/lucebox/tests/test_placement.py b/lucebox/tests/test_placement.py index 43563bb89..6cb99c170 100644 --- a/lucebox/tests/test_placement.py +++ b/lucebox/tests/test_placement.py @@ -95,7 +95,7 @@ def test_strix_128gb_does_not_double_reserve_memory_for_deepseek( amd_gpu_count=1, amd_vram_gb=125, amd_gpu_arch="gfx1151", - amd_gpu_list_csv=_amd_csv(("Radeon 8060S", "gfx1151", 0)), + amd_gpu_list_csv=_amd_csv(("Radeon 8060S", "gfx1151", 512)), ) cfg = Config( variant="rocm", @@ -104,8 +104,13 @@ def test_strix_128gb_does_not_double_reserve_memory_for_deepseek( model=ModelMeta(preset="deepseek-v4-flash"), ) + topology = from_config(cfg) plan = automatic_plan(cfg, optimizer_drafter_available=False) + assert topology.primary is not None + assert topology.primary.unified_memory is True + assert topology.primary.physical_vram_gb == 0 + assert topology.primary.effective_memory_gb == 109 assert plan.placement.runnable is True assert plan.placement.runtime.mode == "single" assert plan.placement.runtime.target_device == "hip:0" From 76c227c3a0b41e6e0defccb649bbf6a9030ecc2e Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:30:14 +0200 Subject: [PATCH 21/25] feat(lucebox): add safe on-machine calibration --- README.md | 17 +- lucebox.sh | 213 +++++++- lucebox/README.md | 10 + lucebox/pyproject.toml | 8 +- lucebox/src/lucebox/__init__.py | 3 +- lucebox/src/lucebox/calibration.py | 748 +++++++++++++++++++++++++++++ lucebox/src/lucebox/cli.py | 142 ++++++ lucebox/tests/test_calibration.py | 368 ++++++++++++++ scripts/test_lucebox_sh.sh | 116 ++++- server/scripts/entrypoint.sh | 22 +- 10 files changed, 1607 insertions(+), 40 deletions(-) create mode 100644 lucebox/src/lucebox/calibration.py create mode 100644 lucebox/tests/test_calibration.py diff --git a/README.md b/README.md index 697f6e827..24bfda16b 100644 --- a/README.md +++ b/README.md @@ -114,7 +114,8 @@ lucebox `lucebox` opens the branded menu. **Quick setup** detects NVIDIA CUDA or AMD ROCm, lets you choose and download a model, applies a model- and hardware-aware -**Automatic** profile, and starts the inference service. Automatic explains its +**Automatic** profile, offers a one-time on-machine calibration, and starts the +inference service. Automatic explains its DFlash, PFlash, KVFlash, and Spark choices, and also prints the resolved prefill, decode, and KV-cache strategy. A typed capability contract prevents unsupported model/backend combinations; preview paths are labeled and stay off @@ -126,6 +127,17 @@ prefill is available there as an approximate, monolithic-HIP preview; exact MLA prefill remains the default. Wi-Fi and device provisioning are intentionally outside this inference-only CLI. +**Calibrate and measure performance** runs a fixed three-turn coding workload on +the selected model and reports measured prefill tok/s, decode tok/s, and warm +prefix-cache reuse. On DDTree backends it tests at most the planner budget and +its two nearest safe neighbours. A candidate is accepted only when its +normalized greedy output and cache behavior are identical and its decode score +is at least 5% faster; failures restore the original config and engine state. +Results are cached against the model files, runtime profile, driver, and GPU. +Spark placement, KVFlash sizing, PFlash request policy, GPU top-K split, and +DSpark width remain engine-owned adaptive policies rather than a second set of +CLI heuristics. + Automatic inventories every NVIDIA and AMD accelerator, then writes one validated execution plan together with the optimization profile. A model that fits stays on the faster primary GPU. When capacity requires it, supported @@ -170,6 +182,7 @@ command, for example: lucebox models select lucebox optimize --yes lucebox optimize --advanced +lucebox calibrate # cached; --force measures again lucebox start ./lucebox.sh build cuda ./lucebox.sh native cuda @@ -358,7 +371,7 @@ Requests that omit `temperature` use the model card's sampling (Qwen3.6: `temper | Flag | Default | Effect | |---|---|---| | `--ddtree` | off (chain) | Enable tree verify | -| `--ddtree-budget N` | `22` | Tree size. 22 on 3090 (default), 40 on 5090, re-sweep on GB10 | +| `--ddtree-budget N` | `22` | Tree size. The planner picks a hardware baseline; `lucebox calibrate` verifies it on the actual machine. | | `--fa-window N` | `0` / `2048` (full attention) | Sliding FA window. Leave at 0: a finite window breaks tool calls (the full-attention layers lose the system prompt/tools). | | `--draft-residency {auto,persistent,request-scoped}` | `auto` | When draft weights are evicted from VRAM. `request-scoped` parks/frees them after each request's draft work (frees VRAM for the target on tight GPUs); `persistent` keeps them resident across requests; `auto` preserves current behavior while honoring the low-VRAM / `--lazy-draft` hint. Reported at `/props.runtime.draft_residency`. | | `--lazy-draft` | off | Legacy alias for `--draft-residency=request-scoped` (defer draft load until first request, release after) | diff --git a/lucebox.sh b/lucebox.sh index 92559d84a..48156dfad 100755 --- a/lucebox.sh +++ b/lucebox.sh @@ -8,7 +8,7 @@ # dispatch into the in-container Python CLI via `docker run`. The # Python CLI lives at /opt/lucebox-hub/lucebox/ inside the image and is # the single source of truth for orchestration logic — TOML config, -# autotune rules + sweep, server bring-up, smoke tests, model downloads. +# optimization planning, calibration probes, and model downloads. # # 2) Manage a user-level systemd unit (~/.config/systemd/user/lucebox.service) # so the server can run as a long-lived service without keeping a shell @@ -1072,12 +1072,19 @@ _set_tty_flags() { # usage: _set_tty_flags arrayname build_orchestrator_argv() { local variant="$1" caller_has_tty="$2"; shift 2 + local cli_group="${1:-}" cli_action="${2:-}" local tty=(-i) [ "$caller_has_tty" = "1" ] && tty=(-it) local argv=(docker run --rm "${tty[@]}") _append_gpu_args argv "$variant" argv+=(--name "${CONTAINER_NAME}-cli-$$") argv+=(--user "$(id -u):$(id -g)") + # Native/hybrid servers listen on the host loopback rather than in a + # Lucebox container. The internal calibration probe needs that namespace + # when it cannot use docker exec against a running inference container. + if [ "$cli_group" = "_calibration" ] && [ "$cli_action" = "probe" ]; then + argv+=(--network host) + fi # Only bind-mount the docker socket when DOCKER_HOST actually points # at a unix socket on this host. With DOCKER_HOST=tcp://… or ssh://… # the path we'd construct is `tcp` or empty, and `docker run -v` would @@ -1588,6 +1595,161 @@ cmd_systemctl_passthrough() { esac } +# One-time machine calibration. The Python package owns candidate generation, +# workload measurement, quality-equivalence checks, winner selection, and the +# result cache. This host layer owns only the lifecycle it alone can control: +# restart the user service for each startup-scoped DDTree budget and restore +# both config and service state on every failure or interrupt. +_calibration_remove_run_dir() { + if [ -n "${LUCEBOX_CALIBRATION_RUN_DIR:-}" ] \ + && [[ "$LUCEBOX_CALIBRATION_RUN_DIR" == "$CONFIG_HOME"/.calibration-run.* ]]; then + rm -rf -- "$LUCEBOX_CALIBRATION_RUN_DIR" + fi +} + +_calibration_on_exit() { + local rc=$? + trap - EXIT HUP INT TERM + if [ "${LUCEBOX_CALIBRATION_COMMITTED:-0}" != "1" ] \ + && [ -f "${LUCEBOX_CALIBRATION_BACKUP:-}" ]; then + cp "$LUCEBOX_CALIBRATION_BACKUP" "$(_lucebox_config_path)" || true + if [ "${LUCEBOX_CALIBRATION_HAD_RECORD:-0}" = "1" ] \ + && [ -f "${LUCEBOX_CALIBRATION_RECORD_BACKUP:-}" ]; then + cp "$LUCEBOX_CALIBRATION_RECORD_BACKUP" \ + "$CONFIG_HOME/calibration.json" || true + else + rm -f "$CONFIG_HOME/calibration.json" + fi + if [ "${LUCEBOX_CALIBRATION_WAS_ACTIVE:-0}" = "1" ]; then + systemctl --user restart "$UNIT_NAME" >/dev/null 2>&1 || true + else + systemctl --user stop "$UNIT_NAME" >/dev/null 2>&1 || true + fi + warn "Calibration did not complete; the original profile and service state were restored." + fi + _calibration_remove_run_dir + exit "$rc" +} + +cmd_calibrate() { + local force=0 + while [ $# -gt 0 ]; do + case "$1" in + --force) force=1 ;; + --help|-h) + cat </dev/null 2>&1 \ + || die "calibration requires flock (provided by util-linux)" + + mkdir -p "$CONFIG_HOME" + exec 9>"$CONFIG_HOME/calibration.lock" + flock -n 9 || die "another Lucebox calibration is already running" + + LUCEBOX_CALIBRATION_RUN_DIR=$(mktemp -d "$CONFIG_HOME/.calibration-run.XXXXXX") \ + || die "could not create calibration workspace" + LUCEBOX_CALIBRATION_BACKUP="$LUCEBOX_CALIBRATION_RUN_DIR/config.toml.backup" + cp "$config_path" "$LUCEBOX_CALIBRATION_BACKUP" + LUCEBOX_CALIBRATION_RECORD_BACKUP="$LUCEBOX_CALIBRATION_RUN_DIR/calibration.json.backup" + LUCEBOX_CALIBRATION_HAD_RECORD=0 + if [ -f "$CONFIG_HOME/calibration.json" ]; then + cp "$CONFIG_HOME/calibration.json" "$LUCEBOX_CALIBRATION_RECORD_BACKUP" + LUCEBOX_CALIBRATION_HAD_RECORD=1 + fi + LUCEBOX_CALIBRATION_WAS_ACTIVE=0 + systemctl --user is-active --quiet "$UNIT_NAME" 2>/dev/null \ + && LUCEBOX_CALIBRATION_WAS_ACTIVE=1 + LUCEBOX_CALIBRATION_COMMITTED=0 + export LUCEBOX_CALIBRATION_RUN_DIR LUCEBOX_CALIBRATION_BACKUP + export LUCEBOX_CALIBRATION_RECORD_BACKUP LUCEBOX_CALIBRATION_HAD_RECORD + export LUCEBOX_CALIBRATION_WAS_ACTIVE LUCEBOX_CALIBRATION_COMMITTED + trap _calibration_on_exit EXIT + trap 'exit 130' HUP INT TERM + + local budgets=() baseline budget result_path successful=0 + mapfile -t budgets < <(bash "$SCRIPT_PATH" _calibration budgets) + [ "${#budgets[@]}" -gt 0 ] || die "calibration planner returned no budget" + [ "${#budgets[@]}" -le 3 ] || die "calibration planner exceeded its three-cell bound" + baseline="${budgets[0]}" + info "Calibrating $model on this machine (${#budgets[@]} server start(s))." + hint "The engine may take several minutes to load each cell. Ctrl-C safely restores the original profile." + + for budget in "${budgets[@]}"; do + [[ "$budget" =~ ^[0-9]+$ ]] || die "invalid calibration budget: $budget" + info "Calibration cell: DDTree budget $budget" + bash "$SCRIPT_PATH" _calibration apply "$budget" \ + || die "could not apply calibration budget $budget" + if ! systemctl --user restart "$UNIT_NAME"; then + if [ "$budget" = "$baseline" ]; then + die "the baseline server failed to start" + fi + warn "Skipping budget $budget because the server failed to start." + continue + fi + result_path="$LUCEBOX_CALIBRATION_RUN_DIR/budget-$budget.json" + if bash "$SCRIPT_PATH" _calibration probe "$budget" "$result_path"; then + successful=$((successful + 1)) + elif [ "$budget" = "$baseline" ]; then + die "the baseline performance probe failed" + else + warn "Skipping budget $budget because its performance probe failed." + fi + done + [ "$successful" -gt 0 ] || die "calibration produced no measurements" + + bash "$SCRIPT_PATH" _calibration finish \ + "$LUCEBOX_CALIBRATION_RUN_DIR" --baseline "$baseline" \ + || die "calibration could not select a safe result" + local winner + winner=$(<"$LUCEBOX_CALIBRATION_RUN_DIR/winner") + [[ "$winner" =~ ^[0-9]+$ ]] || die "calibration returned an invalid winner" + + if [ "$LUCEBOX_CALIBRATION_WAS_ACTIVE" = "1" ]; then + systemctl --user restart "$UNIT_NAME" \ + || die "the calibrated server failed to restart" + ok "Calibrated profile is active (DDTree budget $winner)." + else + systemctl --user stop "$UNIT_NAME" \ + || die "could not restore the engine's stopped state" + ok "Calibrated profile saved (DDTree budget $winner); engine remains stopped." + fi + + LUCEBOX_CALIBRATION_COMMITTED=1 + export LUCEBOX_CALIBRATION_COMMITTED + trap - EXIT HUP INT TERM + _calibration_remove_run_dir + flock -u 9 +} + cmd_logs() { ensure_probed if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ] && [ -f "$UNIT_PATH" ]; then @@ -2795,7 +2957,7 @@ _lucebox_complete() { COMPREPLY=() cur="${COMP_WORDS[COMP_CWORD]}" prev="${COMP_WORDS[COMP_CWORD-1]}" - cmds="menu setup connect install uninstall start stop restart enable disable status logs \ + cmds="menu setup calibrate connect install uninstall start stop restart enable disable status logs \ serve build package-runtime native harness pull update check completion config models \ optimize print-run help version" config_verbs="get set unset" @@ -2839,7 +3001,7 @@ ZSH # lucebox fish completion. Source from ~/.config/fish/config.fish: # lucebox completion fish | source complete -c lucebox -f -set -l __lucebox_cmds menu setup connect install uninstall start stop restart enable disable \ +set -l __lucebox_cmds menu setup calibrate connect install uninstall start stop restart enable disable \ status logs serve build package-runtime native harness pull update check completion config models \ optimize print-run help version for cmd in $__lucebox_cmds @@ -3070,7 +3232,7 @@ _lucebox_container_running() { # server), filesystem, and mounts — no bind mounts needed. # - skips the ~1-3s cold-start cost of a fresh `docker run --rm`. # - only safe for steady-state / read-only / config-only subcommands. Any -# command that restarts the lucebox service (autotune --sweep, serve) +# command that restarts the lucebox service (calibrate, serve) # would kill the very container the exec is in — caller must route those # to cmd_in_container instead. # @@ -3116,6 +3278,12 @@ _lucebox_prefer_exec() { config|models|optimize|check|print-run|print-serve-argv) return 0 ;; + _calibration) + case "${1:-}" in + apply|budgets|finish|probe|status) return 0 ;; + esac + return 1 + ;; *) return 1 ;; @@ -3129,7 +3297,7 @@ _lucebox_prefer_exec() { # 1. LUCEBOX_NO_EXEC=1 / --no-exec was set → always run, never exec. # Useful for debugging the wrapper or when the in-container Python is # stale relative to the image. -# 2. cmd is not in the prefer-exec list → run (sweep, service mutators). +# 2. cmd is not in the prefer-exec list → run (service mutators). # 3. container is running → exec (the fast path, hits the live server). # 4. container is not running → run (fall back so first-run / pre-install # flows still work without a live service). @@ -3210,7 +3378,7 @@ _menu_restart_if_running() { _menu_run restart else _menu_run stop - warn "The foreground engine was stopped. Start it again from menu option 4." + warn "The foreground engine was stopped. Start it again from menu option 5." fi } @@ -3281,8 +3449,10 @@ cmd_setup() { _ensure_configured_image || return $? printf '\n' + local optimized=0 if _confirm "Enable recommended optimizations (may add a ~1.2 GB shared scorer)?" 1; then bash "$SCRIPT_PATH" optimize --yes || return $? + optimized=1 else hint "The safe base profile is active; optional scorer-based features remain off." fi @@ -3294,6 +3464,11 @@ cmd_setup() { bash "$SCRIPT_PATH" install || return $? fi fi + if [ "$optimized" = "1" ] && [ -f "$UNIT_PATH" ] \ + && _confirm "Measure and calibrate this model on this GPU? (one-time; several minutes)" 1; then + bash "$SCRIPT_PATH" calibrate \ + || warn "Calibration was skipped; the Automatic profile is unchanged." + fi if [ -f "$UNIT_PATH" ] && _confirm "Start the inference engine now?" 1; then if [ "$(_engine_state)" = "running" ]; then bash "$SCRIPT_PATH" restart || return $? @@ -3474,11 +3649,12 @@ cmd_menu() { printf ' 1 Quick setup\n' printf ' 2 Choose or download a model\n' printf ' 3 Review optimizations\n' - printf ' 4 Start the inference engine\n' - printf ' 5 Stop the inference engine\n' - printf ' 6 Status\n' - printf ' 7 Recent logs\n' - printf ' 8 Connect or open your harness\n' + printf ' 4 Calibrate and measure performance\n' + printf ' 5 Start the inference engine\n' + printf ' 6 Stop the inference engine\n' + printf ' 7 Status\n' + printf ' 8 Recent logs\n' + printf ' 9 Connect or open your harness\n' printf ' d Developer tools %b(%s)%b\n' "$C_DIM" "$repo_hint" "$C_RST" printf ' q Quit\n\n' printf 'Choose: ' @@ -3496,10 +3672,11 @@ cmd_menu() { if _menu_run optimize; then _menu_restart_if_running; fi _menu_pause ;; - 4) _menu_start; _menu_pause ;; - 5) _menu_run stop; _menu_pause ;; - 6) _menu_run status; _menu_pause ;; - 7) + 4) _menu_run calibrate; _menu_pause ;; + 5) _menu_start; _menu_pause ;; + 6) _menu_run stop; _menu_pause ;; + 7) _menu_run status; _menu_pause ;; + 8) if [ "$LUCEBOX_HOST_HAS_SYSTEMD" = "1" ] && [ -f "$UNIT_PATH" ]; then _menu_run logs -n 80 --no-pager else @@ -3507,10 +3684,10 @@ cmd_menu() { fi _menu_pause ;; - 8) _menu_run connect; _menu_pause ;; + 9) _menu_run connect; _menu_pause ;; d|D) cmd_developer_menu ;; q|Q|quit|exit) return 0 ;; - *) warn "Choose 1–8, d, or q"; _menu_pause ;; + *) warn "Choose 1–9, d, or q"; _menu_pause ;; esac done } @@ -3523,6 +3700,7 @@ Interactive: $SCRIPT_NAME open the branded menu (when run in a terminal) menu open the menu explicitly setup guided backend, model, optimization, and service setup + calibrate [--force] measure this model/GPU and tune safe startup knobs connect [harness] link an installed AI client to the local Lucebox API Service management (via user systemd): @@ -3615,6 +3793,7 @@ main() { # Branded interactive surface / guided first run. menu) cmd_menu "$@" ;; setup) cmd_setup "$@" ;; + calibrate) cmd_calibrate "$@" ;; connect) cmd_connect "$@" ;; # Systemd surface diff --git a/lucebox/README.md b/lucebox/README.md index c8ac35d60..3c6c741f1 100644 --- a/lucebox/README.md +++ b/lucebox/README.md @@ -11,6 +11,7 @@ which invokes the package in the appropriate container: lucebox models select # numbered picker + download + activate lucebox optimize # review/apply the Automatic profile lucebox optimize --advanced # override individual optimizations + lucebox calibrate # measure + tune this model/GPU safely lucebox start The wrapper inventories the host and selects the compatible CUDA image for @@ -21,6 +22,15 @@ download, optimization and placement settings, and construction of the final server command. Host facts are passed through `LUCEBOX_HOST_*` environment variables so the package never has to guess the host configuration. +Calibration measures a real three-turn request and reports prefill, decode, and +warm-prefix performance from the server's own timing data. It brackets only the +DDTree budget, only on architectures where that startup value is consumed, and +requires matching normalized output/cache behavior plus a 5% decode gain before +changing the planner default. The cached record is invalidated by model, +runtime, driver, or GPU changes. Spark, KVFlash, PFlash, GPU split, and DSpark +continue to self-tune inside the engine; the CLI does not duplicate those +policies. + The guided picker leads with Lucebox's four featured model paths: Qwen3.6 27B, Qwen3.6 35B-A3B, Laguna XS.2, and DeepSeek V4 Flash. It checks the resolved placement before downloading, so a machine without enough compatible diff --git a/lucebox/pyproject.toml b/lucebox/pyproject.toml index 8c873e3a1..7939b33cd 100644 --- a/lucebox/pyproject.toml +++ b/lucebox/pyproject.toml @@ -12,7 +12,7 @@ authors = [{ name = "Lucebox" }] license = { text = "Apache-2.0" } # Kept intentionally narrow. typer pulls click+rich; tomli-w gives us TOML -# writes (stdlib tomllib only reads). httpx for the smoke + readiness probes. +# writes (stdlib tomllib only reads). httpx for calibration + readiness probes. # huggingface_hub for download-models — used directly (not via subprocess) # so we can drive a Rich progress bar + verify sha256 against the repo # metadata before re-fetching multi-GB GGUFs. @@ -22,12 +22,6 @@ dependencies = [ "httpx>=0.27", "tomli-w>=1.0", "huggingface_hub>=0.27", - # luce-bench is consumed lazily by the autotune sweep scorer - # (agent_replay_pass_rate in sweep.py does a function-local - # `from lucebench.areas.agent_recorded import ...` wrapped in try/except). - # It's deliberately NOT a hard dep here because the workspace can't lock - # against it until #337 (luce-bench in-tree) lands. Install with - # `uv pip install luce-bench` on the host running the scorer. ] [project.urls] diff --git a/lucebox/src/lucebox/__init__.py b/lucebox/src/lucebox/__init__.py index 79e8f7a8e..f7e92b2f2 100644 --- a/lucebox/src/lucebox/__init__.py +++ b/lucebox/src/lucebox/__init__.py @@ -4,7 +4,8 @@ run` plumbing, systemd integration, and isolated profiles for already-installed AI clients. This package owns: TOML config, the host-derived DFLASH_* serve heuristic, docker daemon calls (via the mounted socket), and model download. -The empirical autotune sweep and profiling remain separate workflows. +Bounded calibration measures the resolved plan without duplicating the +engine's own adaptive policies. """ # Version is generated by hatch-vcs at build time into _version.py. diff --git a/lucebox/src/lucebox/calibration.py b/lucebox/src/lucebox/calibration.py new file mode 100644 index 000000000..8d9e7fdd0 --- /dev/null +++ b/lucebox/src/lucebox/calibration.py @@ -0,0 +1,748 @@ +"""Bounded, quality-safe calibration against a running Lucebox server. + +The automatic planner owns feature selection and placement. Calibration is +deliberately narrower: it measures the resolved plan on the actual machine and +only brackets DDTree's startup budget on backends that consume that value. +Engine-owned adaptive policies (Spark, KVFlash sizing, PFlash request policy, +and DSpark width) are observed, never duplicated here. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import tempfile +import time +from dataclasses import asdict, dataclass, replace +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +import httpx + +import lucebox.autotune as autotune_mod +import lucebox.config as config_mod +import lucebox.download as download_mod +from lucebox import __version__ +from lucebox.types import Config + +SCHEMA_VERSION = 1 +MIN_WINNER_GAIN = 0.05 +_BUDGET_GRID = (4, 8, 16, 22, 32, 40, 64) +_BUDGET_ARCHITECTURES = frozenset({"qwen35", "qwen35moe", "laguna"}) + + +@dataclass(frozen=True, slots=True) +class TurnMeasurement: + """Server-reported timings for one measured request.""" + + decode_tokens_per_sec: float + prefill_tokens_per_sec: float | None + completion_tokens: int + cache_hit: bool + cached_prefix_tokens: int + prefilled_tokens: int + + def as_dict(self) -> dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, raw: Any) -> TurnMeasurement: + if not isinstance(raw, dict) or not isinstance(raw.get("cache_hit"), bool): + raise ValueError("malformed turn measurement") + try: + raw_prefill = raw["prefill_tokens_per_sec"] + prefill = None if raw_prefill is None else float(raw_prefill) + result = cls( + decode_tokens_per_sec=float(raw["decode_tokens_per_sec"]), + prefill_tokens_per_sec=prefill, + completion_tokens=int(raw["completion_tokens"]), + cache_hit=raw["cache_hit"], + cached_prefix_tokens=int(raw["cached_prefix_tokens"]), + prefilled_tokens=int(raw["prefilled_tokens"]), + ) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("malformed turn measurement") from exc + if ( + not math.isfinite(result.decode_tokens_per_sec) + or result.decode_tokens_per_sec <= 0.0 + or ( + result.prefill_tokens_per_sec is not None + and ( + not math.isfinite(result.prefill_tokens_per_sec) + or result.prefill_tokens_per_sec <= 0.0 + ) + ) + or result.completion_tokens <= 0 + or result.cached_prefix_tokens < 0 + or result.prefilled_tokens < 0 + ): + raise ValueError("invalid turn measurement") + return result + + +@dataclass(frozen=True, slots=True) +class ProbeResult: + """Comparable result for one server startup budget.""" + + budget: int + model: str + architecture: str + score: float + response_signature: str + cold: TurnMeasurement + warm: TurnMeasurement + server: dict[str, Any] + + def as_dict(self) -> dict[str, Any]: + return { + "schema": SCHEMA_VERSION, + "budget": self.budget, + "model": self.model, + "architecture": self.architecture, + "score": self.score, + "response_signature": self.response_signature, + "cold": self.cold.as_dict(), + "warm": self.warm.as_dict(), + "server": self.server, + } + + @classmethod + def from_dict(cls, raw: dict[str, Any]) -> ProbeResult: + if raw.get("schema") != SCHEMA_VERSION: + raise ValueError("unsupported calibration result schema") + try: + result = cls( + budget=int(raw["budget"]), + model=str(raw["model"]), + architecture=str(raw["architecture"]), + score=float(raw["score"]), + response_signature=str(raw["response_signature"]), + cold=TurnMeasurement.from_dict(raw["cold"]), + warm=TurnMeasurement.from_dict(raw["warm"]), + server=dict(raw["server"]), + ) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("malformed calibration result") from exc + if ( + not 0 <= result.budget <= 64 + or not result.model + or not result.architecture + or not math.isfinite(result.score) + or result.score <= 0.0 + or not result.response_signature + ): + raise ValueError("invalid calibration result") + return result + + +@dataclass(frozen=True, slots=True) +class CalibrationSummary: + """Winner plus all successfully measured candidates.""" + + baseline_budget: int + winner: ProbeResult + results: tuple[ProbeResult, ...] + rejected_budgets: tuple[int, ...] + + +def _preset(cfg: Config) -> download_mod.ModelPreset | None: + return download_mod.PRESETS.get(cfg.model.preset) + + +def budget_is_tunable(cfg: Config) -> bool: + """Whether changing ``dflash.budget`` can affect this resolved runtime.""" + preset = _preset(cfg) + return bool( + preset is not None + and preset.architecture in _BUDGET_ARCHITECTURES + and cfg.dflash.speculative_decode + and not cfg.dflash.spark + and 4 <= cfg.dflash.budget <= 64 + and autotune_mod.draft_available(cfg, preset) + ) + + +def candidate_budgets(cfg: Config) -> tuple[int, ...]: + """Return baseline first, followed by its two nearest safe neighbours.""" + current = cfg.dflash.budget + if not budget_is_tunable(cfg): + return (current,) + lower = [value for value in _BUDGET_GRID if value < current] + upper = [value for value in _BUDGET_GRID if value > current] + candidates = [current] + if lower: + candidates.append(lower[-1]) + if upper: + candidates.append(upper[0]) + return tuple(candidates) + + +def apply_budget(budget: int, *, final: bool = False) -> None: + """Atomically apply one calibration cell without changing profile ownership.""" + if not 0 <= budget <= 64: + raise ValueError(f"calibration budget must be in [0, 64], got {budget}") + cfg = config_mod.load() + if cfg is None: + raise ValueError("calibration requires an existing config.toml") + mode = config_mod.optimization_mode() + if mode not in {"automatic", "custom"}: + raise ValueError("run `lucebox optimize` before calibration") + config_mod.write_optimization_runtime( + replace(cfg.dflash, budget=budget), + placement=cfg.placement, + mode=mode, + source="calibrated" if final else "calibration-candidate", + ) + + +def _base_urls(cfg: Config) -> tuple[str, ...]: + explicit = os.environ.get("LUCEBOX_CALIBRATION_URL", "").rstrip("/") + if explicit: + return (explicit,) + urls = [f"http://127.0.0.1:{cfg.port}"] + # A probe executed inside the inference container shares its network + # namespace, where the server always listens on the internal port 8080. + if cfg.port != 8080: + urls.append("http://127.0.0.1:8080") + return tuple(urls) + + +def _wait_for_server( + client: httpx.Client, + base_urls: tuple[str, ...], + timeout_s: float, +) -> tuple[str, dict[str, Any]]: + deadline = time.monotonic() + timeout_s + last_error = "server did not respond" + while time.monotonic() < deadline: + for base_url in base_urls: + try: + response = client.get(f"{base_url}/props", timeout=3.0) + response.raise_for_status() + body = response.json() + if isinstance(body, dict): + return base_url, body + last_error = "/props did not return an object" + except (httpx.HTTPError, json.JSONDecodeError, ValueError) as exc: + last_error = str(exc) + time.sleep(1.0) + raise TimeoutError(f"Lucebox was not ready after {timeout_s:g}s: {last_error}") + + +def _representative_context() -> str: + """Compactly generate a repeatable repository-shaped prefill workload.""" + return "\n".join( + f"src/worker_{index:02d}.py: parse request, validate job {index}, " + f"write an atomic result, and preserve cancellation state." + for index in range(12) + ) + + +_TOOLS = [ + { + "type": "function", + "function": { + "name": "read_file", + "description": "Read a UTF-8 source file from the repository.", + "parameters": { + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"], + "additionalProperties": False, + }, + }, + }, + { + "type": "function", + "function": { + "name": "run_tests", + "description": "Run a named, already-installed test target.", + "parameters": { + "type": "object", + "properties": {"target": {"type": "string"}}, + "required": ["target"], + "additionalProperties": False, + }, + }, + }, +] + + +def _request_body( + messages: list[dict[str, Any]], + model: str, + max_tokens: int, + *, + include_tools: bool, +) -> dict[str, Any]: + body: dict[str, Any] = { + "model": model or "lucebox", + "messages": messages, + "max_tokens": max_tokens, + "temperature": 0.0, + "seed": 20260801, + "stream": False, + "chat_template_kwargs": {"enable_thinking": False}, + } + if include_tools: + body["tools"] = _TOOLS + # Keep the representative tool schema in the prompt while making the + # measured completion deterministic and directly comparable across + # candidate runtimes. Tool execution itself is outside calibration. + body["tool_choice"] = "none" + return body + + +def _assistant_turn( + payload: dict[str, Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + """Return a replayable assistant message and a stable quality signature.""" + choices = payload.get("choices") + if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict): + raise ValueError("chat response has no choice") + message = choices[0].get("message") + if not isinstance(message, dict): + raise ValueError("chat response has no assistant message") + + parts = [message.get("reasoning_content"), message.get("content")] + text = "\n".join(part for part in parts if isinstance(part, str) and part.strip()).strip() + replay: dict[str, Any] = {"role": "assistant", "content": text or None} + normalized_calls: list[dict[str, Any]] = [] + raw_calls = message.get("tool_calls", []) + if raw_calls is None: + raw_calls = [] + if not isinstance(raw_calls, list): + raise ValueError("chat response has malformed tool calls") + replay_calls: list[dict[str, Any]] = [] + for index, raw_call in enumerate(raw_calls): + if not isinstance(raw_call, dict): + raise ValueError("chat response has malformed tool calls") + raw_function = raw_call.get("function") + if not isinstance(raw_function, dict): + raise ValueError("chat response has malformed tool calls") + name = raw_function.get("name") + arguments = raw_function.get("arguments", "") + if not isinstance(name, str) or not name: + raise ValueError("chat response has a nameless tool call") + if not isinstance(arguments, str): + arguments = json.dumps(arguments, sort_keys=True, separators=(",", ":")) + try: + normalized_arguments: Any = json.loads(arguments) + except json.JSONDecodeError: + normalized_arguments = arguments.strip() + call_id = raw_call.get("id") + if not isinstance(call_id, str) or not call_id: + call_id = f"calibration_call_{index}" + replay_calls.append( + { + "id": call_id, + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + ) + normalized_calls.append( + { + "type": "function", + "function": {"name": name, "arguments": normalized_arguments}, + } + ) + if replay_calls: + replay["tool_calls"] = replay_calls + if not text and not replay_calls: + raise ValueError("chat response was empty") + return replay, {"text": text, "tool_calls": normalized_calls} + + +def _continue_conversation( + messages: list[dict[str, Any]], + assistant: dict[str, Any], + user_content: str, +) -> list[dict[str, Any]]: + """Append a turn, simulating any requested tool result without executing it.""" + continued = [*messages, assistant] + calls = assistant.get("tool_calls", []) + if isinstance(calls, list): + for call in calls: + if isinstance(call, dict) and isinstance(call.get("id"), str): + continued.append( + { + "role": "tool", + "tool_call_id": call["id"], + "content": '{"status":"skipped","reason":"calibration"}', + } + ) + continued.append({"role": "user", "content": user_content}) + return continued + + +def _measurement(payload: dict[str, Any]) -> TurnMeasurement: + usage = payload.get("usage") + if not isinstance(usage, dict): + raise ValueError("chat response has no usage object") + timings = usage.get("timings") + if not isinstance(timings, dict): + raise ValueError("server does not expose usage.timings") + try: + decode_tps = float(timings["decode_tokens_per_sec"]) + prefill_ms = float(timings["prefill_ms"]) + prefilled = int(timings["prefilled_tokens"]) + completion = int(usage["completion_tokens"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("usage.timings is incomplete") from exc + if not math.isfinite(decode_tps) or decode_tps <= 0.0 or completion <= 0: + raise ValueError("server returned a non-positive decode measurement") + prefill_tps = None + if prefill_ms > 0.0 and prefilled > 0: + prefill_tps = prefilled * 1000.0 / prefill_ms + return TurnMeasurement( + decode_tokens_per_sec=decode_tps, + prefill_tokens_per_sec=prefill_tps, + completion_tokens=completion, + cache_hit=bool(timings.get("cache_hit", False)), + cached_prefix_tokens=int(timings.get("cached_prefix_tokens", 0)), + prefilled_tokens=prefilled, + ) + + +def _chat( + client: httpx.Client, + base_url: str, + body: dict[str, Any], + timeout_s: float, +) -> dict[str, Any]: + response = client.post( + f"{base_url}/v1/chat/completions", + json=body, + timeout=timeout_s, + ) + response.raise_for_status() + payload = response.json() + if not isinstance(payload, dict): + raise ValueError("chat endpoint did not return an object") + return payload + + +def probe( + cfg: Config, + expected_budget: int, + *, + ready_timeout_s: float = 600.0, + request_timeout_s: float = 300.0, + client: httpx.Client | None = None, + base_urls: tuple[str, ...] | None = None, +) -> ProbeResult: + """Measure cold prefill, decode, and warm multi-turn prefix reuse.""" + owned_client = client is None + active_client = client or httpx.Client(trust_env=False) + try: + base_url, props = _wait_for_server( + active_client, + base_urls or _base_urls(cfg), + ready_timeout_s, + ) + speculative = props.get("speculative") + if budget_is_tunable(cfg): + actual = speculative.get("ddtree_budget") if isinstance(speculative, dict) else None + if actual != expected_budget: + raise ValueError( + f"server started with DDTree budget {actual!r}, expected {expected_budget}" + ) + capabilities = props.get("capabilities") + include_tools = bool( + isinstance(capabilities, dict) and capabilities.get("tools_supported") + ) + + # Warm kernels and allocator paths without polluting the measured prefix. + _chat( + active_client, + base_url, + _request_body( + [{"role": "user", "content": "Reply with only the word ready."}], + cfg.model.preset, + 8, + include_tools=include_tools, + ), + request_timeout_s, + ) + + first_messages: list[dict[str, Any]] = [ + { + "role": "system", + "content": ( + "You are a concise senior coding assistant. Do not call tools for this " + "calibration request. Preserve behavior, cancellation, and atomic writes." + ), + }, + { + "role": "user", + "content": ( + "Here is a synthetic repository index:\n" + f"{_representative_context()}\n\n" + "Write a short Python function that atomically records a completed job. " + "Return code plus one sentence." + ), + }, + ] + first = _chat( + active_client, + base_url, + _request_body( + first_messages, + cfg.model.preset, + 64, + include_tools=include_tools, + ), + request_timeout_s, + ) + first_message, first_signature = _assistant_turn(first) + second_messages = _continue_conversation( + first_messages, + first_message, + "Now add idempotency and keep the answer under 12 lines.", + ) + second = _chat( + active_client, + base_url, + _request_body( + second_messages, + cfg.model.preset, + 64, + include_tools=include_tools, + ), + request_timeout_s, + ) + second_message, second_signature = _assistant_turn(second) + third_messages = _continue_conversation( + second_messages, + second_message, + "Finally add type hints without changing the behavior.", + ) + third = _chat( + active_client, + base_url, + _request_body( + third_messages, + cfg.model.preset, + 64, + include_tools=include_tools, + ), + request_timeout_s, + ) + _, third_signature = _assistant_turn(third) + cold = _measurement(first) + warm = _measurement(third) + score = 2.0 / ( + 1.0 / cold.decode_tokens_per_sec + 1.0 / warm.decode_tokens_per_sec + ) + signature_body = json.dumps( + [first_signature, second_signature, third_signature], + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ) + signature = hashlib.sha256(signature_body.encode()).hexdigest() + raw_model_props = props.get("model") + model_props: dict[str, Any] = ( + raw_model_props if isinstance(raw_model_props, dict) else {} + ) + preset = _preset(cfg) + return ProbeResult( + budget=expected_budget, + model=cfg.model.preset, + architecture=str( + model_props.get("arch") or (preset.architecture if preset else "") + ), + score=score, + response_signature=signature, + cold=cold, + warm=warm, + server={ + "build_info": props.get("build_info"), + "runtime": props.get("runtime"), + "speculative": speculative, + "pflash": props.get("pflash"), + "prefix_cache": props.get("prefix_cache"), + }, + ) + except httpx.HTTPError as exc: + raise ValueError(f"server request failed: {exc}") from exc + finally: + if owned_client: + active_client.close() + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile("w", dir=path.parent, delete=False) as handle: + json.dump(payload, handle, indent=2, sort_keys=True) + handle.write("\n") + temporary = Path(handle.name) + try: + os.replace(temporary, path) + except BaseException: + temporary.unlink(missing_ok=True) + raise + + +def write_probe(path: Path, result: ProbeResult) -> None: + _atomic_json(path, result.as_dict()) + + +def read_probe(path: Path) -> ProbeResult: + try: + raw = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"cannot read calibration result {path}") from exc + if not isinstance(raw, dict): + raise ValueError(f"calibration result {path} is not an object") + return ProbeResult.from_dict(raw) + + +def calibration_record_path() -> Path: + return config_mod.default_config_path().with_name("calibration.json") + + +def _artifact_stat(cfg: Config, relative: Path) -> dict[str, Any]: + logical = cfg.models_dir / relative + container = Path("/opt/lucebox-hub/server/models") / relative + stat = None + for candidate in (logical, container): + try: + stat = candidate.stat() + break + except OSError: + continue + if stat is None: + return {"file": relative.as_posix(), "present": False} + return { + "file": relative.as_posix(), + "present": True, + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + } + + +def _fingerprint_payload(cfg: Config) -> dict[str, Any]: + preset = _preset(cfg) + artifacts: list[dict[str, Any]] = [] + target = cfg.model.target_file or (preset.target_file if preset else "") + draft = cfg.model.draft_file or (preset.draft_file if preset else "") + if target: + artifacts.append(_artifact_stat(cfg, Path(target))) + if draft: + artifacts.append(_artifact_stat(cfg, Path("draft") / draft)) + if preset is not None and preset.speculator_dir: + root = Path("draft") / preset.speculator_dir + artifacts.extend(_artifact_stat(cfg, root / name) for name in preset.speculator_files) + host = cfg.host + return { + "cli_version": __version__, + "variant": cfg.variant, + "model": asdict(cfg.model), + "runtime": asdict(cfg.dflash), + "placement": asdict(cfg.placement), + "host": { + "nproc": host.nproc, + "ram_gb": host.ram_gb, + "gpu_name": host.gpu_name, + "gpu_count": host.gpu_count, + "gpu_sm": host.gpu_sm, + "vram_gb": host.vram_gb, + "driver_version": host.driver_version, + "rocm_version": host.rocm_version, + "is_wsl": host.is_wsl, + "nvidia_gpu_name": host.nvidia_gpu_name, + "nvidia_gpu_count": host.nvidia_gpu_count, + "nvidia_vram_gb": host.nvidia_vram_gb, + "nvidia_gpu_arch": host.nvidia_gpu_arch, + "nvidia_gpu_list_csv": host.nvidia_gpu_list_csv, + "amd_gpu_name": host.amd_gpu_name, + "amd_gpu_count": host.amd_gpu_count, + "amd_vram_gb": host.amd_vram_gb, + "amd_gpu_arch": host.amd_gpu_arch, + "amd_gpu_list_csv": host.amd_gpu_list_csv, + "hybrid_runtime": host.hybrid_runtime, + }, + "artifacts": artifacts, + } + + +def fingerprint(cfg: Config) -> str: + payload = json.dumps(_fingerprint_payload(cfg), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode()).hexdigest() + + +def current_record(cfg: Config) -> dict[str, Any] | None: + try: + raw = json.loads(calibration_record_path().read_text()) + except (OSError, json.JSONDecodeError): + return None + if not isinstance(raw, dict) or raw.get("schema") != SCHEMA_VERSION: + return None + return raw if raw.get("fingerprint") == fingerprint(cfg) else None + + +def finish( + result_dir: Path, + baseline_budget: int, + *, + cfg: Config | None = None, +) -> CalibrationSummary: + """Select a quality-equivalent winner, persist it, and cache the evidence.""" + results = tuple( + read_probe(path) for path in sorted(result_dir.glob("budget-*.json")) + ) + if not results: + raise ValueError("calibration produced no successful measurements") + baseline = next((item for item in results if item.budget == baseline_budget), None) + if baseline is None: + raise ValueError("baseline calibration measurement failed") + if not math.isfinite(baseline.score) or baseline.score <= 0.0: + raise ValueError("baseline calibration score is invalid") + + equivalent = tuple( + item + for item in results + if item.model == baseline.model + and item.response_signature == baseline.response_signature + and item.warm.cache_hit == baseline.warm.cache_hit + and item.warm.cached_prefix_tokens == baseline.warm.cached_prefix_tokens + and math.isfinite(item.score) + and item.score > 0.0 + ) + rejected = tuple(item.budget for item in results if item not in equivalent) + fastest = max(equivalent, key=lambda item: item.score) + winner = ( + fastest + if fastest.score >= baseline.score * (1.0 + MIN_WINNER_GAIN) + else baseline + ) + apply_budget(winner.budget, final=True) + live_cfg = cfg or config_mod.load() + if live_cfg is None: # pragma: no cover - apply_budget already proved this + raise RuntimeError("config disappeared while finishing calibration") + final_cfg = replace( + live_cfg, + dflash=replace(live_cfg.dflash, budget=winner.budget), + ) + record = { + "schema": SCHEMA_VERSION, + "captured_at": datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ"), + "fingerprint": fingerprint(final_cfg), + "baseline_budget": baseline_budget, + "winner_budget": winner.budget, + "minimum_gain": MIN_WINNER_GAIN, + "results": [item.as_dict() for item in results], + "rejected_budgets": list(rejected), + } + _atomic_json(calibration_record_path(), record) + (result_dir / "winner").write_text(f"{winner.budget}\n") + return CalibrationSummary( + baseline_budget=baseline_budget, + winner=winner, + results=results, + rejected_budgets=rejected, + ) diff --git a/lucebox/src/lucebox/cli.py b/lucebox/src/lucebox/cli.py index 03044c2da..d4d888cca 100644 --- a/lucebox/src/lucebox/cli.py +++ b/lucebox/src/lucebox/cli.py @@ -18,6 +18,7 @@ import sys from dataclasses import replace +from pathlib import Path from typing import Annotated, Literal import typer @@ -26,6 +27,7 @@ from rich.table import Table import lucebox.autotune as autotune_mod +import lucebox.calibration as calibration_mod import lucebox.capabilities as capabilities_mod import lucebox.config as config_mod import lucebox.docker_run as docker_run @@ -248,6 +250,146 @@ def print_serve_argv() -> None: print(tok) +# ── host-wrapper calibration protocol ───────────────────────────────────── + + +calibration_app = typer.Typer( + no_args_is_help=True, + help="Internal host-wrapper calibration protocol.", +) +app.add_typer(calibration_app, name="_calibration", hidden=True) + + +def _render_measurement(label: str, result: calibration_mod.ProbeResult) -> None: + cold_prefill = result.cold.prefill_tokens_per_sec + prefill = f"{cold_prefill:,.0f} prefill tok/s" if cold_prefill else "prefill n/a" + cache = ( + f"warm cache hit ({result.warm.cached_prefix_tokens:,} tokens)" + if result.warm.cache_hit + else "warm cache not observed" + ) + console.print( + f"{label}: budget [bold]{result.budget}[/bold] · {prefill} · " + f"[bold]{result.cold.decode_tokens_per_sec:.1f}[/bold] cold / " + f"[bold]{result.warm.decode_tokens_per_sec:.1f}[/bold] warm decode tok/s · {cache}" + ) + + +@calibration_app.command("budgets") +def calibration_budgets() -> None: + """Emit the bounded candidate set, baseline first.""" + cfg = _load_or_build() + if config_mod.optimization_mode() not in {"automatic", "custom"}: + error_console.print( + "[red]Calibration needs an optimization profile; run `lucebox optimize` first.[/red]" + ) + raise typer.Exit(code=2) + for budget in calibration_mod.candidate_budgets(cfg): + print(budget) + + +@calibration_app.command("status") +def calibration_status() -> None: + """Show a matching cached result; exit 1 when calibration is stale.""" + cfg = _load_or_build() + record = calibration_mod.current_record(cfg) + if record is None: + raise typer.Exit(code=1) + try: + winner_budget = int(record["winner_budget"]) + except (KeyError, TypeError, ValueError) as exc: + raise typer.Exit(code=1) from exc + raw_results = record.get("results") + if not isinstance(raw_results, list): + raise typer.Exit(code=1) + try: + results = tuple( + calibration_mod.ProbeResult.from_dict(raw) + for raw in raw_results + if isinstance(raw, dict) + ) + except ValueError as exc: + raise typer.Exit(code=1) from exc + winner = next((item for item in results if item.budget == winner_budget), None) + if winner is None: + raise typer.Exit(code=1) + console.print("[green]Calibration is current for this model and machine.[/green]") + _render_measurement("Measured", winner) + + +@calibration_app.command("apply") +def calibration_apply( + budget: Annotated[int, typer.Argument(help="DDTree budget for this cell.")], + final: Annotated[ + bool, + typer.Option("--final", help="Mark this budget as the selected result."), + ] = False, +) -> None: + """Apply one calibration cell while preserving Automatic/Custom ownership.""" + try: + calibration_mod.apply_budget(budget, final=final) + except (OSError, ValueError) as exc: + error_console.print(f"[red]Cannot apply calibration cell:[/red] {escape(str(exc))}") + raise typer.Exit(code=2) from exc + + +@calibration_app.command("probe") +def calibration_probe( + budget: Annotated[int, typer.Argument(help="Expected active DDTree budget.")], + output: Annotated[Path, typer.Argument(help="Result JSON path.")], + ready_timeout: Annotated[ + float, + typer.Option("--ready-timeout", min=1.0, help="Model startup timeout in seconds."), + ] = 600.0, +) -> None: + """Measure one live server cell and atomically write its result.""" + cfg = _load_or_build() + try: + result = calibration_mod.probe( + cfg, + budget, + ready_timeout_s=ready_timeout, + ) + calibration_mod.write_probe(output, result) + except (OSError, TimeoutError, ValueError) as exc: + error_console.print(f"[red]Calibration probe failed:[/red] {escape(str(exc))}") + raise typer.Exit(code=2) from exc + _render_measurement("Measured", result) + + +@calibration_app.command("finish") +def calibration_finish( + result_dir: Annotated[Path, typer.Argument(help="Directory containing cell JSON files.")], + baseline: Annotated[int, typer.Option("--baseline", help="Original DDTree budget.")], +) -> None: + """Select and persist the fastest quality-equivalent result.""" + cfg = _load_or_build() + try: + summary = calibration_mod.finish(result_dir, baseline, cfg=cfg) + except (OSError, ValueError) as exc: + error_console.print(f"[red]Cannot finish calibration:[/red] {escape(str(exc))}") + raise typer.Exit(code=2) from exc + for result in summary.results: + label = "Selected" if result.budget == summary.winner.budget else "Candidate" + _render_measurement(label, result) + if summary.rejected_budgets: + rejected = ", ".join(str(value) for value in summary.rejected_budgets) + console.print( + f"[yellow]Rejected budgets {rejected} because output/cache behavior changed " + "or metrics were invalid.[/yellow]" + ) + if summary.winner.budget == summary.baseline_budget: + console.print( + "[green]Kept the planner budget; no quality-equivalent candidate was " + "at least 5% faster.[/green]" + ) + else: + console.print( + f"[green]Calibrated DDTree budget: {summary.baseline_budget} → " + f"{summary.winner.budget}.[/green]" + ) + + # ── config sub-app ───────────────────────────────────────────────────────── diff --git a/lucebox/tests/test_calibration.py b/lucebox/tests/test_calibration.py new file mode 100644 index 000000000..2573d4f79 --- /dev/null +++ b/lucebox/tests/test_calibration.py @@ -0,0 +1,368 @@ +"""Tests for bounded on-machine calibration.""" + +from __future__ import annotations + +import json +from dataclasses import replace +from pathlib import Path + +import httpx +import pytest +from lucebox.cli import app +from lucebox.types import Config, HostFacts, ModelMeta, PlacementRuntime +from typer.testing import CliRunner + +from lucebox import calibration, config, download + + +def _qwen_config(tmp_path: Path, *, budget: int = 22, with_draft: bool = True) -> Config: + preset = download.PRESETS["qwen3.6-27b"] + cfg = Config( + models_dir=tmp_path / "models", + host=HostFacts( + gpu_vendor="nvidia", + gpu_name="RTX 3090", + gpu_count=1, + gpu_sm="86", + vram_gb=24, + ), + model=ModelMeta( + preset=preset.name, + target_file=preset.target_file, + draft_file=preset.draft_file or "", + ), + ) + cfg = replace(cfg, dflash=replace(cfg.dflash, budget=budget)) + if with_draft and preset.draft_file: + draft = cfg.models_dir / "draft" / preset.draft_file + draft.parent.mkdir(parents=True, exist_ok=True) + draft.write_bytes(b"draft") + return cfg + + +def _turn(tps: float, *, cache_hit: bool = False) -> calibration.TurnMeasurement: + return calibration.TurnMeasurement( + decode_tokens_per_sec=tps, + prefill_tokens_per_sec=1000.0, + completion_tokens=32, + cache_hit=cache_hit, + cached_prefix_tokens=900 if cache_hit else 0, + prefilled_tokens=100 if cache_hit else 1000, + ) + + +def _result(budget: int, score: float, signature: str = "same") -> calibration.ProbeResult: + return calibration.ProbeResult( + budget=budget, + model="qwen3.6-27b", + architecture="qwen35", + score=score, + response_signature=signature, + cold=_turn(score), + warm=_turn(score, cache_hit=True), + server={}, + ) + + +def test_candidate_budgets_are_bounded_and_keep_baseline_first(tmp_path: Path) -> None: + assert calibration.candidate_budgets(_qwen_config(tmp_path, budget=22)) == (22, 16, 32) + assert calibration.candidate_budgets(_qwen_config(tmp_path, budget=8)) == (8, 4, 16) + assert calibration.candidate_budgets(_qwen_config(tmp_path, budget=40)) == (40, 32, 64) + + +def test_candidate_budgets_do_not_sweep_inert_or_unavailable_runtime(tmp_path: Path) -> None: + no_draft = _qwen_config(tmp_path, with_draft=False) + assert calibration.candidate_budgets(no_draft) == (22,) + + deepseek = replace( + no_draft, + model=ModelMeta(preset="deepseek-v4-flash"), + ) + assert calibration.candidate_budgets(deepseek) == (22,) + + spark = _qwen_config(tmp_path) + spark = replace(spark, dflash=replace(spark.dflash, spark=True)) + assert calibration.candidate_budgets(spark) == (22,) + + +def test_internal_budget_protocol_requires_an_optimization_profile( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lucebox.cli as cli + + monkeypatch.setattr(cli, "_load_or_build", lambda: _qwen_config(tmp_path)) + monkeypatch.setattr(config, "optimization_mode", lambda: "unconfigured") + + result = CliRunner().invoke(app, ["_calibration", "budgets"]) + + assert result.exit_code == 2 + assert "run `lucebox optimize` first" in result.output + + +def test_probe_uses_server_timings_and_warm_prefix_cache(tmp_path: Path) -> None: + cfg = _qwen_config(tmp_path) + posts = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal posts + if request.method == "GET" and request.url.path == "/props": + return httpx.Response( + 200, + json={ + "build_info": "test", + "model": {"arch": "qwen35"}, + "runtime": {"backend": "cuda"}, + "speculative": {"enabled": True, "ddtree_budget": 22}, + "capabilities": {"tools_supported": True}, + "pflash": {"enabled": True}, + "prefix_cache": {"capacity": 32}, + }, + ) + assert request.method == "POST" + body = json.loads(request.content) + assert "tools" in body + assert body["tool_choice"] == "none" + posts += 1 + if posts == 4: + assert [message["role"] for message in body["messages"][-3:]] == [ + "assistant", + "tool", + "user", + ] + assert body["messages"][-2]["tool_call_id"] == "call_test" + completion = 8 if posts == 1 else 32 + text = "ready" if posts == 1 else f"answer-{posts}" + message: dict[str, object] = {"role": "assistant", "content": text} + if posts == 3: + message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_test", + "type": "function", + "function": { + "name": "run_tests", + "arguments": '{"target": "unit"}', + }, + } + ], + } + return httpx.Response( + 200, + json={ + "choices": [{"message": message}], + "usage": { + "completion_tokens": completion, + "timings": { + "decode_tokens_per_sec": 20.0 if posts < 4 else 25.0, + "prefill_ms": 500.0 if posts == 2 else 10.0, + "cache_hit": posts == 4, + "cached_prefix_tokens": 900 if posts == 4 else 0, + "prefilled_tokens": 100 if posts == 4 else 1000, + }, + }, + }, + ) + + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + result = calibration.probe( + cfg, + 22, + client=client, + base_urls=("http://test",), + ready_timeout_s=1, + ) + + assert posts == 4 + assert result.cold.prefill_tokens_per_sec == 2000.0 + assert result.cold.decode_tokens_per_sec == 20.0 + assert result.warm.decode_tokens_per_sec == 25.0 + assert result.warm.cache_hit is True + assert result.warm.cached_prefix_tokens == 900 + assert result.score == pytest.approx(22.222222) + + +def test_probe_rejects_a_server_started_with_the_wrong_budget(tmp_path: Path) -> None: + cfg = _qwen_config(tmp_path) + + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "model": {"arch": "qwen35"}, + "speculative": {"enabled": True, "ddtree_budget": 16}, + }, + ) + + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + with pytest.raises(ValueError, match="expected 22"): + calibration.probe( + cfg, + 22, + client=client, + base_urls=("http://test",), + ready_timeout_s=1, + ) + + +def test_tool_call_signature_ignores_transport_ids_and_json_spacing() -> None: + def payload(call_id: str, arguments: str) -> dict[str, object]: + return { + "choices": [ + { + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": { + "name": "run_tests", + "arguments": arguments, + }, + } + ], + } + } + ] + } + + first_message, first_signature = calibration._assistant_turn( + payload("call_one", '{"target":"unit"}') + ) + second_message, second_signature = calibration._assistant_turn( + payload("call_two", '{ "target": "unit" }') + ) + + assert first_message["tool_calls"][0]["id"] == "call_one" + assert second_message["tool_calls"][0]["id"] == "call_two" + assert first_signature == second_signature + + +def test_probe_result_rejects_corrupt_cached_measurements() -> None: + raw = _result(22, 10.0).as_dict() + raw["warm"]["decode_tokens_per_sec"] = "not-a-number" + + with pytest.raises(ValueError, match="malformed calibration result"): + calibration.ProbeResult.from_dict(raw) + + +def test_finish_requires_quality_equivalence_and_material_gain( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cfg = _qwen_config(tmp_path) + result_dir = tmp_path / "results" + for result in ( + _result(22, 10.0), + _result(16, 10.4), + _result(32, 15.0, signature="different"), + ): + calibration.write_probe(result_dir / f"budget-{result.budget}.json", result) + + applied: list[tuple[int, bool]] = [] + monkeypatch.setattr( + calibration, + "apply_budget", + lambda budget, *, final=False: applied.append((budget, final)), + ) + monkeypatch.setattr( + calibration, + "calibration_record_path", + lambda: tmp_path / "calibration.json", + ) + + summary = calibration.finish(result_dir, 22, cfg=cfg) + + assert summary.winner.budget == 22 + assert summary.rejected_budgets == (32,) + assert applied == [(22, True)] + assert (result_dir / "winner").read_text() == "22\n" + + +def test_finish_selects_a_five_percent_faster_equivalent_candidate( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cfg = _qwen_config(tmp_path) + result_dir = tmp_path / "results" + for result in (_result(22, 10.0), _result(16, 11.0)): + calibration.write_probe(result_dir / f"budget-{result.budget}.json", result) + monkeypatch.setattr(calibration, "apply_budget", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + calibration, + "calibration_record_path", + lambda: tmp_path / "calibration.json", + ) + + summary = calibration.finish(result_dir, 22, cfg=cfg) + + assert summary.winner.budget == 16 + + +def test_finish_rejects_changed_warm_cache_behavior( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cfg = _qwen_config(tmp_path) + result_dir = tmp_path / "results" + baseline = _result(22, 10.0) + changed_cache = replace( + _result(16, 15.0), + warm=replace(_turn(15.0, cache_hit=True), cached_prefix_tokens=512), + ) + for result in (baseline, changed_cache): + calibration.write_probe(result_dir / f"budget-{result.budget}.json", result) + monkeypatch.setattr(calibration, "apply_budget", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + calibration, + "calibration_record_path", + lambda: tmp_path / "calibration.json", + ) + + summary = calibration.finish(result_dir, 22, cfg=cfg) + + assert summary.winner.budget == 22 + assert summary.rejected_budgets == (16,) + + +def test_apply_budget_preserves_automatic_profile_ownership( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LUCEBOX_HOME", str(tmp_path)) + config.write_optimization_runtime( + _qwen_config(tmp_path).dflash, + placement=PlacementRuntime(target_device="cuda:0"), + mode="automatic", + ) + + calibration.apply_budget(16) + + loaded = config.load() + assert loaded is not None + assert loaded.dflash.budget == 16 + assert loaded.placement.target_device == "cuda:0" + doc = config.load_doc() + assert doc["autotune"]["mode"] == "automatic" + assert doc["autotune"]["source"] == "calibration-candidate" + + +def test_cached_result_is_invalidated_by_runtime_changes( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LUCEBOX_HOME", str(tmp_path)) + cfg = _qwen_config(tmp_path) + record = { + "schema": calibration.SCHEMA_VERSION, + "fingerprint": calibration.fingerprint(cfg), + } + calibration.calibration_record_path().write_text(json.dumps(record)) + + assert calibration.current_record(cfg) == record + changed = replace(cfg, dflash=replace(cfg.dflash, max_ctx=65536)) + assert calibration.current_record(changed) is None diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index 06679182d..8c678c1f1 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -221,6 +221,7 @@ assert_runs "interactive menu renders and quits" \ "computers for agents" assert_runs "version" "bash '$SCRIPT' version" "" assert_runs "--version" "bash '$SCRIPT' --version" "" +assert_runs "calibrate --help" "bash '$SCRIPT' calibrate --help" "at most three" # ── 4. check — host-only, must run to completion even without docker/nvidia. # This is the path that broke last time (multi-byte glyph + set -u). @@ -235,7 +236,7 @@ assert_runs "check" "bash '$SCRIPT' check" "host readiness report" # On the bare runner there is no user systemd, no installed unit, and no # docker — so every command is expected to exit non-zero with a CLEAN error # message. What we forbid is a raw bash "unbound variable" leak. -for sub in start stop restart enable disable status install uninstall; do +for sub in start stop restart enable disable status install uninstall calibrate; do assert_no_set_u_leak "$sub dispatch (no set -u leak)" "$SCRIPT" "$sub" done # `logs` is special: it execs `journalctl -f` which streams every historical @@ -762,6 +763,119 @@ STUB } test_install_writes_robust_unit +# Calibration owns temporary service restarts, so its rollback contract is a +# production invariant rather than a cosmetic CLI detail. Exercise both the +# successful three-cell path and a failed baseline probe with isolated shims. +test_calibration_lifecycle_and_rollback() { + local label="calibration lifecycle commits success and rolls back failure" + local sandbox state shim_dir out rc restart_count apply_count + sandbox=$(mktemp -d) + state="$sandbox/state" + shim_dir="$sandbox/bin" + mkdir -p "$state" "$shim_dir" + printf 'original\n' > "$state/config.toml" + : > "$sandbox/unit" + sed '$d' "$SCRIPT" > "$sandbox/library.sh" + + cat > "$sandbox/cli-shim" <<'STUB' +#!/usr/bin/env bash +set -euo pipefail +printf '%s\n' "$*" >> "${LUCEBOX_TEST_CLI_LOG:?}" +case "${1:-}:${2:-}" in + _calibration:status) exit 1 ;; + _calibration:budgets) printf '22\n16\n32\n' ;; + _calibration:apply) printf 'budget=%s\n' "$3" > "$LUCEBOX_HOME/config.toml" ;; + _calibration:probe) + if [ "${LUCEBOX_TEST_FAIL_BASELINE:-0}" = "1" ] && [ "$3" = "22" ]; then + exit 2 + fi + printf '{}\n' > "$4" + ;; + _calibration:finish) + printf 'calibrated=22\n' > "$LUCEBOX_HOME/config.toml" + printf '22\n' > "$3/winner" + ;; + *) exit 2 ;; +esac +STUB + cat > "$shim_dir/systemctl" <<'STUB' +#!/usr/bin/env bash +printf '%s\n' "$*" >> "${LUCEBOX_TEST_SYSTEMCTL_LOG:?}" +exit 0 +STUB + cat > "$shim_dir/flock" <<'STUB' +#!/usr/bin/env bash +exit 0 +STUB + chmod +x "$sandbox/cli-shim" "$shim_dir/systemctl" "$shim_dir/flock" + + run_calibration_case() { + local fail_baseline="$1" + HOME="$sandbox/home" \ + XDG_CONFIG_HOME="$sandbox/xdg" \ + LUCEBOX_HOME="$state" \ + LUCEBOX_TEST_CLI_LOG="$sandbox/cli.log" \ + LUCEBOX_TEST_SYSTEMCTL_LOG="$sandbox/systemctl.log" \ + LUCEBOX_TEST_FAIL_BASELINE="$fail_baseline" \ + PATH="$shim_dir:$PATH" \ + NO_COLOR=1 \ + bash -c ' + source "$1/library.sh" + SCRIPT_PATH="$1/cli-shim" + SCRIPT_NAME=lucebox + CONFIG_HOME="$LUCEBOX_HOME" + UNIT_NAME=lucebox.service + UNIT_PATH="$1/unit" + ensure_probed() { :; } + _ensure_configured_image() { :; } + require_systemd() { :; } + _lucebox_config_path() { printf "%s/config.toml" "$CONFIG_HOME"; } + _lucebox_config_get() { printf "qwen3.6-27b"; } + info() { :; } + hint() { :; } + ok() { :; } + warn() { printf "WARN %s\n" "$*" >&2; } + die() { printf "ERROR %s\n" "$*" >&2; exit 1; } + cmd_calibrate + ' bash "$sandbox" + } + + : > "$sandbox/cli.log" + : > "$sandbox/systemctl.log" + rc=0 + out=$(run_calibration_case 0 2>&1) || rc=$? + restart_count=$(grep -cF -- '--user restart lucebox.service' "$sandbox/systemctl.log" || true) + apply_count=$(grep -cF -- '_calibration apply' "$sandbox/cli.log" || true) + if [ "$rc" -ne 0 ] || [ "$(<"$state/config.toml")" != "calibrated=22" ] \ + || [ "$restart_count" -ne 4 ] || [ "$apply_count" -ne 3 ] \ + || compgen -G "$state/.calibration-run.*" >/dev/null; then + report fail "$label" \ + "success case rc=$rc restarts=$restart_count applies=$apply_count output=$(head -5 <<<"$out")" + rm -rf "$sandbox" + return + fi + + printf 'original\n' > "$state/config.toml" + : > "$sandbox/cli.log" + : > "$sandbox/systemctl.log" + rc=0 + out=$(run_calibration_case 1 2>&1) || rc=$? + restart_count=$(grep -cF -- '--user restart lucebox.service' "$sandbox/systemctl.log" || true) + if [ "$rc" -eq 0 ] || [ "$(<"$state/config.toml")" != "original" ] \ + || [ "$restart_count" -ne 2 ] \ + || ! grep -qF 'original profile and service state were restored' <<<"$out" \ + || compgen -G "$state/.calibration-run.*" >/dev/null; then + report fail "$label" \ + "rollback case rc=$rc restarts=$restart_count config=$(<"$state/config.toml") output=$(head -5 <<<"$out")" + rm -rf "$sandbox" + return + fi + + report ok "$label" + rm -rf "$sandbox" +} +test_calibration_lifecycle_and_rollback + # ── 9. entrypoint.sh dispatch — confirm the in-container dispatch routes # trivial subcommands (shell, an unknown passthrough) without firing # `set -u` on DFLASH_* / DRAFT_* vars that only get assigned on the diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index fd1e61295..9c2939501 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -2,13 +2,13 @@ # In-container ENTRYPOINT for lucebox-hub. # # Normal path: the host-side `lucebox` CLI has already populated every -# DFLASH_* env var from its detection / autotune sweep, so this script -# just resolves paths and execs the native dflash_server binary. +# DFLASH_* env var from its model/hardware plan (and optional calibration), +# so this script just resolves paths and execs the native dflash_server binary. # # Fallback path: a user runs the image directly (`docker run --gpus all # ghcr.io/luce-org/lucebox-hub:cuda12`) with no env-var prep. We then do a -# minimal VRAM-tiered autotune — same tiers as `lucebox autotune`, kept in -# sync by hand. NVIDIA and AMD are both supported; elaborate driver/version +# minimal VRAM-tiered fallback — same conservative tiers as the host planner, +# kept in sync by hand. NVIDIA and AMD are both supported; elaborate driver/version # diagnostics and heterogeneous-backend selection stay in the host CLI. set -euo pipefail @@ -27,7 +27,7 @@ die() { printf '\033[1;31m[ERROR]\033[0m %s\n' "$*" >&2; exit 1; } # `shell` — drop into bash inside the container (debug). # `lucebox` — dispatch to the Python CLI. Any subcommand # `lucebox.sh` doesn't handle on the host arrives here -# (check, config, pull, print-run, smoke, …). +# (check, config, pull, print-run, calibration probes, …). # `python` or anything else # — pass through to exec, so `docker run … python -m foo` # still works for dev. @@ -469,7 +469,7 @@ fi # Qwen3.6 DFlash drafters use sliding-window attention in the draft. Some GGUFs # carry this metadata directly; keep the documented env override as the startup -# default so older drafts behave like the autotune-sweep path. +# default so older drafts retain the documented startup behavior. case "$(basename "$DFLASH_TARGET")" in *Qwen3.6*|*qwen3.6*) if [ -z "${DFLASH27B_DRAFT_SWA:-}" ]; then @@ -573,10 +573,8 @@ if [ -d "$DFLASH_DRAFT" ]; then # even if the init on line ~257 was somehow skipped (e.g. a future refactor # that moves the init out of this block, or a partial-rewrite during a # rebase that drops it). Coalesce-to-empty inline so a regression can't - # re-trip the unbound-variable crash that fired on the sindri sweep with - # multiple target GGUFs in models/ (commit a87bb93 was a partial fix — - # the recurrence proved that "initialize once at the top of the block" - # is too easy to undo). Cost: zero bytes at runtime. + # re-trip the unbound-variable crash seen with multiple target GGUFs in + # models/. Coalescing at the read site keeps a future refactor safe. DRAFT_FAMILY_GLOB="${DRAFT_FAMILY_GLOB:-}" if [ -n "$DRAFT_FILE" ] && [ -f "$DRAFT_FILE" ]; then DRAFT_ARG="$DRAFT_FILE" @@ -663,8 +661,8 @@ fi # `--prefill-drafter` and `--draft` are present (look for the runtime # warning `--lazy-draft ignored: requires both --prefill-drafter and # --draft`). Warn loudly here when the operator's config asked for lazy -# but we're about to drop it — sweeping past the silent no-op was the -# fingerprint left in every sindri decode-tuning docker.stderr. +# but we're about to drop it; otherwise performance measurements silently run +# a different profile from the one the operator selected. if [ "$DFLASH_LAZY" = "1" ]; then if [ -z "$DRAFT_ARG" ] || [ -z "$DFLASH_PREFILL_DRAFTER" ]; then warn "DFLASH_LAZY=1 ignored: requires both DFLASH_DRAFT and DFLASH_PREFILL_DRAFTER (see entrypoint.sh comment). Continuing without --lazy-draft." From 87ade52cd8b7ccee88426763cef39c92e65e16c1 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 2 Aug 2026 03:16:26 +0200 Subject: [PATCH 22/25] fix(lucebox): harden long-context optimization safety --- README.md | 4 +- lucebox/src/lucebox/autotune.py | 48 +++++- lucebox/src/lucebox/capabilities.py | 8 + lucebox/src/lucebox/config.py | 2 +- lucebox/src/lucebox/types.py | 7 +- lucebox/tests/test_autotune.py | 77 ++++++++- lucebox/tests/test_capabilities.py | 1 + lucebox/tests/test_config.py | 2 +- lucebox/tests/test_docker_run.py | 2 +- lucebox/tests/test_placement.py | 1 + server/DEVELOPER.md | 2 +- server/docs/DS4.md | 4 + server/docs/ENVIRONMENT.md | 2 +- server/docs/PREFIX_CACHE.md | 15 +- server/docs/SPEC_PREFILL.md | 3 +- server/scripts/entrypoint.sh | 5 +- server/src/common/model_capabilities.h | 31 +++- server/src/common/platform_env.h | 10 ++ server/src/deepseek4/deepseek4_backend.cpp | 160 +++++++++++++++--- server/src/deepseek4/deepseek4_backend.h | 18 +- server/src/flashprefill.cpp | 53 +++++- server/src/flashprefill.h | 27 ++- server/src/server/http_server.cpp | 26 ++- server/src/server/http_server.h | 2 +- server/src/server/prefix_cache.cpp | 69 ++++++-- server/src/server/prefix_cache.h | 16 +- server/src/server/server_main.cpp | 11 +- server/test/test_entrypoint_cache_defaults.sh | 28 +++ server/test/test_feature_gate.cpp | 4 + server/test/test_server_unit.cpp | 62 ++++++- 30 files changed, 608 insertions(+), 92 deletions(-) diff --git a/README.md b/README.md index 24bfda16b..e46d2b73f 100644 --- a/README.md +++ b/README.md @@ -431,7 +431,7 @@ End-to-end repro: `DFLASH_SAMP=0.8,1.0,0,1.1,42 python server/scripts/bench_llm. | `--prefill-drafter ` | required if on | Drafter weights (Qwen3-0.6B BF16 GGUF) | | `--prefill-skip-park` | off | Keep drafter resident across requests (more VRAM, faster) | | `PFLASH_FREEZE_HOT_WINDOW=N` | `2` | FlowKV: how many of the most recent messages stay verbatim. Everything older than this window (but after the system prompt) is compressed once and cached. Larger = more recent context kept uncompressed. | -| `DFLASH_FP_USE_BSA=1` | `0` | Dispatch sparse FA through BSA (sm_80+); required for headline 10.4× | +| `DFLASH_FP_USE_BSA=1` | `0` | Dispatch sparse FA through BSA (qualified sm_80+ except GB10 sm_121, where PFlash safely stays on exact prefill); required for headline 10.4× | | `DFLASH_FP_ALPHA=0.85` | `0.12` | Block-selection threshold; higher = stricter = fewer K-blocks | | `DFLASH_FP_PROFILE=1` | `0` | Per-stage timing log | @@ -445,7 +445,7 @@ When compression is on, the request policy keeps reusable state reusable. Struct | `DFLASH27B_KV_TQ3=1` | (default) | Preset TQ3_0 K+V (3.5 bpv, fits 256K @ 24 GB) | | `DFLASH27B_KV_Q4=1` | off | Q4_0 K+V (4.5 bpv, legacy, ~128K ceiling) | | `--prefix-cache-slots N` | — | Live prefix-cache slot count | -| `DFLASH_PREFIX_CACHE_SLOTS=N` | `32` | Container-entrypoint equivalent of `--prefix-cache-slots`; the native binary itself uses the CLI flag. | +| `DFLASH_PREFIX_CACHE_SLOTS=N` | architecture default (`8`; DeepSeek `4`) | Container-entrypoint equivalent of `--prefix-cache-slots`; the native binary itself uses the CLI flag. | | `--prefill-cache-slots N` / `DFLASH_PREFILL_CACHE_SLOTS=N` | `0` | Exact full-prompt snapshot slots. Automatic assigns four for qualified PFlash models; exact repeats skip scorer and target prefill. | | `--kv-cache-dir ` | — | Persist prefix cache to disk | | `--kv-cache-budget N` | — | On-disk cache size cap | diff --git a/lucebox/src/lucebox/autotune.py b/lucebox/src/lucebox/autotune.py index 991c1fc5b..fb50200e7 100644 --- a/lucebox/src/lucebox/autotune.py +++ b/lucebox/src/lucebox/autotune.py @@ -83,6 +83,31 @@ def _budget_for_host(host: HostFacts) -> int: return 22 +def _pflash_qualified_on_host(host: HostFacts, backend: str) -> bool: + """Whether this host has a qualified local PFlash scorer path.""" + if backend != "cuda": + return True + cuda_sm = host.nvidia_gpu_arch or host.gpu_sm + # The custom CUDA sparse forwards fault on GB10 at long context. The + # server safely retains exact prefill, but must not expose the unavailable + # scorer as an Advanced-mode optimization. + return cuda_sm != "121" + + +def _decode_qualified_on_host( + profile: ModelOptimizationProfile, + host: HostFacts, + backend: str, +) -> bool: + """Apply known device-level exclusions to speculative decode.""" + if profile.architecture != "deepseek4" or backend != "cuda": + return True + cuda_sm = host.nvidia_gpu_arch or host.gpu_sm + # DSpark currently reaches an out-of-bounds tensor read on GB10 even at + # 32K context. Keep exact AR until that CUDA path is requalified. + return cuda_sm != "121" + + def runtime_from_host(host: HostFacts) -> DflashRuntime: """Pick conservative context/cache defaults from the selected GPU. @@ -345,7 +370,11 @@ def automatic_plan( f"but the model catalog declares {preset.architecture!r}" ) backend = _selected_backend(cfg) - runtime = replace(runtime, max_ctx=min(runtime.max_ctx, preset.native_context)) + runtime = replace( + runtime, + max_ctx=min(runtime.max_ctx, preset.native_context), + prefix_cache_slots=profile.prefix_cache_slots, + ) vram_known = cfg.host.vram_gb > 0 headroom = cfg.host.vram_gb - preset.approx_total_gb if vram_known else 0 if vram_known: @@ -360,7 +389,8 @@ def automatic_plan( decode_feature = profile.speculative_decode decode_qualification = _qualification(decode_feature, backend) decode_supported = decode_feature is not None and decode_feature.available_on(backend) - dflash_available = has_draft and decode_supported + decode_hardware_qualified = _decode_qualified_on_host(profile, cfg.host, backend) + dflash_available = has_draft and decode_supported and decode_hardware_qualified dflash_enabled = ( dflash_available and decode_feature is not None @@ -371,6 +401,11 @@ def automatic_plan( dflash_reason = _preview_reason(decode_feature, backend) elif dflash_enabled: dflash_reason = "matching speculative draft is available for this model" + elif has_draft and decode_supported and not decode_hardware_qualified: + dflash_reason = ( + "DSpark is not yet qualified on GB10 CUDA; " + "autoregressive decode stays active" + ) elif has_draft and not decode_supported: dflash_reason = "the installed draft is not supported by the selected backend" elif preset.speculator_dir: @@ -416,11 +451,13 @@ def automatic_plan( pflash_supported = ( pflash_feature is not None and pflash_feature.available_on(backend) ) - pflash_available = ( + pflash_model_available = ( pflash_supported and pflash_capability is not None and runtime.max_ctx >= pflash_capability.minimum_context ) + pflash_hardware_qualified = _pflash_qualified_on_host(cfg.host, backend) + pflash_available = pflash_model_available and pflash_hardware_qualified pflash_profile = ( pflash_available and vram_known @@ -441,6 +478,11 @@ def automatic_plan( ) elif pflash_profile: pflash_reason = "shared 1.2 GB scorer is not installed" + elif pflash_model_available and not pflash_hardware_qualified: + pflash_reason = ( + "GB10's local sparse scorer kernels are not yet qualified; " + "exact prefill and prefix reuse stay active" + ) elif pflash_available and pflash_feature is not None: pflash_reason = _preview_reason(pflash_feature, backend) elif not pflash_supported: diff --git a/lucebox/src/lucebox/capabilities.py b/lucebox/src/lucebox/capabilities.py index ade2c9f99..f2766acf8 100644 --- a/lucebox/src/lucebox/capabilities.py +++ b/lucebox/src/lucebox/capabilities.py @@ -155,6 +155,11 @@ class ModelOptimizationProfile: prefill_baseline: str decode_baseline: str kv_baseline: str + # Turn-prefix snapshots grow with the cached context. Eight slots cover + # the common harness/session set without letting 100K+ contexts reserve an + # unbounded amount of host or unified memory. DeepSeek's allocated MLA + # state has a smaller architecture-specific production default. + prefix_cache_slots: int = 8 # The architecture stores an exact but substantially compressed native KV # representation. It still grows with context, but does not need the blunt # full-KV headroom cap used for ordinary attention caches. @@ -170,6 +175,8 @@ def __post_init__(self) -> None: raise ValueError("model profile preset and architecture must not be empty") if not all(label.strip() for label in self.phase_baselines): raise ValueError("every model profile must name all three baseline strategies") + if self.prefix_cache_slots < 0: + raise ValueError("model prefix-cache slots must not be negative") @property def phase_baselines(self) -> tuple[str, str, str]: @@ -292,6 +299,7 @@ class ArchitecturePlacementCapabilities: prefill_baseline="Exact MLA prefill", decode_baseline="Autoregressive decode", kv_baseline="Native MLA-compressed KV cache", + prefix_cache_slots=4, compact_native_kv=True, speculative_decode=_DSPARK, # Sparse DeepSeek prefill is engine-backed but approximate and diff --git a/lucebox/src/lucebox/config.py b/lucebox/src/lucebox/config.py index dd104e2c6..7ddf134a8 100644 --- a/lucebox/src/lucebox/config.py +++ b/lucebox/src/lucebox/config.py @@ -367,7 +367,7 @@ def _from_dict(raw: dict[str, Any]) -> Config: budget=int(df.get("budget", 22)), max_ctx=int(df.get("max_ctx", 16384)), lazy=_cast_bool(df.get("lazy", False)), - prefix_cache_slots=int(df.get("prefix_cache_slots", 32)), + prefix_cache_slots=int(df.get("prefix_cache_slots", 8)), prefill_cache_slots=int(df.get("prefill_cache_slots", 0)), cache_type_k=str(df.get("cache_type_k", "")), cache_type_v=str(df.get("cache_type_v", "")), diff --git a/lucebox/src/lucebox/types.py b/lucebox/src/lucebox/types.py index 39e9c1a17..497103263 100644 --- a/lucebox/src/lucebox/types.py +++ b/lucebox/src/lucebox/types.py @@ -236,9 +236,10 @@ class DflashRuntime: budget: int = 22 max_ctx: int = 16384 lazy: bool = False - # Exact turn-boundary snapshots. This is safe for every model/backend and - # is the server's production default; 0 remains an explicit opt-out. - prefix_cache_slots: int = 32 + # Exact turn-boundary snapshots. Eight slots bound memory at long context + # while covering the common harness/session set; 0 remains an explicit + # opt-out and Advanced mode can raise the cap deliberately. + prefix_cache_slots: int = 8 # Exact full-prompt snapshots, keyed by the raw prompt. PFlash can use # these to skip both rescoring and target prefill when a request repeats. prefill_cache_slots: int = 0 diff --git a/lucebox/tests/test_autotune.py b/lucebox/tests/test_autotune.py index cb4dea8b5..ff644a742 100644 --- a/lucebox/tests/test_autotune.py +++ b/lucebox/tests/test_autotune.py @@ -22,7 +22,7 @@ def test_wsl_24gb_defaults_leave_cuda_headroom() -> None: # to False makes the host config match runtime behaviour. See the # `entrypoint.sh` warning emitted when the two are out-of-sync. assert runtime.lazy is False - assert runtime.prefix_cache_slots == 32 + assert runtime.prefix_cache_slots == 8 def test_native_24gb_caps_context_below_vmm_failure_boundary() -> None: @@ -31,7 +31,7 @@ def test_native_24gb_caps_context_below_vmm_failure_boundary() -> None: assert runtime.budget == 22 assert runtime.max_ctx == 98304 assert runtime.lazy is False # see WSL test above - assert runtime.prefix_cache_slots == 32 + assert runtime.prefix_cache_slots == 8 def test_no_heuristic_tier_sets_lazy_without_prefill_drafter() -> None: @@ -127,6 +127,40 @@ def test_qwen_24gb_enables_dflash_and_pflash_but_prefers_full_kv( assert "DFLASH_SPARK" not in env +def test_gb10_hides_unqualified_pflash_and_keeps_exact_prefill( + tmp_path: Path, +) -> None: + cfg = _cfg( + tmp_path, + "qwen3.6-27b", + HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + gpu_name="NVIDIA GB10", + gpu_count=1, + vram_gb=105, + gpu_sm="121", + ram_gb=121, + nvidia_gpu_name="NVIDIA GB10", + nvidia_gpu_count=1, + nvidia_gpu_arch="121", + nvidia_unified_memory=True, + ), + ) + _install_optimizer_drafter(cfg) + _install_decode_draft(cfg) + + plan = automatic_plan(cfg) + + assert plan.dflash.enabled is True + assert plan.pflash.available is False + assert plan.pflash.enabled is False + assert plan.runtime.prefill_mode == "off" + assert "GB10" in plan.pflash.reason + assert "prefix reuse" in plan.pflash.reason + assert "DFLASH_PREFILL_MODE" not in _server_env_for_plan(cfg, plan) + + def test_published_draft_is_not_claimed_until_it_is_installed(tmp_path: Path) -> None: cfg = _cfg( tmp_path, @@ -313,6 +347,35 @@ def test_deepseek_compact_mla_cache_avoids_the_generic_full_kv_context_cap( assert plan.runtime.ds4_prefill == "exact" +def test_deepseek_disables_unqualified_dspark_on_gb10(tmp_path: Path) -> None: + cfg = _cfg( + tmp_path, + "deepseek-v4-flash", + HostFacts( + gpu_vendor="nvidia", + has_nvidia_gpu=True, + gpu_name="NVIDIA GB10", + vram_gb=128, + ram_gb=256, + gpu_sm="121", + nvidia_gpu_name="NVIDIA GB10", + nvidia_gpu_count=1, + nvidia_vram_gb=128, + nvidia_gpu_arch="121", + ), + ) + _install_decode_draft(cfg) + + plan = automatic_plan(cfg) + + assert plan.dflash.name == "DSpark" + assert plan.dflash.available is False + assert plan.dflash.enabled is False + assert plan.runtime.speculative_decode is False + assert plan.decode_strategy == "Autoregressive decode" + assert "GB10" in plan.dflash.reason + + def test_spark_stays_off_when_host_ram_cannot_hold_cold_experts( tmp_path: Path, ) -> None: @@ -568,14 +631,15 @@ def test_featured_model_strix_only_contract_reaches_exact_server_environment( assert plan.placement.runtime.target_device == "hip:0" assert plan.placement.runtime.uses_multiple_devices is False assert plan.runtime.max_ctx == 131072 - assert plan.runtime.prefix_cache_slots == 32 + expected_prefix_slots = 4 if preset_name == "deepseek-v4-flash" else 8 + assert plan.runtime.prefix_cache_slots == expected_prefix_slots assert plan.runtime.prefill_cache_slots == (4 if pflash_enabled else 0) assert plan.dflash.enabled is decode_enabled assert plan.pflash.enabled is pflash_enabled assert plan.kvflash.enabled is False assert plan.spark.enabled is False assert env["DFLASH_TARGET_DEVICE"] == "hip:0" - assert env["DFLASH_PREFIX_CACHE_SLOTS"] == "32" + assert env["DFLASH_PREFIX_CACHE_SLOTS"] == str(expected_prefix_slots) assert env["DFLASH_PREFILL_CACHE_SLOTS"] == ("4" if pflash_enabled else "0") assert "DFLASH_TARGET_DEVICES" not in env assert ("DFLASH_PREFILL_MODE" in env) is pflash_enabled @@ -638,7 +702,8 @@ def test_featured_model_roomy_gpu_contract_reaches_exact_server_environment( env = dict(server_run_spec(launch_cfg).env) assert plan.placement.runnable is True - assert plan.runtime.prefix_cache_slots == 32 + expected_prefix_slots = 4 if preset_name == "deepseek-v4-flash" else 8 + assert plan.runtime.prefix_cache_slots == expected_prefix_slots assert plan.runtime.prefill_cache_slots == (4 if pflash_enabled else 0) assert plan.dflash.enabled is bool(decode_name) assert plan.dflash.name == (decode_name or "DFlash") @@ -647,7 +712,7 @@ def test_featured_model_roomy_gpu_contract_reaches_exact_server_environment( assert plan.spark.enabled is False assert env["DFLASH_MODEL_NAME"] == preset_name assert env["DFLASH_TARGET_DEVICE"] == ("cuda:0" if vendor == "nvidia" else "hip:0") - assert env["DFLASH_PREFIX_CACHE_SLOTS"] == "32" + assert env["DFLASH_PREFIX_CACHE_SLOTS"] == str(expected_prefix_slots) assert env["DFLASH_PREFILL_CACHE_SLOTS"] == ("4" if pflash_enabled else "0") assert ("DFLASH_PREFILL_MODE" in env) is pflash_enabled assert "DFLASH_KVFLASH" not in env diff --git a/lucebox/tests/test_capabilities.py b/lucebox/tests/test_capabilities.py index 0b77feb46..5a6bc75f5 100644 --- a/lucebox/tests/test_capabilities.py +++ b/lucebox/tests/test_capabilities.py @@ -32,6 +32,7 @@ def test_deepseek_uses_native_mla_instead_of_claiming_generic_flash_paths() -> N assert profile.pflash is None assert profile.kvflash is None + assert profile.prefix_cache_slots == 4 assert "MLA" in profile.prefill_baseline assert "MLA" in profile.kv_baseline assert profile.deepseek_prefill is not None diff --git a/lucebox/tests/test_config.py b/lucebox/tests/test_config.py index 92f3aea06..f7304eb62 100644 --- a/lucebox/tests/test_config.py +++ b/lucebox/tests/test_config.py @@ -483,7 +483,7 @@ def test_automatic_runtime_preserves_model_card_thinking_budget(tmp_path: Path) loaded = config.load(path) assert loaded is not None assert loaded.dflash.think_max is None - assert loaded.dflash.prefix_cache_slots == 32 + assert loaded.dflash.prefix_cache_slots == 8 def test_explicit_thinking_budget_round_trips(tmp_path: Path) -> None: diff --git a/lucebox/tests/test_docker_run.py b/lucebox/tests/test_docker_run.py index e09672e3b..e179f3f68 100644 --- a/lucebox/tests/test_docker_run.py +++ b/lucebox/tests/test_docker_run.py @@ -377,7 +377,7 @@ def test_server_run_spec_always_emits_core_dflash_env(tmp_path: Path) -> None: env = _env(docker_run.server_run_spec(cfg)) assert env["DFLASH_BUDGET"] == "22" assert env["DFLASH_MAX_CTX"] == "32768" - assert env["DFLASH_PREFIX_CACHE_SLOTS"] == "32" + assert env["DFLASH_PREFIX_CACHE_SLOTS"] == "8" assert env["DFLASH_PREFILL_CACHE_SLOTS"] == "0" assert "DFLASH_THINK_MAX" not in env assert env["DFLASH_PORT"] == "8080" diff --git a/lucebox/tests/test_placement.py b/lucebox/tests/test_placement.py index 6cb99c170..c08223712 100644 --- a/lucebox/tests/test_placement.py +++ b/lucebox/tests/test_placement.py @@ -622,6 +622,7 @@ def test_python_architecture_capabilities_match_engine_table() -> None: r'\{"(?P[^\"]+)",\s*' r"(?Ptrue|false),\s*(?Ptrue|false),\s*" r"(?Ptrue|false),\s*(?Ptrue|false),\s*" + r"(?P[0-9]+),\s*" r"(?PkNever|kMono|kBoth)," ) engine_rows = {match["arch"]: match.groupdict() for match in row_pattern.finditer(header)} diff --git a/server/DEVELOPER.md b/server/DEVELOPER.md index 31148dc74..a1b85335d 100644 --- a/server/DEVELOPER.md +++ b/server/DEVELOPER.md @@ -10,7 +10,7 @@ | VRAM | 22 GB | 24 GB | | OS | Ubuntu 22.04 (jammy) | Ubuntu 24.04 (noble) | -> **Note:** FlashPrefill and BSA (Block-Sparse Attention) require **sm_80+** (Ampere or newer). +> **Note:** FlashPrefill and BSA (Block-Sparse Attention) require **sm_80+** (Ampere or newer). GB10's sm_121 custom sparse kernels are not yet qualified, so PFlash safely stays on exact prefill. > On Turing (sm_75) the drafter falls back to ggml's `flash_attn_ext`. ### System packages diff --git a/server/docs/DS4.md b/server/docs/DS4.md index e0fc2dba7..25c95b1e5 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -207,6 +207,10 @@ GGUF carries its auxiliary projections under the existing `dflash.dspark.*` tensor contract. DeepSeek4/MTP checkpoints store compatible heads under the `mtp.2.*` namespace, which the converter maps as follows. +CUDA sm_121 (GB10) currently stays on autoregressive decode: the DSpark path is +disabled in both the automatic planner and server until its tensor-read failure +is requalified. HIP and other qualified CUDA architectures are unchanged. + Supported DeepSeek4/MTP input tensors: | DeepSeek4/MTP tensor | GGUF tensor | diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index a85e1fa65..d2d553735 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -33,7 +33,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_MOE_TP_*` / `DFLASH_MOE_HYBRID_PREFILL_EAGER` | unset | BURN-IN: model-neutral names for common heterogeneous-MoE scheduling and kernel policy. Existing `DFLASH_DS4_*` names remain compatibility aliases. | | `DFLASH_MMID_TELEMETRY` | unset | DEBUG: report MUL_MAT_ID dispatch, MMVQ variant, and per-node graph compatibility. | | `DFLASH_KVFLASH` | unset | Prefer the CLI: `--kvflash` (token count or `auto`). | -| `DFLASH_PREFIX_CACHE_SLOTS` | 32 | Container-entrypoint equivalent of `--prefix-cache-slots`; not read directly by the native binary. | +| `DFLASH_PREFIX_CACHE_SLOTS` | architecture default (`8`; DeepSeek `4`) | Container-entrypoint equivalent of `--prefix-cache-slots`; not read directly by the native binary. | | `DFLASH_PREFILL_CACHE_SLOTS` | 0 | Container-entrypoint equivalent of `--prefill-cache-slots`; not read directly by the native binary. | | `DFLASH_SPLIT_FAST_ROLLBACK` | unset | OPT-IN: exact F32 checkpoints and replay-free rollback for local qwen35 target layer splits. Prefer `--target-split-fast-rollback`; adds checkpoint VRAM (~1.65 GiB for the measured Qwen3.6-27B q=16 split). | | `DFLASH_STALL_TOOL_PREFIX` | unset | OPT-IN: recover a stalled tool call by injecting the prepared tool prefix when generation stops after an action suffix. | diff --git a/server/docs/PREFIX_CACHE.md b/server/docs/PREFIX_CACHE.md index 392bb1448..72f311094 100644 --- a/server/docs/PREFIX_CACHE.md +++ b/server/docs/PREFIX_CACHE.md @@ -194,21 +194,22 @@ free_snapshot_backend(snap_backend_, compute_backend_); // then backend | Server flag | Default | Description | |-------------|---------|-------------| -| `--prefix-cache-slots N` | 32 | Max turn-boundary prefix cache slots | +| `--prefix-cache-slots N` | 8 (DeepSeek: 4) | Max turn-boundary prefix cache slots | | `--prefill-cache-slots N` | 0 | Max exact full-prompt prefill cache slots | | `--skip-park` | false | Skip parking draft model during compress | ### Choosing `--prefix-cache-slots` -With right-sized, CPU-resident snapshots the limiting resource is **system RAM**, -not VRAM. Each slot costs approximately `cur_pos × 5 KB` (for Qwen3.5-27B Q8_0 KV), -so 32 slots with an average prefix of 2000 tokens ≈ 320 MB of system RAM — negligible -on most workstations. +With right-sized snapshots the limiting resource is **system or unified RAM**, +not only VRAM. Cost grows linearly with the cached prefix and varies by model. +Short chat prefixes are inexpensive, but 100K+ agent/tool prefixes can make one +slot several GiB. The production defaults therefore favor a bounded working +set; raise the cap only after accounting for the selected model and context. | Scenario | Typical prefix length | Recommended cap | |----------|----------------------|-----------------| -| Single-user chat | 200–2000 tokens | 16–32 | -| Multi-session agent | 500–5000 tokens | 32–64 | +| Single-user chat | 200–2000 tokens | 8 | +| Multi-session agent | 500–5000 tokens | 8–16 | | Batch / benchmark | N/A (cold starts) | 4 | The hard limit is `MAX_SLOTS = 64`. Beyond that, increase the constant in diff --git a/server/docs/SPEC_PREFILL.md b/server/docs/SPEC_PREFILL.md index 0859c50ee..82e87a798 100644 --- a/server/docs/SPEC_PREFILL.md +++ b/server/docs/SPEC_PREFILL.md @@ -26,7 +26,8 @@ cmake --build . --target test_dflash test_flashprefill_kernels -- -j8 ``` Required: -- CUDA Toolkit 12.0+ (sm_80+ for BSA path; sm_86 RTX 3090 is the +- CUDA Toolkit 12.0+ (custom BF16 sparse kernels are qualified on sm_80+ + except GB10 sm_121, where PFlash stays on exact prefill; sm_86 RTX 3090 is the reference target). - `git submodule update --init --recursive` to pull `deps/Block-Sparse-Attention` (with cutlass). `deps/llama.cpp` is vendored diff --git a/server/scripts/entrypoint.sh b/server/scripts/entrypoint.sh index 9c2939501..1ea98e281 100755 --- a/server/scripts/entrypoint.sh +++ b/server/scripts/entrypoint.sh @@ -253,7 +253,10 @@ GPU_COUNT=0 if command -v nvidia-smi &>/dev/null; then if mem_mib=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null \ | head -1) && [ -n "$mem_mib" ]; then - GPU_VRAM_GB=$((mem_mib / 1024)) + mem_mib=$(_trim "$mem_mib") + if [[ "$mem_mib" =~ ^[0-9]+$ ]]; then + GPU_VRAM_GB=$((mem_mib / 1024)) + fi fi GPU_COUNT=$(nvidia-smi -L 2>/dev/null | awk '/^GPU /{n++} END{print n+0}') || GPU_COUNT=0 elif command -v amd-smi &>/dev/null; then diff --git a/server/src/common/model_capabilities.h b/server/src/common/model_capabilities.h index 2189fe312..2e1ffe802 100644 --- a/server/src/common/model_capabilities.h +++ b/server/src/common/model_capabilities.h @@ -48,6 +48,7 @@ struct ArchCapabilities { // --collect-routing, --adaptive-experts). Note // deepseek4 is mixture-of-experts but has no such // path, so this is narrower than "is MoE". + int default_prefix_cache_slots; // bounded by architecture snapshot cost // Placement-dependent. FeatureSupport decode_draft; // --draft @@ -62,13 +63,15 @@ inline constexpr FeatureSupport kMono = FeatureSupport::Monolithic; inline constexpr FeatureSupport kBoth = FeatureSupport::Both; inline constexpr ArchCapabilities kArchCapabilities[] = { -// arch split rdraft pflash offload draft ddtree vwidth fa_win dswa - {"qwen35", true, true, true, false, kBoth, kBoth, kNever, kBoth, kBoth}, - {"qwen35moe", false, false, false, true, kMono, kMono, kNever, kMono, kMono}, - {"laguna", true, false, false, true, kMono, kMono, kMono, kNever, kNever}, - {"qwen3", false, false, true, false, kNever, kNever, kNever, kNever, kNever}, - {"gemma4", true, false, false, false, kMono, kNever, kNever, kBoth, kNever}, - {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever}, +// arch split rdraft pflash offload pc draft ddtree vwidth fa_win dswa + {"qwen35", true, true, true, false, 8, kBoth, kBoth, kNever, kBoth, kBoth}, + {"qwen35moe", false, false, false, true, 8, kMono, kMono, kNever, kMono, kMono}, + {"laguna", true, false, false, true, 8, kMono, kMono, kMono, kNever, kNever}, + {"qwen3", false, false, true, false, 8, kNever, kNever, kNever, kNever, kNever}, + {"gemma4", true, false, false, false, 8, kMono, kNever, kNever, kBoth, kNever}, + // DeepSeek snapshots copy its allocated MLA state. At 131K, one slot is + // roughly 0.9 GiB, so the generic 32-slot default could consume ~30 GiB. + {"deepseek4", true, false, false, false, 4, kNever, kNever, kNever, kNever, kNever}, }; inline constexpr std::size_t kArchCount = @@ -116,6 +119,13 @@ constexpr bool table_rows_named() { return true; } +constexpr bool table_cache_defaults_valid() { + for (const ArchCapabilities & c : kArchCapabilities) { + if (c.default_prefix_cache_slots < 0) return false; + } + return true; +} + constexpr bool table_rows_unique() { for (std::size_t i = 0; i < kArchCount; ++i) { for (std::size_t j = i + 1; j < kArchCount; ++j) { @@ -155,6 +165,8 @@ static_assert(detail::table_rows_named(), "every capability row needs a non-empty architecture name"); static_assert(detail::table_rows_unique(), "duplicate architecture row in kArchCapabilities"); +static_assert(detail::table_cache_defaults_valid(), + "architecture prefix-cache defaults must not be negative"); static_assert(detail::table_split_coherent(), "an architecture with no layer-split adapter cannot support an " "option on 'Both' placements; use Monolithic"); @@ -210,6 +222,11 @@ inline bool arch_has_expert_offload(const std::string & arch) { return detail::arch_has(arch, &ArchCapabilities::expert_offload); } +inline int arch_default_prefix_cache_slots(const std::string & arch) { + const ArchCapabilities * caps = find_arch_capabilities(arch); + return caps ? caps->default_prefix_cache_slots : 8; +} + inline bool arch_supports_decode_draft(const std::string & arch, bool is_layer_split) { return detail::arch_has(arch, &ArchCapabilities::decode_draft, is_layer_split); diff --git a/server/src/common/platform_env.h b/server/src/common/platform_env.h index 934d9df8d..92fb22d90 100644 --- a/server/src/common/platform_env.h +++ b/server/src/common/platform_env.h @@ -3,9 +3,19 @@ #pragma once #include +#include namespace dflash::common { +// Treat an unset, empty, or explicit "0" value as disabled. This lets a +// caller override a default with FOO=0 instead of mere variable presence +// accidentally enabling the feature. +inline bool environment_variable_enabled(const char * name) { + const char * value = std::getenv(name); + return value != nullptr && value[0] != '\0' && + std::strcmp(value, "0") != 0; +} + // Match POSIX setenv() semantics on every platform. In particular, // overwrite=false must preserve an existing value; _putenv_s() does not // provide that behavior by itself. diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index e6702903b..4c66b208e 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -24,6 +24,10 @@ namespace dflash::common { +bool deepseek4_dspark_supports_cuda_sm(int sm) { + return sm != 121; +} + namespace { using Clock = std::chrono::steady_clock; @@ -40,6 +44,19 @@ static bool env_flag_enabled(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } +static bool dspark_supported_on_current_device(int gpu, int & cuda_sm) { + cuda_sm = 0; +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + (void) gpu; + return true; +#else + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, gpu) != cudaSuccess) return false; + cuda_sm = prop.major * 10 + prop.minor; + return deepseek4_dspark_supports_cuda_sm(cuda_sm); +#endif +} + static void configure_gfx1151_dspark_mmvq_default(int gpu) { #if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) if (!env_flag_enabled("DFLASH_DS4_SPEC") || @@ -627,7 +644,7 @@ void DeepSeek4Backend::release_spec_drafter(bool mark_parked) { spec_backend_ = nullptr; } spec_enabled_ = false; - spec_feat_window_.clear(); + std::vector().swap(spec_feat_window_); spec_drafter_parked_ = mark_parked && !spec_draft_path_.empty(); } @@ -704,7 +721,20 @@ bool DeepSeek4Backend::init() { prefill_attention_mode_name(cfg_.prefill_mode), moe_hybrid_ ? " [hybrid]" : ""); - if (env_flag_enabled("DFLASH_DS4_SPEC")) { + const bool dspark_requested = env_flag_enabled("DFLASH_DS4_SPEC"); + int dspark_cuda_sm = 0; + if (dspark_requested && + !dspark_supported_on_current_device(cfg_.device.gpu, dspark_cuda_sm)) { + if (dspark_cuda_sm == 121) { + std::fprintf(stderr, + "[deepseek4] DSpark disabled: CUDA sm_121 is not qualified; " + "continuing with autoregressive decode\n"); + } else { + std::fprintf(stderr, + "[deepseek4] DSpark disabled: selected CUDA device could not " + "be qualified; continuing with autoregressive decode\n"); + } + } else if (dspark_requested) { const char * dp = std::getenv("DFLASH_DS4_DRAFT"); if (dp && *dp) { spec_draft_path_ = dp; @@ -941,9 +971,10 @@ bool DeepSeek4Backend::park(ParkTarget target) { maybe_save_routing_stats(); for (int i = 0; i < PREFIX_SLOTS; ++i) { - free_deepseek4_snapshot(snapshots_[i]); + snapshot_free(i); } - last_logits_.clear(); + std::vector().swap(last_logits_); + std::vector().swap(spec_feat_window_); free_deepseek4_cache(cache_); expert_runtime_.reset(); stream_engine_.destroy(); @@ -1044,7 +1075,9 @@ bool DeepSeek4Backend::unpark(ParkTarget target) { int DeepSeek4Backend::do_prefill(const std::vector & tokens, const DaemonIO & io, - int kv_offset) { + int kv_offset, + int snap_pos, + int snap_slot) { // The all-hot layer-range path supports causal chunked prefill. The // optimized graph snapshots the previous raw SWA window, attends over // that snapshot plus the current ubatch, and commits only the final SWA @@ -1108,10 +1141,17 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, DeepSeek4StepTelemetry tel_acc; int steps = 0; - for (int i = 0; i < n_total; i += chunk) { + for (int i = 0; i < n_total;) { if (io.cancelled) return pos; - const int n_tok = std::min(chunk, n_total - i); + int n_tok = std::min(chunk, n_total - i); + // The cache key names an exact token prefix. Split only the chunk + // that crosses that prefix so the saved recurrent/compressor state + // cannot include tokens beyond the key. Exact prefill is already + // tokenwise; preview batched modes keep normal chunking elsewhere. + if (snap_slot >= 0 && snap_pos > pos && snap_pos < pos + n_tok) { + n_tok = snap_pos - pos; + } // Embed tokens std::vector embed(w_.n_embd * n_tok); @@ -1177,6 +1217,19 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, } last_logits_ = std::move(logits); pos += n_tok; + i += n_tok; + + if (snap_slot >= 0 && snap_pos > 0 && pos == snap_pos) { + if (snapshot_save(snap_slot)) { + std::fprintf(stderr, + "[deepseek4] inline-snap slot=%d cur_pos=%d\n", + snap_slot, snapshot_cur_pos(snap_slot)); + } else { + std::fprintf(stderr, + "[deepseek4] inline-snap failed slot=%d requested=%d\n", + snap_slot, snap_pos); + } + } } if (timing) { log_step_tel("prefill", n_total, steps, elapsed_s(phase_t0), tel_acc); @@ -1317,12 +1370,23 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, } // Prefill - int committed = do_prefill(req.prompt, out_io); + int committed = do_prefill( + req.prompt, out_io, /*kv_offset=*/0, req.snap_pos, req.snap_slot); if (committed < 0) { result.fail(GenerateErrorCode::PrefillFailed); return result; } - result.prefill_s = elapsed_s(t0); + return finish_generation(req, out_io, committed, elapsed_s(t0), req.prompt); +} + +GenerateResult DeepSeek4Backend::finish_generation( + const GenerateRequest & req, + const DaemonIO & out_io, + int committed, + double prefill_s, + const std::vector & history_prefix) { + GenerateResult result; + result.prefill_s = prefill_s; if (out_io.cancelled) { result.succeed(); @@ -1396,7 +1460,7 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, gen_tokens.reserve(req.n_gen); bool forced_close = false; - if (!do_decode(committed, req.n_gen, req.prompt, gen_tokens, out_io, + if (!do_decode(committed, req.n_gen, history_prefix, gen_tokens, out_io, req.budget_hook, &forced_close)) { result.fail(GenerateErrorCode::DecodeFailed); return result; @@ -1414,30 +1478,88 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, bool DeepSeek4Backend::snapshot_save(int slot) { if (slot < 0 || slot >= PREFIX_SLOTS) return false; - // TODO: Implement snapshot save (copy KV cache + HC state to CPU) - return false; + auto & snap = snapshots_[slot]; + snapshot_free(slot); + if (!deepseek4_snapshot_save(cache_, snap_backend_, snap.cache)) { + snapshot_free(slot); + return false; + } + snap.last_logits = last_logits_; + snap.spec_feat_window = spec_feat_window_; + return true; } void DeepSeek4Backend::snapshot_free(int slot) { if (slot < 0 || slot >= PREFIX_SLOTS) return; - free_deepseek4_snapshot(snapshots_[slot]); + auto & snap = snapshots_[slot]; + free_deepseek4_snapshot(snap.cache); + std::vector().swap(snap.last_logits); + std::vector().swap(snap.spec_feat_window); } bool DeepSeek4Backend::snapshot_used(int slot) const { if (slot < 0 || slot >= PREFIX_SLOTS) return false; - return snapshots_[slot].ctx != nullptr; + return snapshots_[slot].cache.ctx != nullptr; } int DeepSeek4Backend::snapshot_cur_pos(int slot) const { if (slot < 0 || slot >= PREFIX_SLOTS) return 0; - return snapshots_[slot].cur_pos; + return snapshots_[slot].cache.cur_pos; } GenerateResult DeepSeek4Backend::restore_and_generate_impl( int slot, const GenerateRequest & req, const DaemonIO & io) { - // TODO: Implement snapshot restore + generate - (void)slot; - return generate_impl(req, io); + GenerateResult result; + DaemonIO out_io = io.with_token_callback(req.on_token); + if (!snapshot_used(slot)) { + result.fail(GenerateErrorCode::InvalidSnapshotSlot); + out_io.emit(-1); + return result; + } + + auto & snap = snapshots_[slot]; + if (!deepseek4_snapshot_restore(snap.cache, cache_)) { + result.fail(GenerateErrorCode::InvalidSnapshotSlot, + "DeepSeek cache restore failed"); + out_io.emit(-1); + return result; + } + last_logits_ = snap.last_logits; + spec_feat_window_ = snap.spec_feat_window; + + sampler_ = req.sampler; + if (req.do_sample && sampler_.seed != 0) { + sampler_rng_.seed(sampler_.seed); + } + + const int snap_pos = snap.cache.cur_pos; + const int prompt_len = (int) req.prompt.size(); + if (prompt_len < snap_pos) { + std::fprintf(stderr, + "[pc] snapshot longer than prompt (snap=%d > prompt=%d) — " + "fresh prefill fallback\n", snap_pos, prompt_len); + return generate_impl(req, io); + } + + auto t0 = Clock::now(); + int committed = snap_pos; + if (prompt_len > snap_pos) { + std::vector delta( + req.prompt.begin() + snap_pos, req.prompt.end()); + committed = do_prefill( + delta, out_io, snap_pos, req.snap_pos, req.snap_slot); + if (committed < 0) { + result.fail(GenerateErrorCode::PrefillFailed); + return result; + } + } else if (last_logits_.empty() && req.n_gen > 0) { + result.fail(GenerateErrorCode::DecodeSeedMissing, + "DeepSeek snapshot has no prefill logits"); + return result; + } + + return finish_generation( + req, out_io, committed, elapsed_s(t0), req.prompt); } bool DeepSeek4Backend::handle_compress(const std::string & line, @@ -1466,7 +1588,7 @@ void DeepSeek4Backend::shutdown() { maybe_save_routing_stats(); free_drafter(); for (int i = 0; i < PREFIX_SLOTS; i++) { - free_deepseek4_snapshot(snapshots_[i]); + snapshot_free(i); } free_deepseek4_cache(cache_); expert_runtime_.reset(); diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index c7102ff85..c1d2d5c66 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -26,6 +26,10 @@ namespace dflash::common { +// DSpark's CUDA path is qualified broadly except for the known GB10 sm_121 +// tensor-read failure. Kept pure so the runtime policy has a unit test. +bool deepseek4_dspark_supports_cuda_sm(int sm); + class DeepSeek4Backend : public ModelBackend { public: explicit DeepSeek4Backend(const DeepSeek4BackendConfig & cfg); @@ -80,7 +84,12 @@ class DeepSeek4Backend : public ModelBackend { // Snapshots static constexpr int PREFIX_SLOTS = 64; - DeepSeek4Snapshot snapshots_[PREFIX_SLOTS]; + struct PrefixSnapshot { + DeepSeek4Snapshot cache; + std::vector last_logits; + std::vector spec_feat_window; + }; + PrefixSnapshot snapshots_[PREFIX_SLOTS]; std::vector last_logits_; // DSpark speculative decode (opt-in: DFLASH_DS4_SPEC=1 + DFLASH_DS4_DRAFT=). @@ -96,7 +105,12 @@ class DeepSeek4Backend : public ModelBackend { // Prefill prompt tokens in chunks, return absolute committed position. int do_prefill(const std::vector & tokens, const DaemonIO & io, - int kv_offset = 0); + int kv_offset = 0, int snap_pos = -1, int snap_slot = -1); + + GenerateResult finish_generation( + const GenerateRequest & req, const DaemonIO & io, + int committed, double prefill_s, + const std::vector & history_prefix); // Autoregressive decode loop. bool do_decode(int committed, int n_gen, diff --git a/server/src/flashprefill.cpp b/server/src/flashprefill.cpp index 8502f0059..8d1fc2e44 100644 --- a/server/src/flashprefill.cpp +++ b/server/src/flashprefill.cpp @@ -6,7 +6,9 @@ #include "flashprefill.h" #include "flashprefill_launchers.h" +#include "common/platform_env.h" +#include #include #include #include @@ -209,6 +211,47 @@ namespace { inline int cdiv(int a, int b) { return (a + b - 1) / b; } } +bool custom_bf16_sparse_supported_on_current_device() { +#ifdef DFLASH27B_BACKEND_HIP + // HIP uses Lucebox's rocWMMA implementation; the GB10 restriction is + // specific to the CUDA sm_121 kernels. + return true; +#else + int device = 0; + cudaDeviceProp properties{}; + if (cudaGetDevice(&device) != cudaSuccess || + cudaGetDeviceProperties(&properties, device) != cudaSuccess) { + return false; + } + const int sm = properties.major * 10 + properties.minor; + const bool supported = custom_bf16_sparse_supports_cuda_sm(sm); + if (!supported) { + static std::atomic warned{false}; + if (!warned.exchange(true)) { + std::fprintf(stderr, + "[flashprefill] custom BF16 sparse kernels unavailable on " + "CUDA sm_%d; caller must use a safe fallback\n", sm); + } + } + return supported; +#endif +} + +bool local_pflash_supported_on_current_device() { +#ifdef DFLASH27B_BACKEND_HIP + return true; +#else + int device = 0; + cudaDeviceProp properties{}; + if (cudaGetDevice(&device) != cudaSuccess || + cudaGetDeviceProperties(&properties, device) != cudaSuccess) { + return false; + } + return local_pflash_supports_cuda_sm( + properties.major * 10 + properties.minor); +#endif +} + #if defined(DFLASH27B_HAVE_FLASHPREFILL) || defined(DFLASH27B_HAVE_SM80_FLASHPREFILL) // ── BF16 (sm_80+) dispatch: native BF16 WMMA kernels ── @@ -218,6 +261,12 @@ int flash_prefill_forward_bf16( float scale, const FlashPrefillConfig & cfg) { + // Direct users such as the registered ggml sparse-attention callback do + // not pass a backend and therefore cannot use flash_prefill_forward_q8. + // Reject the unqualified kernel before any launch so the owning request + // path can keep the original prompt and choose exact prefill instead. + if (!custom_bf16_sparse_supported_on_current_device()) return -2; + const int B = batch; const int S = seq_len; const int H = n_q_heads; @@ -328,7 +377,9 @@ int flash_prefill_forward_bf16( } // 4. sparse flash forward (BSA-or-WMMA) #ifdef DFLASH27B_HAVE_BSA - static const bool use_bsa = (std::getenv("DFLASH_FP_USE_BSA") != nullptr); + const bool bsa_requested = + environment_variable_enabled("DFLASH_FP_USE_BSA"); + const bool use_bsa = bsa_requested; if (use_bsa && D == 128 && BLOCK == 128) { launch_bsa_sparse_flash_forward_bf16( Q, K, V, O, dIdx, dCnt, scale, diff --git a/server/src/flashprefill.h b/server/src/flashprefill.h index fd0c64e04..6cc021563 100644 --- a/server/src/flashprefill.h +++ b/server/src/flashprefill.h @@ -11,9 +11,11 @@ // O[B, S, n_q_heads, D] // // Backends: -// - Default: WMMA m16n16k16 sparse forward (sm_70+). Functional everywhere. +// - Default: WMMA m16n16k16 sparse forward (qualified CUDA sm_80+ and HIP). // - Set env DFLASH_FP_USE_BSA=1 to dispatch to the Block-Sparse-Attention -// kernel (FA-2 derived, m16n8k16 PTX, sm_80+ via cuBLAS BF16 GEMM). +// kernel (FA-2 derived, m16n8k16 PTX, qualified CUDA sm_80+ except GB10 +// sm_121). PFlash stays on exact prefill there because neither bundled +// custom BF16 sparse kernel is qualified. // Requires building with -DDFLASH27B_ENABLE_BSA=ON. ~3x faster than WMMA // on RTX 3090 at S=128K. // @@ -35,6 +37,23 @@ namespace dflash::common { namespace flashprefill { +// Keep the low-level BF16-kernel policy separate from the product-level local +// PFlash policy. Legacy CUDA architectures have other compiled fallbacks, but +// GB10's sm_121 is the one device on which the complete local scorer path is +// currently unsafe. +inline constexpr bool custom_bf16_sparse_supports_cuda_sm(int sm) { + return sm >= 80 && sm != 121; +} +inline constexpr bool local_pflash_supports_cuda_sm(int sm) { + return sm != 121; +} + +// Runtime helpers are implemented in flashprefill.cpp, which is present in +// CUDA and custom-kernel HIP builds. Call the PFlash helper only from CUDA +// code; ordinary HIP builds use their independent scorer implementation. +bool custom_bf16_sparse_supported_on_current_device(); +bool local_pflash_supported_on_current_device(); + // Algorithmic parameters for the FlashPrefill selection + sparse forward. struct FlashPrefillConfig { int block_size = 128; // K stride; query block size = K block size @@ -92,7 +111,8 @@ int flash_prefill_forward_q8( // ── Unified dispatch ────────────────────────────────────────────────────────── // Picks the best available kernel at compile time + runtime buffer type: -// BF16 buffers + sm_80 build → flash_prefill_forward_bf16 +// BF16 buffers + qualified sm_80 build → flash_prefill_forward_bf16 +// BF16 buffers + unqualified device → -2 (caller chooses exact fallback) // F16 buffers + Volta build → flash_prefill_forward_f16 // otherwise → flash_prefill_forward_q8 (ggml FA fallback) // @@ -107,6 +127,7 @@ inline int flash_prefill_forward( { #if defined(DFLASH27B_HAVE_FLASHPREFILL) || defined(DFLASH27B_HAVE_SM80_FLASHPREFILL) if (qkv_type == GGML_TYPE_BF16) { + if (!custom_bf16_sparse_supported_on_current_device()) return -2; return flash_prefill_forward_bf16(Q, K, V, O, batch, seq_len, n_q_heads, n_k_heads, head_dim, scale, cfg); } diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 69b34281a..cdf7efc6c 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -23,6 +23,7 @@ #include "tool_hint.h" #include "common/sha1.h" #include "freeze_history.h" +#include "flashprefill.h" #ifdef DFLASH_HAS_CURL #include @@ -1072,7 +1073,7 @@ HttpServer::HttpServer(ModelBackend & backend, , tokenizer_(tokenizer) , config_(config) , chat_format_(ChatFormat::QWEN3) // default, overridden by arch - , prefix_cache_(config.prefix_cache_cap, tokenizer) + , prefix_cache_(config.prefix_cache_cap, tokenizer, config.arch) , disk_cache_({config.disk_cache_dir, config.disk_cache_budget_mb * (size_t)(1024 * 1024), config.disk_cache_min_tokens, @@ -2711,6 +2712,18 @@ HttpServer::PreparedPrompt HttpServer::prepare_prompt( if (config_.pflash_mode != ServerConfig::PflashMode::OFF && drafter_tokenizer_ != nullptr) { + // GB10's local custom sparse scorer kernels are not qualified. Check + // before loading or running the drafter so a stale/manual PFlash + // profile degrades to exact prefill instead of touching that path. + // A remote drafter runs on its own backend and remains independent. +#if defined(DFLASH27B_BACKEND_CUDA) + if (!config_.pflash_remote_drafter && + !flashprefill::local_pflash_supported_on_current_device()) { + std::fprintf(stderr, + "[pflash] local scorer unavailable on this GPU; using exact prefill\n"); + return prepared; + } +#endif const int prompt_tokens = (int) req.prompt_tokens.size(); const bool continuation = is_continuation_request(req.messages); const bool has_tools = req.tools.is_array() && !req.tools.empty(); @@ -2754,8 +2767,15 @@ HttpServer::PreparedPrompt HttpServer::prepare_prompt( } else if (strategy == PflashRequestStrategy::WholePrompt) { prepared.error = apply_pflash_compression(req, prepared); if (!prepared.error.empty()) { - prepared.error_status = 500; - return prepared; + // PFlash is optional acceleration. The backend restores its + // target residency before returning a compression failure, + // so an otherwise valid request can safely continue verbatim. + std::fprintf(stderr, + "[pflash] %s; falling back to exact prefill\n", + prepared.error.c_str()); + prepared.error.clear(); + prepared.tokens = req.prompt_tokens; + prepared.compressed = false; } } } diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index aa3c3aa6a..71bc57842 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -58,7 +58,7 @@ struct ServerConfig { int max_ctx = 0; // 0 = use backend's DevicePlacement default (8192) bool enable_cors = true; std::string model_name = "dflash"; - int prefix_cache_cap = 32; // prefix cache slots (0 disables) + int prefix_cache_cap = 8; // prefix cache slots (0 disables) int prefill_cache_cap = 0; // full-prompt/prefill cache slots (0 disables) // Thinking-budget v2. Applied when a request opts in via diff --git a/server/src/server/prefix_cache.cpp b/server/src/server/prefix_cache.cpp index 4ab20c590..381a13a34 100644 --- a/server/src/server/prefix_cache.cpp +++ b/server/src/server/prefix_cache.cpp @@ -12,7 +12,27 @@ namespace dflash::common { // ─── Chat marker resolution ──────────────────────────────────────────── -bool resolve_chat_markers(const Tokenizer & tok, ChatMarkers & out) { +bool resolve_chat_markers(const Tokenizer & tok, const std::string & arch, + ChatMarkers & out) { + // DeepSeek V4 uses DSML. Unlike ChatML, messages are separated by role + // starts rather than paired start/end delimiters: + // {system}{user}...... + // Resolve this before tokenizer-based probes so the architecture cannot + // be mistaken for another family whose marker strings are tokenizable. + if (arch == "deepseek4") { + auto bos = tok.encode("<｜begin▁of▁sentence｜>"); + auto user = tok.encode("<｜User｜>"); + auto assistant = tok.encode("<｜Assistant｜>"); + if (!bos.empty() && !user.empty() && !assistant.empty()) { + out.family = "deepseek4"; + out.sys_role_prefix = bos; + out.next_role_starts = {user, assistant}; + out.boundary_on_role_start = true; + return true; + } + return false; + } + // Try Qwen family: <|im_end|> and <|im_start|> should be single tokens. auto im_end = tok.encode("<|im_end|>"); auto im_start = tok.encode("<|im_start|>"); @@ -37,20 +57,23 @@ bool resolve_chat_markers(const Tokenizer & tok, ChatMarkers & out) { return true; } - // Try Laguna family: XML-style markers. - auto start_sys = tok.encode(""); - auto end_sys = tok.encode(""); - auto start_usr = tok.encode(""); - auto end_usr = tok.encode(""); - auto start_ast = tok.encode(""); - auto end_ast = tok.encode(""); - if (!start_sys.empty() && !end_sys.empty() && !start_usr.empty() && - !end_usr.empty() && !start_ast.empty() && !end_ast.empty()) { - out.family = "laguna"; - out.sys_role_prefix = start_sys; - out.end_msg_seqs = {end_sys, end_usr, end_ast}; - out.next_role_starts = {start_usr, start_ast, start_sys}; - return true; + // Laguna uses XML strings that every BPE tokenizer can encode, so only + // select this template when the loader identified the architecture. + if (arch == "laguna") { + auto start_sys = tok.encode(""); + auto end_sys = tok.encode(""); + auto start_usr = tok.encode(""); + auto end_usr = tok.encode(""); + auto start_ast = tok.encode(""); + auto end_ast = tok.encode(""); + if (!start_sys.empty() && !end_sys.empty() && !start_usr.empty() && + !end_usr.empty() && !start_ast.empty() && !end_ast.empty()) { + out.family = "laguna"; + out.sys_role_prefix = start_sys; + out.end_msg_seqs = {end_sys, end_usr, end_ast}; + out.next_role_starts = {start_usr, start_ast, start_sys}; + return true; + } } return false; @@ -98,6 +121,17 @@ std::vector find_all_boundaries(const std::vector & ids, if (sys_idx < 0) return out; int cursor = sys_idx + (int)markers.sys_role_prefix.size(); + if (markers.boundary_on_role_start) { + while (true) { + auto [role_idx, role_len] = + find_first_seq_any(ids, markers.next_role_starts, cursor); + if (role_idx < 0) break; + cursor = role_idx + role_len; + out.push_back(cursor); + } + return out; + } + while (true) { auto [end_idx, end_len] = find_first_seq_any(ids, markers.end_msg_seqs, cursor); if (end_idx < 0) break; @@ -183,7 +217,8 @@ int select_inline_snapshot_boundary(const std::vector & boundaries, // ─── PrefixCache ──────────────────────────────────────────────────────── -PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer) +PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer, + const std::string & arch) : cap_(std::min(cap, MAX_SLOTS)) { if (cap_ <= 0) { @@ -191,7 +226,7 @@ PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer) cap_ = 0; return; } - if (!resolve_chat_markers(tokenizer, markers_)) { + if (!resolve_chat_markers(tokenizer, arch, markers_)) { std::fprintf(stderr, "[pc] could not resolve chat markers; prefix cache disabled\n"); disabled_ = true; cap_ = 0; diff --git a/server/src/server/prefix_cache.h b/server/src/server/prefix_cache.h index 2a0515749..3000f7c36 100644 --- a/server/src/server/prefix_cache.h +++ b/server/src/server/prefix_cache.h @@ -27,15 +27,23 @@ namespace dflash::common { // ─── Chat marker detection ────────────────────────────────────────────── struct ChatMarkers { - std::string family; // "qwen", "gemma", or "laguna" + std::string family; // "qwen", "gemma", "laguna", or "deepseek4" // Token sequences for boundary detection std::vector sys_role_prefix; std::vector> end_msg_seqs; std::vector> next_role_starts; + // ChatML-style templates delimit a safe boundary with an end marker + // followed by the next role marker. DeepSeek's DSML template has no + // explicit system/user end marker, so each role start is itself a safe + // cut point after the marker has been consumed. + bool boundary_on_role_start = false; }; -// Resolve chat markers from the tokenizer (detects Qwen, Gemma, or Laguna family). -bool resolve_chat_markers(const Tokenizer & tok, ChatMarkers & out); +// Resolve chat markers from the tokenizer and detected model architecture. +// Architecture is required for templates whose delimiters are ordinary text: +// blindly probing them would classify every tokenizer as that family. +bool resolve_chat_markers(const Tokenizer & tok, const std::string & arch, + ChatMarkers & out); // Find all turn-boundary cut points in a token stream. std::vector find_all_boundaries(const std::vector & ids, @@ -81,7 +89,7 @@ class PrefixCache { static constexpr int MAX_SLOTS = 64; // cap = number of prefix-cache slots (0 disables). - PrefixCache(int cap, const Tokenizer & tokenizer); + PrefixCache(int cap, const Tokenizer & tokenizer, const std::string & arch); bool disabled() const { return disabled_; } diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 9cb7feafc..6cb5f0d83 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -105,7 +105,7 @@ static void print_usage(const char * prog) { " WARNING: >0 drops system prompt / tool definitions\n" " from attention at long contexts. Use 0 for tools.\n" " --model-name Model name for /v1/models (default: dflash)\n" - " --prefix-cache-slots Prefix cache slots (default: 32, 0 disables)\n" + " --prefix-cache-slots Prefix cache slots (default: 8; DeepSeek: 4; 0 disables)\n" " --prefill-cache-slots Full prompt/prefill cache slots (default: 0)\n" " --fast-rollback Enable speculative fast rollback (default: on)\n" " --no-fast-rollback Disable speculative fast rollback, even with --ddtree\n" @@ -222,6 +222,7 @@ int main(int argc, char ** argv) { std::string cache_type_v; // explicit --cache-type-v override bool target_device_seen = false; bool target_devices_seen = false; + bool prefix_cache_slots_set = false; bool fast_rollback_forced_off = false; bool target_split_fast_rollback_cli = false; bool adaptive_experts_set = false; // --adaptive-experts (MoE architectures only) @@ -353,6 +354,7 @@ int main(int argc, char ** argv) { sconfig.model_name = argv[++i]; } else if (std::strcmp(argv[i], "--prefix-cache-slots") == 0 && i + 1 < argc) { sconfig.prefix_cache_cap = std::atoi(argv[++i]); + prefix_cache_slots_set = true; } else if (std::strcmp(argv[i], "--prefill-cache-slots") == 0 && i + 1 < argc) { sconfig.prefill_cache_cap = std::atoi(argv[++i]); } else if (std::strcmp(argv[i], "--fast-rollback") == 0) { @@ -626,6 +628,9 @@ int main(int argc, char ** argv) { } const ResolvedBackendPlan & backend_plan = backend_preparation.plan; const std::string & arch = backend_plan.arch(); + if (!prefix_cache_slots_set) { + sconfig.prefix_cache_cap = arch_default_prefix_cache_slots(arch); + } if (target_split_fast_rollback_cli && arch != "qwen35") { std::fprintf(stderr, "[server] --target-split-fast-rollback is only supported for " @@ -1062,7 +1067,9 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "[server] │ pflash_drafter_exec= %s\n", sconfig.pflash_remote_drafter ? "remote-ipc" : "local"); std::fprintf(stderr, "[server] │ pflash_skip_park= %s\n", sconfig.pflash_skip_park ? "ON" : "off"); - std::fprintf(stderr, "[server] │ fp_use_bsa = %s\n", getenv("DFLASH_FP_USE_BSA") ? "ON" : "off"); + std::fprintf(stderr, "[server] │ fp_use_bsa = %s\n", + environment_variable_enabled("DFLASH_FP_USE_BSA") + ? "ON" : "off"); std::fprintf(stderr, "[server] │ fp_alpha = %s\n", getenv("DFLASH_FP_ALPHA") ? getenv("DFLASH_FP_ALPHA") : "0.12 (default)"); } std::fprintf(stderr, "[server] │ draft_residency = %s\n", diff --git a/server/test/test_entrypoint_cache_defaults.sh b/server/test/test_entrypoint_cache_defaults.sh index 6cd302e11..e680de02e 100755 --- a/server/test/test_entrypoint_cache_defaults.sh +++ b/server/test/test_entrypoint_cache_defaults.sh @@ -65,4 +65,32 @@ assert_arg_pair "$configured_output" --prefix-cache-slots 4 assert_arg_pair "$configured_output" --prefill-cache-slots 2 assert_arg_pair "$configured_output" --think-max-tokens 15488 +# Grace Hopper / GB10 drivers can report `[N/A]` for unified GPU memory. +# The container must leave the fallback at zero instead of evaluating that +# string as a Bash arithmetic expression before the host topology is applied. +FAKE_BIN="$TMP_DIR/bin" +mkdir -p "$FAKE_BIN" +cat >"$FAKE_BIN/nvidia-smi" <<'EOF' +#!/usr/bin/env bash +if [[ " $* " == *" -L "* ]]; then + printf 'GPU 0: NVIDIA GB10 (UUID: GPU-test)\n' +else + printf '[N/A]\n' +fi +EOF +chmod +x "$FAKE_BIN/nvidia-smi" + +gb10_stderr="$TMP_DIR/gb10.stderr" +env \ + PATH="$FAKE_BIN:$PATH" \ + DFLASH_DIR="$TMP_DIR" \ + DFLASH_TARGET="$TARGET" \ + DFLASH_DRAFT="$TMP_DIR/no-draft" \ + DFLASH_SERVER_BIN="$FAKE_SERVER" \ + bash "$ENTRYPOINT" serve >/dev/null 2>"$gb10_stderr" +if grep -Fq "syntax error: operand expected" "$gb10_stderr"; then + echo "entrypoint attempted arithmetic on GB10 [N/A] memory" >&2 + exit 1 +fi + echo "entrypoint cache defaults: PASS" diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index 63091244b..d0be44e98 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -447,6 +447,10 @@ static void test_model_capability_tables() { // deepseek4 is mixture-of-experts but has no hot/cold offload path. TEST_ASSERT(!arch_has_expert_offload("deepseek4")); + TEST_ASSERT(arch_default_prefix_cache_slots("qwen35") == 8); + TEST_ASSERT(arch_default_prefix_cache_slots("deepseek4") == 4); + TEST_ASSERT(arch_default_prefix_cache_slots("unknown") == 8); + // Every capability predicate must be false for an architecture the // factory cannot build, so no rule can admit an unbuildable model. TEST_ASSERT(!arch_supports_layer_split("qwen36")); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index c211e77d8..c5ec337b9 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -21,7 +21,9 @@ #include "common/sampler.h" #include "common/backend_precision.h" #include "common/backend_ipc.h" +#include "common/platform_env.h" #include "common/moe_hybrid_ffn_eval.h" +#include "flashprefill.h" #include "placement/pflash_placement.h" #include "common/io_utils.h" #include "placement/placement_config.h" @@ -34,6 +36,7 @@ #include "ggml-cpu.h" #include "server/prompt_normalize.h" #include "qwen3_drafter_model.h" +#include "deepseek4/deepseek4_backend.h" #include "dflash27b.h" #include "gguf.h" #include @@ -90,6 +93,30 @@ struct ServerUnitFixture {}; } \ } while (0) +TEST_CASE(ServerUnitFixture, test_feature_flag_and_sparse_hardware_policy) { + constexpr const char * flag = "DFLASH_TEST_FEATURE_FLAG"; + dflash_unsetenv(flag); + TEST_ASSERT(!environment_variable_enabled(flag)); + dflash_setenv(flag, "0"); + TEST_ASSERT(!environment_variable_enabled(flag)); + dflash_setenv(flag, "1"); + TEST_ASSERT(environment_variable_enabled(flag)); + dflash_unsetenv(flag); + + TEST_ASSERT(!flashprefill::custom_bf16_sparse_supports_cuda_sm(75)); + TEST_ASSERT(flashprefill::custom_bf16_sparse_supports_cuda_sm(80)); + TEST_ASSERT(flashprefill::custom_bf16_sparse_supports_cuda_sm(86)); + TEST_ASSERT(flashprefill::custom_bf16_sparse_supports_cuda_sm(120)); + TEST_ASSERT(!flashprefill::custom_bf16_sparse_supports_cuda_sm(121)); + TEST_ASSERT(flashprefill::local_pflash_supports_cuda_sm(75)); + TEST_ASSERT(flashprefill::local_pflash_supports_cuda_sm(120)); + TEST_ASSERT(!flashprefill::local_pflash_supports_cuda_sm(121)); + + TEST_ASSERT(deepseek4_dspark_supports_cuda_sm(90)); + TEST_ASSERT(deepseek4_dspark_supports_cuda_sm(120)); + TEST_ASSERT(!deepseek4_dspark_supports_cuda_sm(121)); +} + // ─── Helper: create an SseEmitter with minimal config ────────────────── static json weather_tools() { @@ -1787,6 +1814,31 @@ TEST_CASE(ServerUnitFixture, test_tool_schema_is_part_of_stable_system_boundary) hash_prefix(prompt_new_tools.data(), system_end)); } +TEST_CASE(ServerUnitFixture, test_deepseek_role_starts_are_safe_boundaries) { + // Synthetic DSML-shaped prompt: + // system u1 a1 u2 + // DeepSeek has no explicit system/user end delimiter. The role marker + // itself is the safe cut after all preceding content has been consumed. + ChatMarkers markers; + markers.family = "deepseek4"; + markers.sys_role_prefix = {10}; + markers.next_role_starts = {{20}, {30}}; + markers.boundary_on_role_start = true; + + const std::vector prompt = { + 10, 100, 20, 200, 30, 300, 40, 20, 400, 30, + }; + const auto bounds = find_all_boundaries(prompt, markers); + TEST_ASSERT(bounds == std::vector({3, 5, 8, 10})); + + // The snapshot excludes the current user text while preserving the full + // stable conversation prefix, matching the existing ChatML policy. + TEST_ASSERT(select_inline_snapshot_boundary(bounds) == 8); + + const std::vector missing_bos = {100, 20, 200, 30}; + TEST_ASSERT(find_all_boundaries(missing_bos, markers).empty()); +} + TEST_CASE(ServerUnitFixture, test_inline_snapshot_boundary_advances_past_restore) { const std::vector boundaries = {100, 240, 380, 520}; // Second-to-last is the boundary before the current user turn. @@ -4118,7 +4170,7 @@ TEST_CASE(ServerUnitFixture, test_sampler_needs_logit_processing) { TEST_CASE(ServerUnitFixture, test_server_config_cache_defaults) { ServerConfig cfg; - TEST_ASSERT(cfg.prefix_cache_cap == 32); + TEST_ASSERT(cfg.prefix_cache_cap == 8); TEST_ASSERT(cfg.prefill_cache_cap == 0); } @@ -4172,7 +4224,7 @@ TEST_CASE(ServerUnitFixture, test_props_model_card_wholesale_sidecar) { }; ServerConfig cfg = make_props_config_with_sidecar(sidecar); Tokenizer tok; - PrefixCache pc(0, tok); + PrefixCache pc(0, tok, cfg.arch); ToolMemory tm; json body = build_props_body(cfg, pc, tm); @@ -4205,7 +4257,7 @@ TEST_CASE(ServerUnitFixture, test_props_model_card_null_on_family_fallback) { cfg.hard_limit_reply_budget = 512; cfg.think_max_tokens = 32256; Tokenizer tok; - PrefixCache pc(0, tok); + PrefixCache pc(0, tok, cfg.arch); ToolMemory tm; json body = build_props_body(cfg, pc, tm); @@ -4241,7 +4293,7 @@ TEST_CASE(ServerUnitFixture, test_props_budget_envelope_shape) { cfg.effort_tiers.max = 500; Tokenizer tok; - PrefixCache pc(0, tok); + PrefixCache pc(0, tok, cfg.arch); ToolMemory tm; json body = build_props_body(cfg, pc, tm); @@ -4290,7 +4342,7 @@ TEST_CASE(ServerUnitFixture, test_props_runtime_shape) { cfg.draft_device = "auto:0"; Tokenizer tok; - PrefixCache pc(0, tok); + PrefixCache pc(0, tok, cfg.arch); ToolMemory tm; json body = build_props_body(cfg, pc, tm); From 65887ec883f52d14f47c72f31ed906544f500a3b Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 2 Aug 2026 03:26:01 +0200 Subject: [PATCH 23/25] fix(build): keep BSA flag initialization goto-safe --- server/src/flashprefill.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/server/src/flashprefill.cpp b/server/src/flashprefill.cpp index 8d1fc2e44..86c6cb9ed 100644 --- a/server/src/flashprefill.cpp +++ b/server/src/flashprefill.cpp @@ -377,9 +377,8 @@ int flash_prefill_forward_bf16( } // 4. sparse flash forward (BSA-or-WMMA) #ifdef DFLASH27B_HAVE_BSA - const bool bsa_requested = + static const bool use_bsa = environment_variable_enabled("DFLASH_FP_USE_BSA"); - const bool use_bsa = bsa_requested; if (use_bsa && D == 128 && BLOCK == 128) { launch_bsa_sparse_flash_forward_bf16( Q, K, V, O, dIdx, dCnt, scale, From 1b249425a07cb5b50efdfd57d7bdc9fe90170a0b Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:26:16 +0200 Subject: [PATCH 24/25] fix: address remaining Cubic P1/P2 review findings - flashprefill.h: unqualified BF16 devices (GB10 sm_121) now route through the exact ggml FA path instead of returning -2, which both graph callers treated as a hard scorer failure. - http_server.cpp: the local-PFlash unsupported-GPU fallback no longer returns early from prepare_prompt, so over-context requests still hit effective_prompt_overflows() and get the normal 400. - deepseek4_backend.cpp: qualify the DFLASH_DS4_DRAFT_GPU device before creating the DSpark draft backend; the startup gate only covered the target GPU. - calibration.py: include the image registry in the calibration fingerprint so changing LUCEBOX_IMAGE forces a fresh probe. - test_lucebox_sh.sh: cover the originally-stopped engine in the calibration lifecycle test; commit and rollback must stop the unit rather than restart it. --- lucebox/src/lucebox/calibration.py | 1 + scripts/test_lucebox_sh.sh | 45 ++++++++++++++++++++++ server/src/deepseek4/deepseek4_backend.cpp | 15 ++++++++ server/src/flashprefill.h | 10 ++++- server/src/server/http_server.cpp | 28 ++++++++------ 5 files changed, 85 insertions(+), 14 deletions(-) diff --git a/lucebox/src/lucebox/calibration.py b/lucebox/src/lucebox/calibration.py index 8d9e7fdd0..0c24de378 100644 --- a/lucebox/src/lucebox/calibration.py +++ b/lucebox/src/lucebox/calibration.py @@ -641,6 +641,7 @@ def _fingerprint_payload(cfg: Config) -> dict[str, Any]: return { "cli_version": __version__, "variant": cfg.variant, + "image": cfg.image, "model": asdict(cfg.model), "runtime": asdict(cfg.dflash), "placement": asdict(cfg.placement), diff --git a/scripts/test_lucebox_sh.sh b/scripts/test_lucebox_sh.sh index 8c678c1f1..09d01525f 100755 --- a/scripts/test_lucebox_sh.sh +++ b/scripts/test_lucebox_sh.sh @@ -801,6 +801,11 @@ STUB cat > "$shim_dir/systemctl" <<'STUB' #!/usr/bin/env bash printf '%s\n' "$*" >> "${LUCEBOX_TEST_SYSTEMCTL_LOG:?}" +if [ "${LUCEBOX_TEST_UNIT_INACTIVE:-0}" = "1" ]; then + for arg in "$@"; do + [ "$arg" = "is-active" ] && exit 3 + done +fi exit 0 STUB cat > "$shim_dir/flock" <<'STUB' @@ -811,12 +816,14 @@ STUB run_calibration_case() { local fail_baseline="$1" + local unit_inactive="${2:-0}" HOME="$sandbox/home" \ XDG_CONFIG_HOME="$sandbox/xdg" \ LUCEBOX_HOME="$state" \ LUCEBOX_TEST_CLI_LOG="$sandbox/cli.log" \ LUCEBOX_TEST_SYSTEMCTL_LOG="$sandbox/systemctl.log" \ LUCEBOX_TEST_FAIL_BASELINE="$fail_baseline" \ + LUCEBOX_TEST_UNIT_INACTIVE="$unit_inactive" \ PATH="$shim_dir:$PATH" \ NO_COLOR=1 \ bash -c ' @@ -871,6 +878,44 @@ STUB return fi + # An originally stopped engine must stay stopped: the commit path stops + # the unit instead of restarting it after the last calibration cell. + printf 'original\n' > "$state/config.toml" + : > "$sandbox/cli.log" + : > "$sandbox/systemctl.log" + rc=0 + out=$(run_calibration_case 0 1 2>&1) || rc=$? + restart_count=$(grep -cF -- '--user restart lucebox.service' "$sandbox/systemctl.log" || true) + local stop_count + stop_count=$(grep -cF -- '--user stop lucebox.service' "$sandbox/systemctl.log" || true) + if [ "$rc" -ne 0 ] || [ "$(<"$state/config.toml")" != "calibrated=22" ] \ + || [ "$restart_count" -ne 3 ] || [ "$stop_count" -ne 1 ] \ + || compgen -G "$state/.calibration-run.*" >/dev/null; then + report fail "$label" \ + "inactive success case rc=$rc restarts=$restart_count stops=$stop_count output=$(head -5 <<<"$out")" + rm -rf "$sandbox" + return + fi + + # And a rollback on an originally stopped engine must also stop the unit, + # not restart it. + printf 'original\n' > "$state/config.toml" + : > "$sandbox/cli.log" + : > "$sandbox/systemctl.log" + rc=0 + out=$(run_calibration_case 1 1 2>&1) || rc=$? + restart_count=$(grep -cF -- '--user restart lucebox.service' "$sandbox/systemctl.log" || true) + stop_count=$(grep -cF -- '--user stop lucebox.service' "$sandbox/systemctl.log" || true) + if [ "$rc" -eq 0 ] || [ "$(<"$state/config.toml")" != "original" ] \ + || [ "$restart_count" -ne 1 ] || [ "$stop_count" -ne 1 ] \ + || ! grep -qF 'original profile and service state were restored' <<<"$out" \ + || compgen -G "$state/.calibration-run.*" >/dev/null; then + report fail "$label" \ + "inactive rollback case rc=$rc restarts=$restart_count stops=$stop_count config=$(<"$state/config.toml") output=$(head -5 <<<"$out")" + rm -rf "$sandbox" + return + fi + report ok "$label" rm -rf "$sandbox" } diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index c8306e748..d99cfd9a5 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -699,6 +699,21 @@ bool DeepSeek4Backend::load_spec_drafter() { return false; } const PlacementBackend target_kind = placement_backend_of(backend_); + // Startup qualification covers only the target GPU. DFLASH_DS4_DRAFT_GPU + // can select a different CUDA device, so qualify the resolved draft + // device before creating a backend on it. + if (draft_kind == PlacementBackend::Cuda && + (draft_kind != target_kind || draft_gpu != cfg_.device.gpu)) { + int draft_sm = 0; + if (!dspark_supported_on_current_device(draft_gpu, draft_sm)) { + std::fprintf(stderr, + "[deepseek4] DSpark disabled: draft CUDA device %d " + "(sm_%d) is not qualified; continuing with " + "autoregressive decode\n", + draft_gpu, draft_sm); + return false; + } + } if (draft_kind != target_kind || draft_gpu != cfg_.device.gpu || separate_draft_stream) { std::string backend_error; diff --git a/server/src/flashprefill.h b/server/src/flashprefill.h index 6cc021563..8197e66a3 100644 --- a/server/src/flashprefill.h +++ b/server/src/flashprefill.h @@ -112,7 +112,7 @@ int flash_prefill_forward_q8( // ── Unified dispatch ────────────────────────────────────────────────────────── // Picks the best available kernel at compile time + runtime buffer type: // BF16 buffers + qualified sm_80 build → flash_prefill_forward_bf16 -// BF16 buffers + unqualified device → -2 (caller chooses exact fallback) +// BF16 buffers + unqualified device → flash_prefill_forward_q8 (exact FA) // F16 buffers + Volta build → flash_prefill_forward_f16 // otherwise → flash_prefill_forward_q8 (ggml FA fallback) // @@ -127,7 +127,13 @@ inline int flash_prefill_forward( { #if defined(DFLASH27B_HAVE_FLASHPREFILL) || defined(DFLASH27B_HAVE_SM80_FLASHPREFILL) if (qkv_type == GGML_TYPE_BF16) { - if (!custom_bf16_sparse_supported_on_current_device()) return -2; + if (!custom_bf16_sparse_supported_on_current_device()) { + // Unqualified device (e.g. GB10 sm_121): run the exact ggml FA + // path instead of failing the scorer. + return flash_prefill_forward_q8(backend, Q, K, V, O, + batch, seq_len, n_q_heads, n_k_heads, head_dim, scale, + qkv_type, cfg); + } return flash_prefill_forward_bf16(Q, K, V, O, batch, seq_len, n_q_heads, n_k_heads, head_dim, scale, cfg); } diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 96baafdfd..5895b3ce1 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -2807,20 +2807,24 @@ HttpServer::PreparedPrompt HttpServer::prepare_prompt( PreparedPrompt prepared; prepared.tokens = req.prompt_tokens; - if (config_.pflash_mode != ServerConfig::PflashMode::OFF && - drafter_tokenizer_ != nullptr) { - // GB10's local custom sparse scorer kernels are not qualified. Check - // before loading or running the drafter so a stale/manual PFlash - // profile degrades to exact prefill instead of touching that path. - // A remote drafter runs on its own backend and remains independent. + bool pflash_available = + config_.pflash_mode != ServerConfig::PflashMode::OFF && + drafter_tokenizer_ != nullptr; #if defined(DFLASH27B_BACKEND_CUDA) - if (!config_.pflash_remote_drafter && - !flashprefill::local_pflash_supported_on_current_device()) { - std::fprintf(stderr, - "[pflash] local scorer unavailable on this GPU; using exact prefill\n"); - return prepared; - } + // GB10's local custom sparse scorer kernels are not qualified. Check + // before loading or running the drafter so a stale/manual PFlash + // profile degrades to exact prefill instead of touching that path. + // A remote drafter runs on its own backend and remains independent. + // Fall through to the context check below so an over-length prompt + // still gets the normal 400 instead of reaching exact generation. + if (pflash_available && !config_.pflash_remote_drafter && + !flashprefill::local_pflash_supported_on_current_device()) { + std::fprintf(stderr, + "[pflash] local scorer unavailable on this GPU; using exact prefill\n"); + pflash_available = false; + } #endif + if (pflash_available) { const int prompt_tokens = (int) req.prompt_tokens.size(); const bool continuation = is_continuation_request(req.messages); const bool has_tools = req.tools.is_array() && !req.tools.empty(); From add6cb571165434f21f9b87ccf7953b437478a91 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:10:57 +0200 Subject: [PATCH 25/25] test(harness): non-empty GGUF stubs for timeout regressions The client launchers now reject zero-byte GGUF placeholders, so the model-free timeout tests must stub the target/draft with one byte of content instead of touch. --- harness/tests/test_client_launcher_timeouts.sh | 5 ++++- harness/tests/test_run_pi_timeout.sh | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/harness/tests/test_client_launcher_timeouts.sh b/harness/tests/test_client_launcher_timeouts.sh index 712cfcf55..37cae3ce4 100644 --- a/harness/tests/test_client_launcher_timeouts.sh +++ b/harness/tests/test_client_launcher_timeouts.sh @@ -16,7 +16,10 @@ FAKE_SERVER="$TMP_DIR/dflash_server" FAKE_CLIENT="$TMP_DIR/client" mkdir -p "$FAKE_BIN" mkdir -p "$FAKE_REPO" -touch "$FAKE_TARGET" "$FAKE_DRAFT" +# The launchers reject zero-byte GGUF placeholders, so give the +# model-free stubs one byte of content. +printf stub > "$FAKE_TARGET" +printf stub > "$FAKE_DRAFT" cat >"$FAKE_SERVER" <<'EOF' #!/usr/bin/env bash diff --git a/harness/tests/test_run_pi_timeout.sh b/harness/tests/test_run_pi_timeout.sh index fb38082d1..c3873bfd7 100755 --- a/harness/tests/test_run_pi_timeout.sh +++ b/harness/tests/test_run_pi_timeout.sh @@ -12,7 +12,10 @@ FAKE_DRAFT="$TMP_DIR/draft.gguf" FAKE_SERVER="$TMP_DIR/dflash_server" FAKE_PI="$TMP_DIR/pi" mkdir -p "$FAKE_BIN" -touch "$FAKE_TARGET" "$FAKE_DRAFT" +# The launchers reject zero-byte GGUF placeholders, so give the +# model-free stubs one byte of content. +printf stub > "$FAKE_TARGET" +printf stub > "$FAKE_DRAFT" cat >"$FAKE_SERVER" <<'EOF' #!/usr/bin/env bash