compose.yml

services:
  train-roberta-sand-model:
    environment:
      - HF_TOKEN=${HF_TOKEN}
    build:
      context: .
      dockerfile: Dockerfile
    volumes:
      - ./:/workspace
    working_dir: /workspace
    command: "python -u -m src.model.train_model --roberta --checkpoint dabagyan/roberta-sarcasm-model --sand --push"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0", "1"]
              capabilities: [gpu]
  train-roberta-model:
    environment:
      - HF_TOKEN=${HF_TOKEN}
    build:
      context: .
      dockerfile: Dockerfile
    volumes:
      - ./:/workspace
    working_dir: /workspace
    command: "python -u -m src.model.train_model --roberta --push"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
  train-roberta-context-model:
    environment:
      - HF_TOKEN=${HF_TOKEN}
    build:
      context: .
      dockerfile: Dockerfile
    volumes:
      - ./:/workspace
    working_dir: /workspace
    command: "python -u -m src.model.train_model --roberta --context --push"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["1"]
              capabilities: [gpu]
  train-context-model:
    build:
      context: .
      dockerfile: Dockerfile
    volumes:
      - ./:/workspace
    working_dir: /workspace
    command: "python -u -m src.model.train_model --context"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]

  train-base-model:
    build:
      context: .
      dockerfile: Dockerfile
    volumes:
      - ./:/workspace
    working_dir: /workspace
    command: "python -u -m src.model.train_model"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["1"]
              capabilities: [gpu]

  test-cuda:
    image: nvcr.io/nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
    command: nvidia-smi
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]