name: "${COMPOSE_PROJECT_NAME:-dgx-spark-gguf-template}"

x-common-env: &common-env
  HF_TOKEN: "${HF_TOKEN:-}"
  NVIDIA_VISIBLE_DEVICES: all
  NVIDIA_DRIVER_CAPABILITIES: compute,utility
  HOST: "${HOST:-0.0.0.0}"
  LLAMACPP_PORT: "${LLAMACPP_PORT:-8000}"
  MAX_MODEL_LEN: "${MAX_MODEL_LEN:-32768}"
  LLAMA_THREADS: "${LLAMA_THREADS:-20}"
  BATCH_SIZE: "${BATCH_SIZE:-2048}"
  UBATCH_SIZE: "${UBATCH_SIZE:-1024}"
  CACHE_TYPE_K: "${CACHE_TYPE_K:-q8_0}"
  CACHE_TYPE_V: "${CACHE_TYPE_V:-q8_0}"
  GPU_LAYERS: "${GPU_LAYERS:-99}"
  FLASH_ATTN: "${FLASH_ATTN:-on}"
  MLOCK: "${MLOCK:-0}"
  ENABLE_MMPROJ: "${ENABLE_MMPROJ:-0}"
  AUTO_DOWNLOAD: "${AUTO_DOWNLOAD:-1}"
  GGUF_INCLUDE: "${GGUF_INCLUDE:-}"
  GGUF_REQUIRED_FILES: "${GGUF_REQUIRED_FILES:-}"
  EXTRA_ARGS_PROFILE: "${EXTRA_ARGS_PROFILE:-}"
  LLAMACPP_EXTRA_ARGS: "${LLAMACPP_EXTRA_ARGS:-}"

x-common: &common
  build:
    context: .
    dockerfile: Dockerfile
    args:
      LLAMACPP_REPO: "${LLAMACPP_REPO:-https://github.com/stepfun-ai/llama.cpp.git}"
      LLAMACPP_REF: "${LLAMACPP_REF:-step3.7}"
      CMAKE_CUDA_ARCHITECTURES: "${CMAKE_CUDA_ARCHITECTURES:-121}"
      BUILD_JOBS: "${BUILD_JOBS:-2}"
  image: dgx-spark-gguf-template/llamacpp:gb10
  network_mode: host
  ipc: host
  shm_size: "32gb"
  volumes:
    - ./models:/models
    - ./logs:/logs
    - ./models/hf:/root/.cache/huggingface
  healthcheck:
    test: ["CMD-SHELL", "curl -fsS http://127.0.0.1:${LLAMACPP_PORT:-8000}/v1/models >/dev/null"]
    interval: 30s
    timeout: 5s
    retries: 20
    start_period: 5m
  deploy:
    resources:
      reservations:
        devices:
          - driver: nvidia
            count: all
            capabilities: [gpu]

services:
  step37-flash-q4-k-s:
    <<: *common
    container_name: step37-flash-q4-k-s
    restart: "no"
    environment:
      <<: *common-env
      PROFILE_NAME: step37-flash-q4-k-s
      GGUF_REPO: stepfun-ai/Step-3.7-Flash-GGUF
      GGUF_FILE: Q4_K_S/Step-3.7-flash-Q4_K_S-00001-of-00003.gguf
      GGUF_INCLUDE: Q4_K_S/*.gguf
      GGUF_REQUIRED_FILES: >-
        Q4_K_S/Step-3.7-flash-Q4_K_S-00001-of-00003.gguf
        Q4_K_S/Step-3.7-flash-Q4_K_S-00002-of-00003.gguf
        Q4_K_S/Step-3.7-flash-Q4_K_S-00003-of-00003.gguf
      MMPROJ_FILE: mmproj-step3.7-flash-f16.gguf
      SERVED_MODEL_NAME: step37-flash-q4-k-s