update the docker file

.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
 
-
-

.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
-
-
-

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"]

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"]

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 `<your-dockerhub-username>` with your Docker Hub username before running.
+
+**CPU image:**
+
+```bash
+DOCKER_USER=<your-dockerhub-username>
+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=<your-dockerhub-username>
+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 \
+  .
+```