diff --git a/.github/workflows/docker-image-cpu.yml b/.github/workflows/docker-image-cpu.yml index b0435d9..3366313 100644 --- a/.github/workflows/docker-image-cpu.yml +++ b/.github/workflows/docker-image-cpu.yml @@ -6,7 +6,7 @@ on: pull_request: branches: [ "main" ] schedule: - - cron: '0 1 * * *' + - cron: '0 1 * * *' jobs: @@ -16,25 +16,21 @@ jobs: steps: - uses: actions/checkout@v3 - - name: Build the CPU Docker image + - name: Build and push the CPU Docker image run: | TAG=$(date -u +%Y%m%d%H%M%S) - docker build --no-cache --build-arg VHLS_PATH=/mnt/applications/Xilinx/23.1 \ - --build-arg VHLS_VERSION=2023.1 \ + docker build --no-cache \ -f Dockerfile-cpu \ --tag docker.io/deepwok/mase-docker-cpu:$TAG . || exit 1 - # Push to dockerhub if [ "${{ secrets.DOCKER_HUB_PASSWORD }}" != "" ]; then - echo ${{secrets.DOCKER_HUB_PASSWORD}} | docker login docker.io -u deepwok --password-stdin + echo ${{ secrets.DOCKER_HUB_PASSWORD }} | docker login docker.io -u deepwok --password-stdin docker push docker.io/deepwok/mase-docker-cpu:$TAG docker tag docker.io/deepwok/mase-docker-cpu:$TAG docker.io/deepwok/mase-docker-cpu:latest docker push docker.io/deepwok/mase-docker-cpu:latest - echo "MASE Docker image for CPU pushed." + echo "MASE CPU Docker image pushed." else - echo "Skipped pushing docker image for CPU." + echo "Skipped pushing: DOCKER_HUB_PASSWORD secret not configured." fi - - diff --git a/.github/workflows/docker-image-gpu.yml b/.github/workflows/docker-image-gpu.yml index c9ae749..05428f4 100644 --- a/.github/workflows/docker-image-gpu.yml +++ b/.github/workflows/docker-image-gpu.yml @@ -14,25 +14,20 @@ jobs: steps: - uses: actions/checkout@v3 - - name: Build the GPU Docker image + - name: Build and push the GPU Docker image run: | TAG=$(date -u +%Y%m%d%H%M%S) - docker build --no-cache --build-arg VHLS_PATH=/mnt/applications/Xilinx/23.1 \ - --build-arg VHLS_VERSION=2023.1 \ + docker build --no-cache \ -f Dockerfile-gpu \ --tag docker.io/deepwok/mase-docker-gpu:$TAG . || exit 1 - # Push to dockerhub if [ "${{ secrets.DOCKER_HUB_PASSWORD }}" != "" ]; then - echo ${{secrets.DOCKER_HUB_PASSWORD}} | docker login docker.io -u deepwok --password-stdin + echo ${{ secrets.DOCKER_HUB_PASSWORD }} | docker login docker.io -u deepwok --password-stdin docker push docker.io/deepwok/mase-docker-gpu:$TAG docker tag docker.io/deepwok/mase-docker-gpu:$TAG docker.io/deepwok/mase-docker-gpu:latest docker push docker.io/deepwok/mase-docker-gpu:latest - echo "MASE Docker image for GPU pushed." + echo "MASE GPU Docker image pushed." else - echo "Skipped pushing docker image for GPU." + echo "Skipped pushing: DOCKER_HUB_PASSWORD secret not configured." fi - - - diff --git a/Dockerfile-cpu b/Dockerfile-cpu index ccb2a58..5f10e2f 100644 --- a/Dockerfile-cpu +++ b/Dockerfile-cpu @@ -1,27 +1,20 @@ -# This Dockerfile configures a Docker environment that -# contains all the required packages for the tool FROM ubuntu:22.04 - USER root -# Install apt packages ADD install-pkgs.sh install-pkgs.sh RUN bash install-pkgs.sh -CMD ["bash"] - -# Install PyTorch and Torch-MLIR RUN pip3 install --upgrade pip -RUN pip3 install --pre torch-mlir torchvision \ - -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels \ - --extra-index-url https://download.pytorch.org/whl/nightly/cpu -# Install pip packages +# Pin CPU-only PyTorch before MASE install to prevent pip from pulling in the CUDA variant +RUN pip3 install --no-cache-dir \ + torch==2.6 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cpu + +# Install MASE dependencies, remove any NVIDIA packages pulled in, then re-pin CPU torch ADD install-pips.sh install-pips.sh -RUN bash install-pips.sh +RUN bash install-pips.sh && \ + pip3 list | grep -E '^nvidia-' | awk '{print $1}' | xargs pip3 uninstall -y 2>/dev/null || true && \ + pip3 install --force-reinstall --no-deps --no-cache-dir \ + torch==2.6 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cpu -# Add environment variables -ARG VHLS_PATH -ARG VHLS_VERSION -ADD install-env.sh install-env.sh -RUN bash install-env.sh $VHLS_PATH $VHLS_VERSION +CMD ["bash"] \ No newline at end of file diff --git a/Dockerfile-gpu b/Dockerfile-gpu index 8759223..8a5bf27 100644 --- a/Dockerfile-gpu +++ b/Dockerfile-gpu @@ -1,25 +1,33 @@ -# This Dockerfile configures a Docker environment that -# contains all the required packages for the tool -FROM ubuntu:22.04 +# 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 apt packages +# Install system packages ADD install-pkgs.sh install-pkgs.sh RUN bash install-pkgs.sh -CMD ["bash"] - -# Install PyTorch and Torch-MLIR RUN pip3 install --upgrade pip -RUN pip3 install torch torchvision torchaudio -# Install pip packages +# 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 -# Add environment variables -ARG VHLS_PATH -ARG VHLS_VERSION -ADD install-env.sh install-env.sh -RUN bash install-env.sh $VHLS_PATH $VHLS_VERSION +CMD ["bash"] diff --git a/README.md b/README.md index 2c9ef3e..f1d51f4 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,68 @@ # MASE Docker -The docker container used to run MASE. +Docker images for running MASE. Two variants are provided: -The deps are in `setup.py` of deepwok/mase +- `Dockerfile-cpu` — CPU-only, used for CI and general development +- `Dockerfile-gpu` — CUDA-enabled, for GPU-accelerated training and inference -## How to update dependencies? +## Official Images (DeepWok) -1. Create a PR to update deps in mase-tools +The official MASE Docker images are automatically built and pushed to Docker Hub +by GitHub Actions whenever changes are merged into `main`: -2. Manually trigger the build in mase-docker +| Image | Docker Hub | +|-------|-----------| +| CPU | `deepwok/mase-docker-cpu:latest` | +| GPU | `deepwok/mase-docker-gpu:latest` | + +To use the official images: + +```bash +# CPU +docker pull deepwok/mase-docker-cpu:latest + +# GPU +docker pull deepwok/mase-docker-gpu:latest +``` + +--- + +## Build Your Own Image + +If you need to customise the image (e.g. different CUDA version, private fork), +you can build and push it yourself. + +**Prerequisites:** +- [Docker Desktop](https://docs.docker.com/get-docker/) with `buildx` support +- A [Docker Hub](https://hub.docker.com/) account +- For GPU image: check your host driver version with `nvidia-smi` + +Replace `` with your Docker Hub username before running. + +**CPU image:** + +```bash +DOCKER_USER= +TAG=$(date -u +%Y%m%d%H%M%S) && \ +docker buildx build --no-cache \ + --platform linux/amd64 \ + -f Dockerfile-cpu \ + --tag docker.io/$DOCKER_USER/mase-cpu:$TAG \ + --tag docker.io/$DOCKER_USER/mase-cpu:latest \ + --push \ + . +``` + +**GPU image** (see `Dockerfile-gpu` for available CUDA version options): + +```bash +DOCKER_USER= +TAG=$(date -u +%Y%m%d%H%M%S) && \ +docker buildx build --no-cache \ + --platform linux/amd64 \ + -f Dockerfile-gpu \ + --tag docker.io/$DOCKER_USER/mase-gpu:$TAG \ + --tag docker.io/$DOCKER_USER/mase-gpu:latest \ + --push \ + . +``` \ No newline at end of file