x-default: &default
  restart: unless-stopped
  volumes:
    - &localtime /etc/localtime:/etc/localtime:ro
    - &timezone /etc/timezone:/etc/timezone:ro
  logging:
    driver: json-file
    options:
      max-size: 100m

services:
  pytorch:
    <<: *default
    image: pytorch/pytorch:${PYTORCH_VERSION:-2.6.0-cuda12.6-cudnn9-runtime}
    container_name: pytorch
    ports:
      - "${JUPYTER_PORT_OVERRIDE:-8888}:8888"
      - "${TENSORBOARD_PORT_OVERRIDE:-6006}:6006"
    environment:
      NVIDIA_VISIBLE_DEVICES: ${NVIDIA_VISIBLE_DEVICES:-all}
      NVIDIA_DRIVER_CAPABILITIES: ${NVIDIA_DRIVER_CAPABILITIES:-compute,utility}
      JUPYTER_ENABLE_LAB: ${JUPYTER_ENABLE_LAB:-yes}
    command: >
      bash -c "pip install --no-cache-dir jupyter tensorboard &&
               jupyter lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root
               --NotebookApp.token='${JUPYTER_TOKEN:-pytorch}'"
    volumes:
      - *localtime
      - *timezone
      - pytorch_notebooks:/workspace
      - pytorch_data:/data
    working_dir: /workspace
    deploy:
      resources:
        limits:
          cpus: '4.0'
          memory: 16G
        reservations:
          cpus: '2.0'
          memory: 8G
          devices:
            - driver: nvidia
              count: ${GPU_COUNT:-1}
              capabilities: [gpu]

volumes:
  pytorch_notebooks:
  pytorch_data: