fix(eval): report scores on resumed Evaluation and Eval Set Run spans

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"

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),
+        },
+    )

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,40 +428,18 @@ 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,
             }
             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