services:
  # Main training service
  reasoning-models:
    build:
      context: .
      dockerfile: src/training_pipeline/Dockerfile
    container_name: reasoning-models-training
    
    # GPU support using deploy syntax (working with system Docker daemon)
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    
    volumes:
      # Mount the entire project directory
      - .:/workspace
      # Mount a persistent volume for outputs
      - ./output:/workspace/output
      # Mount text cache
      - ./text_cache:/workspace/text_cache
    
    working_dir: /workspace
    
    # Run the improved training script with config selection
    command: >
      bash -c "
        if [ -n \"$TRAINING_CONFIG\" ]; then
          CONFIG_FILE=\"src/training_pipeline/model_config_${TRAINING_CONFIG}.yml\"
          if [ -f \"$CONFIG_FILE\" ]; then
            echo \"Using configuration: $CONFIG_FILE\"
            python src/training_pipeline/train_model_next_gen.py --config \"$CONFIG_FILE\"
          else
            echo \"Configuration file not found: $CONFIG_FILE\"
            echo \"Available configs: high_performance, low_resource, research, production\"
            exit 1
          fi
        else
          echo \"No TRAINING_CONFIG specified, using default config\"
          python src/training_pipeline/train_model_next_gen.py
        fi
      "
    
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
      - NVIDIA_DRIVER_CAPABILITIES=compute,utility
      - PYTHONPATH=/workspace/src
      - CUDA_VISIBLE_DEVICES=0
      # Configuration override (optional)
      # - TRAINING_CONFIG=high_performance  # Options: high_performance, memory_efficient, high_quality, quick_test
    
    # Restart policy - only restart if there's a failure
    restart: "no"