leoustc/pytorch repository overview

⁠PyTorch Cluster Images: From Data Science to CUDA RDMA Superclusters

This DockerHub repository hosts a series of layered images designed for AI and HPC workloads, ranging from Jupyter-based data science to CUDA/NCCL-enabled multi-node GPU clusters with RDMA. This project provides production-ready Docker containers for PyTorch development and GPU benchmarking, from single-node development to multi-node distributed training.

⁠Docker Hub Repository

All images are available on Docker Hub:

docker pull leoustc/pytorch:<tag>

Available tags:

• base: CPU-only PyTorch with JupyterLab
• cuda12.6: CUDA 12.6 enabled PyTorch
• cuda12.6-mpi: CUDA 12.6 with MPI support for distributed training
• cuda12.6-mpi-rdma: CUDA 12.6 with MPI on RDMA support for distributed training

⁠Features

• Multiple Base Images:
  • Python 3.12 slim for CPU workloads
  • CUDA 12.6 with cuDNN for GPU acceleration
  • MPI-enabled images for distributed computing
• Pre-installed Packages:
  • PyTorch (CPU/GPU versions)
  • Scientific Python stack (NumPy, Pandas, Matplotlib, Seaborn)
  • JupyterLab with widgets support
  • GPU monitoring tools (nvtop)
• Development Tools:
  • Built-in examples and notebooks
  • GPU benchmarking scripts
  • Multi-node MPI support

⁠Quick Start

⁠CPU Development

docker run -p 8888:8888 -v $(pwd)/project:/root/project leoustc/pytorch:base

⁠GPU Development

docker run --gpus all -p 8888:8888 -v $(pwd)/project:/root/project leoustc/pytorch:cuda12.6

⁠Multi-Node MPI Setup

docker run --gpus all -p 8888:8888 \
  -v ./ssh:/root/.ssh \
  -v ./project:/root/project \
  leoustc/pytorch:cuda12.6-mpi

⁠Docker Compose Examples

⁠Basic Development

version: "3.8"
services:
  pytorch:
    image: leoustc/pytorch:base
    ports:
      - "8888:8888"
    environment:
      - TOKEN=abc123
    volumes:
      - ./project:/workspace/project

⁠GPU Development

version: "3.8"
services:
  pytorch-gpu:
    image: leoustc/pytorch:cuda12.6
    runtime: nvidia
    shm_size: '4gb'
    ports:
      - "8888:8888"
    environment:
      - TOKEN=abc123
    deploy:
      resources:
        reservations:
          devices:
            - capabilities: [gpu]
    volumes:
      - ./project:/root/project

⁠Access JupyterLab

Open http://localhost:8888⁠ in your browser.
Default token: abc123 (configurable via TOKEN environment variable)

⁠File Structure

pytorch-container/
├── Dockerfile                 # Base CPU image
├── Dockerfile.cuda12.6        # CUDA 12.6 image
├── Dockerfile.cuda12.6-mpi    # MPI-enabled CUDA image
├── docker-compose.yml         # Basic compose file
├── docker-compose.cuda12.yml  # CUDA compose file
└── docker-compose.cuda12.6-mpi.yml  # MPI cluster compose

⁠License

This project is licensed under the Apache 2.0 License.

Maintained by:
Leo Li (leoustc)

Docker Hub Repository:
leoustc/pytorch⁠