Add load-shedding behavior under pressure for AI service (#621)

app/ai-service/api/v1/inference.py

@@ -9,6 +9,7 @@
 from pydantic import BaseModel
 
 import tasks
+from exceptions import LoadShedError
 from services.cache import cached_response
 from config import settings
 
@@ -64,6 +65,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)}")

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
+
     # Cache TTL settings (in seconds)
     cache_ttl_task_status: int = 30  # Short TTL for responsive polling
     cache_ttl_artifact_access: int = 60  # 1 minute for artifact metadata

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)

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
@@ -271,22 +271,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:
@@ -417,6 +406,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)}")
@@ -589,6 +580,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)

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'])

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 False
+
+        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"

app/ai-service/services/load_shedder.py

@@ -0,0 +1,178 @@
+"""
+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:
+        # Broker unreachable is not a queue-depth overload signal. Let the
+        # request proceed so validation and enqueue logic can handle it.
+        return None
+    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,
+        )

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()

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