diff --git a/perf/throttle.py b/perf/throttle.py new file mode 100644 index 0000000..b909382 --- /dev/null +++ b/perf/throttle.py @@ -0,0 +1,79 @@ +import psutil +import asyncio +import time +import util + +class DynamicThrottle(object): + + def __init__(self, **kwargs): + self.target_cpu_usage = kwargs.get("target_cpu_usage", 50) + self.sleep_interval = 0.0 + + self.sleep_increment = kwargs.get("sleep_increment", 0.01) + self.sleep_decrement = kwargs.get("sleep_decrement", 0.01) + + self.sleep_max = kwargs.get("sleep_max", 0.1) + self.sleep_min = kwargs.get("sleep_min", 0.01) + + self.psutil_interval = kwargs.get("psutil_interval", 0.1) + + self.log = kwargs.get("log", util.get_logger(self.__class__.__name__)) + + self.consecutive_increments = 0 + self.consecutive_decrements = 0 + + self.last_was_increment = False + + if kwargs.get("async"): + self.dynamic_throttle = self.dynamic_throttle_async + else: + self.dynamic_throttle = self.dynamic_throttle + + + async def dynamic_throttle_async(self): + """ + Dynamically sleeps before a request if CPU usage is above our target. + """ + current_cpu_usage = psutil.cpu_percent(interval=self.psutil_interval) + + if current_cpu_usage > self.target_cpu_usage: + self.sleep_interval += self.sleep_increment + if self.sleep_interval > self.sleep_max: + self.sleep_interval = self.sleep_max + self.log.info( + f"CPU {current_cpu_usage}% > {self.target_cpu_usage}%, " + f"=> sleep {self.sleep_interval:.3f}s" + ) + elif current_cpu_usage < self.target_cpu_usage and self.sleep_interval > self.sleep_min: + self.sleep_interval -= self.sleep_decrement + self.log.info( + f"CPU {current_cpu_usage}% < {self.target_cpu_usage}%, " + f"=> sleep {self.sleep_interval:.3f}s" + ) + + if self.sleep_interval > 0: + await asyncio.sleep(self.sleep_interval) + + def dynamic_throttle(self): + """ + Dynamically sleeps before a request if CPU usage is above our target. + """ + current_cpu_usage = psutil.cpu_percent(interval=self.psutil_interval) + + if current_cpu_usage > self.target_cpu_usage: + self.sleep_interval += self.sleep_increment + if self.sleep_interval > self.sleep_max: + self.sleep_interval = self.sleep_max + self.log.info( + f"CPU {current_cpu_usage}% > {self.target_cpu_usage}%, " + f"=> sleep {self.sleep_interval:.3f}s" + ) + elif current_cpu_usage < self.target_cpu_usage and self.sleep_interval > self.sleep_min: + self.sleep_interval -= self.sleep_decrement + self.log.info( + f"CPU {current_cpu_usage}% < {self.target_cpu_usage}%, " + f"=> sleep {self.sleep_interval:.3f}s" + ) + + if self.sleep_interval > 0: + time.sleep(self.sleep_interval) \ No newline at end of file