From c4c3e64ab6ef8826f64043f0b6d50771987d726c Mon Sep 17 00:00:00 2001
From: Ryoichi Ando <ryich.ando@gmail.com>
Date: Fri, 10 Jan 2025 10:57:36 +0900
Subject: [PATCH] two step Docker image build

---
 .github/workflows/build-docker.yml  | 21 ++++++++++++++++++---
 .github/workflows/vast/provision.sh |  2 +-
 Dockerfile                          | 24 +++++++++++++++---------
 README.md                           |  4 ++--
 4 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/.github/workflows/build-docker.yml b/.github/workflows/build-docker.yml
index bc24eb5..9bb08cc 100644
--- a/.github/workflows/build-docker.yml
+++ b/.github/workflows/build-docker.yml
@@ -19,6 +19,9 @@ jobs:
       contents: read
       packages: write
     steps:
+    - name: print mode
+      run: |
+        echo "Mode: ${{ github.event.inputs.target }}" >> $GITHUB_STEP_SUMMARY
     - name: Checkout
       uses: actions/checkout@v3
     - name: Docker Login
@@ -27,9 +30,21 @@ jobs:
         registry: ${{ env.REGISTRY }}
         username: ${{ github.actor }}
         password: ${{ secrets.GITHUB_TOKEN }}
-    - name: Build Image
+    - name: Build Base Image
+      if: github.event.inputs.target == 'base'
       uses: docker/build-push-action@v5
       with:
         push: true
-        build-args: BUILD_MODE=${{ github.event.inputs.target }}
-        tags: ${{ env.REGISTRY }}/${{ github.repository }}/${{ github.event.inputs.target }}:latest
\ No newline at end of file
+        build-args: |
+          MODE=base
+          IMAGE=nvidia/cuda:11.8.0-devel-ubuntu22.04
+        tags: ${{ env.REGISTRY }}/${{ github.repository }}-base:latest
+    - name: Build Compiled Image
+      if: github.event.inputs.target == 'compiled'
+      uses: docker/build-push-action@v5
+      with:
+        push: true
+        build-args: |
+          MODE=compiled
+          IMAGE=ghcr.io/st-tech/ppf-contact-solver-base:latest
+        tags: ${{ env.REGISTRY }}/${{ github.repository }}-compiled:latest
\ No newline at end of file
diff --git a/.github/workflows/vast/provision.sh b/.github/workflows/vast/provision.sh
index 358c599..6bd3f55 100644
--- a/.github/workflows/vast/provision.sh
+++ b/.github/workflows/vast/provision.sh
@@ -18,7 +18,7 @@ DISK_SPACE=32
 GPU_NAME=RTX_4090
 
 # Image
-VAST_IMAGE="ghcr.io/st-tech/ppf-contact-solver/ppf-contact-solver-base:latest"
+VAST_IMAGE="ghcr.io/st-tech/ppf-contact-solver-base:latest"
 
 # Retry interval
 RETRY_INTERVAL=10
diff --git a/Dockerfile b/Dockerfile
index 9a3b14b..3a44057 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,25 +2,31 @@
 # Author: Ryoichi Ando (ryoichi.ando@zozo.com)
 # License: Apache v2.0
 
-FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
+ARG IMAGE
+FROM $IMAGE
 
-ARG BUILD_MODE=base
+ARG MODE
 ENV NVIDIA_DRIVER_CAPABILITIES=utility,compute
 ENV LANG=en_US.UTF-8
 ENV PROJ_NAME=ppf-contact-solver
-ENV BUILT_MODE=${BUILD_MODE}
+ENV BUILT_MODE=$MODE
 
 COPY . /root/${PROJ_NAME}
 WORKDIR /root/${PROJ_NAME}
 
-RUN apt update
-RUN apt install -y git python3 curl
-RUN python3 warmup.py
+RUN echo "building in ${MODE} mode"
+RUN sleep 5
 
-RUN if [ "$BUILD_MODE" = "compiled" ]; then \
+RUN if [ "$MODE" = "compiled" ]; then \
         /root/.cargo/bin/cargo build --release; \
-    else \
+    elif [ "$MODE" = "base" ]; then \
+        apt update; \
+        apt install -y git python3 curl; \
+        python3 warmup.py; \
         cd /root && rm -rf /root/${PROJ_NAME}; \
+    else \
+        echo "unknown build mode ${BUILT_MODE}"; \
+        exit 1; \
     fi
 
 WORKDIR /root
@@ -28,6 +34,6 @@ RUN rm -rf /var/lib/apt/lists/*
 
 CMD if [ "$BUILT_MODE" = "compiled" ]; then \
         cd /root/${PROJ_NAME} && python3 warmup.py jupyter; \
-    else \
+    elif [ "$BUILT_MODE" = "base" ]; then \
         bash; \
     fi
\ No newline at end of file
diff --git a/README.md b/README.md
index c91be07..d04cf28 100644
--- a/README.md
+++ b/README.md
@@ -78,7 +78,7 @@ Next, run the following command to start the 📦 container:
 
 ```
 $MY_WEB_PORT = 8080  # Web port number for web interface
-$IMAGE_NAME = "ghcr.io/st-tech/ppf-contact-solver/ppf-contact-solver:latest"
+$IMAGE_NAME = "ghcr.io/st-tech/ppf-contact-solver-compiled:latest"
 docker run --rm --gpus all -p ${MY_WEB_PORT}:8080 $IMAGE_NAME
 ```
 
@@ -86,7 +86,7 @@ docker run --rm --gpus all -p ${MY_WEB_PORT}:8080 $IMAGE_NAME
 
 ```
 MY_WEB_PORT=8080  # Web port number for web interface
-IMAGE_NAME=ghcr.io/st-tech/ppf-contact-solver/ppf-contact-solver:latest
+IMAGE_NAME=ghcr.io/st-tech/ppf-contact-solver-compiled:latest
 docker run --rm --gpus all -p ${MY_WEB_PORT}:8080 $IMAGE_NAME
 ```