version: "3.8"

services:
  # ── SFT shard encoding (current workload) ──────────────────────────
  sft-worker:
    build:
      context: .
      dockerfile: Dockerfile.sft
    image: bharathkumar192/codecbench-sft:latest
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
      - VAST_OFFER_ID=${VAST_OFFER_ID:-local}
    env_file:
      - .env
    volumes:
      - /tmp/pipeline:/tmp/pipeline
    command: >
      sft-run
      --batch-size=4
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    restart: unless-stopped

  sft-test:
    build:
      context: .
      dockerfile: Dockerfile.sft
    image: bharathkumar192/codecbench-sft:latest
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
      - VAST_OFFER_ID=${VAST_OFFER_ID:-test}
    env_file:
      - .env
    volumes:
      - /tmp/pipeline:/tmp/pipeline
    command: >
      sft-run
      --max-shards=${TEST_SHARDS:-1}
      --batch-size=4
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    profiles: ["test"]

  # ── Legacy pretraining pipeline ────────────────────────────────────
  pipeline-worker:
    build: .
    image: bharathkumar192/codecbench-pipeline:latest
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
      - VAST_OFFER_ID=${VAST_OFFER_ID:-local}
    env_file:
      - .env
    volumes:
      - /tmp/pipeline:/tmp/pipeline
    command: >
      run
      --shard-count=50
      --prefetch=2
      --tmp-dir=/tmp/pipeline
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    restart: unless-stopped
    profiles: ["pretrain"]

  pipeline-test:
    build: .
    image: bharathkumar192/codecbench-pipeline:latest
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
      - VAST_OFFER_ID=${VAST_OFFER_ID:-test}
    env_file:
      - .env
    volumes:
      - /tmp/pipeline:/tmp/pipeline
    command: >
      run
      --max-videos=${TEST_VIDEOS:-3}
      --shard-count=3
      --prefetch=2
      --tmp-dir=/tmp/pipeline
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    profiles: ["pretrain-test"]
