From 538165a6ec69c77bab50180a8d8fc5a5d0f5cf00 Mon Sep 17 00:00:00 2001
From: Josh Bailey <josh@vandervecken.com>
Date: Tue, 10 Oct 2023 04:55:11 +0000
Subject: [PATCH] Build torchserve container with both cpu (default) and cuda
 versions.

---
 .github/workflows/ci-test.yml       |  2 +-
 .github/workflows/docker-extras.yml | 34 +++++++++++
 docker/Dockerfile.cuda-torchserve   | 15 +++++
 docker/Dockerfile.torchserve        | 87 +++--------------------------
 docker/pytorch-patch.txt            | 25 ---------
 docker/torchserve-patch.txt         | 25 ---------
 6 files changed, 58 insertions(+), 130 deletions(-)
 create mode 100644 docker/Dockerfile.cuda-torchserve
 delete mode 100644 docker/pytorch-patch.txt
 delete mode 100644 docker/torchserve-patch.txt

diff --git a/.github/workflows/ci-test.yml b/.github/workflows/ci-test.yml
index 75dd3f0e..42c44395 100644
--- a/.github/workflows/ci-test.yml
+++ b/.github/workflows/ci-test.yml
@@ -40,7 +40,7 @@ jobs:
             sox \
             uhd-host \
             wget \
