refactor(evaluator): trim comments and extract the run session

Address review: the run-level cache is not specific to structured output,
so it moves to its own module as begin_evaluation_session(). Extract the
agent-eval preflight sweep into a function, hoist a function-local import,
and cut comments and docstrings that narrated the edits rather than the
code.

Signed-off-by: Sandy Chapman <schapman@nvidia.com>

Author: SandyChapman

packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/evaluator.py

@@ -51,7 +51,7 @@
     resolve_target_structured_output_mode,
     run_sync,
 )
-from nemo_evaluator_sdk.structured_output import structured_output_mode_session
+from nemo_evaluator_sdk.session import begin_evaluation_session
 from nemo_evaluator_sdk.execution.samples import build_metric_input
 from nemo_evaluator_sdk.inference import InferenceFn
 from nemo_evaluator_sdk.metrics.protocol import Metric, MetricWithPreflight, validate_metric_result
@@ -166,10 +166,7 @@ async def run(
         runtime_config = resolved_config.model_copy(update={"run_id": run_id})
         started_at = datetime.now(UTC)
 
-        # One detection session for the whole run: generation probes the target and scoring probes
-        # any judge model, and imported-trial runs still score, so scoping this to generation alone
-        # would leave judges probing per call.
-        async with structured_output_mode_session():
+        async with begin_evaluation_session():
             # Branch on which seam was supplied so the type checker can narrow ``target`` to a
             # concrete ``AgentEvalTarget`` without a cast.
             if trials is not None:
@@ -248,18 +245,7 @@ async def _score_trials(
             if not task.metrics:
                 raise ValueError(f"task {task.id!r} does not declare any metrics")
 
-        # Agent-eval scores metrics directly rather than through prepare_metric_for_execution, so
-        # nothing else runs their preflight. An LLM judge detects its endpoint's structured-output
-        # encoding there; without this it would score using the provisional guess from new_hooks.
-        # Deduplicated by identity because the same metric object is scored once per trial. The run
-        # session would collapse repeat probes to one request anyway; this just avoids the repeated
-        # awaits. Identity is stable here: `tasks` holds every metric for the duration of the loop.
-        preflighted: set[int] = set()
-        for task in tasks:
-            for metric in task.metrics:
-                if isinstance(metric, MetricWithPreflight) and id(metric) not in preflighted:
-                    preflighted.add(id(metric))
-                    await metric.preflight()
+        await _preflight_task_metrics(tasks)
 
         semaphore = asyncio.Semaphore(config.parallelism)
 
@@ -333,8 +319,6 @@ async def _generate_trials(
         params = _resolve_live_params(config, target)
         prompt_template = config.prompt_template or _default_prompt_template(target)
         semaphore = asyncio.Semaphore(params.parallelism)
-        # Hooks are built per row below; the run-level session opened by run() is what keeps the
-        # endpoint probe to one round trip for the whole pass instead of one per row.
 
         # Use the injected transport client when provided; otherwise build a default for the
         # resolved target type and close it when generation finishes.
@@ -402,6 +386,20 @@ async def generate_one(index: int, task: AgentEvalTask) -> AgentEvalTrial:
                 await close_client()
 
 
+async def _preflight_task_metrics(tasks: Sequence[AgentEvalTask]) -> None:
+    """Run each distinct metric's preflight once.
+
+    Agent-eval scores metrics directly rather than through ``prepare_metric_for_execution``, so
+    nothing else runs their preflight.
+    """
+    preflighted: set[int] = set()
+    for task in tasks:
+        for metric in task.metrics:
+            if isinstance(metric, MetricWithPreflight) and id(metric) not in preflighted:
+                preflighted.add(id(metric))
+                await metric.preflight()
+
+
 async def _generate_sample(
     *,
     target: Model | Agent,
@@ -423,8 +421,6 @@ async def _generate_sample(
         model_inference_fn = (
             cast(InferenceFn, inference_fn) if inference_fn is not None else inference.make_inference_request
         )
-        # Hooks are built per row here, so this relies on detection being cached per endpoint:
-        # without the probe the request would carry whichever encoding new_hooks guessed.
         await resolve_target_structured_output_mode(
             preprocess_hooks=preprocess_hooks,
             model=target,

packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/execution/backends/local/backend.py

@@ -24,7 +24,7 @@
 from nemo_evaluator_sdk.inference import PostprocessResponse, PreprocessRequest
 from nemo_evaluator_sdk.metrics.protocol import Metric
 from nemo_evaluator_sdk.resolvers import LocalModelResolver, LocalSecretResolver
-from nemo_evaluator_sdk.structured_output import structured_output_mode_session
+from nemo_evaluator_sdk.session import begin_evaluation_session
 from nemo_evaluator_sdk.values import Agent, DatasetInput, FieldMapping, Model
 from nemo_evaluator_sdk.values.multi_metric_results import BenchmarkEvaluationResult, namespace_result
 from nemo_evaluator_sdk.values.results import AggregateFieldName, EvaluationResult
@@ -85,10 +85,8 @@ async def _evaluate_one(
         Returns:
             A namespaced single-metric evaluation result.
         """
-        # Wraps preparation as well as execution: prepare_metric_for_execution runs the metric's
-        # preflight, which is where an LLM judge probes its endpoint. evaluate_metric opens a
-        # session too; the session is re-entrant, so both share this one cache.
-        async with structured_output_mode_session():
+        # Preparation runs each metric's preflight, so it belongs inside the session.
+        async with begin_evaluation_session():
             return await self._evaluate_one_in_session(
                 metric=metric,
                 metric_key=metric_key,
@@ -114,7 +112,6 @@ async def _evaluate_one_in_session(
         postprocess_hooks: tuple[PostprocessResponse, ...] | None,
         rows: list[dict[str, Any]],
     ) -> EvaluationResult:
-        """Body of :meth:`_evaluate_one`, inside a structured-output detection session."""
         prepared_metric = await prepare_metric_for_execution(
             metric,
             params=params,
@@ -168,10 +165,8 @@ async def evaluate_dataset(
         """
         rows = _prepare_rows(dataset, params, field_mapping)
         metric_keys = unique_metric_keys(metrics)
-        # Opened before metric preparation on purpose: prepare_metric_for_execution runs each
-        # metric's preflight, which is where an LLM judge probes its endpoint. Starting the session
-        # later would leave those probes uncached, so two judges on one endpoint would each probe.
-        async with structured_output_mode_session():
+        # Preparation runs each metric's preflight, so it belongs inside the session.
+        async with begin_evaluation_session():
             return await self._evaluate_dataset_in_session(
                 metrics=metrics,
                 metric_keys=metric_keys,
@@ -197,7 +192,6 @@ async def _evaluate_dataset_in_session(
         preprocess_hooks: Sequence[PreprocessRequest] | None,
         postprocess_hooks: Sequence[PostprocessResponse] | None,
     ) -> BenchmarkEvaluationResult:
-        """Body of :meth:`evaluate_dataset`, inside a structured-output detection session."""
         prepared_metrics = [
             await prepare_metric_for_execution(
                 metric,
@@ -218,16 +212,13 @@ async def _evaluate_dataset_in_session(
         else:
             merged_preprocess_hooks = tuple(preprocess_hooks or ())
             merged_postprocess_hooks = tuple(postprocess_hooks or ())
-        # This multi-metric path builds target hooks itself and calls evaluate_benchmark directly,
-        # bypassing evaluate_metric, so it has to resolve the hook here or generation would send
-        # whichever encoding new_hooks guessed.
+        # This path builds target hooks itself and bypasses evaluate_metric, so it resolves here.
         if isinstance(target, Model):
             await resolve_target_structured_output_mode(
                 preprocess_hooks=merged_preprocess_hooks,
                 model=target,
-                # Resolved from the benchmark module at call time, not bound here, so the probe
-                # always travels the exact transport that evaluate_benchmark's own default uses --
-                # including when that binding is swapped.
+                # Attribute lookup, not a bound import: the probe must use the same transport
+                # evaluate_benchmark defaults to, including when that binding is swapped.
                 inference_fn=benchmark_execution.make_inference_request,
                 params=params,
             )

packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/execution/metric_execution.py

@@ -48,7 +48,11 @@
 from nemo_evaluator_sdk.metrics.utils import metric_type_name
 from nemo_evaluator_sdk.resilience.api import run_indexed_tasks, use_resilience_session
 from nemo_evaluator_sdk.resilience.errors import get_evaluation_error
-from nemo_evaluator_sdk.structured_output import structured_output_mode_session
+from nemo_evaluator_sdk.session import begin_evaluation_session
+from nemo_evaluator_sdk.structured_output import (
+    InferenceStructuredOutput,
+    detect_structured_output_mode,
+)
 from nemo_evaluator_sdk.templates import render_request
 from nemo_evaluator_sdk.values import (
     Agent,
@@ -265,25 +269,11 @@ async def resolve_target_structured_output_mode(
     inference_fn: inference.InferenceFn,
     params: RunConfig | RunConfigOnline | RunConfigOnlineModel | None,
 ) -> None:
-    """Probe the target endpoint so generation uses an encoding it actually honours.
-
-    The judge has its own preflight; target generation had none, and previously relied on the
-    model's ``format`` label to pick an encoding. That label is deprecated and ignored, so without
-    this probe a target that only accepts ``nvext.guided_json`` would be sent ``response_format``
-    and either reject the request or silently drop the constraint.
-
-    Only the hook built from ``params.structured_output`` is retuned. A caller that hand-builds an
-    InferenceStructuredOutput and splices it in via ``preprocess_hooks`` has chosen a mode
-    deliberately, and silently overriding that choice would be surprising; ``_merge_online_hooks``
-    places the run-level hook ahead of caller-supplied ones, so the first match is ours.
-
-    Like the judge's own preflight, this runs before ``use_resilience_session()`` opens, so probe
-    attempts use the process-global scheduler and do not appear in the session summary. Detection
-    itself is cached per endpoint for the run, so a Model used as both target and judge is probed
-    once, not once per role.
-    """
-    from nemo_evaluator_sdk.structured_output import InferenceStructuredOutput, detect_structured_output_mode
+    """Probe the target endpoint and set the generation hook's encoding.
 
+    Only the hook built from ``params.structured_output`` is retuned; a caller that supplies its own
+    InferenceStructuredOutput has chosen a mode deliberately.
+    """
     if not (isinstance(params, RunConfigOnlineModel) and params.structured_output):
         return
 
@@ -308,13 +298,7 @@ def _maybe_set_default_max_tokens(
     request: dict[str, Any],
     params: RunConfigOnlineModel | None,
 ) -> None:
-    """Apply the default max token cap when neither params nor request set one.
-
-    This was previously applied only to ``nim``-format models. Model format is deprecated and no
-    longer distinguishes endpoints, so the cap is unconditional: an unbounded generation against a
-    reasoning-capable judge is a cost risk on any endpoint, and callers that want more can set
-    ``max_tokens`` or ``max_completion_tokens`` explicitly.
-    """
+    """Apply the default max token cap when neither params nor request set one."""
     inference_params = params.inference if isinstance(params, RunConfigOnlineModel) else None
     if inference_params is not None and (
         inference_params.max_tokens is not None or inference_params.max_completion_tokens is not None
@@ -868,9 +852,7 @@ async def evaluate_metric(
         log.warning("No rows found in dataset, returning empty evaluation result")
         return empty_evaluation_result()
 
-    # Scope endpoint structured-output detection to this run so probes are cached across rows
-    # without leaking a result (or a transient failure) into the next run.
-    async with structured_output_mode_session():
+    async with begin_evaluation_session():
         return await _evaluate_metric_in_session(
             metric=metric,
             rows=rows,
@@ -892,7 +874,6 @@ async def _evaluate_metric_in_session(
     preprocess_hooks: Sequence[inference.PreprocessRequest] | None,
     postprocess_hooks: Sequence[inference.PostprocessResponse] | None,
 ) -> EvaluationResult:
-    """Body of :func:`evaluate_metric`, run inside a structured-output detection session."""
     params = resolve_params(params, target)
 
     client_close_fn = None

packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/metrics/llm_judge.py

@@ -200,14 +200,8 @@ async def resolve_models(self, model_resolver: ModelResolver) -> None:
         await resolve_model_refs(self, model_resolver)
         self._client = None
         self._api_key = None
-        # `_preflight_lock` is deliberately NOT reset here. Clearing it would split preflight
-        # synchronisation while a concurrent compute_scores() is mid-probe: a second scorer would
-        # build a new lock, and an in-flight scorer could then render through the freshly rebuilt
-        # unresolved hook and send a guessed encoding with only a warning. A lock retained from a
-        # previous event loop instead fails loudly on its next contended acquire, and both cases
-        # require resolving models while scoring with the same metric object, which is unsupported:
-        # resolution is preparation-time, before evaluation starts. A loud failure beats a silent
-        # unprobed request in a change whose whole purpose is removing silent degradation.
+        # `_preflight_lock` is deliberately not reset: clearing it mid-run would split preflight
+        # synchronisation and let an in-flight scorer send an unprobed request.
         self._ensure_default_prompt_template()
         preprocess_hooks, postprocess_hooks = new_hooks(self)
         self.with_hooks(preprocess=preprocess_hooks, postprocess=postprocess_hooks)
@@ -250,11 +244,7 @@ async def preflight(self) -> None:
         _logger.info("Structured output mode selected for %s: %s", model.name, mode.value)
 
     def _require_preflight_lock(self) -> asyncio.Lock:
-        """Return this metric's preflight lock, created on first use.
-
-        Built lazily rather than as a field default so the metric stays deep-copyable and does not
-        bind a lock to an event loop at construction time.
-        """
+        """Return this metric's preflight lock, created lazily to keep the metric deep-copyable."""
         if self._preflight_lock is None:
             self._preflight_lock = asyncio.Lock()
         return self._preflight_lock
@@ -357,13 +347,8 @@ def _retry_with_max_completion_tokens(self, request: dict) -> dict:
 
     async def compute_scores(self, input: MetricInput) -> MetricResult:
         """Compute structured score output for one item/sample pair."""
-        # Executors call preflight() before scoring, but this is public API: a caller can build the
-        # metric and score with it directly, and then nothing would have probed the endpoint. Detect
-        # on first use so the encoding is never a guess.
-        #
-        # Locked and re-checked because a direct caller may gather many compute_scores() at once
-        # with no run session to cache detection: without this, every one of them would see an
-        # unresolved hook and probe the same endpoint concurrently.
+        # Public API: a caller may score without an executor having run preflight. Locked because
+        # concurrent direct callers would otherwise each probe the same endpoint.
         if self._has_unresolved_structured_output_hook():
             async with self._require_preflight_lock():
                 if self._has_unresolved_structured_output_hook():

packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/session.py

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Run-scoped cache shared by everything participating in one evaluation."""
+
+import asyncio
+from collections.abc import AsyncIterator, Hashable
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
+from typing import Any
+
+_CACHE: ContextVar[dict[Hashable, Any] | None] = ContextVar("evaluation_session_cache", default=None)
+_LOCK: ContextVar[asyncio.Lock | None] = ContextVar("evaluation_session_lock", default=None)
+
+
+@asynccontextmanager
+async def begin_evaluation_session() -> AsyncIterator[None]:
+    """Scope cached evaluation state to one run.
+
+    Entries live only inside this boundary, which is what makes caching a failed result safe: a
+    transient failure is retried by the next run instead of persisting for the life of the process.
+    Outside a session nothing is cached and every caller recomputes.
+
+    Re-entrant. Entry points nest -- a backend opens a session, then the metric executor opens
+    another -- and a fresh inner cache would recompute what the run already resolved.
+    """
+    if _CACHE.get() is not None:
+        yield
+        return
+
+    cache_token = _CACHE.set({})
+    lock_token = _LOCK.set(asyncio.Lock())
+    try:
+        yield
+    finally:
+        _CACHE.reset(cache_token)
+        _LOCK.reset(lock_token)
+
+
+def session_cache() -> dict[Hashable, Any] | None:
+    """Return the active run's cache, or None outside a session."""
+    return _CACHE.get()
+
+
+def session_lock() -> asyncio.Lock:
+    """Return the active run's lock, or a throwaway one outside a session."""
+    return _LOCK.get() or asyncio.Lock()