Dockerfile-gpu

# GPU Dockerfile for MASE
#
# Uses a plain Ubuntu base image. CUDA support is provided by the PyTorch
# wheel which bundles the necessary CUDA runtime libraries.
#
# Requirements:
#   - NVIDIA GPU with drivers installed on the host
#   - Run with: docker run --gpus all ...
#
# Example build command:
#   docker buildx build --platform linux/amd64 \
#     -f Dockerfile-gpu \
#     --tag aaronyirenzhao/mase-gpu:latest \
#     --push .

FROM ubuntu:22.04
USER root

# Install system packages
ADD install-pkgs.sh install-pkgs.sh
RUN bash install-pkgs.sh

RUN pip3 install --upgrade pip

# Install PyTorch with CUDA support.
# pip installs the CUDA-enabled wheel by default on Linux (includes bundled CUDA runtime).
RUN pip3 install --no-cache-dir torch torchvision torchaudio

# Install MASE and its dependencies
ADD install-pips.sh install-pips.sh
RUN bash install-pips.sh

CMD ["bash"]