-          && grep -h 'git clone' docker/*|sed -E 's/RUN\s+git clone\s+//g'|sort|uniq|xargs -L1 git clone
+          && grep -h 'git clone' docker/*|grep -v pytorch|sed -E 's/RUN\s+git clone\s+//g'|sort|uniq|xargs -L1 git clone
       - name: Install dependencies
         run: |
           for repodir in flatbuffers json libsigmf gr-iqtlabs ; do \
diff --git a/.github/workflows/docker-extras.yml b/.github/workflows/docker-extras.yml
index 3652731f..024814fc 100644
--- a/.github/workflows/docker-extras.yml
+++ b/.github/workflows/docker-extras.yml
@@ -7,6 +7,40 @@ on:
 
 jobs:
   buildx-torchserve:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+    - name: Get the version
+      id: get_version
+      run: echo ::set-output name=VERSION::$(echo $GITHUB_REF | cut -d / -f 3)
+    - name: Change for main
+      id: change_version
+      run: if [ "${{ steps.get_version.outputs.VERSION }}" == "main" ]; then echo ::set-output name=VERSION::latest; else echo ::set-output name=VERSION::${{ steps.get_version.outputs.VERSION }}; fi
+    - name: Set up qemu
+      uses: docker/setup-qemu-action@v3
+      with:
+        platforms: all
+    - name: Set up Docker Buildx
+      id: buildx
+      uses: docker/setup-buildx-action@v3
+    - name: Docker Login
+      env:
+        DOCKER_PASSWORD: ${{ secrets.DOCKER_TOKEN }}
+      run: |
+        echo "${DOCKER_PASSWORD}" | docker login --username "${{ secrets.DOCKER_USERNAME }}" --password-stdin
+      if: github.repository == 'iqtlabs/gamutrf' && github.event_name == 'push'
+    - name: Build and push platforms
+      uses: docker/build-push-action@v5
+      with:
+        context: docker
+        file: docker/Dockerfile.torchserve
+        platforms: linux/amd64,linux/arm64
+        push: true
+        tags: iqtlabs/gamutrf-torchserve:${{ steps.change_version.outputs.VERSION }}
+      if: github.repository == 'iqtlabs/gamutrf' && github.event_name == 'push'
+  buildx-cuda-torchserve:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
diff --git a/docker/Dockerfile.cuda-torchserve b/docker/Dockerfile.cuda-torchserve
new file mode 100644
index 00000000..efb54cc0
--- /dev/null
+++ b/docker/Dockerfile.cuda-torchserve
@@ -0,0 +1,15 @@
+FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
+ENV DEBIAN_FRONTEND=noninteractive
+WORKDIR /root
+RUN apt-get update && \
+    apt-get install -y \
+      git \
+      python3-pip
+RUN pip config set global.no-cache-dir false && \
+  git clone https://github.com/pytorch/serve -b v0.8.2 && \
+  cd serve && \
+  python3 ./ts_scripts/install_dependencies.py --cuda cu118 --environment prod && \
+  pip3 install . && \
+  cd .. && \
+  rm -rf serve
+RUN /usr/local/bin/torchserve --help
diff --git a/docker/Dockerfile.torchserve b/docker/Dockerfile.torchserve
index 9b11c004..46640b90 100644
--- a/docker/Dockerfile.torchserve
+++ b/docker/Dockerfile.torchserve
@@ -1,86 +1,15 @@
-FROM --platform=$BUILDPLATFORM ubuntu:22.04 as git-cloner
-WORKDIR /root
-RUN apt-get update && \
-    apt-get install -y \
-      ca-certificates \
-      git \
-      python3
-RUN git clone https://github.com/google/shaderc -b v2023.6
-RUN git clone https://github.com/pytorch/pytorch -b v2.1.0
-RUN git clone https://github.com/pytorch/serve -b v0.8.2
-WORKDIR /root/shaderc
-RUN ./utils/git-sync-deps
-WORKDIR /root/pytorch
-RUN git submodule update --init --recursive
-
-FROM ubuntu:22.04 as shader-compiler
-WORKDIR /root
-RUN apt-get update && \
-    apt-get install -y \
-      build-essential \
-      cmake \
-      git \
-      ninja-build \
-      python3-dev \
-      python3-pip
-COPY --from=git-cloner /root/shaderc /root/shaderc
-WORKDIR /root/shaderc/build
-RUN cmake -GNinja -DCMAKE_BUILD_TYPE=Release -DSHADERC_SKIP_TESTS=on .. && ninja && ninja install
-
-FROM ubuntu:22.04 as pytorch-compiler
+FROM ubuntu:22.04
 ENV DEBIAN_FRONTEND=noninteractive
-ENV USE_CUDA=0
-ENV USE_VULKAN=1
-ENV USE_VULKAN_SHADERC_RUNTIME=1
-ENV USE_VULKAN_WRAPPER=0
 WORKDIR /root
-COPY --from=shader-compiler /usr/local /usr/local
-COPY --from=git-cloner /root/pytorch /root/pytorch
 RUN apt-get update && \
     apt-get install -y \
-      build-essential \
-      cmake \
       git \
-      libvulkan-dev \
-      ninja-build \
-      python3-dev \
       python3-pip
-WORKDIR /root/pytorch
-COPY pytorch-patch.txt /root/pytorch
-RUN patch -p1 < pytorch-patch.txt
-RUN pip3 install -U pyyaml
-RUN python3 setup.py install && python3 setup.py develop
-RUN python3 -c "import torch ; assert(torch.is_vulkan_available())"
-
-FROM ubuntu:22.04 as torchserve-builder
-RUN apt-get update && \
-    apt-get install -y \
-      build-essential \
-      cmake \
-      git \
-      libvulkan-dev \
-      ninja-build \
-      openjdk-17-jdk \
-      python3-dev \
-      python3-pip
-WORKDIR /root
-COPY --from=git-cloner /root/serve /root/serve
-WORKDIR /root/serve
-COPY torchserve-patch.txt /root/serve
-RUN patch -p1 < torchserve-patch.txt
-COPY --from=pytorch-compiler /usr/local /usr/local
-RUN python3 -c "import torch ; assert(torch.is_vulkan_available())"
-RUN python3 ./ts_scripts/install_dependencies.py --environment prod
-RUN pip3 install .
-
-FROM ubuntu:22.04
-RUN apt-get update && \
-    apt-get install -y \
-      libvulkan1 \
-      openjdk-17-jre \
-      python3 \
-      python3-pip
-WORKDIR /root
-COPY --from=torchserve-builder /usr/local /usr/local
-RUN python3 -c "import torch ; assert(torch.is_vulkan_available())"
+RUN pip config set global.no-cache-dir false && \
+  git clone https://github.com/pytorch/serve -b v0.8.2 && \
+  cd serve && \
+  python3 ./ts_scripts/install_dependencies.py --environment prod && \
+  pip3 install . && \
+  cd .. && \
+  rm -rf serve
 RUN /usr/local/bin/torchserve --help
diff --git a/docker/pytorch-patch.txt b/docker/pytorch-patch.txt
deleted file mode 100644
index d4546756..00000000
--- a/docker/pytorch-patch.txt
+++ /dev/null
@@ -1,25 +0,0 @@
-diff --git a/aten/src/ATen/native/vulkan/impl/Arithmetic.cpp b/aten/src/ATen/native/vulkan/impl/Arithmetic.cpp
-index ef6b794fbd9..900dd8ddf4c 100644
---- a/aten/src/ATen/native/vulkan/impl/Arithmetic.cpp
-+++ b/aten/src/ATen/native/vulkan/impl/Arithmetic.cpp
-@@ -17,6 +17,7 @@ api::ShaderInfo get_shader(const OpType type) {
-     case OpType::DIV:
-       return VK_KERNEL(div);
-   }
-+  return VK_KERNEL(add);
- }
- 
- struct Params final {
-diff --git a/aten/src/ATen/native/vulkan/impl/Registry.cpp b/aten/src/ATen/native/vulkan/impl/Registry.cpp
-index 3cf3148c874..b750914dfeb 100644
---- a/aten/src/ATen/native/vulkan/impl/Registry.cpp
-+++ b/aten/src/ATen/native/vulkan/impl/Registry.cpp
-@@ -33,7 +33,7 @@ const api::ShaderInfo& look_up_shader_info(const std::string& op_name) {
-   const RegistryKeyMap& registry_key_map = registry_iterator->second;
- 
-   // Look for "override" and "catchall" keys
--  for (const std::string& key : {"override", "catchall"}) {
-+  for (const std::string key : {"override", "catchall"}) {
-     const RegistryKeyMap::const_iterator registry_key_iterator =
-         registry_key_map.find(key);
-     if (registry_key_iterator != registry_key_map.end()) {
diff --git a/docker/torchserve-patch.txt b/docker/torchserve-patch.txt
deleted file mode 100644
index 0c99606f..00000000
--- a/docker/torchserve-patch.txt
+++ /dev/null
@@ -1,25 +0,0 @@
-diff --git a/requirements/torch_cu118_linux.txt b/requirements/torch_cu118_linux.txt
-index d34969ef..fe823925 100644
---- a/requirements/torch_cu118_linux.txt
-+++ b/requirements/torch_cu118_linux.txt
-@@ -1,7 +1,7 @@
- #pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
- --extra-index-url https://download.pytorch.org/whl/cu118
- -r torch_common.txt
--torch==2.0.1+cu118; sys_platform == 'linux'
-+# torch==2.0.1+cu118; sys_platform == 'linux'
- torchvision==0.15.2+cu118; sys_platform == 'linux'
- torchtext==0.15.2; sys_platform == 'linux'
- torchaudio==2.0.2+cu118; sys_platform == 'linux'
-diff --git a/ts_scripts/install_dependencies.py b/ts_scripts/install_dependencies.py
-index 6aef56db..0875d4c6 100644
---- a/ts_scripts/install_dependencies.py
-+++ b/ts_scripts/install_dependencies.py
-@@ -23,6 +23,7 @@ class Common:
-         pass
- 
-     def install_torch_packages(self, cuda_version):
-+        return
-         if cuda_version:
-             if platform.system() == "Darwin":
-                 print(