refactor to better handle creating profile generator

neuralmagic · Jul 2, 2024 · fb1ebcb · fb1ebcb
1 parent 9a0c582
commit fb1ebcb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/src/guidellm/executor/executor.py b/src/guidellm/executor/executor.py
@@ -14,13 +14,16 @@ def __init__(
         self,
         request_generator: RequestGenerator,
         backend: Backend,
-        profile_mode: Union[str, ProfileGenerationModes] = "fixed_rate",
+        rate_type: str = "sweep",
         profile_args: Optional[Dict[str, Any]] = None,
         max_requests: Optional[int] = None,
         max_duration: Optional[float] = None,
     ):
         self.request_generator = request_generator
         self.backend = backend
+        profile_mode = "sweep"
+        if rate_type in {"synchronous", "constant", "poisson"}:
+            profile_mode = "fixed_rate"
         self.profile = ProfileGenerator.create_generator(
             profile_mode, **(profile_args or {})
         )

diff --git a/src/guidellm/main.py b/src/guidellm/main.py
@@ -107,11 +107,10 @@ def main(
         raise ValueError(f"Unknown data type: {data_type}")
 
     # Create executor
-    profile_mode = rate_type if rate_type == "sweep" else "fixed_rate"
     executor = Executor(
         request_generator=request_generator,
         backend=backend,
-        profile_mode=profile_mode,
+        rate_type=rate_type,
         profile_args={"rate_type": rate_type, "rate": rate},
         max_requests=num_requests,
         max_duration=num_seconds,