# Nemotron on a single DGX Spark (GB10, 128 GB unified memory)
#
# Run ONE model at a time — at the recipe memory settings they don't co-reside:
#   docker compose --profile pull run --rm pull      # stage weights once (~85 GB)
#   docker compose --profile super     up -d         # Nemotron 3 Super 120B-A12B  -> :8000
#   docker compose --profile lightning up -d         # Nemotron 3.5 Lightning 30B-A3B -> :8001
#   docker compose --profile super     down          # before switching
#
# Prereqs:
#   - NVIDIA driver 580.x  (590.x has a CUDA-graph capture deadlock on GB10)
#   - nvidia-container-toolkit, Docker Compose v2
#   - HF_TOKEN in a .env file next to this compose (weights are gated)
#
# Sources (flag sets copied verbatim, only host/port/name added):
#   Super:     recipes.vllm.ai/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16  "DGX Spark (GB10)" (2026-07-31)
#   Lightning: huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4  "1x DGX Spark (GB10)"

x-vllm-common: &vllm-common
  restart: unless-stopped
  ipc: host
  shm_size: "32g"
  ulimits:
    memlock: -1
    stack: 67108864
  volumes:
    - ./hf-cache:/root/.cache/huggingface       # download once, mount everywhere
  environment:
    HF_TOKEN: ${HF_TOKEN:-}
    VLLM_FLOAT32_MATMUL_PRECISION: "high"      # keeps the Mamba scan numerically stable
    VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
  deploy:
    resources:
      reservations:
        devices:
          - driver: nvidia
            count: all
            capabilities: [gpu]
  healthcheck:
    test: ["CMD-SHELL", "curl -fsS http://localhost:8000/health || exit 1"]
    interval: 30s
    timeout: 10s
    retries: 40
    start_period: 900s                          # Super takes several minutes to load + capture graphs

services:

  # ───────────── Nemotron 3 Super 120B-A12B (NVFP4) — batch / long-context ─────────────
  # Expect ~23 tok/s single-stream decode after warm-up. Marlin/unfused MoE fallback
  # messages at load are normal on GB10. 256K context is the single-node ceiling;
  # 1M needs a second Spark (TP=2 over the QSFP link).
  super:
    <<: *vllm-common
    profiles: ["super"]
    image: vllm/vllm-openai:v0.28.0-ubuntu2404
    container_name: nemotron-super
    ports:
      - "8000:8000"
    command:
      - nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4
      - --served-model-name
      - nemotron-super
      - --host
      - 0.0.0.0
      - --port
      - "8000"
      - --max-model-len
      - "262144"
      - --gpu-memory-utilization
      - "0.8"
      - --max-num-seqs
      - "8"
      - --max-num-batched-tokens
      - "16384"
      - --load-format
      - fastsafetensors
      - --attention-backend
      - flashinfer
      - --kv-cache-dtype
      - fp8
      - --mamba-cache-mode
      - align
      - --mamba-ssm-cache-dtype
      - float32
      - --speculative-config
      - '{"method":"mtp","num_speculative_tokens":3}'   # Super uses built-in MTP, NOT DSpark
      - --enable-prefix-caching
      - --reasoning-parser
      - nemotron_v3
      - --enable-auto-tool-choice
      - --tool-call-parser
      - qwen3_xml

  # ───────────── Nemotron 3.5 Lightning 30B-A3B (NVFP4) + DSpark — interactive ─────────────
  # ~17 GB weights + ~1 GB drafter. Marlin is the intended MoE path on GB10 (no native
  # FP4 tensor-core path). 1M context is validated by default on a single Spark.
  lightning:
    <<: *vllm-common
    profiles: ["lightning"]
    image: vllm/vllm-openai:v0.27.1
    container_name: nemotron-lightning
    ports:
      - "8001:8000"
    command:
      - nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4
      - --served-model-name
      - nemotron-lightning
      - --host
      - 0.0.0.0
      - --port
      - "8000"
      - --moe-backend
      - marlin
      - --kv-cache-dtype
      - fp8
      - --enable-prefix-caching
      - --gpu-memory-utilization
      - "0.85"
      - --speculative_config.model
      - nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DSpark
      - --speculative_config.num_speculative_tokens
      - "3"
      - --mamba-backend
      - flashinfer
      - --mamba-cache-mode
      - align
      - --reasoning-parser
      - nemotron_v3
      - --tool-call-parser
      - qwen3_coder
      - --enable-auto-tool-choice

  # ───────────── One-shot weight pre-stage ─────────────
  pull:
    profiles: ["pull"]
    image: vllm/vllm-openai:v0.28.0-ubuntu2404
    entrypoint: ["/bin/sh", "-c"]
    environment:
      HF_TOKEN: ${HF_TOKEN:-}
    volumes:
      - ./hf-cache:/root/.cache/huggingface
    command:
      - >
        hf download nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 &&
        hf download nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4 &&
        hf download nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DSpark

# ── Running both at once (optional, not recipe-validated) ──
# Unified memory is shared, so co-residency means shrinking both:
#   super:     --gpu-memory-utilization 0.60  --max-model-len 32768
#   lightning: --gpu-memory-utilization 0.25
# That leaves Super ~8 GB of KV — fine for short batch jobs, not for long charts.
# Drop the `profiles:` keys on both services to launch them together.