1212 "ProfileGenerationModes" ,
1313 "Profile" ,
1414 "ProfileGenerator" ,
15- "MultiProfileGenerator " ,
15+ "FixedRateProfileGenerator " ,
1616 "SweepProfileGenerator" ,
1717]
1818
1919
2020class ProfileGenerationModes (Enum ):
21- MULTI = "multi "
21+ FIXED_RATE = "fixed_rate "
2222 SWEEP = "sweep"
2323
2424
@@ -27,6 +27,12 @@ class Profile:
2727 load_gen_mode : LoadGenerationModes
2828 load_gen_rate : Optional [float ]
2929
30+ RateTypeLoadGenModeMap = {
31+ "constant" : LoadGenerationModes .CONSTANT ,
32+ "synchronous" : LoadGenerationModes .SYNCHRONOUS ,
33+ "poisson" : LoadGenerationModes .POISSON ,
34+ }
35+
3036
3137class ProfileGenerator (ABC ):
3238 _registry = {}
@@ -61,10 +67,12 @@ def next_profile(
6167 pass
6268
6369
64- @ProfileGenerator .register_generator (ProfileGenerationModes .MULTI )
65- class MultiProfileGenerator (ProfileGenerator ):
66- def __init__ (self , rate : List [float ], rate_type : str , ** kwargs ):
67- super ().__init__ (ProfileGenerationModes .MULTI )
70+ @ProfileGenerator .register_generator (ProfileGenerationModes .FIXED_RATE )
71+ class FixedRateProfileGenerator (ProfileGenerator ):
72+ def __init__ (self , rate : List [float ], rate_type : LoadGenerationModes , ** kwargs ):
73+ super ().__init__ (ProfileGenerationModes .FIXED_RATE )
74+ if rate_type == "synchronous" and rate .length > 0 :
75+ raise ValueError ("custom rates are not supported in synchronous mode" )
6876 self ._rates = rate
6977 self ._rate_index = 0
7078 self ._rate_type = rate_type
@@ -75,22 +83,18 @@ def next_profile(
7583 ) -> Optional [Profile ]:
7684 if self ._rate_index >= len (self ._rates ):
7785 return None
78-
79- if self ._rate_type == "constant" :
80- return Profile (
81- load_gen_mode = LoadGenerationModes .CONSTANT , load_gen_rate = self ._rates [self ._rate_index ]
82- )
83-
86+ current_rate = self ._rates [self ._rate_index ]
87+ self ._rate_index += 1
88+
8489 if self ._rate_type == "synchronous" :
8590 return Profile (
8691 load_gen_mode = LoadGenerationModes .SYNCHRONOUS , load_gen_rate = None
8792 )
88-
89- if self . _rate_type == "poisson" :
93+ if self . _rate_type in { "constant" , "poisson" }:
94+ load_gen_mode = RateTypeLoadGenModeMap [ self . _rate_type ]
9095 return Profile (
91- load_gen_mode = LoadGenerationModes . POISSON , load_gen_rate = self . _rates [ self . _rate_index ]
96+ load_gen_mode = load_gen_mode , load_gen_rate = current_rate
9297 )
93-
9498 raise ValueError (f"Invalid rate type: { self ._rate_type } " )
9599
96100
0 commit comments