From 3f78b09df73a266ec9b09fbb9fd3bbc456a90cb5 Mon Sep 17 00:00:00 2001
From: jianfengjiang <jiangjianfeng.jjf@antgroup.com>
Date: Sun, 26 Sep 2021 11:18:51 +0800
Subject: [PATCH] Update pytorch demo with copy_bom

---
 .github/workflows/demo_test.yml            |  2 +-
 demos/pytorch/install_python_with_conda.sh |  2 +-
 demos/pytorch/pytorch.yaml                 | 17 +++++++++++++++
 demos/pytorch/run_pytorch_on_occlum.sh     | 24 ++++++++--------------
 4 files changed, 28 insertions(+), 17 deletions(-)
 create mode 100644 demos/pytorch/pytorch.yaml

diff --git a/.github/workflows/demo_test.yml b/.github/workflows/demo_test.yml
index 77836ed7..23ff09a1 100644
--- a/.github/workflows/demo_test.yml
+++ b/.github/workflows/demo_test.yml
@@ -346,7 +346,7 @@ jobs:
     # FIXME: PyTorch can't exit normally in SIM mode
     - name: Kill the container
       run: |
-        sleep 180;
+        sleep 240;
         cat "$GITHUB_WORKSPACE/log";
         if grep -q Done "$GITHUB_WORKSPACE/log"; then
           docker stop pytorch_test
diff --git a/demos/pytorch/install_python_with_conda.sh b/demos/pytorch/install_python_with_conda.sh
index aa1766b0..011b1864 100755
--- a/demos/pytorch/install_python_with_conda.sh
+++ b/demos/pytorch/install_python_with_conda.sh
@@ -8,4 +8,4 @@ script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}"  )" >/dev/null 2>&1 && pwd )"
 # 2. Install python and dependencies to specified position
 [ -d Miniconda3-latest-Linux-x86_64.sh ] || wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
 [ -d miniconda ] || bash ./Miniconda3-latest-Linux-x86_64.sh -b -p $script_dir/miniconda
-$script_dir/miniconda/bin/conda create --prefix $script_dir/occlum_instance/image/opt/python-occlum -y  python=3.7 pytorch numpy torchvision -c pytorch
+$script_dir/miniconda/bin/conda create --prefix $script_dir/python-occlum -y  python=3.7 pytorch numpy torchvision -c pytorch
diff --git a/demos/pytorch/pytorch.yaml b/demos/pytorch/pytorch.yaml
new file mode 100644
index 00000000..fdce3e6b
--- /dev/null
+++ b/demos/pytorch/pytorch.yaml
@@ -0,0 +1,17 @@
+includes:
+  - base.yaml
+targets:
+  - target: /bin
+    createlinks:
+      - src: /opt/python-occlum/bin/python3
+        linkname: python3
+  # python packages
+  - target: /opt
+    copy: 
+      - dirs:
+          - ../python-occlum
+  # python code
+  - target: /
+    copy:
+      - files: 
+          - ../demo.py
\ No newline at end of file
diff --git a/demos/pytorch/run_pytorch_on_occlum.sh b/demos/pytorch/run_pytorch_on_occlum.sh
index 21f9e9ed..f701209b 100755
--- a/demos/pytorch/run_pytorch_on_occlum.sh
+++ b/demos/pytorch/run_pytorch_on_occlum.sh
@@ -7,26 +7,20 @@ NC='\033[0m'
 script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}"  )" >/dev/null 2>&1 && pwd )"
 python_dir="$script_dir/occlum_instance/image/opt/python-occlum"
 
+cd occlum_instance && rm -rf image
+copy_bom -f ../pytorch.yaml --root image --include-dir /opt/occlum/etc/template
+
 if [ ! -d $python_dir ];then
     echo "Error: cannot stat '$python_dir' directory"
     exit 1
 fi
 
-cd occlum_instance
-# Copy files into Occlum Workspace and build
-if [ ! -L "image/bin/python3" ];then
-    ln -s /opt/python-occlum/bin/python3 image/bin/python3
-    cp -f /opt/occlum/glibc/lib/libdl.so.2 image/opt/occlum/glibc/lib/
-    cp -f /opt/occlum/glibc/lib/libutil.so.1 image/opt/occlum/glibc/lib/
-    cp -f /opt/occlum/glibc/lib/librt.so.1 image/opt/occlum/glibc/lib/
-    cp -f ../demo.py image
-    new_json="$(jq '.resource_limits.user_space_size = "6000MB" |
-                    .resource_limits.kernel_space_heap_size = "256MB" |
-                    .process.default_mmap_size = "4000MB" |
-                    .env.default += ["PYTHONHOME=/opt/python-occlum"]' Occlum.json)" && \
-    echo "${new_json}" > Occlum.json
-    occlum build
-fi
+new_json="$(jq '.resource_limits.user_space_size = "6000MB" |
+                .resource_limits.kernel_space_heap_size = "256MB" |
+                .process.default_mmap_size = "4000MB" |
+                .env.default += ["PYTHONHOME=/opt/python-occlum"]' Occlum.json)" && \
+echo "${new_json}" > Occlum.json
+occlum build
 
 # Run the python demo
 echo -e "${BLUE}occlum run /bin/python3 demo.py${NC}"