From accf70ea453487df46dd0a1408b4b2911fca1ea3 Mon Sep 17 00:00:00 2001
From: Sherif Akoush <sa@seldon.io>
Date: Fri, 21 Jul 2023 20:46:34 +0100
Subject: [PATCH] add replica

---
 tests/k6/components/settings.js | 8 ++++++++
 tests/k6/components/utils.js    | 2 +-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/tests/k6/components/settings.js b/tests/k6/components/settings.js
index e7f8689b96..dae5ff41cc 100644
--- a/tests/k6/components/settings.js
+++ b/tests/k6/components/settings.js
@@ -110,6 +110,13 @@ function inferBatchSize() {
     return [1]
 }
 
+function modelReplicas() {
+    if (__ENV.MODEL_NUM_REPLICAS) {
+        return __ENV.MODEL_NUM_REPLICAS.split(",").map( s => parseInt(s))
+    }
+    return [1]
+}
+
 function modelStartIdx() {
     if (__ENV.MODEL_START_IDX) {
         return parseInt(__ENV.MODEL_START_IDX)
@@ -217,5 +224,6 @@ export function getConfig() {
         "doWarmup": doWarmup(),
         "requestRate": requestRate(),
         "constantRateDurationSeconds": constantRateDurationSeconds(),
+        "modelReplicas": modelReplicas(),
     }
 }
diff --git a/tests/k6/components/utils.js b/tests/k6/components/utils.js
index 25a2a4963b..d938013c26 100644
--- a/tests/k6/components/utils.js
+++ b/tests/k6/components/utils.js
@@ -22,7 +22,7 @@ export function setupBase(config ) {
         for (let j = 0; j < config.maxNumModels.length; j++) {        
             for (let i = 0; i < config.maxNumModels[j]; i++) {
                 const modelName = config.modelNamePrefix[j] + i.toString()
-                const model = generateModel(config.modelType[j], modelName, 1, 1, config.isSchedulerProxy, config.modelMemoryBytes[j], config.inferBatchSize[j])
+                const model = generateModel(config.modelType[j], modelName, 1, config.modelReplicas[j], config.isSchedulerProxy, config.modelMemoryBytes[j], config.inferBatchSize[j])
                 const modelDefn = model.modelDefn
                 const pipelineDefn = model.pipelineDefn