version: '3.8'

services:
  train:
    build:
      context: .
      dockerfile: Dockerfile.train
    ipc: private 
    shm_size: '2g'
    ulimits:
      memlock: -1
    deploy:
      resources:
        reservations:
          devices:
            - capabilities: [gpu]
    network_mode: host
    environment:
      - PATH=/app/venv/bin/:$PATH
      - WANDB_API_KEY=de8d037779af70199106db0710f417f5157c1818
      - RAG_SERVER_URL=http://localhost:2223
    volumes:
      - ./train.sh:/app/train.sh
      - ./outputs:/app/outputs/
      - ./wandb:/app/wandb/
      - ~/.cache:/root/.cache/
    command: bash train.sh
    # restart: unless-stopped

  inference:
    build:
      context: .
      dockerfile: Dockerfile.inference
    # ports:
    #   - "8000:8000"
    network_mode: host
    ipc: private 
    shm_size: '2g'
    ulimits:
      memlock: -1
    deploy:
      resources:
        reservations:
          devices:
            - capabilities: [gpu]
    environment:
      - PATH=/app/venv/bin/:$PATH
    #   - THREADS=${INFERENCE_THREADS:-4}
    volumes:
      - ./serve.sh:/app/serve.sh
      - ~/.cache:/root/.cache/
    command: bash serve.sh
    restart: always