refactor setting of profile mode, load_gen_mode

neuralmagic · Jul 2, 2024 · 45f6a1b · 45f6a1b
1 parent bfcc329
commit 45f6a1b
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 21 deletions.
diff --git a/src/guidellm/executor/__init__.py b/src/guidellm/executor/__init__.py
@@ -3,7 +3,7 @@
     Profile,
     ProfileGenerationModes,
     ProfileGenerator,
-    MultiProfileGenerator,
+    FixedRateProfileGenerator,
     SweepProfileGenerator,
 )
 
@@ -12,6 +12,6 @@
     "ProfileGenerationModes",
     "Profile",
     "ProfileGenerator",
-    "MultiProfileGenerator",
+    "FixedRateProfileGenerator",
     "SweepProfileGenerator",
 ]
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] = "multi",
+        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/executor/profile_generator.py b/src/guidellm/executor/profile_generator.py
@@ -12,13 +12,13 @@
     "ProfileGenerationModes",
     "Profile",
     "ProfileGenerator",
-    "MultiProfileGenerator",
+    "FixedRateProfileGenerator",
     "SweepProfileGenerator",
 ]
 
 
 class ProfileGenerationModes(Enum):
-    MULTI = "multi"
+    FIXED_RATE = "fixed_rate"
     SWEEP = "sweep"
 
 
@@ -27,6 +27,12 @@ class Profile:
     load_gen_mode: LoadGenerationModes
     load_gen_rate: Optional[float]
 
+RateTypeLoadGenModeMap = {
+    "constant": LoadGenerationModes.CONSTANT,
+    "synchronous": LoadGenerationModes.SYNCHRONOUS,
+    "poisson": LoadGenerationModes.POISSON,
+}
+
 
 class ProfileGenerator(ABC):
     _registry = {}
@@ -61,10 +67,12 @@ def next_profile(
         pass
 
 
-@ProfileGenerator.register_generator(ProfileGenerationModes.MULTI)
-class MultiProfileGenerator(ProfileGenerator):
-    def __init__(self, rate: List[float], rate_type: str, **kwargs):
-        super().__init__(ProfileGenerationModes.MULTI)
+@ProfileGenerator.register_generator(ProfileGenerationModes.FIXED_RATE)
+class FixedRateProfileGenerator(ProfileGenerator):
+    def __init__(self, rate: List[float], rate_type: LoadGenerationModes, **kwargs):
+        super().__init__(ProfileGenerationModes.FIXED_RATE)
+        if rate_type == "synchronous" and rate.length > 0:
+            raise ValueError("custom rates are not supported in synchronous mode")
         self._rates = rate
         self._rate_index = 0
         self._rate_type = rate_type
@@ -75,22 +83,18 @@ def next_profile(
     ) -> Optional[Profile]:
         if self._rate_index >= len(self._rates):
             return None
-
-        if self._rate_type == "constant":
-            return Profile(
-                load_gen_mode=LoadGenerationModes.CONSTANT, load_gen_rate=self._rates[self._rate_index]
-            )
-
+        current_rate = self._rates[self._rate_index]
+        self._rate_index += 1
+
         if self._rate_type == "synchronous":
             return Profile(
                 load_gen_mode=LoadGenerationModes.SYNCHRONOUS, load_gen_rate=None
             )
-
-        if self._rate_type == "poisson":
+        if self._rate_type in {"constant", "poisson"}:
+            load_gen_mode = RateTypeLoadGenModeMap[self._rate_type]
             return Profile(
-                load_gen_mode=LoadGenerationModes.POISSON, load_gen_rate=self._rates[self._rate_index]
+                load_gen_mode=load_gen_mode, load_gen_rate=current_rate
             )
-
         raise ValueError(f"Invalid rate type: {self._rate_type}")
 
 

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 "multi"
     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,