From 136c2cca9561c189f8b0918efa59c6d8c06fa0b9 Mon Sep 17 00:00:00 2001
From: ClawSeven <zehuan97@gmail.com>
Date: Wed, 25 Aug 2021 14:20:47 +0800
Subject: [PATCH] Seperate ssl generation with environment preparation

---
 .github/workflows/build_and_push_ci_image.yml       | 12 ++++++++++--
 demos/tensorflow/tensorflow_serving/README.md       |  8 ++++++--
 .../tensorflow_serving/prepare_model_and_env.sh     |  2 --
 tools/docker/ci/Dockerfile.tf_serving               | 13 +++++++++++++
 4 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/build_and_push_ci_image.yml b/.github/workflows/build_and_push_ci_image.yml
index 2c45b76a..4bb54ce0 100644
--- a/.github/workflows/build_and_push_ci_image.yml
+++ b/.github/workflows/build_and_push_ci_image.yml
@@ -176,6 +176,7 @@ jobs:
         push: true
         tags: occlumbackup/occlum:${{ github.event.inputs.tag }}-ubuntu18.04-python
 
+
   Build_tf_serving_image:
     runs-on: ubuntu-18.04
     if: github.event.inputs.image_name == 'tf_serving'
@@ -187,12 +188,17 @@ jobs:
     - name: Get occlum version
       run: echo "OCCLUM_VERSION=$(grep 'Version =' src/pal/include/occlum_version.h | awk '{print $4}')" >> $GITHUB_ENV;
 
+    - name: Get branch name
+      # GITHUB_REF=refs/heads/branch-name
+      # ${GITHUB_REF##*/} == branch-name
+      run: echo "OCCLUM_BRANCH=$(echo ${GITHUB_REF##*/})" >> $GITHUB_ENV;
+
     # Because "Build and push" step `context` field can't be subdir,
     # we need to copy files needed by dockerfile to root dir of the project
     - name: Copy context for docker build
       run: |
         cp -r tools/docker .
-        
+
     - name: Set up QEMU
       uses: docker/setup-qemu-action@v1
 
@@ -211,6 +217,8 @@ jobs:
         context: .
         file: ./tools/docker/ci/Dockerfile.tf_serving
         platforms: linux/amd64
-        build-args: OCCLUM_VERSION=${{ env.OCCLUM_VERSION }}
+        build-args: |
+          "OCCLUM_VERSION=${{ env.OCCLUM_VERSION }}"
+          "OCCLUM_BRANCH=${{ env.OCCLUM_BRANCH }}"
         push: true
         tags: occlumbackup/occlum:${{ github.event.inputs.tag }}-ubuntu18.04-tf_serving
diff --git a/demos/tensorflow/tensorflow_serving/README.md b/demos/tensorflow/tensorflow_serving/README.md
index 5265164a..a00e7659 100644
--- a/demos/tensorflow/tensorflow_serving/README.md
+++ b/demos/tensorflow/tensorflow_serving/README.md
@@ -10,12 +10,16 @@ TensorFlow Serving is a flexible, high-performance serving system for machine le
 
 #### Executing the Tensorflow serving in Occlum
 
-The following command will download the Resnet50 model and convert the model format. It also will generate the TLS key and certificates for localhost( server domain name). The server.crt will be used by client. The sever.key and ssl.cfg is used by TF serving.
-
+The following command will download the Resnet50 model and convert the model format.
 ```
 ./prepare_model_and_env.sh 
 ```
 
+The following command will generate the TLS key and certificates for localhost( server domain name). The server.crt will be used by client. The sever.key and ssl.cfg is used by TF serving.
+```
+./generate_ssl_config.sh localhost
+```
+
 Run the Tensorflow Serving in occlum.
 
 ```
diff --git a/demos/tensorflow/tensorflow_serving/prepare_model_and_env.sh b/demos/tensorflow/tensorflow_serving/prepare_model_and_env.sh
index f22b5b33..a1a9e458 100755
--- a/demos/tensorflow/tensorflow_serving/prepare_model_and_env.sh
+++ b/demos/tensorflow/tensorflow_serving/prepare_model_and_env.sh
@@ -4,5 +4,3 @@ pip3 install --upgrade pip
 pip3 install --upgrade tensorflow==2.4
 ./download_model.sh
 python3 ./model_graph_to_saved_model.py --import_path ./models/resnet50-v15-fp32/resnet50-v15-fp32.pb --export_dir ./resnet50-v15-fp32 --model_version 1 --inputs input --outputs predict
-./generate_ssl_config.sh localhost
-
diff --git a/tools/docker/ci/Dockerfile.tf_serving b/tools/docker/ci/Dockerfile.tf_serving
index b473c09b..925fbd3e 100644
--- a/tools/docker/ci/Dockerfile.tf_serving
+++ b/tools/docker/ci/Dockerfile.tf_serving
@@ -38,3 +38,16 @@ LABEL maintainer="Zehuan Li <lizehuan.lzh@antgroup.com>"
 WORKDIR /root
 COPY --from=binary_build /usr/local/bin/tensorflow_model_server /root/tensorflow_model_server
 
+ARG OCCLUM_BRANCH
+RUN apt-get update && apt-get install git && \
+    git clone -b $OCCLUM_BRANCH git://github.com/occlum/occlum.git && \
+    cd occlum && \
+    cd demos/tensorflow/tensorflow_serving && \
+    ./prepare_model_and_env.sh && \
+    cd client && \
+    ./prepare_client_env.sh && \
+    mv ../resnet50-v15-fp32/ /root && \
+    rm -rf /root/occlum
+
+WORKDIR /root
+