79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
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) |