From ca73028a3ae63cccc0914f6b229a699b24561025 Mon Sep 17 00:00:00 2001
From: Frank Lee <somerlee.9@gmail.com>
Date: Thu, 23 Jun 2022 14:51:59 +0800
Subject: [PATCH] [workflow] auto-publish docker image upon release (#1164)

---
 .../workflows/draft_github_release_post.yml   |  1 +
 .github/workflows/release_docker.yml          | 40 +++++++++++++++++++
 docker/Dockerfile                             | 10 ++++-
 3 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 .github/workflows/release_docker.yml

diff --git a/.github/workflows/draft_github_release_post.yml b/.github/workflows/draft_github_release_post.yml
index 234fdb847..f970a9091 100644
--- a/.github/workflows/draft_github_release_post.yml
+++ b/.github/workflows/draft_github_release_post.yml
@@ -12,6 +12,7 @@ on:
 jobs:
   release:
     name: Draft Release Post
+    if: github.repository == 'hpcaitech/ColossalAI'
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
diff --git a/.github/workflows/release_docker.yml b/.github/workflows/release_docker.yml
new file mode 100644
index 000000000..d79f2ec3d
--- /dev/null
+++ b/.github/workflows/release_docker.yml
@@ -0,0 +1,40 @@
+name: Publish Docker Image to DockerHub
+
+on:
+  workflow_dispatch:
+  release:
+    types: [published]
+
+jobs:
+  release:
+    name: Publish Docker Image to DockerHub
+    if: github.repository == 'hpcaitech/ColossalAI'
+    runs-on: [self-hosted, gpu]
+    container:
+      image: "hpcaitech/docker-in-docker:latest"
+      options: --gpus all --rm -v /var/run/docker.sock:/var/run/docker.sock
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+      - name: Build Docker
+        run: |
+          version=$(cat version.txt)
+          docker build -t hpcaitech/colossalai:$version ./docker
+      - name: Log in to Docker Hub
+        uses: docker/login-action@f054a8b539a109f9f41c372932f1ae047eff08c9
+        with:
+          username: ${{ secrets.DOCKER_USERNAME }}
+          password: ${{ secrets.DOCKER_PASSWORD }}
+      - name: Extract metadata (tags, labels) for Docker
+        id: meta
+        uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
+        with:
+          images: hpcaitech/colossalai
+      - name: Build and push Docker image
+        uses: docker/build-push-action@ad44023a93711e3deb337508980b4b5e9bcdc5dc
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
\ No newline at end of file
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 123c7235b..7dd82fc93 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,4 +1,12 @@
-FROM hpcaitech/pytorch-cuda:1.10.1-11.3.0
+FROM hpcaitech/cuda-conda:11.3
+
+# install torch
+RUN conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
+
+# install apex
+RUN git clone https://github.com/NVIDIA/apex && \
+    cd apex && \
+    pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
 
 # install colossalai
 RUN git clone https://github.com/hpcaitech/ColossalAI.git \