From cd1e2cd2cf171afe9303bbb0e3d0402e26cd77d5 Mon Sep 17 00:00:00 2001
From: Xuan Son Nguyen <son@huggingface.co>
Date: Fri, 9 Aug 2024 13:16:49 +0200
Subject: [PATCH] add docker load_tests

---
 load_tests/common.js          |  7 +++++--
 load_tests/docker-compose.yml | 12 ++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)
 create mode 100644 load_tests/docker-compose.yml

diff --git a/load_tests/common.js b/load_tests/common.js
index e0a10595..1c451818 100644
--- a/load_tests/common.js
+++ b/load_tests/common.js
@@ -36,7 +36,7 @@ export function get_options() {
             load_test: {
                 executor: 'constant-arrival-rate',
                 duration: '60s',
-                preAllocatedVUs: 100,
+                preAllocatedVUs: 50, // not enough RAM for 100 VUs
                 rate: 1,
                 timeUnit: '1s',
             },
@@ -68,13 +68,16 @@ export default function run() {
     const headers = { 'Content-Type': 'application/json' };
     const query = shareGPT[scenario.iterationInTest % shareGPT.length];
     const payload = JSON.stringify(generate_payload(query, max_new_tokens));
-    const res = http.post(`http://${host}/v1/chat/completions`, payload, {
+    const res = http.post(`https://${host}/v1/chat/completions`, payload, {
         headers,
     });
     if (res.status >= 400 && res.status < 500) {
         return;
     }
 
+    if (res.status !== 200) {
+        console.error(res.body);
+    }
 
     check(res, {
         'Post status is 200': (res) => res.status === 200,
diff --git a/load_tests/docker-compose.yml b/load_tests/docker-compose.yml
new file mode 100644
index 00000000..edac91e0
--- /dev/null
+++ b/load_tests/docker-compose.yml
@@ -0,0 +1,12 @@
+services:
+
+  # How to run:
+  # HOST=..... docker compose up
+  tgi_load_test:
+    image: grafana/k6
+    command: run /load_tests/common.js
+    environment:
+      - HOST=${HOST}
+      - MODEL_ID=${MODEL_ID}
+    volumes:
+      - ./:/load_tests:Z