[ci][lcnc] add trt-llm container to no code tests

deepjavalibrary · Mar 28, 2024 · 06a9535 · 06a9535
1 parent cce574e
commit 06a9535
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 20 deletions.
diff --git a/.github/workflows/lmi-no-code.yml b/.github/workflows/lmi-no-code.yml
@@ -7,6 +7,8 @@ on:
         description: 'The released version of DJL'
         required: false
         default: ''
+  schedule:
+    - cron: '0 17 * * *'
 
 jobs:
   create-runners:
@@ -52,7 +54,7 @@ jobs:
       # Limit to 1 so we don't steal a p4d from another test that may be running
       max-parallel: 1
       matrix:
-        container: [deepspeed]
+        container: [tensorrt-llm]
     steps:
       - uses: actions/checkout@v4
       - name: Clean env
@@ -81,7 +83,7 @@ jobs:
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=s3://djl-llm/llama-2-70b-hf/" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code llama-70b
           docker rm -f $(docker ps -aq)
@@ -90,7 +92,7 @@ jobs:
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=codellama/CodeLlama-34b-hf" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code codellama
           docker rm -f $(docker ps -aq)
@@ -99,7 +101,7 @@ jobs:
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=s3://djl-llm/mixtral-8x7b/" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code mixtral-8x7b
           docker rm -f $(docker ps -aq)
@@ -108,7 +110,7 @@ jobs:
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=s3://djl-llm/falcon-40b/" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code falcon-40b 
           docker rm -f $(docker ps -aq) 
@@ -129,7 +131,7 @@ jobs:
     needs: create-runners
     strategy:
       matrix:
-        container: [deepspeed]
+        container: [tensorrt-llm]
     steps:
       - uses: actions/checkout@v4
       - name: Clean env
@@ -158,7 +160,7 @@ jobs:
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=s3://djl-llm/llama-2-7b-hf/" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code llama-7b
           docker rm -f $(docker ps -aq)
@@ -167,17 +169,17 @@ jobs:
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=s3://djl-llm/llama-2-13b-hf/" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code llama-13b
           docker rm -f $(docker ps -aq)
       - name: Gemma-7b lmi container
-        if: ${{ matrix.container }} == "deepspeed"
+        if: ${{ matrix.container  == 'deepspeed' }}
         working-directory: tests/integration
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=s3://djl-llm/gemma-7b/" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code gemma-7b
           docker rm -f $(docker ps -aq)
@@ -186,27 +188,27 @@ jobs:
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=s3://djl-llm/mistral-7b/" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code mistral-7b
           docker rm -f $(docker ps -aq)
       - name: GPTNeox lmi container
-        if: ${{ matrix.container }} == "deepspeed"
+        if: ${{ matrix.container  == 'deepspeed' }}
         working-directory: tests/integration
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=s3://djl-llm/gpt-neox-20b/" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code gpt-neox
           docker rm -f $(docker ps -aq)
       - name: Phi2 lmi container
-        if: ${{ matrix.container }} == "deepspeed"
+        if: ${{ matrix.container  == 'deepspeed' }}
         working-directory: tests/integration
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=microsoft/phi-2" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code phi-2
           docker rm -f $(docker ps -aq)
@@ -215,7 +217,7 @@ jobs:
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=s3://djl-llm/baichuan2-13b/\nHF_TRUST_REMOTE_CODE=true" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code baichuan-13b
           docker rm -f $(docker ps -aq)
@@ -224,7 +226,7 @@ jobs:
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=Qwen/Qwen1.5-14B" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code qwen-1.5-14b
           docker rm -f $(docker ps -aq)
@@ -233,7 +235,7 @@ jobs:
         run: |
           rm -rf models
           echo -en "HF_MODEL_ID=s3://djl-llm/bigcode-starcoder/" > docker_env
-          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models deepspeed \
+          ./launch_container.sh deepjavalibrary/djl-serving:$DJLSERVING_DOCKER_TAG $PWD/models ${{ matrix.container }} \
           serve
           python3 llm/client.py no_code starcoder
           docker rm -f $(docker ps -aq)

diff --git a/tests/integration/launch_container.sh b/tests/integration/launch_container.sh
@@ -23,7 +23,7 @@ fi
 is_llm=false
 if [[ "$platform" == *"-gpu"* ]]; then # if the platform has cuda capabilities
   runtime="nvidia"
-elif [[ "$platform" == *"deepspeed"* || "$platform" == *"trtllm"* ]]; then # Runs multi-gpu
+elif [[ "$platform" == *"deepspeed"* || "$platform" == *"trtllm"*  || "$platform" == *"tensorrt-llm"* ]]; then # Runs multi-gpu
   runtime="nvidia"
   is_llm=true
   shm="12gb"
@@ -117,7 +117,7 @@ if $is_llm; then
   if [[ "$platform" == *"inf2"* ]]; then
     total_retries=80
   fi
-  if [[ "$platform" == *"trtllm"* ]]; then
+  if [[ "$platform" == *"trtllm"* || "$platform" == *"tensorrt-llm"* ]]; then
     total_retries=100
     echo "extra sleep of 10 min for trtllm compilation"
   fi