nemo-platform-plugin 0.1.0__py3-none-any.whl

nemo_platform/beta/evaluator/resilience/api.py

@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Public API boundary for evaluator resilience scheduling."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
+from typing import ParamSpec, TypeVar, cast
+
+from nemo_platform.beta.evaluator.resilience.config import ResilienceConfig
+from nemo_platform.beta.evaluator.resilience.scheduler import ResilienceScheduler
+
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+_logger = logging.getLogger(__name__)
+
+_current_scheduler: ContextVar[ResilienceScheduler | None] = ContextVar("resilience_scheduler", default=None)
+_default_scheduler = ResilienceScheduler(ResilienceConfig())
+
+
+def _active_scheduler() -> ResilienceScheduler:
+    scheduler = _current_scheduler.get()
+    if scheduler is not None:
+        return scheduler
+    return _default_scheduler
+
+
+@asynccontextmanager
+async def use_resilience_session(
+    *,
+    global_limit: int | None = None,
+    endpoint_max_limit: int | None = None,
+) -> AsyncIterator[None]:
+    """Override active scheduler for the current async context.
+
+    This uses a ContextVar-backed session boundary so concurrent requests/jobs can
+    each have their own scheduler instance and concurrency cap.
+
+    Note:
+        Limits are global only within this session. They are not process-wide caps
+        shared across concurrent evaluator jobs. Service-level coordination remains
+        out of scope for this process-local V2 design.
+    """
+    base = ResilienceConfig()
+    scheduler = ResilienceScheduler(
+        ResilienceConfig(
+            global_limit=max(1, global_limit) if global_limit is not None else base.global_limit,
+            endpoint_max_limit=max(1, endpoint_max_limit)
+            if endpoint_max_limit is not None
+            else base.endpoint_max_limit,
+        )
+    )
+    token = _current_scheduler.set(scheduler)
+    try:
+        yield
+    finally:
+        summary = await scheduler.summary()
+        _logger.info("Resilience session summary", extra=summary)
+        await scheduler.shutdown()
+        _current_scheduler.reset(token)
+
+
+async def run_with_resilience(
+    endpoint_key: str,
+    operation: Callable[_P, Awaitable[_T]],
+    *args: _P.args,
+    max_attempts: int,
+    deadline_at: float | None = None,
+    **kwargs: _P.kwargs,
+) -> _T:
+    """Execute an outbound attempt with scheduler-managed retries/admission."""
+    scheduler = _active_scheduler()
+    effective_attempts = max(1, max_attempts)
+    return await scheduler.run_with_resilience(
+        endpoint_key,
+        operation,
+        *args,
+        max_attempts=effective_attempts,
+        deadline_at=deadline_at,
+        **kwargs,
+    )
+
+
+async def run_indexed_tasks(
+    indices: Sequence[int],
+    worker: Callable[[int], Awaitable[_T]],
+    *,
+    parallelism: int,
+) -> list[_T]:
+    """Execute index-keyed work with bounded in-flight task dispatch.
+
+    Callers provide item indices and an async worker; this helper schedules up to
+    `parallelism` tasks at a time, preserving result ordering by input position.
+
+    Notes:
+        `parallelism` is a hard cap on active row-level tasks for a session.
+        Scheduler admission still applies inside each task, so endpoint/global
+        shedding can reduce actual outbound call concurrency below this value.
+        Keeping this outer cap bounds task memory/CPU overhead while retaining
+        adaptive network pressure control in the scheduler.
+    """
+    if not indices:
+        return []
+
+    results: list[_T | None] = [None] * len(indices)
+    max_inflight = min(len(indices), max(1, parallelism))
+    _logger.debug(
+        "Resilience indexed task execution started",
+        extra={"task_count": len(indices), "worker_parallelism": parallelism, "max_inflight_workers": max_inflight},
+    )
+    next_position = 0
+    in_flight: dict[asyncio.Task[_T], int] = {}
+
+    while next_position < len(indices) or in_flight:
+        try:
+            while next_position < len(indices) and len(in_flight) < max_inflight:
+                position = next_position
+                index = indices[position]
+                in_flight[asyncio.create_task(worker(index))] = position
+                next_position += 1
+
+            done, _ = await asyncio.wait(in_flight.keys(), return_when=asyncio.FIRST_COMPLETED)
+            for task in done:
+                position = in_flight.pop(task)
+                results[position] = task.result()
+        except Exception:
+            for task in in_flight:
+                task.cancel()
+            if in_flight:
+                await asyncio.gather(*in_flight.keys(), return_exceptions=True)
+            raise
+
+    _logger.info(
+        "Resilience indexed task execution completed",
+        extra={"task_count": len(indices), "worker_parallelism": parallelism, "max_inflight_workers": max_inflight},
+    )
+    return cast(list[_T], results)

