Bert base batch test

frontend/server/src/main/java/org/pytorch/serve/wlm/BatchAggregator.java

@@ -63,12 +63,13 @@ public void sendResponse(ModelWorkerResponse message) {
                 // this is from initial load.
                 return;
             }
-
             for (Predictions prediction : message.getPredictions()) {
                 String jobId = prediction.getRequestId();
-                Job job = jobs.remove(jobId);
+                Job job = jobs.get(jobId);
+
                 if (job == null) {
-                    throw new IllegalStateException("Unexpected job: " + jobId);
+                    throw new IllegalStateException(
+                            "Unexpected job in sendResponse() with 200 status code: " + jobId);
                 }
                 job.response(
                         prediction.getResp(),
@@ -77,18 +78,19 @@ public void sendResponse(ModelWorkerResponse message) {
                         prediction.getReasonPhrase(),
                         prediction.getHeaders());
             }
+
         } else {
-            for (String reqId : jobs.keySet()) {
-                Job j = jobs.remove(reqId);
-                if (j == null) {
-                    throw new IllegalStateException("Unexpected job: " + reqId);
+            for (Map.Entry<String, Job> j : jobs.entrySet()) {
+
+                if (j.getValue() == null) {
+                    throw new IllegalStateException(
+                            "Unexpected job in sendResponse() with non 200 status code: "
+                                    + j.getKey());
                 }
-                j.sendError(message.getCode(), message.getMessage());
-            }
-            if (!jobs.isEmpty()) {
-                throw new IllegalStateException("Not all jobs get response.");
+                j.getValue().sendError(message.getCode(), message.getMessage());
             }
         }
+        jobs.clear();
     }
 
     public void sendError(BaseModelRequest message, String error, int status) {
@@ -103,20 +105,20 @@ public void sendError(BaseModelRequest message, String error, int status) {
                 String requestId = req.getRequestId();
                 Job job = jobs.remove(requestId);
                 if (job == null) {
-                    logger.error("Unexpected job: " + requestId);
+                    logger.error("Unexpected job in sendError(): " + requestId);
                 } else {
                     job.sendError(status, error);
                 }
             }
             if (!jobs.isEmpty()) {
                 jobs.clear();
-                logger.error("Not all jobs get response.");
+                logger.error("Not all jobs got an error response.");
             }
         } else {
             // Send the error message to all the jobs
             for (Map.Entry<String, Job> j : jobs.entrySet()) {
                 String jobsId = j.getValue().getJobId();
-                Job job = jobs.remove(jobsId);
+                Job job = jobs.get(jobsId);
 
                 if (job.isControlCmd()) {
                     job.sendError(status, error);
@@ -127,5 +129,6 @@ public void sendError(BaseModelRequest message, String error, int status) {
                 }
             }
         }
+        jobs.clear();
     }
 }

frontend/server/src/test/java/org/pytorch/serve/TestUtils.java

@@ -174,6 +174,39 @@ public static void registerModel(
         }
     }
 
+    public static void registerModel(
+            Channel channel,
+            String url,
+            String modelName,
+            boolean withInitialWorkers,
+            boolean syncChannel,
+            int batchSize,
+            int maxBatchDelay)
+            throws InterruptedException {
+        String requestURL =
+                "/models?url="
+                        + url
+                        + "&model_name="
+                        + modelName
+                        + "&runtime=python"
+                        + "&batch_size="
+                        + batchSize
+                        + "&max_batch_delay="
+                        + maxBatchDelay;
+        if (withInitialWorkers) {
+            requestURL += "&initial_workers=1&synchronous=true";
+        }
+
+        HttpRequest req =
+                new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.POST, requestURL);
+        if (syncChannel) {
+            channel.writeAndFlush(req).sync();
+            channel.closeFuture().sync();
+        } else {
+            channel.writeAndFlush(req);
+        }
+    }
+
     public static void registerWorkflow(
             Channel channel, String url, String workflowName, boolean syncChannel)
             throws InterruptedException {

test/pytest/test_handler.py

@@ -224,7 +224,30 @@ def test_kfserving_mnist_model_register_and_inference_on_valid_model_explain():
 
     assert np.array(json.loads(response.content)['explanations']).shape == (1, 1, 28, 28)
     test_utils.unregister_model("mnist")
+
+def test_huggingface_bert_batch_inference():
+    batch_size = 2
+    batch_delay = 10000 # 10 seconds
+    params = (
+        ('model_name', 'BERTSeqClassification'),
+        ('url', 'https://torchserve.pytorch.org/mar_files/BERTSeqClassification.mar'),
+        ('initial_workers', '1'),
+        ('batch_size', str(batch_size)),
+        ('max_batch_delay', str(batch_delay))
+    )
+    test_utils.start_torchserve(no_config_snapshots=True)
+    test_utils.register_model_with_params(params)
+    input_text = os.path.join(REPO_ROOT, 'examples/Huggingface_Transformers/Seq_classification_artifacts/sample_text.txt')
 
+    # Make 2 curl requests in parallel with &
+    # curl --header \"X-Forwarded-For: 1.2.3.4\" won't work since you can't access local host anymore
+    response = os.popen(f"curl http://127.0.0.1:8080/predictions/BERTSeqClassification -T {input_text} & curl http://127.0.0.1:8080/predictions/BERTSeqClassification -T {input_text}")
+    response = response.read()
+
+
+    ## Assert that 2 responses are returned from the same batch
+    assert response == 'Not AcceptedNot Accepted'
+    test_utils.unregister_model('BERTSeqClassification')
 
 def test_MMF_activity_recognition_model_register_and_inference_on_valid_model():
 
@@ -245,4 +268,4 @@ def test_MMF_activity_recognition_model_register_and_inference_on_valid_model():
     response = ast.literal_eval(response)
     response = [n.strip() for n in response]
     assert response == ['Sitting at a table','Someone is sneezing','Watching a laptop or something on a laptop']
-    test_utils.unregister_model("MMF_activity_recognition_v2")
+    test_utils.unregister_model("MMF_activity_recognition_v2")