version: "3.9"

services:
  pvmd-training:
    build:
      context: ..
      dockerfile: docker/Dockerfile
    image: pvmd-training:latest
    container_name: pvmd-training
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
      - CUDA_VISIBLE_DEVICES=0
      - PYSPARK_PYTHON=/opt/conda/envs/pvmd/bin/python
      - PYSPARK_DRIVER_PYTHON=/opt/conda/envs/pvmd/bin/python
      - HF_HOME=/workspace/data/cache/huggingface
      - TRANSFORMERS_CACHE=/workspace/data/cache/huggingface
      - HF_DATASETS_CACHE=/workspace/data/cache/datasets
    volumes:
      - /mnt/parakeet-training:/workspace/data
      - /root/proiect-PVMD:/workspace
      - /tmp:/tmp
    working_dir: /workspace
    stdin_open: true
    tty: true
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    command: ["/bin/bash"]