111 lines
3.2 KiB
Python
111 lines
3.2 KiB
Python
"""
|
|
Lightweight in-process cooldown for noisy Ollama endpoint failures.
|
|
|
|
This does not change ADR-110 policy order. It suppresses endpoints that just
|
|
failed for short-lived callers while still returning explicit offline/cooldown
|
|
state to health and route-status surfaces.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from collections.abc import Iterable
|
|
|
|
from src.services.ollama_endpoint_resolver import (
|
|
OllamaEndpointSelection,
|
|
OllamaWorkloadType,
|
|
resolve_ollama_order,
|
|
)
|
|
|
|
DEFAULT_OLLAMA_ENDPOINT_COOLDOWN_SECONDS = 60.0
|
|
|
|
_blocked_until_by_url: dict[str, float] = {}
|
|
|
|
|
|
def _normalize_url(url: str) -> str:
|
|
return url.rstrip("/")
|
|
|
|
|
|
def record_ollama_endpoint_failure(
|
|
url: str,
|
|
*,
|
|
cooldown_seconds: float = DEFAULT_OLLAMA_ENDPOINT_COOLDOWN_SECONDS,
|
|
now: float | None = None,
|
|
) -> None:
|
|
"""Temporarily suppress an endpoint after network/5xx failure."""
|
|
if not url:
|
|
return
|
|
current_time = time.monotonic() if now is None else now
|
|
_blocked_until_by_url[_normalize_url(url)] = current_time + cooldown_seconds
|
|
|
|
|
|
def record_ollama_endpoint_success(url: str) -> None:
|
|
"""Clear cooldown when an endpoint succeeds again."""
|
|
if not url:
|
|
return
|
|
_blocked_until_by_url.pop(_normalize_url(url), None)
|
|
|
|
|
|
def is_ollama_endpoint_blocked(url: str, *, now: float | None = None) -> bool:
|
|
if not url:
|
|
return False
|
|
current_time = time.monotonic() if now is None else now
|
|
normalized = _normalize_url(url)
|
|
blocked_until = _blocked_until_by_url.get(normalized)
|
|
if blocked_until is None:
|
|
return False
|
|
if blocked_until <= current_time:
|
|
_blocked_until_by_url.pop(normalized, None)
|
|
return False
|
|
return True
|
|
|
|
|
|
def get_ollama_endpoint_cooldown_remaining_seconds(
|
|
url: str,
|
|
*,
|
|
now: float | None = None,
|
|
) -> float:
|
|
"""Return remaining cooldown seconds for display/debug surfaces."""
|
|
if not url:
|
|
return 0.0
|
|
current_time = time.monotonic() if now is None else now
|
|
normalized = _normalize_url(url)
|
|
blocked_until = _blocked_until_by_url.get(normalized)
|
|
if blocked_until is None:
|
|
return 0.0
|
|
remaining = blocked_until - current_time
|
|
if remaining <= 0:
|
|
_blocked_until_by_url.pop(normalized, None)
|
|
return 0.0
|
|
return remaining
|
|
|
|
|
|
def filter_ollama_urls_with_cooldown(
|
|
urls: Iterable[str],
|
|
*,
|
|
now: float | None = None,
|
|
) -> tuple[str, ...]:
|
|
"""Return non-cooldown URLs, preserving original order and all-blocked recovery."""
|
|
ordered_urls = tuple(url for url in urls if url)
|
|
available = tuple(url for url in ordered_urls if not is_ollama_endpoint_blocked(url, now=now))
|
|
return available or ordered_urls
|
|
|
|
|
|
def resolve_ollama_order_with_cooldown(
|
|
workload_type: OllamaWorkloadType = "interactive",
|
|
*,
|
|
now: float | None = None,
|
|
) -> tuple[OllamaEndpointSelection, ...]:
|
|
"""Return resolver order with short-lived noisy-failure cooldown applied."""
|
|
order = resolve_ollama_order(workload_type)
|
|
available = tuple(
|
|
endpoint
|
|
for endpoint in order
|
|
if endpoint.url and not is_ollama_endpoint_blocked(endpoint.url, now=now)
|
|
)
|
|
return available or order
|
|
|
|
|
|
def reset_ollama_endpoint_cooldown_for_tests() -> None:
|
|
_blocked_until_by_url.clear()
|