826c859b32
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
432 lines
18 KiB
Python
432 lines
18 KiB
Python
import asyncio
|
|
import json
|
|
import time
|
|
from asyncio import AbstractEventLoop
|
|
from threading import Thread
|
|
|
|
import psutil
|
|
from aio_pika import Message
|
|
from aio_pika.abc import AbstractRobustChannel
|
|
|
|
from amqp.adapter.logging_utils import logging_info, logging_warning
|
|
from amqp.config.amq_configuration import AMQConfiguration
|
|
from amqp.model.model import ScalingRequestAlert
|
|
from amqp.router.utils import await_future, await_result
|
|
|
|
CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2
|
|
IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 0.15, 0.90
|
|
RESOURCE_UNSET, RESOURCE_IDLE, RESOURCE_ACTIVE, RESOURCE_OVERLOAD = -1, 0, 1, 2
|
|
|
|
|
|
class ScaleRequestV1:
|
|
|
|
def __init__(
|
|
self,
|
|
serviceId: str,
|
|
taskId: str,
|
|
maxAvailability: int,
|
|
currentAvailability: int,
|
|
requestType: ScalingRequestAlert,
|
|
):
|
|
self.version = 1 # Version of the request, currently 1
|
|
self.serviceId = serviceId
|
|
self.taskId = taskId
|
|
self.maxAvailability = maxAvailability
|
|
self.currentAvailability = currentAvailability
|
|
self.requestType = requestType
|
|
self.thread = None
|
|
|
|
def to_json(self) -> str:
|
|
"""Converts the object to a JSON string matching the Java structure"""
|
|
return json.dumps(
|
|
{
|
|
"version": self.version,
|
|
"serviceId": self.serviceId,
|
|
"taskId": self.taskId,
|
|
"maxAvailability": self.maxAvailability,
|
|
"currentAvailability": self.currentAvailability,
|
|
"requestType": self.requestType.value, # Using .value for Enum serialization
|
|
},
|
|
indent=2,
|
|
)
|
|
|
|
|
|
class RunningAverage:
|
|
"""
|
|
A class to maintain a running average of the last N values.
|
|
The intended use is to track CPU usage over time window, that is directly related to sample rate and window size.
|
|
This is the reason why the constructor takes the number of items to store, calculated as window size / sample rate
|
|
Initial_value is for unit test to emulate OVERLOAD or IDLE condition.
|
|
"""
|
|
|
|
def __init__(self, num_items, initial_value=0):
|
|
self.buffer = [initial_value] * num_items # Fixed-size array
|
|
self.pointer = 0 # Current position in array
|
|
self.num_items = num_items # Maximum number of items to store
|
|
self.is_filled = False # Track if buffer is fully filled
|
|
|
|
def add(self, value):
|
|
# Store the value at current pointer position
|
|
self.buffer[self.pointer] = value
|
|
|
|
# Move pointer to next position with wrap-around
|
|
self.pointer += 1
|
|
if self.pointer >= self.num_items:
|
|
self.pointer = 0
|
|
self.is_filled = True
|
|
|
|
def get_average(self):
|
|
# Determine how many items we should average
|
|
count = self.num_items if self.is_filled else self.pointer
|
|
|
|
if count == 0:
|
|
return 0.0 # Avoid division by zero
|
|
|
|
return sum(self.buffer) / count
|
|
|
|
def get_current_values(self):
|
|
"""Returns the values in chronological order (oldest first)"""
|
|
if not self.is_filled:
|
|
return self.buffer[: self.pointer]
|
|
|
|
return self.buffer[self.pointer :] + self.buffer[: self.pointer]
|
|
|
|
|
|
class BackpressureHandler:
|
|
# Track the number of messages currently being processed
|
|
current_parallel_executions = 0
|
|
# helps detect IDLE condition
|
|
last_data_message_time = 0
|
|
# helps prevent flooding the system with backpressure events
|
|
last_backpressure_event_time = 0
|
|
last_backpressure_event = ScalingRequestAlert.UPDATE
|
|
# _callback_list is used in unit tests to record the invoked callbacks
|
|
_callback_list = None
|
|
|
|
def __init__(
|
|
self,
|
|
channel: AbstractRobustChannel,
|
|
loop: AbstractEventLoop,
|
|
config: AMQConfiguration,
|
|
):
|
|
self.lock = None
|
|
self.channel = channel
|
|
self.loop = loop
|
|
self.config = config
|
|
self.exchange = None
|
|
self.swarm_service_id = self.config.amq_adapter.swarm_service_id
|
|
self.swarm_task_id = self.config.amq_adapter.swarm_task_id
|
|
self.parallel_workers = self.config.backpressure.threshold_threads
|
|
self.average_cpu_usage = None
|
|
self.do_loop = -1
|
|
self._resource_usage_state = RESOURCE_UNSET
|
|
self._resource_usage_changed = 0
|
|
self._resource_average_value = 0
|
|
self._last_resource_max_value = 100 # Default max value for CPU usage
|
|
|
|
def _do_loop(self) -> bool:
|
|
"""
|
|
Check if the loop should continue running.
|
|
Positive value of do_loop indicates the number of iterations left.
|
|
Negative value indicates the loop should run indefinitely.
|
|
"""
|
|
_val = self.do_loop != 0
|
|
if self.do_loop > 0:
|
|
self.do_loop -= 1
|
|
return _val
|
|
|
|
def increase_parallel_executions(self):
|
|
"""Increase the number of parallel executions"""
|
|
logging_info(
|
|
"Backpressure: Increase parallel executions, current=%s",
|
|
self.current_parallel_executions,
|
|
)
|
|
self.current_parallel_executions += 1
|
|
|
|
def decrease_parallel_executions(self):
|
|
"""Decrease the number of parallel executions"""
|
|
logging_info(
|
|
"Backpressure: Decrease parallel executions, current=%s",
|
|
self.current_parallel_executions,
|
|
)
|
|
if self.current_parallel_executions > 0:
|
|
self.current_parallel_executions -= 1
|
|
|
|
def update_parallel_executions(self, count: int):
|
|
"""Update the number of parallel executions"""
|
|
self.current_parallel_executions = count
|
|
|
|
def update_last_data_message_time(self):
|
|
logging_info("Backpressure: Update last data message time")
|
|
"""Update the last data message time"""
|
|
self.last_data_message_time = time.time()
|
|
|
|
async def update_backpressure_value(self, current: int, maximum: int):
|
|
"""
|
|
Update the current backpressure value and check for overload conditions.
|
|
This method is called by the AMQService.backpressure() method to update
|
|
the current parallel executions count and check for overload conditions.
|
|
|
|
Args:
|
|
current: Current value of the backpressure metric
|
|
maximum: Maximum value of the backpressure metric
|
|
"""
|
|
logging_info(
|
|
"Backpressure: Updating backpressure value, current=%s, maximum=%s",
|
|
current,
|
|
maximum
|
|
)
|
|
|
|
# Update the current parallel executions count
|
|
self.update_parallel_executions(current)
|
|
|
|
# Update the last data message time to indicate activity
|
|
self.update_last_data_message_time()
|
|
|
|
# Check for overload conditions
|
|
await self.check_overload_condition()
|
|
|
|
# If maximum has changed, update the parallel workers threshold
|
|
if maximum != self.parallel_workers and maximum > 0:
|
|
logging_info(
|
|
"Backpressure: Updating maximum parallel workers from %s to %s",
|
|
self.parallel_workers,
|
|
maximum
|
|
)
|
|
self.parallel_workers = maximum
|
|
|
|
def start_backpressure_monitor(self) -> Thread:
|
|
# Check if CPU monitoring is enabled in configuration
|
|
cpu_monitoring_enabled = self.config.backpressure.cpu_monitoring_enabled
|
|
|
|
if not cpu_monitoring_enabled:
|
|
logging_info("Backpressure CPU monitoring is disabled by configuration")
|
|
return None
|
|
|
|
# Start the Backpressure monitor loop
|
|
self.thread = Thread(target=self.backpressure_monitor_loop)
|
|
self.thread.daemon = True # This makes it a daemon thread
|
|
self.thread.start()
|
|
return self.thread
|
|
|
|
async def check_overload_condition(self):
|
|
"""
|
|
Check if the current parallel executions exceed the limit (OVERLOAD)
|
|
or if the service has been idle for too long (IDLE).
|
|
"""
|
|
current_time = time.time()
|
|
self.update_last_data_message_time()
|
|
|
|
logging_info(
|
|
"Backpressure: Check conditions, current=%s, max=%s, last_activity=%s",
|
|
self.current_parallel_executions,
|
|
self.parallel_workers,
|
|
current_time - self.last_data_message_time,
|
|
)
|
|
|
|
# Check for OVERLOAD condition
|
|
if self.current_parallel_executions >= self.parallel_workers - 1:
|
|
# Check if the last backpressure event was not an overload
|
|
if self.last_backpressure_event != ScalingRequestAlert.OVERLOAD:
|
|
# Trigger the overload event
|
|
await self.handle_backpressure_overload_event()
|
|
self.last_backpressure_event_time = current_time
|
|
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
|
|
|
|
# Check for IDLE condition
|
|
elif self.current_parallel_executions == 0:
|
|
idle_duration = self.config.backpressure.idle_duration
|
|
# Check if service has been idle for longer than the configured duration
|
|
if (current_time - self.last_data_message_time > idle_duration and
|
|
(current_time - self.last_backpressure_event_time > self.config.backpressure.time_window or
|
|
self.last_backpressure_event != ScalingRequestAlert.IDLE)):
|
|
|
|
logging_info(
|
|
"Backpressure: Service has been idle for %s seconds (threshold: %s seconds)",
|
|
current_time - self.last_data_message_time,
|
|
idle_duration
|
|
)
|
|
|
|
# Trigger the idle event
|
|
await self._handle_backpressure_idle_event()
|
|
self.last_backpressure_event_time = current_time
|
|
self.last_backpressure_event = ScalingRequestAlert.IDLE
|
|
|
|
# If neither OVERLOAD nor IDLE, and it's time for an update, send UPDATE event
|
|
elif (current_time - self.last_backpressure_event_time > self.config.backpressure.time_window):
|
|
_scaling_request: ScaleRequestV1 = ScaleRequestV1(
|
|
self.swarm_service_id,
|
|
self.swarm_task_id,
|
|
self.parallel_workers,
|
|
self.parallel_workers - self.current_parallel_executions,
|
|
ScalingRequestAlert.UPDATE,
|
|
)
|
|
# Address the message to any adapter capable of supporting BACKPRESSURE request
|
|
await self.publish_backpressure_request(_scaling_request)
|
|
self.last_backpressure_event_time = current_time
|
|
self.last_backpressure_event = ScalingRequestAlert.UPDATE
|
|
|
|
async def _backpressure_monitor(self):
|
|
"""Monitor the backpressure conditions and trigger events accordingly"""
|
|
_time_window_millis = self.config.backpressure.time_window
|
|
_idle_duration_millis = self.config.backpressure.idle_duration
|
|
_overload_duration_millis = self.config.backpressure.cpu_overload_duration
|
|
_monitor_interval = 0.5 # Monitor every 500ms
|
|
CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2
|
|
IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 15, 90
|
|
if self.average_cpu_usage is None:
|
|
self.average_cpu_usage = RunningAverage(
|
|
int(max(_overload_duration_millis, _idle_duration_millis) // _monitor_interval)
|
|
)
|
|
_cpu_usage_state = CPU_ACTIVE
|
|
_cpu_usage_changed = time.time()
|
|
while self._do_loop():
|
|
_current_cpu_usage = psutil.cpu_percent(interval=None)
|
|
self.average_cpu_usage.add(_current_cpu_usage)
|
|
_average_cpu_usage = self.average_cpu_usage.get_average()
|
|
_now = time.time()
|
|
# logging_info(
|
|
# "Backpressure: current CPU[pct]=%s, avg=%s, state=%s, last dataMsg=%s, last event=%s, last eventTime=%s",
|
|
# _current_cpu_usage,
|
|
# _average_cpu_usage,
|
|
# _cpu_usage_state,
|
|
# self.last_data_message_time,
|
|
# self.last_backpressure_event,
|
|
# self.last_backpressure_event_time,
|
|
# )
|
|
if _average_cpu_usage < IDLE_THRESHOLD:
|
|
if _cpu_usage_state != CPU_IDLE:
|
|
_cpu_usage_changed = _now
|
|
_cpu_usage_state = CPU_IDLE
|
|
elif _average_cpu_usage > OVERLOAD_THRESHOLD:
|
|
if _cpu_usage_state != CPU_OVERLOAD:
|
|
_cpu_usage_changed = _now
|
|
_cpu_usage_state = CPU_OVERLOAD
|
|
else:
|
|
if _cpu_usage_state != CPU_ACTIVE:
|
|
_cpu_usage_changed = _now
|
|
_cpu_usage_state = CPU_ACTIVE
|
|
|
|
# Check if the current time exceeds the last backpressure event time
|
|
if (
|
|
_cpu_usage_state == CPU_OVERLOAD
|
|
and _now - _cpu_usage_changed > _overload_duration_millis
|
|
and (
|
|
_now - self.last_backpressure_event_time > _time_window_millis
|
|
or self.last_backpressure_event != ScalingRequestAlert.OVERLOAD
|
|
)
|
|
):
|
|
# Trigger the overload event
|
|
await self.handle_backpressure_overload_event()
|
|
self.last_backpressure_event_time = _now
|
|
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
|
|
elif (
|
|
_cpu_usage_state == CPU_IDLE
|
|
and _now - _cpu_usage_changed > _idle_duration_millis
|
|
and _now - self.last_data_message_time > _idle_duration_millis
|
|
and (
|
|
_now - self.last_backpressure_event_time > _time_window_millis
|
|
or self.last_backpressure_event != ScalingRequestAlert.IDLE
|
|
)
|
|
):
|
|
# Trigger the idle event
|
|
await self._handle_backpressure_idle_event()
|
|
self.last_backpressure_event = ScalingRequestAlert.IDLE
|
|
self.last_backpressure_event_time = _now
|
|
else:
|
|
# Trigger update event
|
|
if _now - self.last_backpressure_event_time > _time_window_millis:
|
|
# Trigger the update event
|
|
_scaling_request: ScaleRequestV1 = ScaleRequestV1(
|
|
self.swarm_service_id,
|
|
self.swarm_task_id,
|
|
100,
|
|
_average_cpu_usage,
|
|
ScalingRequestAlert.UPDATE,
|
|
)
|
|
# Address the message to any adapter capable of supporting BACKPRESSURE request
|
|
await self.publish_backpressure_request(_scaling_request)
|
|
self.last_backpressure_event_time = _now
|
|
self.last_backpressure_event = ScalingRequestAlert.UPDATE
|
|
await asyncio.sleep(_monitor_interval)
|
|
|
|
def backpressure_monitor_loop(self):
|
|
_loop = asyncio.new_event_loop()
|
|
_loop.run_until_complete(self._backpressure_monitor())
|
|
|
|
async def handle_backpressure_overload_event(self):
|
|
logging_warning("Backpressure: Capacity close to depleted!")
|
|
# Send an Overload event to the Management service
|
|
scaling_request: ScaleRequestV1 = ScaleRequestV1(
|
|
self.swarm_service_id,
|
|
self.swarm_task_id,
|
|
self.parallel_workers,
|
|
self.parallel_workers - self.current_parallel_executions,
|
|
ScalingRequestAlert.OVERLOAD,
|
|
)
|
|
# Address the message to any management service instance capable of supporting BACKPRESSURE request
|
|
await self.publish_backpressure_request(scaling_request)
|
|
# update / reset time-window so that the OVERLOAD is not sent too often
|
|
self.last_data_message_time = time.time()
|
|
|
|
async def _handle_backpressure_idle_event(self):
|
|
logging_warning("Backpressure: Service is idle.")
|
|
# Send an Idle event to the Management service
|
|
scaling_request: ScaleRequestV1 = ScaleRequestV1(
|
|
self.swarm_service_id,
|
|
self.swarm_task_id,
|
|
self.parallel_workers,
|
|
self.parallel_workers - self.current_parallel_executions,
|
|
ScalingRequestAlert.IDLE,
|
|
)
|
|
# Address the message to any adapter capable of supporting BACKPRESSURE request
|
|
await self.publish_backpressure_request(scaling_request)
|
|
# update / reset time-window so that the IDLE is not sent too often
|
|
self.last_data_message_time = time.time()
|
|
|
|
async def publish_backpressure_request(self, scaling_request: ScaleRequestV1):
|
|
# Publish the backpressure request to the management service
|
|
logging_info(
|
|
f"Publishing backpressure for {scaling_request.serviceId} with request type {scaling_request.requestType}"
|
|
)
|
|
if not self.exchange:
|
|
|
|
async def _wrap_rabbit_mq_api_init(channel):
|
|
_exchange = await channel.get_exchange(name="cleverthis.clevermicro.management")
|
|
return _exchange
|
|
|
|
if BackpressureHandler._callback_list is not None:
|
|
BackpressureHandler._callback_list["_wrap_rabbit_mq_api_init"] = 1
|
|
self.exchange = await await_result(
|
|
asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api_init(self.channel), self.loop)
|
|
)
|
|
|
|
if self.exchange:
|
|
|
|
async def _wrap_rabbit_mq_api():
|
|
if not self.channel.is_closed:
|
|
binary_content: bytes = scaling_request.to_json().encode("utf-8")
|
|
pika_message: Message = Message(
|
|
body=binary_content,
|
|
content_encoding="utf-8",
|
|
delivery_mode=2,
|
|
content_type="application/octet-stream",
|
|
headers=None,
|
|
priority=0,
|
|
correlation_id=None,
|
|
)
|
|
await self.exchange.publish(
|
|
message=pika_message, routing_key="backpressure-scaling-v1"
|
|
)
|
|
logging_info(
|
|
"Service Message Published to %s, msg: %s",
|
|
self.exchange.name,
|
|
str(binary_content),
|
|
)
|
|
return True
|
|
return False
|
|
|
|
if BackpressureHandler._callback_list is not None:
|
|
BackpressureHandler._callback_list["_wrap_rabbit_mq_api"] = 1
|
|
await await_future(asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop))
|