When max_live_batches <= 2, don't prepare the next

batch until the only previous one is done.

PiperOrigin-RevId: 620893396
Change-Id: I917aea93b016eb4eaee70084929c45d83987a1b9

saxml/server/model_service_base.py

@@ -431,8 +431,12 @@ def register_method(
 
     # Start the batching loop.
     def _batching():
-      # Keeps at most 2 active batches in the rest of pipeline.
-      batch_sem = threading.Semaphore(value=2)
+      # Keeps 1 or 2 active batches in the rest of pipeline.
+      if max_live_batches <= 2:
+        sem_limit = 1
+      else:
+        sem_limit = 2
+      batch_sem = threading.Semaphore(value=sem_limit)
 
       def _finish_batch():
         batch_sem.release()