|
1 | 1 | from abc import ABC, abstractmethod |
2 | 2 | from dataclasses import dataclass |
3 | 3 | from enum import Enum |
4 | | -from typing import Dict, Optional, Type, Union |
| 4 | +from typing import Dict, List, Optional, Type, Union |
5 | 5 |
|
6 | 6 | import numpy |
7 | 7 |
|
|
12 | 12 | "ProfileGenerationMode", |
13 | 13 | "Profile", |
14 | 14 | "ProfileGenerator", |
15 | | - "SingleProfileGenerator", |
| 15 | + "FixedRateProfileGenerator", |
16 | 16 | "SweepProfileGenerator", |
17 | 17 | ] |
18 | 18 |
|
| 19 | +rate_type_to_load_gen_mode = { |
| 20 | + "synchronous": LoadGenerationMode.SYNCHRONOUS, |
| 21 | + "constant": LoadGenerationMode.CONSTANT, |
| 22 | + "poisson": LoadGenerationMode.POISSON, |
| 23 | +} |
| 24 | + |
19 | 25 |
|
20 | 26 | class ProfileGenerationMode(Enum): |
21 | | - SINGLE = "single" |
| 27 | + FIXED_RATE = "fixed_rate" |
22 | 28 | SWEEP = "sweep" |
23 | 29 |
|
24 | 30 |
|
| 31 | +rate_type_to_profile_mode = { |
| 32 | + "synchronous": ProfileGenerationMode.FIXED_RATE, |
| 33 | + "constant": ProfileGenerationMode.FIXED_RATE, |
| 34 | + "poisson": ProfileGenerationMode.FIXED_RATE, |
| 35 | + "sweep": ProfileGenerationMode.SWEEP, |
| 36 | +} |
| 37 | + |
| 38 | + |
25 | 39 | @dataclass |
26 | 40 | class Profile: |
27 | 41 | load_gen_mode: LoadGenerationMode |
@@ -55,34 +69,44 @@ def next(self, current_report: TextGenerationBenchmarkReport) -> Optional[Profil |
55 | 69 | pass |
56 | 70 |
|
57 | 71 |
|
58 | | -@ProfileGenerator.register(ProfileGenerationMode.SINGLE) |
59 | | -class SingleProfileGenerator(ProfileGenerator): |
60 | | - def __init__(self, rate: float, rate_type: LoadGenerationMode): |
61 | | - super().__init__(ProfileGenerationMode.SINGLE) |
62 | | - self._rate: float = rate |
63 | | - self._rate_type: LoadGenerationMode = rate_type |
| 72 | +@ProfileGenerator.register(ProfileGenerationMode.FIXED_RATE) |
| 73 | +class FixedRateProfileGenerator(ProfileGenerator): |
| 74 | + def __init__( |
| 75 | + self, |
| 76 | + load_gen_mode: Optional[LoadGenerationMode], |
| 77 | + rates: Optional[List[float]] = None, |
| 78 | + **kwargs, |
| 79 | + ): |
| 80 | + super().__init__(ProfileGenerationMode.FIXED_RATE) |
| 81 | + if load_gen_mode == LoadGenerationMode.SYNCHRONOUS and rates and len(rates) > 0: |
| 82 | + raise ValueError("custom rates are not supported in synchronous mode") |
| 83 | + self._rates: Optional[List[float]] = rates |
| 84 | + self._load_gen_mode = load_gen_mode |
64 | 85 | self._generated: bool = False |
| 86 | + self._rate_index: int = 0 |
65 | 87 |
|
66 | 88 | def next(self, current_report: TextGenerationBenchmarkReport) -> Optional[Profile]: |
67 | | - if self._generated: |
68 | | - return None |
69 | | - |
70 | | - self._generated = True |
71 | | - |
72 | | - if self._rate_type == LoadGenerationMode.CONSTANT: |
73 | | - return Profile( |
74 | | - load_gen_mode=LoadGenerationMode.CONSTANT, load_gen_rate=self._rate |
75 | | - ) |
76 | | - elif self._rate_type == LoadGenerationMode.SYNCHRONOUS: |
| 89 | + if self._load_gen_mode == LoadGenerationMode.SYNCHRONOUS: |
| 90 | + if self._generated: |
| 91 | + return None |
| 92 | + self._generated = True |
77 | 93 | return Profile( |
78 | 94 | load_gen_mode=LoadGenerationMode.SYNCHRONOUS, load_gen_rate=None |
79 | 95 | ) |
80 | | - elif self._rate_type == LoadGenerationMode.POISSON: |
81 | | - return Profile( |
82 | | - load_gen_mode=LoadGenerationMode.POISSON, load_gen_rate=self._rate |
83 | | - ) |
| 96 | + elif self._load_gen_mode in { |
| 97 | + LoadGenerationMode.CONSTANT, |
| 98 | + LoadGenerationMode.POISSON, |
| 99 | + }: |
| 100 | + if self._rates: |
| 101 | + if self._rate_index >= len(self._rates): |
| 102 | + return None |
| 103 | + current_rate = self._rates[self._rate_index] |
| 104 | + self._rate_index += 1 |
| 105 | + return Profile( |
| 106 | + load_gen_mode=self._load_gen_mode, load_gen_rate=current_rate |
| 107 | + ) |
84 | 108 |
|
85 | | - raise ValueError(f"Invalid rate type: {self._rate_type}") |
| 109 | + raise ValueError(f"Invalid rate type: {self._load_gen_mode}") |
86 | 110 |
|
87 | 111 |
|
88 | 112 | @ProfileGenerator.register(ProfileGenerationMode.SWEEP) |
|
0 commit comments