# PII masker on vLLM. Run from this `serving/` directory after cloning the repo —
# the model files in the repo root are mounted as the model dir.
#
#   docker compose up -d          # vllm, then the /v1/scoring adapter
#   curl -s localhost:8002/health
#
# Ports: 8002 = /v1/scoring adapter (batch scoring contract)
#        8003 = vLLM /pooling (raw 38-column logits)
#
# Verified on an NVIDIA DGX Spark (GB10, aarch64) with vLLM 0.26.0.

services:
  vllm:
    image: ${VLLM_IMAGE:-vllm/vllm-openai:v0.26.0}
    container_name: pii-masker-vllm
    ipc: host
    network_mode: host
    volumes:
      - ..:/models/pii-masker:ro
    command:
      - /models/pii-masker
      - --served-model-name=pii-masker
      - --host=127.0.0.1
      - --port=8003
      - --runner=pooling
      - --convert=classify
      # Raw logits, no softmax — required: the client runs Viterbi on raw logits,
      # and the folded sensitivity column (37) must stay exact.
      - '--pooler-config={"use_activation": false}'
      - --dtype=bfloat16
      # Training max_seq_len (the backbone itself supports 32768).
      - --max-model-len=4096
      - --max-num-seqs=16
      # 0.6B model, ~1.8 GB in practice. Raise if this is the only tenant.
      - --gpu-memory-utilization=0.05
      # Safe default on GB10/sm_121 for encoder-only attention; remove to try
      # CUDA graphs.
      - --enforce-eager
    gpus: all
    healthcheck:
      test: ["CMD", "python3", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8003/health')"]
      interval: 15s
      timeout: 5s
      retries: 20
      start_period: 120s
    restart: unless-stopped

  scoring:
    image: ${VLLM_IMAGE:-vllm/vllm-openai:v0.26.0}
    container_name: pii-masker-scoring
    network_mode: host
    entrypoint: []
    command: ["python3", "/adapter/scoring_adapter.py"]
    volumes:
      - ..:/models/pii-masker:ro
      - ./scoring_adapter.py:/adapter/scoring_adapter.py:ro
    environment:
      VLLM_URL: http://127.0.0.1:8003
      SERVED_MODEL_NAME: pii-masker
      TOKENIZER_PATH: /models/pii-masker/tokenizer.json
      # Neither service has auth; bind loopback unless you add some.
      HOST: 127.0.0.1
      PORT: "8002"
    depends_on:
      vllm:
        condition: service_healthy
    healthcheck:
      test: ["CMD", "python3", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8002/health')"]
      interval: 15s
      timeout: 5s
      retries: 5
      start_period: 30s
    restart: unless-stopped