From 506ade2500150d0185d8e3c9866cb21ad073a034 Mon Sep 17 00:00:00 2001 From: Chibi Date: Mon, 19 Oct 2026 10:50:49 -0700 Subject: [PATCH] fix(eval): report scores on resumed Evaluation and Eval Set Run spans On resume the runtime re-entered both spans as NonRecordingSpans, so the scores set on them were dropped and the evaluator span lost execution.id. Re-open them as recording spans under the saved ids and start time so the trace upsert replaces the suspended job's empty output. Co-Authored-By: Claude Opus 5.5 --- packages/uipath/pyproject.toml | 2 +- .../uipath/eval/runtime/_resumable_spans.py | 83 ++++++ .../uipath/src/uipath/eval/runtime/runtime.py | 246 +++--------------- .../tests/cli/eval/test_eval_resume_spans.py | 207 +++++++++++++++ .../cli/eval/test_eval_tracing_integration.py | 2 +- packages/uipath/uv.lock | 2 +- 6 files changed, 323 insertions(+), 219 deletions(-) create mode 100644 packages/uipath/src/uipath/eval/runtime/_resumable_spans.py create mode 100644 packages/uipath/tests/cli/eval/test_eval_resume_spans.py diff --git a/packages/uipath/pyproject.toml b/packages/uipath/pyproject.toml index e0b493573..b95894996 100644 --- a/packages/uipath/pyproject.toml +++ b/packages/uipath/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath" -version = "2.14.41" +version = "2.14.42" description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools." readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" diff --git a/packages/uipath/src/uipath/eval/runtime/_resumable_spans.py b/packages/uipath/src/uipath/eval/runtime/_resumable_spans.py new file mode 100644 index 000000000..a13fbdab0 --- /dev/null +++ b/packages/uipath/src/uipath/eval/runtime/_resumable_spans.py @@ -0,0 +1,83 @@ +"""Eval spans that outlive a suspended job and are finished by the resuming one.""" + +from contextlib import asynccontextmanager +from typing import Any, AsyncIterator, cast + +from opentelemetry.sdk.trace import ReadableSpan, Tracer, TracerProvider +from opentelemetry.sdk.trace.id_generator import IdGenerator +from opentelemetry.trace import Span + +from uipath.runtime import UiPathRuntimeStorageProtocol + +_NAMESPACE = "eval_parent_span" + + +class _SavedIdGenerator(IdGenerator): + def __init__(self, trace_id: int, span_id: int) -> None: + self._trace_id = trace_id + self._span_id = span_id + + def generate_trace_id(self) -> int: + return self._trace_id + + def generate_span_id(self) -> int: + return self._span_id + + +@asynccontextmanager +async def start_resumable_span( + tracer_provider: TracerProvider, + storage: UiPathRuntimeStorageProtocol | None, + runtime_id: str, + key: str, + *, + resume: bool, + name: str, + attributes: dict[str, Any], +) -> AsyncIterator[Span]: + """Start a span, or on resume re-open the one the suspended job saved under `key`.""" + tracer = tracer_provider.get_tracer(__name__) + start_time = None + saved = await _load(storage, runtime_id, key) if resume else None + # The trace backend upserts by span id, so this replaces the suspended job's output. + if saved: + cast(Tracer, tracer).id_generator = _SavedIdGenerator( + int(saved["trace_id"], 16), int(saved["span_id"], 16) + ) + start_time = saved.get("start_time") + + with tracer.start_as_current_span( + name, attributes=attributes, start_time=start_time + ) as span: + await _save(storage, runtime_id, key, span) + yield span + + +async def _load( + storage: UiPathRuntimeStorageProtocol | None, runtime_id: str, key: str +) -> dict[str, Any] | None: + if storage is None: + return None + return await storage.get_value(runtime_id=runtime_id, namespace=_NAMESPACE, key=key) + + +async def _save( + storage: UiPathRuntimeStorageProtocol | None, + runtime_id: str, + key: str, + span: Span, +) -> None: + if storage is None: + return + span_context = span.get_span_context() + await storage.set_value( + runtime_id=runtime_id, + namespace=_NAMESPACE, + key=key, + value={ + "trace_id": format(span_context.trace_id, "032x"), + "span_id": format(span_context.span_id, "016x"), + # Absent on a NonRecordingSpan (SDK disabled). + "start_time": getattr(cast(ReadableSpan, span), "start_time", None), + }, + ) diff --git a/packages/uipath/src/uipath/eval/runtime/runtime.py b/packages/uipath/src/uipath/eval/runtime/runtime.py index ff69d300c..b08d5751e 100644 --- a/packages/uipath/src/uipath/eval/runtime/runtime.py +++ b/packages/uipath/src/uipath/eval/runtime/runtime.py @@ -5,7 +5,7 @@ import json import logging from collections import defaultdict -from contextlib import contextmanager +from contextlib import AbstractAsyncContextManager, contextmanager from time import time from typing import ( Any, @@ -17,14 +17,7 @@ import coverage from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.trace import ( - NonRecordingSpan, - SpanContext, - Status, - StatusCode, - TraceFlags, - use_span, -) +from opentelemetry.trace import Span, Status, StatusCode from pydantic import BaseModel, ValidationError from uipath.core.events import EventBus @@ -78,6 +71,7 @@ ExecutionSpanProcessor, ) from ._parallelization import execute_parallel +from ._resumable_spans import start_resumable_span from ._spans import ( configure_eval_set_run_span, configure_evaluation_span, @@ -434,18 +428,6 @@ async def execute(self) -> UiPathRuntimeResult: logger.debug("=" * 80) with self._mocker_cache(): - tracer = self.trace_manager.tracer_provider.get_tracer(__name__) - - # During resume, restore the parent "Evaluation Set Run" span context - # This prevents creating duplicate eval set run spans across jobs - eval_set_parent_span = await self._restore_parent_span( - "eval_set_run", "Evaluation Set Run" - ) - - # Create "Evaluation Set Run" span or use restored parent context - # NOTE: Do NOT set execution.id on this parent span, as the mixin in - # UiPathExecutionBatchTraceProcessor propagates execution.id from parent - # to child spans, which would overwrite the per-eval execution.id span_attributes: dict[str, str | bool] = { "span_type": "eval_set_run", "uipath.custom_instrumentation": True, @@ -453,21 +435,11 @@ async def execute(self) -> UiPathRuntimeResult: if self.context.eval_set_run_id: span_attributes["eval_set_run_id"] = self.context.eval_set_run_id - eval_set_span_context_manager = ( - use_span( - eval_set_parent_span, end_on_exit=False - ) # Don't end the remote span - if eval_set_parent_span - else tracer.start_as_current_span( - "Evaluation Set Run", attributes=span_attributes - ) - ) - - with eval_set_span_context_manager as span: - await self._save_span_context_for_resume( - span, "eval_set_run", "Evaluation Set Run" - ) - + # No execution.id here: the trace processor would copy it onto every + # child and overwrite each Evaluation span's own id. + async with self._resumable_span( + "eval_set_run", "Evaluation Set Run", span_attributes + ) as span: try: ( evaluation_set, @@ -605,32 +577,17 @@ async def _execute_eval( ) -> UiPathEvalRunResult: execution_id = str(eval_item.id) - tracer = self.trace_manager.tracer_provider.get_tracer(__name__) - - # During resume, restore the parent span context from the previous execution - # This allows evaluators to be properly parented to the original "Evaluation" span - parent_span = await self._restore_parent_span(eval_item.id, "Evaluation") - - # Create "Evaluation" span or use restored parent context - # use_span() handles context management automatically (no manual attach/detach) - span_context_manager = ( - use_span(parent_span, end_on_exit=False) # Don't end the remote span - if parent_span - else tracer.start_as_current_span( - "Evaluation", - attributes={ - "execution.id": execution_id, - "span_type": "evaluation", - "eval_item_id": eval_item.id, - "eval_item_name": eval_item.name, - "uipath.custom_instrumentation": True, - }, - ) - ) - - with span_context_manager as span: - await self._save_span_context_for_resume(span, eval_item.id, "Evaluation") - + async with self._resumable_span( + eval_item.id, + "Evaluation", + { + "execution.id": execution_id, + "span_type": "evaluation", + "eval_item_id": eval_item.id, + "eval_item_name": eval_item.name, + "uipath.custom_instrumentation": True, + }, + ) as span: evaluation_run_results = UiPathEvalRunResult( evaluation_name=eval_item.name, evaluation_run_results=[] ) @@ -1203,162 +1160,19 @@ async def run_evaluator( return result - async def _restore_parent_span( - self, span_key: str, span_type: str - ) -> NonRecordingSpan | None: - """Restore parent span from storage during resume. - - Creates a NonRecordingSpan from saved span context to continue the trace - across job boundaries without creating duplicate spans. - - Args: - span_key: Storage key for the span. Examples: - - "eval_set_run" (string literal) for Evaluation Set Run span - - eval_item.id (e.g., "eval-001") for individual Evaluation span - span_type: Human-readable span type for logging (e.g., "Evaluation Set Run") - - Returns: - NonRecordingSpan if context was restored successfully, None otherwise - """ - if not self.context.resume: - return None - - saved_context = await self._get_saved_parent_span_context(span_key) - if not saved_context: - return None - - try: - trace_id = int(saved_context["trace_id"], 16) - span_id = int(saved_context["span_id"], 16) - span_context = SpanContext( - trace_id=trace_id, - span_id=span_id, - is_remote=True, - trace_flags=TraceFlags(0x01), # Sampled - ) - parent_span = NonRecordingSpan(span_context) - logger.debug( - f"EVAL RUNTIME: Restored {span_type} span context for resume - " - f"trace_id={saved_context['trace_id']}, span_id={saved_context['span_id']}" - ) - return parent_span - except Exception as e: - logger.warning( - f"EVAL RUNTIME: Failed to restore {span_type} span context: {e}" - ) - return None - - async def _save_span_context_for_resume( - self, span: Any, span_key: str, span_type: str - ) -> None: - """Save span context for retrieval during resume. - - Extracts trace_id and span_id from the span and persists them to storage - so they can be restored after suspend/resume across job boundaries. - - Args: - span: The OpenTelemetry span to save context from - span_key: Storage key for the span. Examples: - - "eval_set_run" (string literal) for Evaluation Set Run span - - eval_item.id (e.g., "eval-001") for individual Evaluation span - span_type: Human-readable span type for logging (e.g., "Evaluation") - """ - if span is None or not hasattr(span, "get_span_context"): - return - - span_context = span.get_span_context() - span_id_hex = format(span_context.span_id, "016x") - trace_id_hex = format(span_context.trace_id, "032x") - - await self._save_parent_span_context( - span_key, - { - "span_id": span_id_hex, - "trace_id": trace_id_hex, - }, - ) - - logger.debug( - f"EVAL RUNTIME: Saved {span_type} span context for resume - " - f"trace_id={trace_id_hex}, span_id={span_id_hex}" + def _resumable_span( + self, key: str, name: str, attributes: dict[str, Any] + ) -> AbstractAsyncContextManager[Span]: + return start_resumable_span( + self.trace_manager.tracer_provider, + self._storage, + self.execution_id, + key, + resume=self.context.resume, + name=name, + attributes=attributes, ) - async def _save_parent_span_context( - self, span_key: str, span_context: dict[str, str] - ) -> None: - """Save parent span context for retrieval during resume. - - Uses storage protocol from runtime factory to persist span context - across job boundaries (suspend/resume). - - Storage structure: - - runtime_id: self.execution_id (eval set run ID) - - namespace: "eval_parent_span" - - key: span_key parameter (span-specific identifier) - - Args: - span_key: Storage key for the span. Can be: - - "eval_set_run" for Evaluation Set Run span - - eval_item.id for individual Evaluation spans - span_context: Dictionary with 'span_id' and 'trace_id' keys (hex strings) - """ - if self._storage is not None: - await self._storage.set_value( - runtime_id=self.execution_id, - namespace="eval_parent_span", - key=span_key, - value=span_context, - ) - logger.debug( - f"Saved parent span context to storage for span_key={span_key}: {span_context}" - ) - else: - logger.warning( - f"No storage available, cannot persist parent span context for span_key={span_key}" - ) - - async def _get_saved_parent_span_context( - self, span_key: str - ) -> dict[str, str] | None: - """Retrieve saved parent span context for resume. - - Uses storage protocol from runtime factory to retrieve span context - persisted during suspend. - - Storage lookup: - - runtime_id: self.execution_id (eval set run ID) - - namespace: "eval_parent_span" - - key: span_key parameter (span-specific identifier) - - Args: - span_key: Storage key for the span. Can be: - - "eval_set_run" for Evaluation Set Run span - - eval_item.id for individual Evaluation spans - - Returns: - Dictionary with 'span_id' and 'trace_id' keys (hex strings), or None if not found - """ - if self._storage is not None: - context = await self._storage.get_value( - runtime_id=self.execution_id, - namespace="eval_parent_span", - key=span_key, - ) - if context: - logger.debug( - f"Retrieved parent span context from storage for span_key={span_key}: {context}" - ) - else: - logger.debug( - f"No saved parent span context found in storage for span_key={span_key}" - ) - return context - else: - logger.warning( - f"No storage available, cannot retrieve parent span context for span_key={span_key}" - ) - return None - async def dispose(self) -> None: """Cleanup runtime resources.""" pass diff --git a/packages/uipath/tests/cli/eval/test_eval_resume_spans.py b/packages/uipath/tests/cli/eval/test_eval_resume_spans.py new file mode 100644 index 000000000..27bc985e9 --- /dev/null +++ b/packages/uipath/tests/cli/eval/test_eval_resume_spans.py @@ -0,0 +1,207 @@ +"""Span output across an eval item that suspends in one job and resumes in another.""" + +import json +import uuid +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock + +import pytest +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + +from uipath.core.events import EventBus +from uipath.core.tracing import UiPathTraceManager +from uipath.eval.helpers import EvalHelpers +from uipath.eval.runtime import UiPathEvalContext, evaluate +from uipath.runtime import ( + UiPathExecuteOptions, + UiPathRuntimeFactorySettings, + UiPathRuntimeProtocol, + UiPathRuntimeResult, + UiPathRuntimeStatus, + UiPathRuntimeStorageProtocol, +) +from uipath.runtime.schema import UiPathRuntimeSchema + +EVAL_SET_PATH = str(Path(__file__).parent / "evals" / "eval-sets" / "default.json") + + +class InMemoryStorage: + def __init__(self) -> None: + self._values: dict[tuple[str, str, str], Any] = {} + + async def set_value( + self, runtime_id: str, namespace: str, key: str, value: Any + ) -> None: + self._values[(runtime_id, namespace, key)] = value + + async def get_value(self, runtime_id: str, namespace: str, key: str) -> Any: + return self._values.get((runtime_id, namespace, key)) + + +class StubRuntime: + def __init__(self, result: UiPathRuntimeResult) -> None: + self._result = result + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + return self._result + + async def stream(self, *args: Any, **kwargs: Any) -> Any: + yield self._result + + async def get_schema(self) -> UiPathRuntimeSchema: + return _schema() + + async def dispose(self) -> None: + pass + + +class StubFactory: + def __init__(self, storage: InMemoryStorage, result: UiPathRuntimeResult) -> None: + self._storage = storage + self._result = result + + def discover_entrypoints(self) -> list[str]: + return ["test"] + + async def get_storage(self) -> UiPathRuntimeStorageProtocol | None: + return self._storage + + async def get_settings(self) -> UiPathRuntimeFactorySettings | None: + return None + + async def new_runtime( + self, entrypoint: str, runtime_id: str, **kwargs: Any + ) -> UiPathRuntimeProtocol: + return StubRuntime(self._result) + + async def dispose(self) -> None: + pass + + +def _schema() -> UiPathRuntimeSchema: + return UiPathRuntimeSchema( + filePath="test.py", + uniqueId="test", + type="agent", + input={"type": "object", "properties": {}}, + output={"type": "object", "properties": {}}, + ) + + +SUSPENDED = UiPathRuntimeResult(output={}, status=UiPathRuntimeStatus.SUSPENDED) +COMPLETED = UiPathRuntimeResult( + output={"foo": "bar"}, status=UiPathRuntimeStatus.SUCCESSFUL +) + + +async def _context(resume: bool, execution_id: str) -> UiPathEvalContext: + evaluation_set, _ = EvalHelpers.load_eval_set(EVAL_SET_PATH) + context = UiPathEvalContext() + context.execution_id = execution_id + context.evaluation_set = evaluation_set + context.runtime_schema = _schema() + context.evaluators = await EvalHelpers.load_evaluators( + EVAL_SET_PATH, evaluation_set, agent_model=None + ) + context.resume = resume + return context + + +async def _run_job( + storage: InMemoryStorage, + result: UiPathRuntimeResult, + *, + resume: bool, + execution_id: str, +) -> list[ReadableSpan]: + exporter = InMemorySpanExporter() + trace_manager = UiPathTraceManager() + trace_manager.tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + event_bus = EventBus() + event_bus.publish = AsyncMock() # type: ignore[method-assign] + + await evaluate( + StubFactory(storage, result), + trace_manager, + await _context(resume, execution_id), + event_bus, + ) + return list(exporter.get_finished_spans()) + + +def _one(spans: list[ReadableSpan], span_type: str) -> ReadableSpan: + matches = [ + s for s in spans if s.attributes and s.attributes.get("span_type") == span_type + ] + assert len(matches) == 1, f"expected one {span_type} span, got {len(matches)}" + return matches[0] + + +def _output(span: ReadableSpan) -> dict[str, Any]: + assert span.attributes is not None + return json.loads(str(span.attributes["output"])) + + +@pytest.fixture +async def suspended_then_resumed() -> tuple[list[ReadableSpan], list[ReadableSpan]]: + storage = InMemoryStorage() + execution_id = str(uuid.uuid4()) + suspend_spans = await _run_job( + storage, SUSPENDED, resume=False, execution_id=execution_id + ) + resume_spans = await _run_job( + storage, COMPLETED, resume=True, execution_id=execution_id + ) + return suspend_spans, resume_spans + + +class TestResumedEvaluationSpans: + async def test_evaluation_span_reports_scores(self, suspended_then_resumed): + _, resume_spans = suspended_then_resumed + + evaluation = _one(resume_spans, "evaluation") + + assert _output(evaluation)["scores"] == {"ExactMatchEvaluator": 100.0} + + async def test_eval_set_run_span_reports_scores(self, suspended_then_resumed): + _, resume_spans = suspended_then_resumed + + eval_set_run = _one(resume_spans, "eval_set_run") + + assert _output(eval_set_run)["scores"] == {"ExactMatchEvaluator": 100.0} + + async def test_resumed_spans_keep_original_ids_and_start( + self, suspended_then_resumed + ): + suspend_spans, resume_spans = suspended_then_resumed + + for span_type in ("eval_set_run", "evaluation"): + before = _one(suspend_spans, span_type) + after = _one(resume_spans, span_type) + assert before.context is not None and after.context is not None + assert after.context.span_id == before.context.span_id + assert after.context.trace_id == before.context.trace_id + assert after.parent == before.parent + assert after.start_time == before.start_time + + async def test_evaluator_span_keeps_execution_id(self, suspended_then_resumed): + _, resume_spans = suspended_then_resumed + + evaluation = _one(resume_spans, "evaluation") + evaluator = _one(resume_spans, "evaluator") + + assert evaluator.attributes is not None and evaluation.attributes is not None + assert evaluator.attributes.get("execution.id") == evaluation.attributes.get( + "execution.id" + ) + assert evaluator.parent is not None and evaluation.context is not None + assert evaluator.parent.span_id == evaluation.context.span_id diff --git a/packages/uipath/tests/cli/eval/test_eval_tracing_integration.py b/packages/uipath/tests/cli/eval/test_eval_tracing_integration.py index b106d52ef..13c8f81a7 100644 --- a/packages/uipath/tests/cli/eval/test_eval_tracing_integration.py +++ b/packages/uipath/tests/cli/eval/test_eval_tracing_integration.py @@ -47,7 +47,7 @@ def __init__(self) -> None: @contextmanager def start_as_current_span( - self, name: str, attributes: dict[str, Any] | None = None + self, name: str, attributes: dict[str, Any] | None = None, **kwargs: Any ): """Capture span creation and yield a mock span.""" # Create MockSpan first so we can reference its attributes diff --git a/packages/uipath/uv.lock b/packages/uipath/uv.lock index 1be243190..19132255f 100644 --- a/packages/uipath/uv.lock +++ b/packages/uipath/uv.lock @@ -2599,7 +2599,7 @@ wheels = [ [[package]] name = "uipath" -version = "2.14.41" +version = "2.14.42" source = { editable = "." } dependencies = [ { name = "applicationinsights" },