Files
amq-adapter-python/amqp/adapter/backpressure_handler.py
T
Stanislav Hejny 5bb6597b57
Unit test coverage / pytest (push) Failing after 1m1s
Unit test coverage / pytest (pull_request) Failing after 1m18s
#40 - ensure the JSON format is handled correctly at all levels
2025-05-07 13:05:04 +01:00

325 lines
13 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
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:
def __init__(self, num_items):
self.buffer = [0] * 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
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
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()
def start_backpressure_monitor(self) -> Thread:
# Start the Backpressure monitor loop
self.thread = Thread(target=self.backpressure_monitor_loop)
self.thread.start()
return self.thread
async def check_overload_condition(self):
"""Check if the current parallel executions exceed the limit"""
self.update_last_data_message_time()
logging_info(
"Backpressure: Check overload condition, current=%s, max=%s",
self.current_parallel_executions,
self.parallel_workers,
)
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 = time.time()
self.last_backpressure_event = ScalingRequestAlert.OVERLOAD
def backpressure_monitor_loop(self):
_time_window_millis = self.config.backpressure.time_window
_idle_duration_millis = self.config.backpressure.idle_duration
_overload_duration_millis = self.config.backpressure.idle_duration
_loop = asyncio.new_event_loop()
_monitor_interval = 0.2 # Monitor every 200ms
CPU_IDLE, CPU_ACTIVE, CPU_OVERLOAD = 0, 1, 2
IDLE_THRESHOLD, OVERLOAD_THRESHOLD = 15, 90
self.average_cpu_usage = RunningAverage(
int(
max(_overload_duration_millis, _idle_duration_millis)
// _monitor_interval
)
)
async def _backpressure_monitor():
"""Monitor the backpressure conditions and trigger events accordingly"""
_cpu_usage_state = CPU_ACTIVE
_cpu_usage_changed = time.time()
while True:
_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)
_loop.run_until_complete(_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
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
await await_future(
asyncio.run_coroutine_threadsafe(_wrap_rabbit_mq_api(), self.loop)
)