feat: multi rate generator (#297)

In order to test scale_in properly, we need a way to dynamically increase or decrease sending rate. For that, we added multi rate generator schedule. The key is the replay on which the rate update should happen and the value is the new rate. Added controller configuration as example on how this can be used.

Author: raresgaia123

examples/configs/controller_multi_rate.yaml

@@ -0,0 +1,22 @@
+---
+# api_port: 8080
+# ctrl_port: 31310
+autoscale: true
+deploy_policy: static # options: even, packing, random (default), static
+                      # if autoscale is enabled, deploy_policy is static
+
+# a directory path to job deployment templates/plans
+# pick a directory where plan files are located
+job_plans: ~/projects/infscale/examples/configs/plans
+
+reqgen:
+  sort: multirate_exponential
+  params:
+    in_memory: true
+    replay: 20
+    rate: 1600.0
+    schedule:
+      - replay_index: 3
+        rate: 1800.0
+      - replay_index: 10
+        rate: 300.0

infscale/common/exceptions.py

@@ -43,6 +43,14 @@ def __init__(self, err_msg: str):
         super().__init__(err_msg)
 
 
+class InvalidGenConfig(InfScaleException):
+    """Exception for invalid generator configuration."""
+
+    def __init__(self, err_msg: str):
+        """Initialize InvalidGenConfig exception instance."""
+        super().__init__(err_msg)
+
+
 class InsufficientResources(InfScaleException):
     """Exception for insufficient agent resources."""
 

infscale/configs/controller.py

@@ -20,10 +20,8 @@
 from dataclasses import dataclass, field
 from enum import Enum
 
-from infscale.common.constants import (
-    APISERVER_PORT,
-    CONTROLLER_PORT,
-)
+from infscale.common.constants import APISERVER_PORT, CONTROLLER_PORT
+from infscale.common.exceptions import InvalidGenConfig
 
 
 class DeploymentPolicyEnum(Enum):
@@ -47,6 +45,7 @@ class ReqGenEnum(str, Enum):
 
     DEFAULT = "default"
     EXP = "exponential"
+    MULTIRATE_EXP = "multirate_exponential"
 
 
 @dataclass
@@ -67,7 +66,20 @@ class ExponentialParams(DefaultParams):
     rate: float = 1.0  # rate is per-second
 
 
-GenParams = DefaultParams | ExponentialParams
+@dataclass
+class RateScheduleItem:
+    replay_index: int
+    rate: float
+
+
+@dataclass
+class MultiRateExponentialParams(ExponentialParams):
+    """Multi-rate exponential generator parameters."""
+
+    schedule: list[RateScheduleItem] = field(default_factory=list)
+
+
+GenParams = DefaultParams | ExponentialParams | MultiRateExponentialParams
 
 
 @dataclass
@@ -94,6 +106,32 @@ def __post_init__(self):
             case ReqGenEnum.EXP:
                 self.params = ExponentialParams(**self.params)
 
+            case ReqGenEnum.MULTIRATE_EXP:
+                self.params = MultiRateExponentialParams(**self.params)
+
+                if self.params.replay is None:
+                    raise InvalidGenConfig(f"Replay param is required.")
+
+                if len(self.params.schedule) > 0:
+                    self.params.schedule = [
+                        RateScheduleItem(**item) if isinstance(item, dict) else item
+                        for item in self.params.schedule
+                    ]
+                    replay_indexes = [
+                        item.replay_index for item in self.params.schedule
+                    ]
+                    min_key = min(replay_indexes)
+                    max_key = max(replay_indexes)
+
+                    if min_key <= 0:
+                        raise InvalidGenConfig(
+                            f"invalid schedule: iteration {min_key} must be positive"
+                        )
+                    if max_key > self.params.replay:
+                        msg = "invalid schedule:"
+                        msg += f" iteration {max_key} exceeds replay limit {self.params.replay}"
+                        raise InvalidGenConfig(msg)
+
 
 @dataclass
 class CtrlConfig:

infscale/request/generator.py

@@ -17,16 +17,21 @@
 """generator.py."""
 
 import asyncio
+import os
 from abc import ABC, abstractmethod
 
 import numpy as np
 from torch import Tensor
 
-from infscale.configs.controller import GenParams, ReqGenEnum
+from infscale import get_logger
+from infscale.configs.controller import GenParams, RateScheduleItem, ReqGenEnum
 from infscale.execution.metrics_collector import MetricsCollector
 from infscale.module.dataset import HuggingFaceDataset
 
 
+logger = None
+
+
 class Generator(ABC):
     """Abstact Generator class."""
 
@@ -44,6 +49,9 @@ def initialize(
         self._mc = mc
         self._seqno = 0
 
+        global logger
+        logger = get_logger(f"{os.getpid()}")
+
     @abstractmethod
     async def get(self) -> list[Tensor | None]:
         """Return generated requests as batch."""
@@ -100,9 +108,12 @@ async def _generate(self) -> None:
             self._mc.update(self._seqno)
             self._seqno += 1
 
-            iat = np.random.exponential(scale=1 / self._batch_rate)
+            iat = self._compute_iat()
             await asyncio.sleep(iat)
 
+    def _compute_iat(self):
+        return np.random.exponential(scale=1 / self._batch_rate)
+
     async def get(self) -> list[Tensor | None]:
         """Return one batch of requests.
 
@@ -122,6 +133,74 @@ async def get(self) -> list[Tensor | None]:
         return batches
 
 
+class MultiRateExponentialGenerator(ExponentialGenerator):
+    """Exponential generator with replay-dependent rate schedule."""
+
+    def initialize(
+        self,
+        dataset,
+        params,
+        batch_size,
+        mc,
+    ) -> None:
+        assert params is not None
+        # intentionally bypassing super().initialize
+        # for properly setting up queue and event and to avoid duplicating
+        # asyncio task creation for _generate method
+        Generator.initialize(self, dataset, params, batch_size, mc)
+
+        self.range_list = self._prepare_schedule(
+            self._params.rate, self._params.schedule, self._params.replay
+        )
+
+        self._range_index = 0
+        rate = self.range_list[0][2]
+        self._batch_rate = rate / self._batch_size
+
+        self._queue = asyncio.Queue()
+        self._gen_evt = asyncio.Event()
+        _ = asyncio.create_task(self._generate())
+
+        msg = f"generator initialized with rate={rate}"
+        msg += f" replay rate update schedule={self._params.schedule}"
+        logger.info(msg)
+
+    def _prepare_schedule(
+        self, base_rate: float, schedule: list[RateScheduleItem], max_replay: int
+    ) -> list[tuple[int, int, float]]:
+        """Convert replay-based schedule into continuous replay ranges."""
+        schedule_sorted = sorted(schedule, key=lambda s: s.replay_index)
+
+        rate_schedule_ranges = []
+        prev_replay = 0
+        prev_rate = base_rate
+
+        for item in schedule_sorted:
+            # range [prev_replay, item.replay_index - 1] uses prev_rate
+            rate_schedule_ranges.append((prev_replay, item.replay_index - 1, prev_rate))
+            prev_replay = item.replay_index
+            prev_rate = item.rate
+
+        # last range goes until max_replay
+        rate_schedule_ranges.append((prev_replay, max_replay, prev_rate))
+        return rate_schedule_ranges
+
+    def _compute_iat(self):
+        current_replay = self._params.replay - self._dataset._replay
+        range_info = self.range_list[self._range_index]
+
+        if not range_info[0] <= current_replay <= range_info[1]:
+            self._range_index += 1
+
+            range_info = self.range_list[self._range_index]
+            rate = range_info[2]
+            self._batch_rate = rate / self._batch_size
+
+            logger.info(f"sending rate updated to {rate}")
+
+        return np.random.exponential(scale=1 / self._batch_rate)
+
+
 class GeneratorFactory:
     """Request generator factory class."""
 
@@ -131,6 +210,7 @@ def get(sort: ReqGenEnum) -> Generator:
         generators = {
             ReqGenEnum.DEFAULT: DefaultGenerator(),
             ReqGenEnum.EXP: ExponentialGenerator(),
+            ReqGenEnum.MULTIRATE_EXP: MultiRateExponentialGenerator(),
         }
 
         return generators[sort]