From fa374c27af33d80f2fde476cad784cf8ecd28c08 Mon Sep 17 00:00:00 2001
From: Yue Zhu <16687552+yuezhu1@users.noreply.github.com>
Date: Fri, 10 Oct 2025 15:04:08 +0000
Subject: [PATCH] enable --enforce-concurrent-user-cnt

---
 synthetic-multi-round-qa/multi-round-qa.py | 30 ++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/synthetic-multi-round-qa/multi-round-qa.py b/synthetic-multi-round-qa/multi-round-qa.py
index fb3c4c6..8010610 100644
--- a/synthetic-multi-round-qa/multi-round-qa.py
+++ b/synthetic-multi-round-qa/multi-round-qa.py
@@ -42,6 +42,9 @@ class WorkloadConfig:
     # Whether to include user id in request header
     enable_user_id: bool
 
+    # Whether strictly cap active sessions at num_users
+    enforce_strict_concurrent_users: bool = False
+
 
 @dataclass
 class UserConfig:
@@ -457,6 +460,10 @@ def __init__(
         if self.use_sharegpt:
             self._load_sharegpt_data()
 
+        self.enforce_strict_concurrent_users = (
+            workload_config.enforce_strict_concurrent_users
+        )
+
     def _load_sharegpt_data(self):
         with open("ShareGPT.json", "r", encoding="utf-8") as file:
             self.sharegpt_data = json.load(file)
@@ -499,6 +506,19 @@ def _remove_finished_sessions(self):
                 self.session_summaries.append(session.summary())
         self.sessions = [s for s in self.sessions if not s.finished]
 
+    def _can_join_user(self, timestamp: float) -> bool:
+        # No new user session if gap_between_users time interval not meets
+        if timestamp - self.last_user_join <= self.gap_between_users:
+            return False
+
+        # No user seession if active user count is less than configured
+        if (
+            self.enforce_strict_concurrent_users
+            and len(self.sessions) >= self.workload_config.num_users
+        ):
+            return False
+        return True
+
     def step(self, timestamp: float, executor: RequestExecutor):
         if self.need_ramp_up:
             self._ramp_up(timestamp, self.ramp_up_time)
@@ -506,8 +526,8 @@ def step(self, timestamp: float, executor: RequestExecutor):
         if self.start_time is None:
             self.start_time = timestamp
         
-        # New user session only joins when meets time interval and active user count is less than configured
-        if timestamp - self.last_user_join > self.gap_between_users and len(self.sessions) < self.workload_config.num_users:
+        # Check if can join new user session
+        if self._can_join_user(timestamp):
             new_session = self._create_user_session()
             if new_session is not None:
                 self.last_user_join = timestamp
@@ -704,6 +724,11 @@ def parse_arguments() -> WorkloadConfig:
     parser.add_argument(
         "--sharegpt", action="store_true", help="Whether to use ShareGPT dataset"
     )
+    parser.add_argument(
+        "--enforce-strict-concurrent-users",
+        action="store_true",
+        help="Strictly enforce concurrent users count to match --num-users",
+    )
     args = parser.parse_args()
     return args
 
@@ -750,6 +775,7 @@ def main():
         qps=args.qps,
         model=args.model,
         enable_user_id=args.request_with_user_id,
+        enforce_strict_concurrent_users=args.enforce_strict_concurrent_users,
     )
 
     manager = UserSessionManager(