Notebook changes to pass CI build

aws · Oct 5, 2022 · 1d3f798 · 1d3f798
1 parent fc2a765
commit 1d3f798
Showing 1 changed file with 20 additions and 25 deletions.
diff --git a/...ggingface/sentence-transformers-triton-ensemble/examples/triton_sentence_embeddings.ipynb b/...ggingface/sentence-transformers-triton-ensemble/examples/triton_sentence_embeddings.ipynb
@@ -198,7 +198,17 @@
    "id": "629bfade-8f14-44b6-9be7-88a2fdb84ba9",
    "metadata": {},
    "source": [
-    "> **WARNING**: The next cell will only work if you have first created a custom Studio image, described in Step 2 of this repository's README."
+    "> **WARNING**: The next cell will only work if you have first created a custom Studio image, described in Step 2 of this repository's README. Change the `RUNNING_IN_STUDIO` to `True` if this is the case."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6df2bf7d-a595-46b1-8d1a-43ab946bc858",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "RUNNING_IN_STUDIO = False"
    ]
   },
   {
@@ -216,7 +226,8 @@
    },
    "outputs": [],
    "source": [
-    "! /bin/bash ./workspace/generate_model_trt.sh $model_id && rm -rf ensemble_hf/bert-trt/1 && mkdir -p ensemble_hf/bert-trt/1 && cp ./model.plan ensemble_hf/bert-trt/1/model.plan && rm -rf ./model.plan ./conversion_bs16_dy.txt ./model.onnx"
+    "if RUNNING_IN_STUDIO:\n",
+    "    !/bin/bash ./workspace/generate_model_trt.sh $model_id && rm -rf ensemble_hf/bert-trt/1 && mkdir -p ensemble_hf/bert-trt/1 && cp ./model.plan ensemble_hf/bert-trt/1/model.plan && rm -rf ./model.plan ./conversion_bs16_dy.txt ./model.onnx"
    ]
   },
   {
@@ -350,7 +361,8 @@
    },
    "outputs": [],
    "source": [
-    "!docker run --gpus=all -d --shm-size=4G --rm -p8000:8000 -p8001:8001 -p8002:8002 -v$(pwd)/ensemble_hf:/model_repository nvcr.io/nvidia/tritonserver:21.08-py3 tritonserver --model-repository=/model_repository --exit-on-error=false --strict-model-config=false"
+    "!docker run --gpus=all -d --shm-size=4G --rm -p8000:8000 -p8001:8001 -p8002:8002 -v$(pwd)/ensemble_hf:/model_repository nvcr.io/nvidia/tritonserver:21.08-py3 tritonserver --model-repository=/model_repository --exit-on-error=false --strict-model-config=false\n",
+    "time.sleep(20)"
    ]
   },
   {
@@ -368,15 +380,16 @@
    },
    "outputs": [],
    "source": [
-    "!docker container ls"
+    "CONTAINER_ID=!docker container ls -q\n",
+    "FIRST_CONTAINER_ID = CONTAINER_ID[0]"
    ]
   },
   {
    "cell_type": "markdown",
    "id": "3f903432-6e87-449c-84ea-4c3ed2fa445b",
    "metadata": {},
    "source": [
-    "Paste the container id obtained from the cell above, replacing it by `container_id` in the next cell."
+    "Uncomment the next cell and run it to view the container logs and understand Triton model loading."
    ]
   },
   {
@@ -394,7 +407,7 @@
    },
    "outputs": [],
    "source": [
-    "! docker logs <container_id> -f"
+    "# !docker logs $FIRST_CONTAINER_ID -f"
    ]
   },
   {
@@ -523,24 +536,6 @@
     "    print(output)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "94cc9b2e",
-   "metadata": {
-    "collapsed": false,
-    "jupyter": {
-     "outputs_hidden": false
-    },
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "!docker container ls"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -557,7 +552,7 @@
    "outputs": [],
    "source": [
     "# Use this to stop the container that was started in detached mode\n",
-    "! docker kill <container_id>"
+    "!docker kill $FIRST_CONTAINER_ID"
    ]
   },
   {