nemo_platform/beta/evaluator/resilience/classifier.py

@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Exception classification and endpoint identity helpers for resilience."""
+
+from __future__ import annotations
+
+import hashlib
+from datetime import UTC, datetime
+from email.utils import parsedate_to_datetime
+
+import httpx
+import openai
+from nemo_platform.beta.evaluator.resilience.types import ClassifierResult, FailureClass
+
+_HARD_OVERLOAD_STATUS_CODES = frozenset({429, 503})
+_TRANSIENT_STATUS_CODES = frozenset({408, 500, 502, 504})
+
+
+def endpoint_identity(base_url: str, model_id: str | None = None, auth_identity: str | None = None) -> str:
+    """Build a stable endpoint key for scheduler state and accounting."""
+    auth_hash = ""
+    if auth_identity:
+        auth_hash = hashlib.sha256(auth_identity.encode("utf-8")).hexdigest()[:16]
+    return f"{base_url}|{model_id or '_'}|{auth_hash}"
+
+
+def _status_code_from_exception(exc: Exception) -> int | None:
+    """Return HTTP status code if present on a known exception type."""
+    if isinstance(exc, openai.APIStatusError):
+        return exc.status_code
+    if isinstance(exc, httpx.HTTPStatusError):
+        return exc.response.status_code
+    return None
+
+
+def _parse_retry_after_value(value: str | None) -> float | None:
+    """Parse `Retry-After` header values into seconds."""
+    if value is None:
+        return None
+    raw = value.strip()
+    try:
+        parsed = float(raw)
+    except ValueError:
+        try:
+            when = parsedate_to_datetime(raw)
+        except (TypeError, ValueError):
+            return None
+        if when.tzinfo is None:
+            when = when.replace(tzinfo=UTC)
+        parsed = (when - datetime.now(tz=UTC)).total_seconds()
+    if parsed < 0:
+        return 0.0
+    return parsed
+
+
+def _retry_after_from_headers(headers: httpx.Headers | None) -> float | None:
+    """Extract and parse `Retry-After` from an HTTP header mapping."""
+    if not headers:
+        return None
+    return _parse_retry_after_value(headers.get("Retry-After") or headers.get("retry-after"))
+
+
+def _retry_after_from_exception(exc: Exception) -> float | None:
+    """Extract `Retry-After` seconds from supported exception types."""
+    if isinstance(exc, httpx.HTTPStatusError):
+        return _retry_after_from_headers(exc.response.headers)
+    if isinstance(exc, openai.APIStatusError):
+        response = getattr(exc, "response", None)
+        headers = getattr(response, "headers", None)
+        if isinstance(headers, httpx.Headers):
+            return _retry_after_from_headers(headers)
+        if isinstance(headers, dict):
+            return _parse_retry_after_value(headers.get("Retry-After") or headers.get("retry-after"))
+    return None
+
+
+def classify_exception(exc: Exception) -> ClassifierResult:
+    """Classify errors into retry/failure policy buckets."""
+    retry_after = _retry_after_from_exception(exc)
+    status_code = _status_code_from_exception(exc)
+    error_type = type(exc).__name__
+
+    if isinstance(exc, (httpx.ReadTimeout, httpx.WriteTimeout, httpx.PoolTimeout, openai.APITimeoutError)):
+        return ClassifierResult(
+            failure_class=FailureClass.SOFT_OVERLOAD,
+            retryable=True,
+            retry_after_seconds=retry_after,
+            status_code=status_code,
+            error_type=error_type,
+        )
+    if isinstance(exc, httpx.ConnectTimeout):
+        return ClassifierResult(
+            failure_class=FailureClass.TRANSIENT,
+            retryable=True,
+            retry_after_seconds=retry_after,
+            status_code=status_code,
+            error_type=error_type,
+        )
+    if isinstance(exc, (httpx.NetworkError, openai.APIConnectionError)):
+        return ClassifierResult(
+            failure_class=FailureClass.TRANSIENT,
+            retryable=True,
+            retry_after_seconds=retry_after,
+            status_code=status_code,
+            error_type=error_type,
+        )
+
+    if status_code in _HARD_OVERLOAD_STATUS_CODES:
+        return ClassifierResult(
+            failure_class=FailureClass.HARD_OVERLOAD,
+            retryable=True,
+            retry_after_seconds=retry_after,
+            status_code=status_code,
+            error_type=error_type,
+        )
+    if status_code in _TRANSIENT_STATUS_CODES:
+        return ClassifierResult(
+            failure_class=FailureClass.TRANSIENT,
+            retryable=True,
+            retry_after_seconds=retry_after,
+            status_code=status_code,
+            error_type=error_type,
+        )
+
+    # Rate-limit style errors without explicit status code handling above.
+    if isinstance(exc, (openai.RateLimitError,)):
+        return ClassifierResult(
+            failure_class=FailureClass.HARD_OVERLOAD,
+            retryable=True,
+            retry_after_seconds=retry_after,
+            status_code=status_code,
+            error_type=error_type,
+        )
+
+    return ClassifierResult(
+        failure_class=FailureClass.FATAL,
+        retryable=False,
+        retry_after_seconds=retry_after,
+        status_code=status_code,
+        error_type=error_type,
+    )

nemo_platform/beta/evaluator/resilience/config.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Internal resilience policy configuration for evaluator outbound calls."""
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ResilienceConfig:
+    """Tuning knobs for retry + admission control.
+
+    These settings are internal to evaluator and are not user-facing API.
+
+    Attributes:
+        global_limit: Hard cap on concurrent outbound attempts for one resilience
+            session (job/request context), not a process-wide service cap.
+        endpoint_initial_limit: Starting per-endpoint concurrency limit.
+        endpoint_min_limit: Lower bound for per-endpoint adaptive concurrency.
+        endpoint_max_limit: Upper bound for per-endpoint adaptive concurrency.
+        success_window: Number of successes needed before additive increase (+1) is applied.
+        beta_hard_overload: Multiplicative decrease factor for hard overload failures.
+        beta_soft_overload: Multiplicative decrease factor for soft overload failures (timeouts).
+        cooldown_seconds_hard: Cooldown duration after hard overload feedback.
+        cooldown_seconds_soft: Cooldown duration after soft overload feedback.
+        timeout_soft_overload_escalation_count: Soft overload events needed to escalate to hard behavior.
+        escalation_window_seconds: Sliding window used for soft-overload escalation counting.
+        backoff_initial_ms: Initial retry backoff in milliseconds.
+        backoff_cap_ms: Maximum retry backoff in milliseconds.
+        pressure_gain_k: Pressure multiplier gain for retry wait inflation.
+        max_attempts_default: Default maximum attempts if callsites do not override.
+        task_deadline_seconds_default: Default per-task retry deadline.
+        global_max_queued: Max queued operations across all endpoints before rejecting work.
+        endpoint_max_queued: Max queued operations per endpoint before rejecting work.
+        retry_budget_tokens_per_sec: Retry token refill rate per endpoint.
+        retry_budget_burst: Retry token burst capacity per endpoint.
+        shutdown_grace_seconds: Grace period for graceful scheduler shutdown.
+        endpoint_state_max_entries: Maximum endpoint states retained in memory.
+        endpoint_state_ttl_seconds: TTL for endpoint states before eviction.
+    """
+
+    global_limit: int = 64
+    endpoint_initial_limit: int = 4
+    endpoint_min_limit: int = 1
+    endpoint_max_limit: int = 64
+
+    success_window: int = 20
+    beta_hard_overload: float = 0.5
+    beta_soft_overload: float = 0.7
+    cooldown_seconds_hard: float = 3.0
+    cooldown_seconds_soft: float = 1.5
+    timeout_soft_overload_escalation_count: int = 3
+    escalation_window_seconds: float = 10.0
+
+    backoff_initial_ms: float = 250.0
+    backoff_cap_ms: float = 15_000.0
+    pressure_gain_k: float = 1.0
+    max_attempts_default: int = 3
+    task_deadline_seconds_default: float = 60.0
+
+    global_max_queued: int = 8_192
+    endpoint_max_queued: int = 1_024
+    retry_budget_tokens_per_sec: float = 5.0
+    retry_budget_burst: float = 10.0
+
+    shutdown_grace_seconds: float = 15.0
+
+    endpoint_state_max_entries: int = 4_096
+    endpoint_state_ttl_seconds: float = 300.0

