From 3db79be13de9d474276901f9e5c22f791371ca5d Mon Sep 17 00:00:00 2001 From: "Abdullateef O.G" Date: Wed, 24 Jun 2026 23:55:37 +0100 Subject: [PATCH 1/2] Add load-shedding under pressure for AI service (Issue #621) Return standardized 503 error envelopes and emit requests_shed_total metrics when memory, queue, or provider capacity is exhausted. --- app/ai-service/api/v1/inference.py | 3 + app/ai-service/config.py | 4 + app/ai-service/exceptions.py | 15 ++ app/ai-service/main.py | 39 ++-- app/ai-service/metrics.py | 6 + .../services/humanitarian_verification.py | 18 ++ app/ai-service/services/load_shedder.py | 176 ++++++++++++++++++ app/ai-service/tasks.py | 3 + app/ai-service/tests/test_error_envelope.py | 22 ++- app/ai-service/tests/test_load_shedder.py | 148 +++++++++++++++ app/ai-service/tests/test_versioned_routes.py | 3 + 11 files changed, 419 insertions(+), 18 deletions(-) create mode 100644 app/ai-service/services/load_shedder.py create mode 100644 app/ai-service/tests/test_load_shedder.py diff --git a/app/ai-service/api/v1/inference.py b/app/ai-service/api/v1/inference.py index c3bf21ab..08d566b9 100644 --- a/app/ai-service/api/v1/inference.py +++ b/app/ai-service/api/v1/inference.py @@ -9,6 +9,7 @@ from pydantic import BaseModel import tasks +from exceptions import LoadShedError logger = logging.getLogger(__name__) @@ -62,6 +63,8 @@ async def create_inference_task( "status_url": f"/v1/ai/status/{task_id}", } + except LoadShedError: + raise except Exception as e: logger.error(f"Failed to create inference task: {str(e)}") raise HTTPException(status_code=500, detail=f"Failed to create task: {str(e)}") diff --git a/app/ai-service/config.py b/app/ai-service/config.py index a590acaa..753d2bfb 100644 --- a/app/ai-service/config.py +++ b/app/ai-service/config.py @@ -51,6 +51,10 @@ class Settings(BaseSettings): circuit_breaker_failure_threshold: int = 3 circuit_breaker_recovery_timeout_seconds: float = 30.0 + # Load shedding settings + load_shed_memory_threshold_percent: float = 90.0 + load_shed_max_celery_queue_depth: int = 100 + # Application settings app_env: Literal["development", "staging", "production", "test"] = "development" log_level: str = "INFO" diff --git a/app/ai-service/exceptions.py b/app/ai-service/exceptions.py index ff15ecd4..27d45c57 100644 --- a/app/ai-service/exceptions.py +++ b/app/ai-service/exceptions.py @@ -12,3 +12,18 @@ def __init__(self, message: str, code: str = "AI_SERVICE_ERROR", details: Option def __str__(self) -> str: return f"[{self.code}] {self.message}" + + +class LoadShedError(Exception): + """Raised when the service must reject work due to overload.""" + + def __init__( + self, + reason: str, + message: str, + details: Optional[Any] = None, + ): + self.reason = reason + self.message = message + self.details = details or {} + super().__init__(message) diff --git a/app/ai-service/main.py b/app/ai-service/main.py index 48414a10..ebf44734 100644 --- a/app/ai-service/main.py +++ b/app/ai-service/main.py @@ -16,7 +16,7 @@ from fastapi.exceptions import RequestValidationError from starlette.exceptions import HTTPException as StarletteHTTPException from fastapi.responses import JSONResponse, RedirectResponse -from exceptions import AIServiceError +from exceptions import AIServiceError, LoadShedError from schemas.errors import ErrorDetail, ErrorEnvelope import time import metrics @@ -263,22 +263,11 @@ async def monitor_requests(request: Request, call_next): if path in _NEVER_THROTTLE or is_redirect_path: return await call_next(request) - # Gracefully throttle if memory pressure is critical. - if not metrics.check_system_resources(memory_threshold_percent=90.0): - metrics.REQUEST_COUNT.labels( - method=request.method, - endpoint=path, - http_status=503, - ).inc() - return JSONResponse( - status_code=503, - content={ - "detail": ( - "Service unavailable: System resources (RAM/VRAM) exhausted, " - "gracefully throttling." - ) - }, - ) + from services.load_shedder import evaluate_load_shed + + shed_response = evaluate_load_shed(request) + if shed_response is not None: + return shed_response start_time = time.time() try: @@ -409,6 +398,8 @@ async def _legacy_create_inference_task( "message": "Task queued for processing", "status_url": f"/v1/ai/status/{task_id}", } + except LoadShedError: + raise except Exception as e: logger.error(f"Failed to create inference task: {str(e)}") raise HTTPException(status_code=500, detail=f"Failed to create task: {str(e)}") @@ -581,6 +572,20 @@ async def validation_exception_handler(request, exc: RequestValidationError): ) +@app.exception_handler(LoadShedError) +async def load_shed_exception_handler(request, exc: LoadShedError): + from services.load_shedder import build_shed_response + + path = request.url.path + logger.warning("Load shedding request to %s due to %s", path, exc.reason) + return build_shed_response( + exc.reason, + request.method, + path, + details=exc.details, + ) + + @app.exception_handler(AIServiceError) async def ai_service_exception_handler(request, exc: AIServiceError): logger.error(f"AI service error: {exc.message}", exc_info=True) diff --git a/app/ai-service/metrics.py b/app/ai-service/metrics.py index c41e7c18..50ac796a 100644 --- a/app/ai-service/metrics.py +++ b/app/ai-service/metrics.py @@ -11,6 +11,12 @@ # API metrics REQUEST_COUNT = Counter('api_request_count', 'Total API request count', ['method', 'endpoint', 'http_status']) REQUEST_LATENCY = Histogram('api_request_latency_seconds', 'API request latency', ['method', 'endpoint']) +REQUESTS_SHED_TOTAL = Counter( + 'requests_shed_total', + 'Requests rejected due to overload (load shedding)', + ['reason', 'method', 'endpoint'], +) +CELERY_QUEUE_DEPTH = Gauge('celery_queue_depth', 'Pending tasks in the Celery default queue') # AI Model metrics MODEL_LOAD_TIME = Histogram('model_load_time_seconds', 'Model load time in seconds', ['model_name']) diff --git a/app/ai-service/services/humanitarian_verification.py b/app/ai-service/services/humanitarian_verification.py index 1c3c595b..69fe7579 100644 --- a/app/ai-service/services/humanitarian_verification.py +++ b/app/ai-service/services/humanitarian_verification.py @@ -127,6 +127,24 @@ def _provider_attempt_order(self, provider_preference: str) -> List[str]: return [preference] + [provider for provider in available if provider != preference] return available + def all_providers_unavailable(self) -> bool: + """Return True when every configured LLM provider circuit is open.""" + if settings.test_provider_mode: + return False + + providers = [] + if settings.openai_api_key: + providers.append("openai") + if settings.groq_api_key: + providers.append("groq") + if not providers: + return True + + return all( + provider in self.breakers and not self.breakers[provider].allow_request() + for provider in providers + ) + def _get_model_for_provider(self, provider: str) -> str: if provider == "test": return "test-provider/fixture" diff --git a/app/ai-service/services/load_shedder.py b/app/ai-service/services/load_shedder.py new file mode 100644 index 00000000..dba7c43b --- /dev/null +++ b/app/ai-service/services/load_shedder.py @@ -0,0 +1,176 @@ +""" +Load-shedding for the AI service under pressure (Issue #621). + +Rejects incoming work with HTTP 503 and a standardized error envelope when +system memory, the Celery queue, or configured LLM providers are overloaded. +""" + +import logging +from typing import Any, Dict, Optional, Tuple + +from fastapi import Request +from fastapi.responses import JSONResponse + +import metrics +from config import settings +from exceptions import LoadShedError +from schemas.errors import ErrorDetail, ErrorEnvelope + +logger = logging.getLogger(__name__) + +CELERY_QUEUE_NAME = "celery" +RETRY_AFTER_SECONDS = 30 + +REASON_MESSAGES = { + "memory": "Service temporarily unavailable due to high memory pressure", + "queue_full": "Service temporarily unavailable: task queue is at capacity", + "broker_unavailable": "Service temporarily unavailable: task broker is unreachable", + "provider_down": "Service temporarily unavailable: AI providers are currently down", +} + + +def record_shed_request(reason: str, method: str, endpoint: str) -> None: + metrics.REQUESTS_SHED_TOTAL.labels( + reason=reason, method=method, endpoint=endpoint + ).inc() + metrics.REQUEST_COUNT.labels( + method=method, endpoint=endpoint, http_status=503 + ).inc() + + +def build_shed_response( + reason: str, + method: str, + endpoint: str, + details: Optional[Dict[str, Any]] = None, +) -> JSONResponse: + record_shed_request(reason, method, endpoint) + payload_details: Dict[str, Any] = {"reason": reason, **(details or {})} + return JSONResponse( + status_code=503, + headers={"Retry-After": str(RETRY_AFTER_SECONDS)}, + content=ErrorEnvelope( + error=ErrorDetail( + code="SERVICE_OVERLOADED", + message=REASON_MESSAGES.get( + reason, "Service temporarily unavailable due to high load" + ), + details=payload_details, + ) + ).model_dump(), + ) + + +def get_celery_queue_depth() -> Optional[int]: + """Return pending Celery queue depth, or None when the broker is unreachable.""" + try: + import redis + + client = redis.from_url( + settings.redis_url, + socket_connect_timeout=1.0, + socket_timeout=1.0, + ) + client.ping() + depth = client.llen(CELERY_QUEUE_NAME) + if not isinstance(depth, int): + return 0 + metrics.CELERY_QUEUE_DEPTH.set(depth) + return depth + except Exception as exc: + logger.warning("Failed to check Celery queue depth: %s", exc) + return None + + +def check_memory_pressure() -> Optional[str]: + if not metrics.check_system_resources( + memory_threshold_percent=settings.load_shed_memory_threshold_percent + ): + return "memory" + return None + + +def check_queue_pressure() -> Optional[Tuple[str, Dict[str, Any]]]: + if settings.app_env == "test": + return None + + depth = get_celery_queue_depth() + if depth is None: + return "broker_unavailable", {} + if depth >= settings.load_shed_max_celery_queue_depth: + return "queue_full", { + "queue_depth": depth, + "max_queue_depth": settings.load_shed_max_celery_queue_depth, + } + return None + + +def are_llm_providers_down() -> bool: + if settings.app_env == "test" or settings.test_provider_mode: + return False + + try: + import main as _main + + return _main.humanitarian_verification_service.all_providers_unavailable() + except Exception as exc: + logger.warning("Failed to evaluate LLM provider health: %s", exc) + return False + + +def check_provider_pressure() -> Optional[str]: + if are_llm_providers_down(): + return "provider_down" + return None + + +def _is_job_creation_route(path: str, method: str) -> bool: + if method.upper() != "POST": + return False + return path.endswith("/ai/inference") or path.endswith("/ai/ocr/jobs") + + +def _is_llm_route(path: str, method: str) -> bool: + if method.upper() != "POST": + return False + return path.endswith("/ai/humanitarian/verify") + + +def evaluate_load_shed(request: Request) -> Optional[JSONResponse]: + path = request.url.path + method = request.method + + memory_reason = check_memory_pressure() + if memory_reason: + return build_shed_response( + memory_reason, + method, + path, + details={ + "threshold_percent": settings.load_shed_memory_threshold_percent, + }, + ) + + if _is_job_creation_route(path, method): + queue_result = check_queue_pressure() + if queue_result: + reason, details = queue_result + return build_shed_response(reason, method, path, details=details) + + if _is_llm_route(path, method): + provider_reason = check_provider_pressure() + if provider_reason: + return build_shed_response(provider_reason, method, path) + + return None + + +def ensure_queue_capacity() -> None: + queue_result = check_queue_pressure() + if queue_result: + reason, details = queue_result + raise LoadShedError( + reason, + REASON_MESSAGES.get(reason, "Service temporarily unavailable due to high load"), + details=details, + ) diff --git a/app/ai-service/tasks.py b/app/ai-service/tasks.py index 84fa2736..1d5c00b8 100644 --- a/app/ai-service/tasks.py +++ b/app/ai-service/tasks.py @@ -13,6 +13,7 @@ import metrics from config import settings +from services.load_shedder import ensure_queue_capacity from services.pii_scrubber import PIIScrubberService from services.humanitarian_verification import HumanitarianVerificationService from services.ocr_job import run_ocr_from_base64 @@ -454,6 +455,8 @@ def create_task(task_type: str, payload: Dict[str, Any]) -> str: # Initialize task status update_task_status(task_id, 'pending') + ensure_queue_capacity() + try: # Queue the task using the lazy-registered task task = get_process_heavy_inference_task() diff --git a/app/ai-service/tests/test_error_envelope.py b/app/ai-service/tests/test_error_envelope.py index faeb74d4..6cd8ca2d 100644 --- a/app/ai-service/tests/test_error_envelope.py +++ b/app/ai-service/tests/test_error_envelope.py @@ -7,6 +7,7 @@ import pytest from fastapi import HTTPException from fastapi.testclient import TestClient +import main from main import app from exceptions import AIServiceError @@ -93,8 +94,27 @@ def test_ai_error_details(self): # --------------------------------------------------------------------------- -# 5. Generic 500 Internal Server Error +# 6. 503 Service Overloaded (load shedding) # --------------------------------------------------------------------------- +class TestServiceOverloaded: + def test_503_shape(self): + @main.app.get("/_test/503") + async def _raise_503(): + from exceptions import LoadShedError + + raise LoadShedError( + "memory", + "Service temporarily unavailable due to high memory pressure", + details={"threshold_percent": 90.0}, + ) + + r = client.get("/_test/503") + assert r.status_code == 503 + assert_envelope(r.json(), "SERVICE_OVERLOADED") + assert r.json()["error"]["details"]["reason"] == "memory" + assert r.headers.get("retry-after") == "30" + + class TestInternalServerError: def test_500_shape(self): @app.get("/_test/500") diff --git a/app/ai-service/tests/test_load_shedder.py b/app/ai-service/tests/test_load_shedder.py new file mode 100644 index 00000000..0ac36b70 --- /dev/null +++ b/app/ai-service/tests/test_load_shedder.py @@ -0,0 +1,148 @@ +""" +Tests for load-shedding behavior (Issue #621). +""" + +from unittest.mock import patch + +import metrics +import pytest +from fastapi.testclient import TestClient + +import main +from exceptions import LoadShedError +from services.load_shedder import ( + build_shed_response, + check_memory_pressure, + check_queue_pressure, + evaluate_load_shed, + ensure_queue_capacity, + record_shed_request, +) + + +@pytest.fixture +def client(): + return TestClient(main.app, follow_redirects=True) + + +def assert_shed_envelope(data: dict, expected_reason: str): + assert "error" in data + err = data["error"] + assert err["code"] == "SERVICE_OVERLOADED" + assert isinstance(err["message"], str) + assert err["details"]["reason"] == expected_reason + + +class TestLoadShedResponse: + def test_build_shed_response_shape(self): + import json + + response = build_shed_response("memory", "POST", "/v1/ai/anonymize") + assert response.status_code == 503 + assert response.headers["retry-after"] == "30" + assert_shed_envelope(json.loads(response.body.decode()), "memory") + + def test_record_shed_request_increments_metric(self): + before = metrics.REQUESTS_SHED_TOTAL.labels( + reason="queue_full", method="POST", endpoint="/v1/ai/inference" + )._value.get() + record_shed_request("queue_full", "POST", "/v1/ai/inference") + after = metrics.REQUESTS_SHED_TOTAL.labels( + reason="queue_full", method="POST", endpoint="/v1/ai/inference" + )._value.get() + assert after == before + 1 + + +class TestMemoryPressure: + def test_memory_pressure_detected(self): + with patch.object(metrics, "check_system_resources", return_value=False): + assert check_memory_pressure() == "memory" + + def test_memory_pressure_healthy(self): + with patch.object(metrics, "check_system_resources", return_value=True): + assert check_memory_pressure() is None + + +class TestQueuePressure: + def test_queue_full(self): + with patch( + "services.load_shedder.get_celery_queue_depth", return_value=150 + ), patch("services.load_shedder.settings") as mock_settings: + mock_settings.app_env = "production" + mock_settings.load_shed_max_celery_queue_depth = 100 + result = check_queue_pressure() + assert result is not None + reason, details = result + assert reason == "queue_full" + assert details["queue_depth"] == 150 + + def test_broker_unavailable(self): + with patch( + "services.load_shedder.get_celery_queue_depth", return_value=None + ), patch("services.load_shedder.settings") as mock_settings: + mock_settings.app_env = "production" + result = check_queue_pressure() + assert result == ("broker_unavailable", {}) + + def test_ensure_queue_capacity_raises(self): + with patch( + "services.load_shedder.check_queue_pressure", + return_value=("queue_full", {"queue_depth": 120}), + ): + with pytest.raises(LoadShedError) as exc_info: + ensure_queue_capacity() + assert exc_info.value.reason == "queue_full" + + +class TestMiddlewareLoadShedding: + def test_v1_endpoint_shed_on_memory_pressure(self, client): + with patch.object(metrics, "check_system_resources", return_value=False): + response = client.post( + "/v1/ai/anonymize", + json={"text": "Some text with Jane Smith in Lagos."}, + ) + assert response.status_code == 503 + assert_shed_envelope(response.json(), "memory") + + def test_health_never_shed(self, client): + with patch.object(metrics, "check_system_resources", return_value=False): + response = client.get("/health") + assert response.status_code == 200 + + def test_inference_shed_when_queue_full(self, client): + with patch( + "services.load_shedder.check_queue_pressure", + return_value=("queue_full", {"queue_depth": 200}), + ): + response = client.post("/v1/ai/inference", json={"type": "inference"}) + assert response.status_code == 503 + assert_shed_envelope(response.json(), "queue_full") + + def test_humanitarian_shed_when_providers_down(self, client): + with patch("services.load_shedder.are_llm_providers_down", return_value=True): + response = client.post( + "/v1/ai/humanitarian/verify", + json={"aid_claim": "Need food assistance"}, + ) + assert response.status_code == 503 + assert_shed_envelope(response.json(), "provider_down") + + def test_metrics_endpoint_exposes_shed_counter(self, client): + record_shed_request("memory", "POST", "/v1/ai/anonymize") + response = client.get("/ai/metrics") + assert response.status_code == 200 + assert "requests_shed_total" in response.text + + +class TestLoadShedExceptionHandler: + def test_handler_returns_envelope(self, client): + @main.app.get("/_test/load-shed") + async def _raise_load_shed(): + raise LoadShedError( + "broker_unavailable", + "Service temporarily unavailable: task broker is unreachable", + ) + + response = client.get("/_test/load-shed") + assert response.status_code == 503 + assert_shed_envelope(response.json(), "broker_unavailable") diff --git a/app/ai-service/tests/test_versioned_routes.py b/app/ai-service/tests/test_versioned_routes.py index 1f8d2c85..7a1e9db8 100644 --- a/app/ai-service/tests/test_versioned_routes.py +++ b/app/ai-service/tests/test_versioned_routes.py @@ -461,6 +461,9 @@ def test_v1_endpoint_throttled_when_resources_exhausted(self, following_client): json={"text": "Some text with Jane Smith in Lagos."}, ) assert response.status_code == 503 + body = response.json() + assert body["error"]["code"] == "SERVICE_OVERLOADED" + assert body["error"]["details"]["reason"] == "memory" def test_health_never_throttled(self, client): """Health endpoint must respond even under resource pressure.""" From f60dff045e516ac42ba9950cc32cd8efe32f0644 Mon Sep 17 00:00:00 2001 From: "Abdullateef O.G" Date: Thu, 25 Jun 2026 23:13:38 +0100 Subject: [PATCH 2/2] Fix false provider-down load shedding without API keys Only shed humanitarian routes when configured providers have open circuit breakers, not when no providers are configured at all. --- .../services/humanitarian_verification.py | 2 +- app/ai-service/tests/test_circuit_breaker.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/app/ai-service/services/humanitarian_verification.py b/app/ai-service/services/humanitarian_verification.py index 69fe7579..9f6b44ae 100644 --- a/app/ai-service/services/humanitarian_verification.py +++ b/app/ai-service/services/humanitarian_verification.py @@ -138,7 +138,7 @@ def all_providers_unavailable(self) -> bool: if settings.groq_api_key: providers.append("groq") if not providers: - return True + return False return all( provider in self.breakers and not self.breakers[provider].allow_request() diff --git a/app/ai-service/tests/test_circuit_breaker.py b/app/ai-service/tests/test_circuit_breaker.py index 5206b3dd..bc9784b8 100644 --- a/app/ai-service/tests/test_circuit_breaker.py +++ b/app/ai-service/tests/test_circuit_breaker.py @@ -129,3 +129,21 @@ def test_request_timeout_raises_ai_timeout(self, mock_post, monkeypatch): # The breaker for openai should have recorded the failure assert self.service.breakers["openai"].failure_count == 2 # Primary & fallback attempts both failed + + def test_all_providers_unavailable_false_without_configured_providers(self, monkeypatch): + monkeypatch.setattr(settings, "openai_api_key", None) + monkeypatch.setattr(settings, "groq_api_key", None) + monkeypatch.setattr(settings, "test_provider_mode", False) + + assert self.service.all_providers_unavailable() is False + + def test_all_providers_unavailable_true_when_all_breakers_open(self, monkeypatch): + monkeypatch.setattr(settings, "openai_api_key", "test-key") + monkeypatch.setattr(settings, "groq_api_key", "test-key") + monkeypatch.setattr(settings, "test_provider_mode", False) + + for breaker in self.service.breakers.values(): + breaker.state = "OPEN" + breaker.last_state_change = time.time() + + assert self.service.all_providers_unavailable() is True