# Training + vLLM 0.26.0 rollout server: two containers on one host, HTTP for rollouts and a shared
# NCCL group for weight sync. `docker compose -f docker-compose.vllm.yml up [vllm-server]` builds
# `vllm-server:0.26.0` from Dockerfile.vllm when the tag is absent; the trainer image comes from
# `make build-blackwell` (TRAIN_IMAGE=halo:hopper on Hopper). The sync needs network_mode + ipc host
# and all GPUs visible, VLLM_ and TRAINER_CUDA_DEVICES disjoint; EFA layers docker-compose.vllm.efa.yml.

services:
  vllm-server:
    build:
      context: .
      dockerfile: Dockerfile.vllm
      args:
        VLLM_VERSION: "0.26.0"
    image: vllm-server:0.26.0
    ipc: host
    network_mode: host
    environment:
      - CUDA_VISIBLE_DEVICES=${VLLM_CUDA_DEVICES:-7}
      - VLLM_SERVER_DEV_MODE=1
      # Unset by default. Thinking budgets need VLLM_USE_V2_MODEL_RUNNER=0 (V2 answers them with
      # 400). 0 or 1 only — vLLM parses it with int(), so an empty value kills the server at startup.
      - VLLM_USE_V2_MODEL_RUNNER
      # Off for RL: strict mode constrains every tool-bearing request with a structural-tag grammar,
      # so the served distribution is no longer the policy's, and it costs the engine core CPU per
      # step. The tool parser still parses the calls.
      - VLLM_ENFORCE_STRICT_TOOL_CALLING=${VLLM_ENFORCE_STRICT_TOOL_CALLING:-0}
      # Directory of tuned Triton MoE tile configs, visible inside the container; unset stays unset.
      - VLLM_TUNED_CONFIG_FOLDER
      - NCCL_DEBUG=WARN
      # The no-fabric recipe (InfiniBand off, socket net), spelled the same on every service; the
      # EFA overlay overrides both.
      - NCCL_IB_DISABLE=${NCCL_IB_DISABLE:-1}
      - NCCL_NET=${NCCL_NET:-Socket}
      - NCCL_P2P_LEVEL=${NCCL_P2P_LEVEL:-NVL}
      # Passed through bare from the shell or .env (unset stays unset): a protocol table set on the
      # trainer alone hangs the first collective, so a trainer that pins NCCL_PROTO needs the server
      # pinned the same way, on every recipe.
      - NCCL_PROTO
      # Keep the socket transport off Docker's bridge and per-container veth pairs.
      - NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-^docker,veth}
      - HF_TOKEN=${HF_TOKEN:-}
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    # `vllm serve` args. Three are load-bearing and silent when wrong: --moe-backend triton (the auto
    # Blackwell backends repack expert weights, so a synced write corrupts the experts),
    # --return-tokens-as-token-ids (train_on_sampled_tokens recovers the sampled ids) and
    # --logprobs-mode processed_logprobs (the IS ratio divides by them; raw_logprobs are
    # pre-temperature). Parsers, prefix caching, chat template: agent-docs/infrastructure/rollout-servers.md
    command: >
      ${VLLM_MODEL:-Qwen/Qwen3-0.6B}
      --port ${VLLM_PORT:-8000}
      --tensor-parallel-size ${VLLM_TP:-1}
      --gpu-memory-utilization ${VLLM_GPU_MEM:-0.85}
      --dtype auto
      --weight-transfer-config '{"backend": "nccl"}'
      --moe-backend ${VLLM_MOE_BACKEND:-triton}
      ${VLLM_ATTENTION_BACKEND:+--attention-backend ${VLLM_ATTENTION_BACKEND}}
      ${VLLM_ENABLE_R3:+--enable-return-routed-experts}
      --return-tokens-as-token-ids
      --logprobs-mode processed_logprobs
      ${VLLM_PREFIX_CACHING_FLAG:---enable-prefix-caching}
      ${VLLM_SPECULATIVE_CONFIG:+--speculative-config '${VLLM_SPECULATIVE_CONFIG}'}
      --enable-auto-tool-choice
      ${VLLM_TOOL_PARSER_PLUGIN:+--tool-parser-plugin ${VLLM_TOOL_PARSER_PLUGIN}}
      --tool-call-parser ${VLLM_TOOL_PARSER:-hermes}
      ${VLLM_CHAT_TEMPLATE:+--chat-template ${VLLM_CHAT_TEMPLATE}}
      ${VLLM_REASONING_PARSER_PLUGIN:+--reasoning-parser-plugin ${VLLM_REASONING_PARSER_PLUGIN}}
      ${VLLM_REASONING_PARSER:+--reasoning-parser ${VLLM_REASONING_PARSER}}
    healthcheck:
      # python3 because the base image does not provide a bare `python`.
      test: ["CMD-SHELL", "python3 -c \"import urllib.request; urllib.request.urlopen('http://localhost:${VLLM_PORT:-8000}/health')\""]
      interval: 15s
      timeout: 10s
      retries: 40
      start_period: 120s
    restart: unless-stopped
    volumes:
      - ${HF_HOME:-~/.cache/huggingface}:/root/.cache/huggingface

  training:
    image: ${TRAIN_IMAGE:-halo:blackwell}
    ipc: host
    network_mode: host
    # The RL environments' external API keys live in the repo-root .env, which compose otherwise
    # reads for ${VAR} interpolation only.
    env_file:
      - path: .env
        required: false
    environment:
      - CUDA_VISIBLE_DEVICES=${TRAINER_CUDA_DEVICES:-0,1,2,3,4,5,6}
      # GPU test harness only; production trainers take the server URL from the YAML config.
      - VLLM_SERVER_URL=http://localhost:${VLLM_PORT:-8000}
      - NCCL_DEBUG=WARN
      # Trainer half of the no-fabric recipe. NCCL_NET is process-global: left at Socket on an EFA
      # host it forces every trainer collective onto TCP and breaks DeepEP, so the EFA overlay
      # overrides both.
      - NCCL_IB_DISABLE=${NCCL_IB_DISABLE:-1}
      - NCCL_NET=${NCCL_NET:-Socket}
      - NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-^docker,veth}
      - NCCL_PROTO
      - WANDB_API_KEY=${WANDB_API_KEY:-}
      - HF_TOKEN=${HF_TOKEN:-}
      # Keeps the Triton and FA4 JIT caches on the mounted volume instead of recompiling every run.
      - HF_HOME=/root/.cache/huggingface
      - HF_DATASETS_CACHE=/root/.cache/huggingface/datasets
      # Toolkit scratch on the volume the Makefile targets use, so the S3 dataset cache and profiler
      # traces are shared with `make test-gpu-*`.
      - HALO_DATA_ROOT=${HALO_SCRATCH:-/mnt}
      - TMPDIR=${HALO_SCRATCH:-/mnt}/tmp
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    depends_on:
      vllm-server:
        condition: service_healthy
    volumes:
      - .:/workspace
      - ${HF_HOME:-~/.cache/huggingface}:/root/.cache/huggingface
      # Same path inside and out, so HALO_DATA_ROOT resolves either side of the boundary.
      - ${HALO_SCRATCH:-/mnt}:${HALO_SCRATCH:-/mnt}
      # Optional S3 dataset sources — bring your own credentials (read-write so refreshes persist).
      - ~/.aws:/root/.aws
    working_dir: /workspace
    command: >
      bash -c "echo 'vLLM server ready at http://localhost:${VLLM_PORT:-8000}' && bash"
    stdin_open: true
    tty: true