nemo_platform/beta/evaluator/resilience/errors.py

@@ -0,0 +1,117 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Exception helpers shared by evaluator resilience/task flows."""
+
+from collections.abc import Iterator
+from typing import TypeVar
+
+from jinja2.exceptions import UndefinedError as JinjaUndefinedError
+from nemo_platform.beta.evaluator.execution.values import EvaluationError
+
+E = TypeVar("E", bound=BaseException)
+
+
+def iter_leaf_causes(exc: BaseException) -> Iterator[BaseException]:
+    """Yield all leaf exceptions from an exception-group tree in left-to-right DFS order.
+
+    Recurses into any exception exposing a non-empty tuple ``.exceptions``
+    attribute — matches ``BaseExceptionGroup`` on Python >= 3.11 and the
+    ``exceptiongroup`` backport. Group nesting reflects how tasks were grouped
+    (e.g. nested ``TaskGroup``s), not causation — leaves are independent,
+    concurrent failures.
+
+    Example:
+        Nested group — each branch is fully traversed (children before
+        siblings) before moving to the next sibling at the parent level::
+
+            group = ExceptionGroup(
+                "outer",
+                [ValueError("a"),
+                 ExceptionGroup("inner", [RuntimeError("b"), KeyError("c")]),
+                 TypeError("d")],
+            )
+            list(iter_leaf_causes(group))
+            # [ValueError("a"), RuntimeError("b"), KeyError("c"), TypeError("d")]
+    """
+    stack: list[BaseException] = [exc]
+    while stack:
+        current = stack.pop()
+        exceptions = getattr(current, "exceptions", None)
+        if not isinstance(exceptions, tuple) or not exceptions:
+            yield current
+            continue
+        if any(not isinstance(child, BaseException) for child in exceptions):
+            # Defensive: a non-exception child means this node is effectively a leaf.
+            yield current
+            continue
+        # Push reversed so the leftmost child is popped first (preserves DFS left→right).
+        stack.extend(reversed(exceptions))
+
+
+def first_failure_cause(exc: BaseException) -> BaseException:
+    """Return the first leaf failure from an exception/exception-group tree.
+
+    Use case: surface *a* concrete failure in a log line or wrapped message,
+    without caring which sibling comes first.
+    """
+    return next(iter_leaf_causes(exc), exc)
+
+
+def find_cause(exc: BaseException, cls: type[E]) -> E | None:
+    """Return the first leaf of type ``cls`` anywhere in the exception tree, or ``None``.
+
+    Unlike :func:`first_failure_cause`, which returns only the first leaf in
+    traversal order, ``find_cause`` walks every leaf of any ``ExceptionGroup``
+    via :func:`iter_leaf_causes`. That means a matching exception is found even
+    when it is not the leading sibling or sits under a non-leading branch. The
+    return type is narrowed to ``cls | None`` via ``TypeVar``, avoiding a
+    duplicate ``isinstance`` check at the call site. Common with
+    ``asyncio.TaskGroup``, where sibling order depends on task scheduling and
+    is not deterministic.
+
+    Example:
+        Nested group — the ``EvaluationError`` is buried under a non-leading
+        branch, so ``first_failure_cause`` misses it::
+
+            group = ExceptionGroup("outer", [
+                RuntimeError("sibling"),
+                ExceptionGroup("inner", [
+                    RuntimeError("noise"),
+                    EvaluationError(index=5, message="template error"),
+                ]),
+            ])
+            first_failure_cause(group)          # RuntimeError("sibling")
+            find_cause(group, EvaluationError)  # EvaluationError(index=5, ...)
+    """
+    for leaf in iter_leaf_causes(exc):
+        if isinstance(leaf, cls):
+            return leaf
+    return None
+
+
+def normalize_evaluation_failure(
+    exc: BaseException,
+    *,
+    prefix: str = "Metric evaluation has failed",
+) -> RuntimeError:
+    """Convert queue/task execution failures into the public evaluator error shape."""
+    root = first_failure_cause(exc)
+    if isinstance(root, JinjaUndefinedError):
+        return RuntimeError(f"{prefix} due to templating error: {str(root)}")
+    if isinstance(exc, JinjaUndefinedError):
+        return RuntimeError(f"{prefix} due to templating error: {str(exc)}")
+    return RuntimeError(f"{prefix} with error: {str(root) or root.__class__.__name__}")
+
+
+def get_evaluation_error(exc: BaseException) -> EvaluationError | RuntimeError:
+    """Classify a pipeline failure into an ``EvaluationError`` or normalized ``RuntimeError``.
+
+    Returns the first ``EvaluationError`` found anywhere in the exception tree
+    (via :func:`find_cause`) if present; otherwise falls back to
+    :func:`normalize_evaluation_failure`. Callers decide how to re-raise.
+    """
+    evaluation_error = find_cause(exc, EvaluationError)
+    if evaluation_error is not None:
+        return evaluation_error
+    return normalize_evaluation_failure(exc)

nemo_platform/beta/evaluator/resilience/policy.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Policy math for resilience retry waits and endpoint adaptation."""
+
+from __future__ import annotations
+
+import math
+import random
+
+from nemo_platform.beta.evaluator.resilience.config import ResilienceConfig
+from nemo_platform.beta.evaluator.resilience.types import FailureClass, RetryContext
+
+
+def pressure(limit: int, max_limit: int) -> float:
+    """Compute normalized pressure (0..1) from endpoint capacity state."""
+    if max_limit <= 0:
+        return 1.0
+    return max(0.0, min(1.0, 1.0 - (limit / max_limit)))
+
+
+def retry_wait_seconds(context: RetryContext, config: ResilienceConfig) -> float:
+    """Compute retry delay with jitter, pressure scaling, and overload floors."""
+    exp = min(
+        config.backoff_cap_ms / 1000.0, (config.backoff_initial_ms / 1000.0) * (2 ** max(0, context.attempt_number - 1))
+    )
+    jittered = random.uniform(0.0, exp)
+    pressure_mult = 1.0 + (config.pressure_gain_k * context.pressure)
+    computed = jittered * pressure_mult
+    server_floor = context.retry_after_seconds or 0.0
+    return max(server_floor, context.cooldown_remaining_seconds, computed)
+
+
+def additive_increase(limit: int, max_limit: int) -> int:
+    """Increase endpoint limit by one within configured bounds."""
+    return min(max_limit, limit + 1)
+
+
+def multiplicative_decrease(limit: int, min_limit: int, beta: float) -> int:
+    """Reduce endpoint limit by multiplicative factor within bounds."""
+    return max(min_limit, math.ceil(limit * beta))
+
+
+def cooldown_for_failure_class(failure_class: FailureClass, config: ResilienceConfig) -> float:
+    """Return cooldown period in seconds for a failure class."""
+    if failure_class == FailureClass.HARD_OVERLOAD:
+        return config.cooldown_seconds_hard
+    if failure_class == FailureClass.SOFT_OVERLOAD:
+        return config.cooldown_seconds_soft
+    return 0.0