[CI][IB] Benchmark TGI Models (#1529)

.github/workflows/instant_benchmark.yml

@@ -136,7 +136,7 @@ jobs:
         working-directory: tests/integration
         run: |
           docker rm -f $(docker ps -aq) || true
-          cat logs/serving.log
+          cat logs/serving.log || true
       - name: Upload test artifacts
         uses: actions/upload-artifact@v3
         if: always()

tests/integration/instant_benchmark.py

@@ -106,6 +106,7 @@ def parse_raw_template(url, override_container):
     name = ''
     container = None
     properties = []
+    env = []
     commandline = []
     requirements = []
     vars = []
@@ -123,6 +124,12 @@ def parse_raw_template(url, override_container):
                     lines[iterator]):
                 properties.append(lines[iterator])
                 iterator += 1
+        elif '[env]' == lines[iterator]:
+            iterator += 1
+            while iterator < len(lines) and not is_square_bracket(
+                    lines[iterator]):
+                env.append(lines[iterator])
+                iterator += 1
         elif '[requirements]' == lines[iterator]:
             iterator += 1
             while iterator < len(lines) and not is_square_bracket(
@@ -150,9 +157,10 @@ def parse_raw_template(url, override_container):
                 iterator += 1
         else:
             iterator += 1
-        if name and properties and commandline:
+        if name and commandline:
             cur_result = {
                 "properties": properties,
+                "env": env,
                 "awscurl": ' '.join(commandline),
                 "requirements": requirements
             }
@@ -172,23 +180,28 @@ def parse_raw_template(url, override_container):
             name = ''
             container = None
             properties = []
+            env = []
             commandline = []
             requirements = []
             vars = []
             info = None
     return final_result
 
 
-def write_model_artifacts(properties, requirements=None):
+def write_model_artifacts(properties, requirements=None, env=None):
     model_path = "models/test"
     if os.path.exists(model_path):
         shutil.rmtree(model_path)
     os.makedirs(model_path, exist_ok=True)
-    with open(os.path.join(model_path, "serving.properties"), "w") as f:
-        f.write('\n'.join(properties) + '\n')
+    if properties and len(properties) > 0:
+        with open(os.path.join(model_path, "serving.properties"), "w") as f:
+            f.write('\n'.join(properties) + '\n')
     if requirements:
         with open(os.path.join(model_path, "requirements.txt"), "w") as f:
             f.write('\n'.join(requirements) + '\n')
+    if env and len(env) > 0:
+        with open(os.path.join(model_path, "docker_env"), "w") as f:
+            f.write('\n'.join(env) + '\n')
 
 
 def machine_translation(machine_name: str):
@@ -213,7 +226,7 @@ def build_running_script(template, job, instance):
     job_template['awscurl'] = bytes.fromhex(
         job_template['awscurl']).decode("utf-8")
     write_model_artifacts(job_template['properties'],
-                          job_template['requirements'])
+                          job_template['requirements'], job_template['env'])
 
     container = job_template['container']
 

tests/integration/launch_container.sh

@@ -72,6 +72,18 @@ if $is_partition; then
     ${args}
 
   exit 0
+elif [[ "$docker_image" == *"text-generation-inference"* ]]; then
+  container_id=$(docker run \
+    -itd \
+    --rm \
+    -p 8080:80 \
+    ${model_path:+-v ${model_path}:/opt/ml/model:ro} \
+    -v ~/sagemaker_infra/:/opt/ml/.sagemaker_infra/:ro \
+    ${env_file} \
+    ${runtime:+--runtime="${runtime}"} \
+    ${shm:+--shm-size="${shm}"} \
+    "${docker_image}" \
+    ${args})
 else
   echo "$(whoami), UID: $UID"
   if [[ "$UID" == "1000" ]]; then

tests/integration/record_benchmark.py

@@ -145,6 +145,11 @@ def data_container():
                 data["image"] = split[1]
             else:
                 data["image"] = "cpu"
+        if "text-generation-inference" in container:
+            data["modelServer"] = "TGI"
+            version = container.split(":")[1]
+            if not version.startswith("sha"):
+                data["tgiVersion"] = version
 
 
 def data_from_model_files():
@@ -155,7 +160,7 @@ def data_from_model_files():
                 properties = {}
                 for line in f.readlines():
                     line = line.strip()
-                    if line[0] == "#":
+                    if len(line) == 0 or line[0] == "#":
                         continue
                     if "=" in line:
                         split = line.split("=", 1)
@@ -177,13 +182,28 @@ def data_from_model_files():
                 req = {}
                 for line in f.readlines():
                     line = line.strip()
-                    if line[0] == "#":
+                    if len(line) == 0 or line[0] == "#":
                         continue
                     if "=" in line:
                         split = line.split("=", 1)
                         req[split[0]] = split[1]
                 data["requirements_txt"] = req
 
+        envPath = os.path.join(args.model, "docker_env")
+        if os.path.isfile(envPath):
+            with open(envPath, "r") as f:
+                envs = {}
+                for line in f.readlines():
+                    line = line.strip()
+                    if len(line) == 0 or line[0] == "#":
+                        continue
+                    if "=" in line:
+                        split = line.split("=", 1)
+                        envs[split[0]] = split[1]
+                data["env"] = envs
+                if "modelId" not in data and "MODEL_ID" in envs:
+                    data["modelId"] = envs["MODEL_ID"]
+
 
 def data_from_template():
     if args.template: