Files
amq-adapter-python/amqp/adapter/backpressure_handler.py
T
Stanislav Hejny 9c0cb3a334 refactor: update backpressure logic to handle availability correctly
Co-authored-by: aider (openrouter/openai/o3-mini-high) <aider@aider.chat>
2025-07-16 18:04:09 +01:00

442 lines
19 KiB
Python

import asyncio
import json
import time
from asyncio import AbstractEventLoop
from threading import Thread
from typing import Optional
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_load = 0
# helps detect IDLE condition
# 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.maximum_load = self.config.backpressure.threshold_load
self.average_cpu_usage = None
self.do_loop = -1
self._resource_usage_changed = 0
self._resource_average_value = 0
self._last_resource_max_value = 100 # Default max value for CPU usage
def increase_current_load(self):
"""Increase the number of parallel executions, or current load"""
logging_info(
"Backpressure: Increase parallel executions, current=%s",
self.current_load,
)
self.current_load += 1
def decrease_current_load(self):
"""Decrease the number of parallel executions, or current load"""
logging_info(
"Backpressure: Decrease parallel executions, current=%s",
self.current_load,
)
if self.current_load > 0:
self.current_load -= 1
def update_current_load(self, count: int):
"""Update the number of parallel executions, or current load"""
self.current_load = count
def update_last_data_message_time(self):
# logging_info("Backpressure: Update last data message time")
"""Update the last data message time"""
self.last_backpressure_event_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
)
if maximum > 0 and maximum > self.maximum_load:
self.maximum_load = maximum
# Update the current parallel executions count
self.update_current_load(current)
# Check for overload conditions
await self.check_overload_condition()
# Update the last data message time to indicate activity
self.update_last_data_message_time()
# If maximum has changed, update the parallel workers threshold
if maximum != self.maximum_load and maximum > 0:
logging_info(
"Backpressure: Updating maximum parallel workers from %s to %s",
self.maximum_load,
maximum,
)
self.maximum_load = maximum
def start_backpressure_monitor(self) -> Optional[Thread]:
# 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 availability is too low (OVERLOAD)
or if the service has been idle for too long with high availability (IDLE).
Note: current_availability represents available capacity, not used capacity.
- Low availability (close to 0) means OVERLOAD
- High availability (close to maximum) with no activity means IDLE
"""
current_time = time.time()
last_event_delta = current_time - self.last_backpressure_event_time
# If maximum is not set or invalid, use the highest seen value
if self.maximum_load <= 0 and self.current_load > 0:
self.maximum_load = self.current_load
elif self.current_load > self.maximum_load:
# Update maximum if we see a higher value
self.maximum_load = self.current_load
logging_info(
"Backpressure: Check conditions, current_availability=%s, max=%s, last_activity=%s",
self.current_load,
self.maximum_load,
last_event_delta,
)
# Check for OVERLOAD condition - low availability (less than 20% of maximum)
if self.current_load <= round(0.2 * self.maximum_load):
# 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()
# update / reset time-window so that the OVERLOAD is not sent too often
self.last_backpressure_event_time = current_time
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
# Check for IDLE condition - high availability (more than 80% of maximum) with no activity
elif self.current_load >= round(0.8 * self.maximum_load):
idle_duration = self.config.backpressure.idle_duration
# Check if service has been idle for longer than the configured duration
if (
last_event_delta > idle_duration
and self.last_backpressure_event == ScalingRequestAlert.IDLE
):
logging_info(
"Backpressure: Service has been idle for %s seconds (threshold: %s seconds)",
last_event_delta,
idle_duration,
)
# Trigger the idle event
await self._handle_backpressure_idle_event()
# update / reset time-window so that the IDLE is not sent too often
self.last_backpressure_event_time = current_time
else:
# Don't send IDLE right away when the availability increases, but wait for a while
self.last_backpressure_event = ScalingRequestAlert.IDLE
# If neither OVERLOAD nor IDLE, and it's time for an update, send UPDATE event
elif last_event_delta > self.config.backpressure.time_window:
_scaling_request: ScaleRequestV1 = ScaleRequestV1(
self.swarm_service_id,
self.swarm_task_id,
self.maximum_load,
self.current_load, # Current availability is passed directly
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
self.update_last_data_message_time()
async def _backpressure_monitor(self):
"""Monitor the backpressure conditions and trigger events accordingly"""
# Check if CPU monitoring is enabled in configuration
cpu_monitoring_enabled = self.config.backpressure.cpu_monitoring_enabled
_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():
if cpu_monitoring_enabled:
_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_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_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
else:
# If CPU monitoring is disabled, just check the current load
await self.check_overload_condition()
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.maximum_load,
self.current_load, # Current availability is passed directly
ScalingRequestAlert.OVERLOAD,
)
# Address the message to any management service instance capable of supporting BACKPRESSURE request
await self.publish_backpressure_request(scaling_request)
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.maximum_load,
self.current_load, # Current availability is passed directly
ScalingRequestAlert.IDLE,
)
# Address the message to any adapter capable of supporting BACKPRESSURE request
await self.publish_backpressure_request(scaling_request)
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))
def _do_loop(self) -> bool:
"""
Helper function for unit tests to perform several loops only. 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