diff --git a/fern/versions/latest/pages/concepts/architecture-and-performance.mdx b/fern/versions/latest/pages/concepts/architecture-and-performance.mdx index 9fccfc5a4..b3106be80 100644 --- a/fern/versions/latest/pages/concepts/architecture-and-performance.mdx +++ b/fern/versions/latest/pages/concepts/architecture-and-performance.mdx @@ -399,6 +399,55 @@ except DataDesignerEarlyShutdownError: ... ``` +## Local OpenTelemetry Metrics + +`DataDesigner.create()` starts or reuses a process-wide OpenTelemetry Prometheus endpoint on `http://127.0.0.1:9464/metrics` by default. Use `RunConfig` to disable metrics for an invocation or choose another port: + +```python +import data_designer.config as dd +from data_designer.interface import DataDesigner + +designer = DataDesigner() +designer.set_run_config(dd.RunConfig(otel_metrics_port=None)) + +# Or keep metrics enabled on a different loopback port. +designer.set_run_config(dd.RunConfig(otel_metrics_port=9465)) +``` + +Configure a Prometheus server on the same host to scrape the endpoint: + +```yaml +scrape_configs: + - job_name: data-designer + metrics_path: /metrics + static_configs: + - targets: ["127.0.0.1:9464"] +``` + +The endpoint exposes seven instruments. Prometheus derives throughput and request-completion rates from their counters and histogram counts. + +| Instrument | Type | Meaning | +|------------|------|---------| +| `data_designer.create.duration` | Histogram | Create-job duration in seconds, with failure-only `error.type`. | +| `data_designer.dataset.records` | Counter | Checkpointed generated and dropped records, distinguished by `record.result`; derive record throughput from this counter. | +| `data_designer.dataset.progress` | Gauge | Fraction of scheduled records processed by the current create job, or across all active jobs when creates overlap. The most recently finished value remains visible until the next job starts. | +| `data_designer.scheduler.events` | Counter | Job start/completion, non-retryable task errors, worker spawn failures, deferred retries, and task cancellations, identified by `event.name` with bounded `outcome` and failure-only `error.type` where applicable. | +| `data_designer.model.request.active` | UpDownCounter | Model requests currently in progress, identified by `gen_ai.operation.name`, `gen_ai.provider.name`, and `gen_ai.request.model`. | +| `gen_ai.client.operation.duration` | Histogram | Completed model-request duration in seconds, identified by `gen_ai.operation.name`, `gen_ai.provider.name`, and `gen_ai.request.model`, with failure-only `error.type`. | +| `data_designer.log.records` | Counter | Data Designer log-record counts by normalized `log.severity`. | + +The listener and cumulative metrics remain available until process shutdown so a scrape can collect results after a job finishes. A later enabled job can rebind an idle listener to a different configured port without resetting those metrics. Concurrent jobs that request different ports share the active listener and emit a warning instead of replacing it. + +Generated and dropped record counts advance when a row group is durably checkpointed, so `data_designer.dataset.progress` moves in `buffer_size`-sized steps. Use `data_designer.model.request.active` and the `gen_ai.client.operation.duration` histogram count for live request activity between checkpoints. + +Setting `otel_metrics_port=None` records nothing for that invocation. If an earlier enabled job already started the process listener, disabling a later job does not stop the listener or remove the earlier cumulative metrics. + + +The endpoint binds only to loopback and has no authentication. It is intended for a Prometheus server running on the same host, not for remote exposure. + + +Prometheus pull is metrics-only: `data_designer.log.records` counts safe log metadata, not raw log bodies. Existing raw logs continue to use the configured stdout and file handlers and should be collected from those outputs. The Python Prometheus exporter does not support multiprocessing, so this endpoint is unsupported for multiprocessing-based collection. + ## Common Problems | Problem | Symptom | Solution | diff --git a/packages/data-designer-config/src/data_designer/config/run_config.py b/packages/data-designer-config/src/data_designer/config/run_config.py index 455d27b28..b2585657b 100644 --- a/packages/data-designer-config/src/data_designer/config/run_config.py +++ b/packages/data-designer-config/src/data_designer/config/run_config.py @@ -159,6 +159,10 @@ class RunConfig(ConfigBase): non-TTY environments. Default is False. progress_interval: How often (in seconds) the async progress reporter emits a consolidated log block. Must be > 0. Default is 5.0. + otel_metrics_port: Loopback port for the pull-based OpenTelemetry metrics endpoint. + Set to None to disable instrumentation for this create invocation. An endpoint + already opened by the process may remain available for metrics from prior runs. + The endpoint exposes metrics, not raw log records. Default is 9464. preserve_dropped_columns: If True, write columns removed by drop processors to separate dropped-column parquet files. Set to False to omit those artifacts while still removing dropped columns from the final dataset. Default is True. @@ -200,6 +204,16 @@ class RunConfig(ConfigBase): write_scheduler_events: bool = False display_tui: bool = False progress_interval: float = Field(default=5.0, gt=0.0) + otel_metrics_port: int | None = Field( + default=9464, + ge=1, + le=65535, + description=( + "Loopback port for the pull-based OpenTelemetry metrics endpoint. " + "None disables instrumentation for this create invocation; an existing process endpoint may remain " + "available for prior metrics. Raw log records are not exposed." + ), + ) preserve_dropped_columns: bool = Field( default=True, description=( diff --git a/packages/data-designer-config/tests/config/test_run_config.py b/packages/data-designer-config/tests/config/test_run_config.py index 0d4368c61..daf1a30c6 100644 --- a/packages/data-designer-config/tests/config/test_run_config.py +++ b/packages/data-designer-config/tests/config/test_run_config.py @@ -40,6 +40,36 @@ def test_run_config_accepts_scheduler_event_writes() -> None: assert RunConfig(write_scheduler_events=True).write_scheduler_events is True +def test_run_config_defaults_otel_metrics_port_to_9464() -> None: + assert RunConfig().otel_metrics_port == 9464 + + +def test_run_config_accepts_custom_otel_metrics_port() -> None: + assert RunConfig(otel_metrics_port=4318).otel_metrics_port == 4318 + + +def test_run_config_accepts_disabled_otel_metrics() -> None: + assert RunConfig(otel_metrics_port=None).otel_metrics_port is None + + +@pytest.mark.parametrize("otel_metrics_port", [1, 65535]) +def test_run_config_accepts_otel_metrics_port_bounds(otel_metrics_port: int) -> None: + assert RunConfig(otel_metrics_port=otel_metrics_port).otel_metrics_port == otel_metrics_port + + +@pytest.mark.parametrize("otel_metrics_port", [0, 65536]) +def test_run_config_rejects_otel_metrics_port_outside_bounds(otel_metrics_port: int) -> None: + with pytest.raises(ValidationError, match="otel_metrics_port"): + RunConfig(otel_metrics_port=otel_metrics_port) + + +def test_run_config_preserves_otel_metrics_port_when_serialized() -> None: + serialized = RunConfig(otel_metrics_port=4318).model_dump() + + assert serialized["otel_metrics_port"] == 4318 + assert RunConfig.model_validate(serialized).otel_metrics_port == 4318 + + def test_run_config_progress_bar_shim_translates_to_display_tui() -> None: with pytest.warns(DeprecationWarning, match="RunConfig.progress_bar.*RunConfig.display_tui") as caught: run_config = RunConfig(progress_bar=False) diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py index c22b6dbbb..8b6be7889 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py @@ -77,8 +77,10 @@ from data_designer.engine.observability import ( RuntimeCorrelation, SchedulerAdmissionEvent, + SchedulerAdmissionEventKind, SchedulerAdmissionEventSink, runtime_correlation_provider, + scheduler_event_sink_accepts, ) from data_designer.engine.processing.ginja.exceptions import UserTemplateError from data_designer.engine.progress.reporter import ( @@ -235,7 +237,14 @@ def __init__( self._in_flight: set[Task] = set() self._worker_tasks: set[asyncio.Task] = set() self._wake_event = asyncio.Event() - self._run_id = run_id or f"run-{uuid.uuid4().hex}" + ambient_correlation = runtime_correlation_provider.current() + self._run_id = ( + run_id + if run_id is not None + else ambient_correlation.run_id + if ambient_correlation is not None + else f"run-{uuid.uuid4().hex}" + ) self._scheduler_event_sink = scheduler_event_sink self._scheduler_event_sequence = 0 if salvage_max_rounds < 1: @@ -448,7 +457,7 @@ def _spawn_worker(self, coro: Coroutine[Any, Any, None]) -> asyncio.Task: def _emit_scheduler_event( self, - event_kind: str, + event_kind: SchedulerAdmissionEventKind, *, task: Task | None = None, lease: TaskAdmissionLease | None = None, @@ -457,10 +466,18 @@ def _emit_scheduler_event( reason_or_result: str | None = None, diagnostics: dict[str, object] | None = None, ) -> None: - if self._scheduler_event_sink is None: + if not scheduler_event_sink_accepts(self._scheduler_event_sink, event_kind): return self._scheduler_event_sequence += 1 - correlation = None + correlation = RuntimeCorrelation( + run_id=self._run_id, + row_group=None, + task_column=None, + task_type=None, + scheduling_group_kind=None, + scheduling_group_identity_hash=None, + task_execution_id=None, + ) event_diagnostics = dict(diagnostics or {}) if task is not None: schedulable = lease.item if lease is not None else self._schedulable_task(task) @@ -480,7 +497,7 @@ def _emit_scheduler_event( try: self._scheduler_event_sink.emit_scheduler_event( SchedulerAdmissionEvent.capture( - event_kind, # type: ignore[arg-type] + event_kind, sequence=self._scheduler_event_sequence, correlation=correlation, task_id=stable_task_id(task) if task is not None else None, @@ -530,7 +547,7 @@ def _record_observed_task_state(self) -> None: ) def _emit_scheduler_health_snapshot(self, reason: str) -> None: - if self._scheduler_event_sink is None: + if not scheduler_event_sink_accepts(self._scheduler_event_sink, "scheduler_health_snapshot"): return self._emit_scheduler_event( "scheduler_health_snapshot", @@ -1244,9 +1261,11 @@ async def run(self) -> None: num_rgs = len(self._row_groups) self._run_loop = asyncio.get_running_loop() + outcome = "success" + error_type: str | None = None - with self._progress_bar or contextlib.nullcontext(): - try: + try: + with self._progress_bar or contextlib.nullcontext(): if self._reporter: self._reporter.log_start(num_row_groups=num_rgs, scheduled_records=self._scheduled_records) @@ -1284,9 +1303,6 @@ async def run(self) -> None: self._reporter.log_final() self._emit_scheduler_health_snapshot("completed") - self._emit_scheduler_event( - "scheduler_job_completed", diagnostics=self._scheduler_health_diagnostics(reason="completed") - ) if self._rg_states: incomplete = list(self._rg_states) @@ -1294,10 +1310,31 @@ async def run(self) -> None: f"Scheduler exited with {len(self._rg_states)} unfinished row group(s): {incomplete}. " "These row groups were not checkpointed." ) - finally: - if self._reporter: - self._reporter.close() - self._run_loop = None + except asyncio.CancelledError: + outcome = "cancelled" + raise + except BaseException as exc: + outcome = "error" + error_type = type(exc).__name__ + raise + finally: + self._run_loop = None + if outcome == "success" and self._early_shutdown: + outcome = "early_shutdown" + terminal_diagnostics: dict[str, object] = {"outcome": outcome} + if outcome in {"success", "early_shutdown"} and scheduler_event_sink_accepts( + self._scheduler_event_sink, "scheduler_job_completed" + ): + terminal_diagnostics = self._scheduler_health_diagnostics(reason="completed") | terminal_diagnostics + if error_type is not None: + terminal_diagnostics["error_type"] = error_type + self._emit_scheduler_event( + "scheduler_job_completed", + reason_or_result=outcome, + diagnostics=terminal_diagnostics, + ) + if self._reporter: + self._reporter.close() async def _main_dispatch_loop( self, diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py index 60c6e49b8..13d23e383 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py @@ -50,7 +50,11 @@ from data_designer.engine.dataset_builders.utils.processor_runner import ProcessorRunner, ProcessorStage from data_designer.engine.dataset_builders.utils.row_group_buffer import RowGroupBufferManager from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler -from data_designer.engine.observability import JsonlSchedulerEventSink, SchedulerAdmissionEventSink +from data_designer.engine.observability import ( + JsonlSchedulerEventSink, + SchedulerAdmissionEventSink, + fanout_scheduler_event_sinks, +) from data_designer.engine.processing.processors.base import Processor from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor from data_designer.engine.readiness import run_readiness_check @@ -984,7 +988,10 @@ def on_before_checkpoint(rg_id: int, rg_size: int) -> None: initial_completed_records=initial_actual_num_records, progress_interval=self._resource_provider.run_config.progress_interval, display_tui=self._resource_provider.run_config.display_tui, - scheduler_event_sink=scheduler_event_sink, + scheduler_event_sink=fanout_scheduler_event_sinks( + scheduler_event_sink, + self._resource_provider.scheduler_event_sink, + ), request_pressure_provider=self._resource_provider.model_registry.request_admission, request_pressure_advisory=True, ) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/factory.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/factory.py index acdbe24dc..ddbac6c15 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/factory.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/factory.py @@ -54,6 +54,7 @@ def create_model_client( setup, which happens before any generation task invokes the client. Direct callers of this factory must ensure registration happens before use. + request_event_sink: Optional direct sink for model-request events. Returns: A ``ModelClient`` instance routed by provider type. diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/model_request_executor.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/model_request_executor.py index 13bc1e71c..af4771563 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/model_request_executor.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/model_request_executor.py @@ -35,7 +35,9 @@ ) from data_designer.engine.observability import ( RequestAdmissionEvent, + RequestAdmissionEventKind, RequestAdmissionEventSink, + request_event_sink_accepts, runtime_correlation_provider, ) @@ -122,32 +124,39 @@ def _execute_sync_attempt(self, domain: RequestDomain, call: Callable[[], _T]) - lease = self._request_admission.acquire_sync(item) except RequestAdmissionError as exc: raise self._provider_error_from_request_admission(exc) from exc + started_at = 0.0 + duration_seconds = 0.0 + outcome = "unexpected_exception" try: self._emit_model_event("model_request_started", item=item, lease=lease) - result = call() + started_at = time.perf_counter() + try: + result = call() + finally: + duration_seconds = time.perf_counter() - started_at except ProviderError as exc: + outcome = exc.kind.value self._release_provider_error(lease, exc) - self._emit_model_event( - "model_request_completed", item=item, lease=lease, diagnostics={"outcome": exc.kind.value} - ) raise except TimeoutError: + outcome = "provider_timeout" self._request_admission.release(lease, RequestReleaseOutcome(kind="provider_timeout")) - self._emit_model_event( - "model_request_completed", item=item, lease=lease, diagnostics={"outcome": "provider_timeout"} - ) raise except BaseException as exc: outcome = "local_cancelled" if isinstance(exc, KeyboardInterrupt) else "unexpected_exception" self._request_admission.release(lease, RequestReleaseOutcome(kind=outcome)) - self._emit_model_event("model_request_completed", item=item, lease=lease, diagnostics={"outcome": outcome}) raise else: + outcome = "success" self._request_admission.release(lease, RequestReleaseOutcome(kind="success")) + return result + finally: self._emit_model_event( - "model_request_completed", item=item, lease=lease, diagnostics={"outcome": "success"} + "model_request_completed", + item=item, + lease=lease, + diagnostics={"outcome": outcome, "duration_seconds": duration_seconds}, ) - return result async def _execute_async(self, domain: RequestDomain, call: Callable[[], Awaitable[_T]]) -> _T: for attempt in range(self._max_attempts()): @@ -167,38 +176,43 @@ async def _execute_async_attempt(self, domain: RequestDomain, call: Callable[[], raise self._provider_error_from_request_admission(exc) from exc except asyncio.CancelledError: raise + started_at = 0.0 + duration_seconds = 0.0 + outcome = "unexpected_exception" try: self._emit_model_event("model_request_started", item=item, lease=lease) - result = await call() + started_at = time.perf_counter() + try: + result = await call() + finally: + duration_seconds = time.perf_counter() - started_at except asyncio.CancelledError: + outcome = "local_cancelled" self._request_admission.release(lease, RequestReleaseOutcome(kind="local_cancelled")) - self._emit_model_event( - "model_request_completed", item=item, lease=lease, diagnostics={"outcome": "local_cancelled"} - ) raise except ProviderError as exc: + outcome = exc.kind.value self._release_provider_error(lease, exc) - self._emit_model_event( - "model_request_completed", item=item, lease=lease, diagnostics={"outcome": exc.kind.value} - ) raise except TimeoutError: + outcome = "provider_timeout" self._request_admission.release(lease, RequestReleaseOutcome(kind="provider_timeout")) - self._emit_model_event( - "model_request_completed", item=item, lease=lease, diagnostics={"outcome": "provider_timeout"} - ) raise except BaseException as exc: outcome = "local_cancelled" if isinstance(exc, KeyboardInterrupt) else "unexpected_exception" self._request_admission.release(lease, RequestReleaseOutcome(kind=outcome)) - self._emit_model_event("model_request_completed", item=item, lease=lease, diagnostics={"outcome": outcome}) raise else: + outcome = "success" self._request_admission.release(lease, RequestReleaseOutcome(kind="success")) + return result + finally: self._emit_model_event( - "model_request_completed", item=item, lease=lease, diagnostics={"outcome": "success"} + "model_request_completed", + item=item, + lease=lease, + diagnostics={"outcome": outcome, "duration_seconds": duration_seconds}, ) - return result def _max_attempts(self) -> int: return max(1, self._retry_config.max_retries + 1) @@ -276,20 +290,21 @@ def _image_domain(request: ImageGenerationRequest) -> RequestDomain: def _emit_model_event( self, - event_kind: str, + event_kind: RequestAdmissionEventKind, *, item: RequestAdmissionItem, lease: RequestAdmissionLease, diagnostics: dict[str, object] | None = None, ) -> None: - if self._event_sink is None: + sink = self._event_sink + if sink is None or not request_event_sink_accepts(sink, event_kind): return self._event_sequence += 1 context = item.event_context try: - self._event_sink.emit_request_event( + sink.emit_request_event( RequestAdmissionEvent.capture( - event_kind, # type: ignore[arg-type] + event_kind, sequence=self._event_sequence, correlation=context.captured_correlation if context is not None diff --git a/packages/data-designer-engine/src/data_designer/engine/models/factory.py b/packages/data-designer-engine/src/data_designer/engine/models/factory.py index 61e1e70b9..6104ca00d 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/factory.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/factory.py @@ -16,6 +16,7 @@ from data_designer.engine.models.registry import ModelRegistry from data_designer.engine.models.request_admission.config import RequestAdmissionConfig from data_designer.engine.models.request_admission.controller import AdaptiveRequestAdmissionController + from data_designer.engine.observability import RequestAdmissionEventSink def create_model_registry( @@ -27,6 +28,7 @@ def create_model_registry( client_concurrency_mode: ClientConcurrencyMode = ClientConcurrencyMode.SYNC, run_config: RunConfig | None = None, request_admission: AdaptiveRequestAdmissionController | None = None, + request_event_sink: RequestAdmissionEventSink | None = None, ) -> ModelRegistry: """Factory function for creating a ModelRegistry instance. @@ -47,6 +49,7 @@ def create_model_registry( tuning is translated to the engine-internal request-admission config. request_admission: Optional shared request-admission controller. When omitted, a new controller is created from ``run_config``. + request_event_sink: Optional direct sink for request-admission and model-request events. Returns: A configured ModelRegistry instance. @@ -57,7 +60,7 @@ def create_model_registry( from data_designer.engine.models.registry import ModelRegistry if request_admission is None: - request_admission = create_request_admission_controller(run_config) + request_admission = create_request_admission_controller(run_config, request_event_sink=request_event_sink) def model_facade_factory( model_config: ModelConfig, @@ -72,6 +75,7 @@ def model_facade_factory( retry_config=retry_config, client_concurrency_mode=client_concurrency_mode, request_admission=request_admission, + request_event_sink=request_event_sink, ) return ModelFacade( model_config, @@ -92,6 +96,8 @@ def model_facade_factory( def create_request_admission_controller( run_config: RunConfig | None = None, + *, + request_event_sink: RequestAdmissionEventSink | None = None, ) -> AdaptiveRequestAdmissionController: """Create a request-admission controller from public runtime tuning.""" from data_designer.config.run_config import RunConfig @@ -100,7 +106,8 @@ def create_request_admission_controller( resolved_run_config = run_config or RunConfig() return AdaptiveRequestAdmissionController( - _request_admission_config_from_run_config(resolved_run_config, RequestAdmissionConfig) + _request_admission_config_from_run_config(resolved_run_config, RequestAdmissionConfig), + event_sink=request_event_sink, ) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/request_admission/controller.py b/packages/data-designer-engine/src/data_designer/engine/models/request_admission/controller.py index dadae6be8..e65e9fc56 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/request_admission/controller.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/request_admission/controller.py @@ -33,6 +33,7 @@ RequestAdmissionEvent, RequestAdmissionEventSink, emit_request_admission_event, + request_event_sink_accepts, runtime_correlation_provider, ) @@ -768,10 +769,11 @@ def _request_event_locked( ) def _emit_events(self, events: list[RequestAdmissionEvent]) -> None: + sink = self._event_sink for event in events: - if self._event_sink is not None: + if sink is not None and request_event_sink_accepts(sink, event.event_kind): try: - self._event_sink.emit_request_event(event) + sink.emit_request_event(event) except Exception: logger.warning("Request admission event sink raised; dropping event.", exc_info=True) emit_request_admission_event(event) diff --git a/packages/data-designer-engine/src/data_designer/engine/observability.py b/packages/data-designer-engine/src/data_designer/engine/observability.py index 3d422e994..b77c88418 100644 --- a/packages/data-designer-engine/src/data_designer/engine/observability.py +++ b/packages/data-designer-engine/src/data_designer/engine/observability.py @@ -14,7 +14,7 @@ from enum import Enum from os import PathLike from threading import Lock -from typing import Callable, Literal, Protocol, TextIO +from typing import Callable, Literal, Protocol, TextIO, runtime_checkable logger = logging.getLogger(__name__) @@ -214,6 +214,7 @@ def capture( ) +@runtime_checkable class SchedulerAdmissionEventSink(Protocol): def emit_scheduler_event(self, event: SchedulerAdmissionEvent) -> None: ... @@ -265,10 +266,71 @@ def close(self) -> None: logger.warning("Failed to close scheduler event file %s", self._path, exc_info=True) +def scheduler_event_sink_accepts( + sink: SchedulerAdmissionEventSink | None, + event_kind: SchedulerAdmissionEventKind, +) -> bool: + """Return whether a scheduler sink wants an event, defaulting to all events.""" + if sink is None: + return False + accepts = getattr(sink, "accepts_scheduler_event", None) + if accepts is None: + return True + try: + return bool(accepts(event_kind)) + except Exception: + logger.warning("Scheduler admission event interest check raised; dropping event.", exc_info=True) + return False + + +class _FanoutSchedulerEventSink: + def __init__(self, sinks: tuple[SchedulerAdmissionEventSink, ...]) -> None: + self._sinks = sinks + + def accepts_scheduler_event(self, event_kind: SchedulerAdmissionEventKind) -> bool: + return any(scheduler_event_sink_accepts(sink, event_kind) for sink in self._sinks) + + def emit_scheduler_event(self, event: SchedulerAdmissionEvent) -> None: + for sink in self._sinks: + if not scheduler_event_sink_accepts(sink, event.event_kind): + continue + try: + sink.emit_scheduler_event(event) + except Exception: + logger.warning("Scheduler admission event sink raised; dropping event.", exc_info=True) + + +def fanout_scheduler_event_sinks( + *sinks: SchedulerAdmissionEventSink | None, +) -> SchedulerAdmissionEventSink | None: + """Combine scheduler event sinks while isolating individual sink failures.""" + active_sinks = tuple(sink for sink in sinks if sink is not None) + if len(active_sinks) < 2: + return active_sinks[0] if active_sinks else None + return _FanoutSchedulerEventSink(active_sinks) + + class RequestAdmissionEventSink(Protocol): def emit_request_event(self, event: RequestAdmissionEvent) -> None: ... +def request_event_sink_accepts( + sink: RequestAdmissionEventSink | None, + event_kind: RequestAdmissionEventKind, +) -> bool: + """Return whether a request sink wants an event, defaulting to all events.""" + if sink is None: + return False + accepts = getattr(sink, "accepts_request_event", None) + if accepts is None: + return True + try: + return bool(accepts(event_kind)) + except Exception: + logger.warning("Request admission event interest check raised; dropping event.", exc_info=True) + return False + + RequestAdmissionEventCallback = Callable[[RequestAdmissionEvent], None] _request_event_callback_lock = Lock() diff --git a/packages/data-designer-engine/src/data_designer/engine/resources/resource_provider.py b/packages/data-designer-engine/src/data_designer/engine/resources/resource_provider.py index b6fa89411..9baed6135 100644 --- a/packages/data-designer-engine/src/data_designer/engine/resources/resource_provider.py +++ b/packages/data-designer-engine/src/data_designer/engine/resources/resource_provider.py @@ -21,6 +21,7 @@ from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode from data_designer.engine.models.factory import create_model_registry from data_designer.engine.models.registry import ModelRegistry +from data_designer.engine.observability import RequestAdmissionEventSink, SchedulerAdmissionEventSink from data_designer.engine.resources.person_reader import PersonReader from data_designer.engine.resources.seed_reader import SeedReader, SeedReaderRegistry from data_designer.engine.secret_resolver import SecretResolver @@ -43,6 +44,7 @@ class ResourceProvider(ConfigBase): mcp_registry: MCPRegistry | None = None run_config: RunConfig = RunConfig() seed_reader: SeedReader | None = None + scheduler_event_sink: SchedulerAdmissionEventSink | None = None def get_dataset_metadata(self) -> DatasetMetadata: """Get metadata about the dataset being generated. @@ -70,6 +72,8 @@ def create_resource_provider( tool_configs: list[ToolConfig] | None = None, client_concurrency_mode: ClientConcurrencyMode | None = None, request_admission: AdaptiveRequestAdmissionController | None = None, + scheduler_event_sink: SchedulerAdmissionEventSink | None = None, + request_event_sink: RequestAdmissionEventSink | None = None, ) -> ResourceProvider: """Factory function for creating a ResourceProvider instance. @@ -91,6 +95,8 @@ def create_resource_provider( mcp_providers: Optional list of MCP provider configurations. tool_configs: Optional list of tool configurations. request_admission: Optional shared request-admission controller for model clients. + scheduler_event_sink: Optional direct sink for scheduler events. + request_event_sink: Optional direct sink for request-admission and model-request events. Returns: A configured ResourceProvider instance. @@ -130,9 +136,11 @@ def create_resource_provider( client_concurrency_mode=client_concurrency_mode, run_config=effective_run_config, request_admission=request_admission, + request_event_sink=request_event_sink, ), person_reader=person_reader, mcp_registry=mcp_registry, seed_reader=seed_reader, run_config=effective_run_config, + scheduler_event_sink=scheduler_event_sink, ) diff --git a/packages/data-designer-engine/tests/engine/column_generators/generators/test_custom.py b/packages/data-designer-engine/tests/engine/column_generators/generators/test_custom.py index c50032296..0f9ad9441 100644 --- a/packages/data-designer-engine/tests/engine/column_generators/generators/test_custom.py +++ b/packages/data-designer-engine/tests/engine/column_generators/generators/test_custom.py @@ -480,7 +480,8 @@ async def fake_agenerate(*args: Any, **kwargs: Any) -> tuple: engine_thread.join(timeout=5) -def test_async_bridge_propagates_runtime_correlation() -> None: +@pytest.mark.asyncio(loop_scope="session") +async def test_async_bridge_propagates_runtime_correlation() -> None: facade = Mock() facade.generate.side_effect = SyncClientUnavailableError( "Sync methods are not available on an async-mode HttpModelClient." @@ -512,7 +513,7 @@ async def fake_agenerate(*args: Any, **kwargs: Any) -> tuple: "data_designer.engine.dataset_builders.utils.async_concurrency.ensure_async_engine_loop", return_value=engine_loop, ): - result = proxy.generate("hello", parser=str) + result = await asyncio.to_thread(proxy.generate, "hello", parser=str) finally: runtime_correlation_provider.reset(correlation_token) engine_loop.call_soon_threadsafe(engine_loop.stop) diff --git a/packages/data-designer-engine/tests/engine/conftest.py b/packages/data-designer-engine/tests/engine/conftest.py index cd4b53b48..f76352549 100644 --- a/packages/data-designer-engine/tests/engine/conftest.py +++ b/packages/data-designer-engine/tests/engine/conftest.py @@ -42,6 +42,7 @@ def stub_resource_provider(tmp_path, stub_model_facade): mock_provider.seed_reader = Mock() mock_provider.seed_reader.get_column_names.return_value = [] mock_provider.run_config = RunConfig() + mock_provider.scheduler_event_sink = None return mock_provider diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/test_async_builder_integration.py b/packages/data-designer-engine/tests/engine/dataset_builders/test_async_builder_integration.py index d3d746fb4..164817709 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/test_async_builder_integration.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/test_async_builder_integration.py @@ -409,6 +409,7 @@ def __init__(self, **kwargs: object) -> None: progress_interval=5.0, display_tui=False, ), + scheduler_event_sink=None, ) processor_runner = MagicMock() processor_runner.has_processors_for.return_value = False @@ -448,6 +449,7 @@ def __init__(self, **kwargs: object) -> None: progress_interval=5.0, display_tui=False, ), + scheduler_event_sink=None, ) processor_runner = MagicMock() processor_runner.has_processors_for.return_value = False diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py b/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py index ee0e2a290..8de0e4cee 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py @@ -66,6 +66,11 @@ RequestResourceKey, ) from data_designer.engine.models.resources import ProviderModelKey +from data_designer.engine.observability import ( + RuntimeCorrelation, + SchedulerAdmissionEventKind, + runtime_correlation_provider, +) from data_designer.engine.progress.reporter import AsyncProgressReporter from data_designer.engine.progress.tracker import ProgressTracker from data_designer.engine.resources.resource_provider import ResourceProvider @@ -364,6 +369,15 @@ def emit_scheduler_event(self, _event: object) -> None: raise RuntimeError("sink boom") +class _FilteringSchedulerSink(InMemoryAdmissionEventSink): + def __init__(self, accepted: set[SchedulerAdmissionEventKind]) -> None: + super().__init__() + self._accepted = accepted + + def accepts_scheduler_event(self, event_kind: SchedulerAdmissionEventKind) -> bool: + return event_kind in self._accepted + + # -- Helper to build graph + scheduler ---------------------------------------- @@ -1421,6 +1435,7 @@ async def test_zero_survivor_shutdown_does_not_raise() -> None: ) generators, graph, row_groups, tracker, buffer_mgr, storage = _seed_plus_cell_setup(cell, num_records=5) finalized: list[int] = [] + sink = _FilteringSchedulerSink({"scheduler_job_completed"}) def on_finalize(rg_id: int) -> None: buffer_mgr.checkpoint_row_group(rg_id) @@ -1435,6 +1450,7 @@ def on_finalize(rg_id: int) -> None: on_finalize_row_group=on_finalize, shutdown_error_rate=0.5, shutdown_error_window=2, + scheduler_event_sink=sink, ) # Must not raise (no FileNotFoundError, no DataDesignerGenerationError). await scheduler.run() @@ -1446,6 +1462,10 @@ def on_finalize(rg_id: int) -> None: assert finalized == [] assert scheduler.partial_row_groups == () storage.write_batch_to_parquet_file.assert_not_called() + [completed] = sink.scheduler_events + assert completed.reason_or_result == "early_shutdown" + assert completed.diagnostics["outcome"] == "early_shutdown" + assert completed.diagnostics["reason"] == "completed" @pytest.mark.asyncio(loop_scope="session") @@ -3716,6 +3736,91 @@ async def test_scheduler_emits_job_health_and_row_group_telemetry() -> None: assert checkpointed.diagnostics["row_group_size"] == 2 assert checkpointed.diagnostics["surviving_rows"] == 2 + completed = [event for event in sink.scheduler_events if event.event_kind == "scheduler_job_completed"] + assert len(completed) == 1 + assert completed[0].reason_or_result == "success" + assert completed[0].diagnostics["outcome"] == "success" + assert completed[0].diagnostics["reason"] == "completed" + assert "queued_total" in completed[0].diagnostics + assert "request_pressure" in completed[0].diagnostics + assert "error_type" not in completed[0].diagnostics + + +def test_scheduler_filter_skips_snapshot_and_health_diagnostics() -> None: + sink = _FilteringSchedulerSink({"row_group_checkpointed"}) + scheduler, _tracker = _build_simple_pipeline(scheduler_event_sink=sink) + scheduler.task_admission_snapshot = MagicMock(wraps=scheduler.task_admission_snapshot) + scheduler._scheduler_health_diagnostics = MagicMock(wraps=scheduler._scheduler_health_diagnostics) + + scheduler._emit_scheduler_event("selected") + scheduler._emit_scheduler_health_snapshot("test") + + scheduler.task_admission_snapshot.assert_not_called() + scheduler._scheduler_health_diagnostics.assert_not_called() + assert sink.scheduler_events == [] + + scheduler._emit_scheduler_event("row_group_checkpointed") + scheduler.task_admission_snapshot.assert_called_once_with() + assert [event.event_kind for event in sink.scheduler_events] == ["row_group_checkpointed"] + + +@pytest.mark.asyncio(loop_scope="session") +async def test_scheduler_inherits_ambient_run_correlation_for_taskless_events() -> None: + sink = _FilteringSchedulerSink({"scheduler_job_completed"}) + ambient = RuntimeCorrelation( + run_id="ambient-run", + row_group=None, + task_column=None, + task_type=None, + scheduling_group_kind=None, + scheduling_group_identity_hash=None, + task_execution_id=None, + ) + token = runtime_correlation_provider.set(ambient) + try: + scheduler, _tracker = _build_simple_pipeline(num_records=1, scheduler_event_sink=sink) + finally: + runtime_correlation_provider.reset(token) + + await scheduler.run() + + [completed] = sink.scheduler_events + assert completed.task_id is None + assert completed.captured_correlation["run_id"] == "ambient-run" + assert completed.captured_correlation["row_group"] is None + + +@pytest.mark.asyncio(loop_scope="session") +@pytest.mark.parametrize( + ("raised", "expected_outcome", "expected_error_type"), + [ + (RuntimeError("scheduler failed"), "error", "RuntimeError"), + (asyncio.CancelledError(), "cancelled", None), + ], +) +async def test_scheduler_emits_one_terminal_event_without_replacing_failure( + monkeypatch: pytest.MonkeyPatch, + raised: BaseException, + expected_outcome: str, + expected_error_type: str | None, +) -> None: + sink = _FilteringSchedulerSink({"scheduler_job_completed"}) + scheduler, _tracker = _build_simple_pipeline(num_records=1, scheduler_event_sink=sink) + + async def fail(*_args: object) -> None: + raise raised + + monkeypatch.setattr(scheduler, "_main_dispatch_loop", fail) + + with pytest.raises(type(raised)) as exc_info: + await scheduler.run() + + assert exc_info.value is raised + [completed] = sink.scheduler_events + assert completed.reason_or_result == expected_outcome + assert completed.diagnostics["outcome"] == expected_outcome + assert completed.diagnostics.get("error_type") == expected_error_type + @pytest.mark.asyncio(loop_scope="session") async def test_scheduler_adaptive_row_group_admission_expands_target_for_horizon_idle() -> None: diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py b/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py index 4d625243f..baca1ef2c 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py @@ -37,10 +37,12 @@ from data_designer.engine.dataset_builders.utils.processor_runner import ProcessorRunner from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum from data_designer.engine.models.usage import ModelUsageStats, TokenUsageStats +from data_designer.engine.observability import SchedulerAdmissionEvent from data_designer.engine.processing.processors.base import Processor from data_designer.engine.registry.data_designer_registry import DataDesignerRegistry from data_designer.engine.resources.seed_reader import DataFrameSeedReader from data_designer.engine.storage.artifact_storage import ArtifactStorage, ResumeMode +from data_designer.engine.testing import InMemoryAdmissionEventSink if TYPE_CHECKING: import pandas as pd @@ -150,6 +152,33 @@ def test_dataset_builder_artifact_storage_property(stub_dataset_builder, stub_re assert stub_dataset_builder.artifact_storage == stub_resource_provider.artifact_storage +def test_dataset_builder_combines_scheduler_event_sinks( + stub_resource_provider, + stub_test_config_builder, +) -> None: + provider_sink = InMemoryAdmissionEventSink() + local_sink = InMemoryAdmissionEventSink() + stub_resource_provider.scheduler_event_sink = provider_sink + builder = DatasetBuilder( + data_designer_config=stub_test_config_builder.build(), + resource_provider=stub_resource_provider, + ) + generators, _graph = builder._initialize_generators_and_graph() + for generator in generators: + generator.log_pre_generation = Mock() + + with patch.object(builder_mod, "AsyncTaskScheduler", return_value=Mock()) as scheduler_factory: + builder._prepare_async_run(generators, num_records=1, buffer_size=1, scheduler_event_sink=local_sink) + + sink = scheduler_factory.call_args.kwargs["scheduler_event_sink"] + assert sink is not None + event = SchedulerAdmissionEvent.capture("scheduler_job_started", sequence=1) + sink.emit_scheduler_event(event) + + assert local_sink.scheduler_events == [event] + assert provider_sink.scheduler_events == [event] + + @pytest.mark.parametrize( "config_type,expected_single_configs", [ diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_factory.py b/packages/data-designer-engine/tests/engine/models/clients/test_factory.py index f809db8be..1a64be125 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_factory.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_factory.py @@ -22,6 +22,7 @@ from data_designer.engine.models.clients.retry import RetryConfig from data_designer.engine.models.request_admission.controller import AdaptiveRequestAdmissionController from data_designer.engine.secret_resolver import SecretResolver +from data_designer.engine.testing import InMemoryAdmissionEventSink @pytest.fixture @@ -201,6 +202,26 @@ def test_request_admission_wraps_openai_client( assert client._inner._retry_config.max_retries == 0 +def test_request_event_sink_is_forwarded_to_executor( + openai_model_config: ModelConfig, + secret_resolver: SecretResolver, + openai_registry: ModelProviderRegistry, +) -> None: + controller = AdaptiveRequestAdmissionController() + sink = InMemoryAdmissionEventSink() + + client = create_model_client( + openai_model_config, + secret_resolver, + openai_registry, + request_admission=controller, + request_event_sink=sink, + ) + + assert isinstance(client, ModelRequestExecutor) + assert client._event_sink is sink + + def test_request_admission_wraps_anthropic_client( anthropic_model_config: ModelConfig, secret_resolver: SecretResolver, diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_model_request_executor.py b/packages/data-designer-engine/tests/engine/models/clients/test_model_request_executor.py index e193617ed..84aba60d2 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_model_request_executor.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_model_request_executor.py @@ -5,6 +5,7 @@ import asyncio import logging +from unittest.mock import patch import pytest @@ -28,6 +29,7 @@ RequestAdmissionError, ) from data_designer.engine.models.request_admission.resources import RequestAdmissionItem, RequestDomain +from data_designer.engine.observability import RequestAdmissionEvent, RequestAdmissionEventKind from data_designer.engine.testing import InMemoryAdmissionEventSink @@ -80,6 +82,21 @@ def emit_request_event(self, _event: object) -> None: raise RuntimeError("sink boom") +class _RejectingSink: + def accepts_request_event(self, _event_kind: RequestAdmissionEventKind) -> bool: + return False + + def emit_request_event(self, _event: RequestAdmissionEvent) -> None: + raise AssertionError("rejected request event was emitted") + + +class _InterruptingStartedSink(InMemoryAdmissionEventSink): + def emit_request_event(self, event: RequestAdmissionEvent) -> None: + if event.event_kind == "model_request_started": + raise KeyboardInterrupt + super().emit_request_event(event) + + class _GatedAsyncClient(_Client): def __init__(self) -> None: super().__init__() @@ -169,6 +186,34 @@ def test_model_request_executor_releases_successful_request() -> None: assert snapshot.last_outcome == "success" +def test_model_request_executor_skips_unaccepted_attempt_events() -> None: + controller = AdaptiveRequestAdmissionController() + controller.register(provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=1) + executor = ModelRequestExecutor(_Client(), controller, "nvidia", "nemotron", event_sink=_RejectingSink()) + + with patch.object(controller, "snapshot", side_effect=AssertionError("pressure snapshot should not be captured")): + response = executor.completion(ChatCompletionRequest(model="nemotron", messages=[])) + + assert response.message.content == "ok" + + +def test_model_request_executor_releases_when_started_event_is_interrupted() -> None: + sink = _InterruptingStartedSink() + controller = AdaptiveRequestAdmissionController(event_sink=sink) + controller.register(provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=1) + executor = ModelRequestExecutor(_Client(), controller, "nvidia", "nemotron", event_sink=sink) + + with pytest.raises(KeyboardInterrupt): + executor.completion(ChatCompletionRequest(model="nemotron", messages=[])) + + snapshot = controller.pressure.snapshot(next(iter(controller.pressure.snapshots()))) + assert snapshot is not None + assert snapshot.active_lease_count == 0 + assert snapshot.last_outcome == "local_cancelled" + completed = [event for event in sink.request_events if event.event_kind == "model_request_completed"] + assert completed[-1].diagnostics == {"outcome": "local_cancelled", "duration_seconds": 0.0} + + def test_model_request_executor_classifies_rate_limit() -> None: executor, controller, client = _executor() client.error = ProviderError( @@ -211,6 +256,9 @@ def test_model_request_executor_retries_provider_503_with_fresh_leases() -> None assert len(acquired) == 2 assert len(released) == 2 assert {event.request_lease_id for event in acquired} == {event.request_lease_id for event in released} + completed = [event for event in sink.request_events if event.event_kind == "model_request_completed"] + assert len(completed) == 2 + assert all(event.diagnostics["duration_seconds"] >= 0 for event in completed) def test_model_request_executor_does_not_retry_provider_timeout_without_status() -> None: @@ -284,6 +332,9 @@ async def test_model_request_executor_retries_async_provider_503_with_fresh_leas assert len(acquired) == 2 assert len(released) == 2 assert {event.request_lease_id for event in acquired} == {event.request_lease_id for event in released} + completed = [event for event in sink.request_events if event.event_kind == "model_request_completed"] + assert len(completed) == 2 + assert all(event.diagnostics["duration_seconds"] >= 0 for event in completed) @pytest.mark.asyncio(loop_scope="session") @@ -341,6 +392,25 @@ async def acompletion(self, request: ChatCompletionRequest) -> ChatCompletionRes assert snapshot.last_outcome == "local_cancelled" completed = [event for event in sink.request_events if event.event_kind == "model_request_completed"] assert completed[-1].diagnostics["outcome"] == "local_cancelled" + assert completed[-1].diagnostics["duration_seconds"] >= 0 + + +@pytest.mark.asyncio(loop_scope="session") +async def test_model_request_executor_releases_async_when_started_event_is_interrupted() -> None: + sink = _InterruptingStartedSink() + controller = AdaptiveRequestAdmissionController(event_sink=sink) + controller.register(provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=1) + executor = ModelRequestExecutor(_Client(), controller, "nvidia", "nemotron", event_sink=sink) + + with pytest.raises(KeyboardInterrupt): + await executor.acompletion(ChatCompletionRequest(model="nemotron", messages=[])) + + snapshot = controller.pressure.snapshot(next(iter(controller.pressure.snapshots()))) + assert snapshot is not None + assert snapshot.active_lease_count == 0 + assert snapshot.last_outcome == "local_cancelled" + completed = [event for event in sink.request_events if event.event_kind == "model_request_completed"] + assert completed[-1].diagnostics == {"outcome": "local_cancelled", "duration_seconds": 0.0} def test_model_request_executor_maps_image_chat_domain() -> None: @@ -423,6 +493,10 @@ def test_model_request_executor_emits_attempt_events_with_correlation_fields() - for event in attempt_events: assert isinstance(event.pressure_snapshot, dict) assert event.pressure_snapshot["resource"] == event.request_resource_key + completed = [event for event in attempt_events if event.event_kind == "model_request_completed"] + assert len(completed) == 1 + assert completed[0].diagnostics["outcome"] == "success" + assert completed[0].diagnostics["duration_seconds"] >= 0 def test_model_request_executor_logs_sink_failures(caplog: pytest.LogCaptureFixture) -> None: diff --git a/packages/data-designer-engine/tests/engine/models/request_admission/test_controller.py b/packages/data-designer-engine/tests/engine/models/request_admission/test_controller.py index 82614d1cd..4967cdb1b 100644 --- a/packages/data-designer-engine/tests/engine/models/request_admission/test_controller.py +++ b/packages/data-designer-engine/tests/engine/models/request_admission/test_controller.py @@ -25,7 +25,7 @@ RequestGroupSpec, RequestResourceKey, ) -from data_designer.engine.observability import subscribe_request_admission_events +from data_designer.engine.observability import RequestAdmissionEventKind, subscribe_request_admission_events from data_designer.engine.testing import InMemoryAdmissionEventSink @@ -49,6 +49,11 @@ def emit_request_event(self, _event: object) -> None: raise RuntimeError("sink boom") +class _RejectingRequestSink(InMemoryAdmissionEventSink): + def accepts_request_event(self, _event_kind: RequestAdmissionEventKind) -> bool: + return False + + def test_request_admission_acquires_and_releases_lease() -> None: controller = _controller(cap=1) item = _item() @@ -350,6 +355,15 @@ def test_request_admission_logs_sink_failures(caplog: pytest.LogCaptureFixture) assert "Request admission event sink raised; dropping event." in caplog.text +def test_request_admission_honors_sink_interest_filter() -> None: + sink = _RejectingRequestSink() + controller = AdaptiveRequestAdmissionController(event_sink=sink) + + controller.register(provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=1) + + assert sink.request_events == [] + + def test_request_lease_released_event_records_release_outcome() -> None: sink = InMemoryAdmissionEventSink() controller = AdaptiveRequestAdmissionController(event_sink=sink) diff --git a/packages/data-designer-engine/tests/engine/models/test_model_registry.py b/packages/data-designer-engine/tests/engine/models/test_model_registry.py index 65d4a897f..c92cbcf68 100644 --- a/packages/data-designer-engine/tests/engine/models/test_model_registry.py +++ b/packages/data-designer-engine/tests/engine/models/test_model_registry.py @@ -7,12 +7,13 @@ from data_designer.config.models import ChatCompletionInferenceParams, ModelConfig from data_designer.config.run_config import RequestAdmissionTuningConfig, RunConfig +from data_designer.engine.models.clients.model_request_executor import ModelRequestExecutor from data_designer.engine.models.errors import ModelAuthenticationError, ModelGenerationValidationFailureError from data_designer.engine.models.facade import ModelFacade -from data_designer.engine.models.factory import create_model_registry +from data_designer.engine.models.factory import create_model_registry, create_request_admission_controller from data_designer.engine.models.registry import ModelRegistry from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenCountSource, TokenUsageStats -from data_designer.engine.testing import make_stub_completion_response +from data_designer.engine.testing import InMemoryAdmissionEventSink, make_stub_completion_response from data_designer.logging import LOG_INDENT @@ -85,6 +86,33 @@ def test_create_model_registry_maps_request_admission_tuning_config( assert request_config.startup_ramp_seconds == 30.0 +def test_create_model_registry_shares_request_sink_and_controller( + stub_model_configs: list[ModelConfig], + stub_secrets_resolver: object, + stub_model_provider_registry: object, +) -> None: + sink = InMemoryAdmissionEventSink() + controller = create_request_admission_controller(request_event_sink=sink) + model_registry = create_model_registry( + model_configs=stub_model_configs, + secret_resolver=stub_secrets_resolver, + model_provider_registry=stub_model_provider_registry, + request_admission=controller, + request_event_sink=sink, + ) + + model = model_registry.get_model(model_alias=stub_model_configs[0].alias) + + assert model_registry.request_admission is controller + assert isinstance(model._client, ModelRequestExecutor) + assert model._client._request_admission is controller + assert model._client._event_sink is sink + kinds = [event.event_kind for event in sink.request_events] + assert kinds.count("request_resource_registered") == 1 + assert kinds.count("request_effective_cap_changed") == 1 + assert len({id(event) for event in sink.request_events}) == len(sink.request_events) + + def test_public_props(stub_model_configs, stub_model_registry): assert stub_model_registry.model_configs == { model_config.alias: model_config for model_config in stub_model_configs diff --git a/packages/data-designer-engine/tests/engine/resources/test_resource_provider.py b/packages/data-designer-engine/tests/engine/resources/test_resource_provider.py index 1d9c86530..2f8d2f297 100644 --- a/packages/data-designer-engine/tests/engine/resources/test_resource_provider.py +++ b/packages/data-designer-engine/tests/engine/resources/test_resource_provider.py @@ -12,6 +12,7 @@ create_resource_provider, ) from data_designer.engine.storage.artifact_storage import ArtifactStorage +from data_designer.engine.testing import InMemoryAdmissionEventSink def _stub_model_registry() -> ModelRegistry: @@ -54,6 +55,28 @@ def test_create_resource_provider_error_cases(test_case, expected_error, tmp_pat ) +def test_create_resource_provider_forwards_observability_sinks(tmp_path) -> None: + artifact_storage = ArtifactStorage(artifact_path=str(tmp_path), dataset_name="test") + sink = InMemoryAdmissionEventSink() + + with patch( + "data_designer.engine.resources.resource_provider.create_model_registry", + return_value=_stub_model_registry(), + ) as create_registry: + provider = create_resource_provider( + artifact_storage=artifact_storage, + model_configs=[], + secret_resolver=Mock(), + model_provider_registry=Mock(), + seed_reader_registry=Mock(), + scheduler_event_sink=sink, + request_event_sink=sink, + ) + + assert provider.scheduler_event_sink is sink + assert create_registry.call_args.kwargs["request_event_sink"] is sink + + class TestToolConfigValidation: """Tests for ToolConfig validation against MCP providers.""" diff --git a/packages/data-designer-engine/tests/engine/test_observability.py b/packages/data-designer-engine/tests/engine/test_observability.py index 5e67d3935..2b802480a 100644 --- a/packages/data-designer-engine/tests/engine/test_observability.py +++ b/packages/data-designer-engine/tests/engine/test_observability.py @@ -15,9 +15,14 @@ from data_designer.engine.observability import ( JsonlSchedulerEventSink, RequestAdmissionEvent, + RequestAdmissionEventKind, RuntimeCorrelation, RuntimeCorrelationProvider, SchedulerAdmissionEvent, + SchedulerAdmissionEventKind, + fanout_scheduler_event_sinks, + request_event_sink_accepts, + scheduler_event_sink_accepts, ) from data_designer.engine.testing import CorrelatedRuntimeView, InMemoryAdmissionEventSink @@ -187,6 +192,76 @@ def test_jsonl_scheduler_event_sink_disables_after_write_failure(monkeypatch: py file.close.assert_called_once_with() +def test_scheduler_event_sink_interest_defaults_to_all_and_honors_filter() -> None: + sink = InMemoryAdmissionEventSink() + + assert scheduler_event_sink_accepts(sink, "selected") + + class SelectiveSink(InMemoryAdmissionEventSink): + def accepts_scheduler_event(self, event_kind: SchedulerAdmissionEventKind) -> bool: + return event_kind == "scheduler_job_completed" + + selective_sink = SelectiveSink() + assert scheduler_event_sink_accepts(selective_sink, "scheduler_job_completed") + assert not scheduler_event_sink_accepts(selective_sink, "selected") + + class BrokenInterestSink(InMemoryAdmissionEventSink): + def accepts_scheduler_event(self, event_kind: SchedulerAdmissionEventKind) -> bool: + raise RuntimeError(event_kind) + + assert not scheduler_event_sink_accepts(BrokenInterestSink(), "selected") + + +def test_scheduler_event_sink_fanout_filters_and_isolates_failures( + caplog: pytest.LogCaptureFixture, +) -> None: + all_events = InMemoryAdmissionEventSink() + + class CompletedEvents(InMemoryAdmissionEventSink): + def accepts_scheduler_event(self, event_kind: SchedulerAdmissionEventKind) -> bool: + return event_kind == "scheduler_job_completed" + + class BrokenSink(InMemoryAdmissionEventSink): + def emit_scheduler_event(self, event: SchedulerAdmissionEvent) -> None: + raise RuntimeError(event.event_kind) + + completed_events = CompletedEvents() + sink = fanout_scheduler_event_sinks(all_events, BrokenSink(), completed_events) + assert sink is not None + selected = SchedulerAdmissionEvent.capture("selected", sequence=1) + completed = SchedulerAdmissionEvent.capture("scheduler_job_completed", sequence=2) + + with caplog.at_level(logging.WARNING): + sink.emit_scheduler_event(selected) + sink.emit_scheduler_event(completed) + + assert all_events.scheduler_events == [selected, completed] + assert completed_events.scheduler_events == [completed] + assert "Scheduler admission event sink raised; dropping event." in caplog.text + assert fanout_scheduler_event_sinks() is None + assert fanout_scheduler_event_sinks(None, all_events) is all_events + + +def test_request_event_sink_interest_defaults_to_all_and_honors_filter() -> None: + sink = InMemoryAdmissionEventSink() + + assert request_event_sink_accepts(sink, "model_request_started") + + class SelectiveSink(InMemoryAdmissionEventSink): + def accepts_request_event(self, event_kind: RequestAdmissionEventKind) -> bool: + return event_kind == "model_request_completed" + + selective_sink = SelectiveSink() + assert request_event_sink_accepts(selective_sink, "model_request_completed") + assert not request_event_sink_accepts(selective_sink, "model_request_started") + + class BrokenInterestSink(InMemoryAdmissionEventSink): + def accepts_request_event(self, event_kind: RequestAdmissionEventKind) -> bool: + raise RuntimeError(event_kind) + + assert not request_event_sink_accepts(BrokenInterestSink(), "model_request_started") + + def test_correlated_runtime_view_timeline_sorts_events() -> None: scheduler_event = SchedulerAdmissionEvent(event_kind="selected", captured_at_monotonic=2.0, sequence=1) first_request_event = RequestAdmissionEvent( diff --git a/packages/data-designer/pyproject.toml b/packages/data-designer/pyproject.toml index bbf655c4d..266f3a76c 100644 --- a/packages/data-designer/pyproject.toml +++ b/packages/data-designer/pyproject.toml @@ -42,7 +42,11 @@ path = "dev-tools/hatch_build.py" dependencies = [ "data-designer-config=={{ version }}", "data-designer-engine=={{ version }}", + "opentelemetry-api>=1.43,<1.44", + "opentelemetry-exporter-prometheus>=0.64b0,<0.65", + "opentelemetry-sdk>=1.43,<1.44", "packaging>=25,<27", + "prometheus-client>=0.24,<1", "prompt-toolkit>=3.0.0,<4", "typer>=0.12.0,<1", ] diff --git a/packages/data-designer/src/data_designer/integrations/opentelemetry.py b/packages/data-designer/src/data_designer/integrations/opentelemetry.py new file mode 100644 index 000000000..f1a6f27c6 --- /dev/null +++ b/packages/data-designer/src/data_designer/integrations/opentelemetry.py @@ -0,0 +1,526 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import atexit +import contextlib +import hashlib +import logging +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from enum import Enum +from typing import Any + +from data_designer.config.version import get_library_version +from data_designer.engine.observability import ( + RequestAdmissionEvent, + RequestAdmissionEventKind, + RuntimeCorrelation, + SchedulerAdmissionEvent, + SchedulerAdmissionEventKind, + runtime_correlation_provider, +) + +logger = logging.getLogger(__name__) + +_HOST = "127.0.0.1" +_MAX_ATTRIBUTE_LENGTH = 128 +_SCHEDULER_EVENT_NAMES: dict[SchedulerAdmissionEventKind, str | None] = { + "scheduler_job_started": "job.started", + "scheduler_job_completed": "job.completed", + "row_group_checkpointed": None, + "non_retryable_dropped": "task.error", + "worker_spawn_failed": "worker.spawn_failed", + "retry_deferred": "task.retry_deferred", + "cancelled": "task.cancelled", +} +_OPERATION_NAMES = { + "chat": "chat", + "embedding": "embeddings", + "image": "generate_content", + "healthcheck": "healthcheck", +} +_GEN_AI_DURATION_BUCKETS = ( + 0.01, + 0.02, + 0.04, + 0.08, + 0.16, + 0.32, + 0.64, + 1.28, + 2.56, + 5.12, + 10.24, + 20.48, + 40.96, + 81.92, +) + + +class _MetricLogHandler(logging.Handler): + def __init__(self, runtime: OpenTelemetryRuntime) -> None: + super().__init__() + self._runtime = runtime + + def emit(self, record: logging.LogRecord) -> None: + try: + self._runtime.record_log(record.levelno) + except Exception: + pass + + +class OpenTelemetryRuntime: + """Process-owned Prometheus exporter and adapter for runtime events.""" + + def __init__(self) -> None: + self._lock = threading.RLock() + self._active_run_ids: set[str] = set() + self._active_dataset_runs: dict[str, tuple[int, int]] = {} + self._dataset_progress_value: float | None = None + self._active_request_attempts: dict[tuple[str, str], dict[str, str]] = {} + self._initialized = False + self._port: int | None = None + self._registry: Any = None + self._meter_provider: Any = None + self._server: Any = None + self._server_thread: threading.Thread | None = None + self._start_http_server: Callable[..., Any] | None = None + self._create_duration: Any = None + self._dataset_records: Any = None + self._dataset_progress: Any = None + self._scheduler_events: Any = None + self._active_model_requests: Any = None + self._request_duration: Any = None + self._log_records: Any = None + self._log_handler: logging.Handler | None = None + + @property + def metrics_url(self) -> str | None: + with self._lock: + return f"http://{_HOST}:{self._port}/metrics" if self._server is not None else None + + @contextlib.contextmanager + def observe_create(self, port: int | None) -> Iterator[None]: + run_id = f"run-{uuid.uuid4().hex}" + token = runtime_correlation_provider.set( + RuntimeCorrelation( + run_id=run_id, + row_group=None, + task_column=None, + task_type=None, + scheduling_group_kind=None, + scheduling_group_identity_hash=None, + task_execution_id=None, + ) + ) + started_at = time.perf_counter() + enabled = False + try: + if port is not None: + enabled = self._activate_run(run_id, port) + try: + yield + except BaseException as exc: + if enabled: + self._record_create_duration(time.perf_counter() - started_at, type(exc).__name__) + raise + else: + if enabled: + self._record_create_duration(time.perf_counter() - started_at) + finally: + reconciliation_failures: list[Exception] = [] + if enabled: + with self._lock: + self._active_run_ids.discard(run_id) + self._finish_active_dataset(run_id) + reconciliation_failures = self._finish_active_requests(run_id) + runtime_correlation_provider.reset(token) + for failure in reconciliation_failures: + logger.warning("Failed to reconcile an active OpenTelemetry model request.", exc_info=failure) + + def accepts_scheduler_event(self, event_kind: SchedulerAdmissionEventKind) -> bool: + with self._lock: + return bool(self._active_run_ids) and event_kind in _SCHEDULER_EVENT_NAMES + + def accepts_request_event(self, event_kind: RequestAdmissionEventKind) -> bool: + with self._lock: + return bool(self._active_run_ids) and event_kind in {"model_request_started", "model_request_completed"} + + def emit_scheduler_event(self, event: SchedulerAdmissionEvent) -> None: + run_id = self._event_run_id(event.captured_correlation) + if run_id is None: + return + with self._lock: + if run_id not in self._active_run_ids: + return + if event.event_kind == "row_group_checkpointed": + diagnostics = event.diagnostics + generated = _non_negative_int(diagnostics.get("surviving_rows")) + dropped = _non_negative_int(diagnostics.get("dropped_rows")) + self._dataset_records.add( + generated, + {"record.result": "generated"}, + ) + self._dataset_records.add( + dropped, + {"record.result": "dropped"}, + ) + self._advance_dataset_progress(run_id, generated + dropped) + return + if event.event_kind == "scheduler_job_started": + self._start_dataset_progress(run_id, _non_negative_int(event.diagnostics.get("row_group_total_rows"))) + + event_name = _SCHEDULER_EVENT_NAMES.get(event.event_kind) + if event_name is None: + return + attributes: dict[str, str] = {"event.name": event_name} + if event.event_kind == "scheduler_job_completed": + attributes["outcome"] = _bounded_attribute(event.reason_or_result or "success") + if event.event_kind in {"scheduler_job_completed", "non_retryable_dropped"}: + error_type = event.diagnostics.get("error_type") + elif event.event_kind == "worker_spawn_failed": + error_type = "worker_spawn_failed" + else: + error_type = None + if error_type is not None: + attributes["error.type"] = _bounded_attribute(error_type) + self._scheduler_events.add(1, attributes) + if event.event_kind == "scheduler_job_completed": + self._finish_active_dataset(run_id) + + def emit_request_event(self, event: RequestAdmissionEvent) -> None: + run_id = self._event_run_id(event.captured_correlation) + if run_id is None or event.event_kind not in {"model_request_started", "model_request_completed"}: + return + resource = event.request_resource_key if isinstance(event.request_resource_key, dict) else {} + raw_domain = resource.get("domain") + if isinstance(raw_domain, Enum): + raw_domain = raw_domain.value + domain = raw_domain if isinstance(raw_domain, str) else "" + attributes = { + "gen_ai.operation.name": _OPERATION_NAMES.get(domain, "_OTHER"), + } + provider_name = resource.get("provider_name") + if provider_name: + attributes["gen_ai.provider.name"] = _bounded_identity_attribute(provider_name) + model_id = resource.get("model_id") + if model_id: + attributes["gen_ai.request.model"] = _bounded_identity_attribute(model_id) + attempt_id = event.request_attempt_id or event.request_lease_id + attempt_key = (run_id, attempt_id) if isinstance(attempt_id, str) and attempt_id else None + with self._lock: + if run_id not in self._active_run_ids: + return + if event.event_kind == "model_request_started": + if attempt_key is not None: + if attempt_key not in self._active_request_attempts: + self._active_model_requests.add(1, attributes) + self._active_request_attempts[attempt_key] = dict(attributes) + return + + if attempt_key is not None: + active_attributes = self._active_request_attempts.get(attempt_key) + if active_attributes is not None: + self._active_model_requests.add(-1, active_attributes) + self._active_request_attempts.pop(attempt_key, None) + duration = event.diagnostics.get("duration_seconds") + if not isinstance(duration, int | float) or duration < 0: + return + outcome = event.diagnostics.get("outcome") + if outcome not in (None, "success"): + attributes = dict(attributes) + attributes["error.type"] = _bounded_attribute(outcome) + self._request_duration.record(float(duration), attributes) + + def record_log(self, level: int) -> None: + correlation = runtime_correlation_provider.current() + if correlation is None: + return + with self._lock: + if correlation.run_id not in self._active_run_ids: + return + self._log_records.add(1, {"log.severity": _log_severity(level)}) + + def shutdown(self) -> None: + with self._lock: + if not self._initialized: + return + self._active_run_ids.clear() + self._active_dataset_runs.clear() + self._dataset_progress_value = None + self._active_request_attempts.clear() + server, thread = self._server, self._server_thread + provider, handler = self._meter_provider, self._log_handler + self._server = None + self._server_thread = None + self._port = None + self._meter_provider = None + self._registry = None + self._start_http_server = None + self._create_duration = None + self._dataset_records = None + self._dataset_progress = None + self._scheduler_events = None + self._active_model_requests = None + self._request_duration = None + self._log_records = None + self._log_handler = None + self._initialized = False + + if handler is not None: + logging.getLogger("data_designer").removeHandler(handler) + handler.close() + _stop_server(server, thread) + if provider is not None: + with contextlib.suppress(Exception): + provider.shutdown() + + def _activate_run(self, run_id: str, port: int) -> bool: + old_server: Any = None + old_thread: threading.Thread | None = None + bound_url: str | None = None + reused_url: str | None = None + try: + with self._lock: + if not self._initialized: + self._initialize() + if self._server is None: + old_server, old_thread = self._rebind(port) + bound_url = f"http://{_HOST}:{port}/metrics" + elif self._port != port and not self._active_run_ids: + old_server, old_thread = self._rebind(port) + bound_url = f"http://{_HOST}:{port}/metrics" + elif self._port != port: + reused_url = f"http://{_HOST}:{self._port}/metrics" + if self._server is None: + return False + self._active_run_ids.add(run_id) + except Exception: + logger.warning("OpenTelemetry metrics are unavailable; continuing without them.", exc_info=True) + return False + + _stop_server(old_server, old_thread) + if bound_url is not None: + logger.info("OpenTelemetry metrics available at %s", bound_url) + if reused_url is not None: + logger.warning( + "OpenTelemetry metrics already active at %s; reusing it instead of requested port %d.", + reused_url, + port, + ) + return True + + def _initialize(self) -> None: + from opentelemetry.exporter.prometheus import PrometheusMetricReader + from opentelemetry.metrics import Observation + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.view import ExplicitBucketHistogramAggregation, View + from opentelemetry.sdk.resources import Resource + from prometheus_client import CollectorRegistry, start_http_server + + provider: Any = None + handler: logging.Handler | None = None + try: + version = get_library_version() + registry = CollectorRegistry() + reader = PrometheusMetricReader(registry=registry) + provider = MeterProvider( + resource=Resource({"service.name": "data-designer", "service.version": version}), + metric_readers=[reader], + shutdown_on_exit=False, + views=[ + View( + instrument_name="gen_ai.client.operation.duration", + aggregation=ExplicitBucketHistogramAggregation(boundaries=_GEN_AI_DURATION_BUCKETS), + ) + ], + ) + meter = provider.get_meter("data_designer", version) + create_duration = meter.create_histogram( + "data_designer.create.duration", + unit="s", + description="Duration of DataDesigner create operations.", + ) + dataset_records = meter.create_counter( + "data_designer.dataset.records", + unit="{record}", + description="Dataset records generated or dropped by DataDesigner.", + ) + + def observe_dataset_progress(_options: object) -> list[Observation]: + with self._lock: + value = self._dataset_progress_value + return [] if value is None else [Observation(value)] + + dataset_progress = meter.create_observable_gauge( + "data_designer.dataset.progress", + callbacks=[observe_dataset_progress], + unit="1", + description="Fraction of records processed by current DataDesigner create jobs.", + ) + scheduler_events = meter.create_counter( + "data_designer.scheduler.events", + unit="{event}", + description="Job lifecycle and selected error events from the DataDesigner scheduler.", + ) + active_model_requests = meter.create_up_down_counter( + "data_designer.model.request.active", + unit="{request}", + description="Model requests currently in progress.", + ) + request_duration = meter.create_histogram( + "gen_ai.client.operation.duration", + unit="s", + description="GenAI operation duration.", + ) + log_records = meter.create_counter( + "data_designer.log.records", + unit="{record}", + description="DataDesigner log records by severity.", + ) + handler = _MetricLogHandler(self) + logging.getLogger("data_designer").addHandler(handler) + except Exception: + if handler is not None: + logging.getLogger("data_designer").removeHandler(handler) + handler.close() + if provider is not None: + with contextlib.suppress(Exception): + provider.shutdown() + raise + + self._registry = registry + self._meter_provider = provider + self._start_http_server = start_http_server + self._create_duration = create_duration + self._dataset_records = dataset_records + self._dataset_progress = dataset_progress + self._scheduler_events = scheduler_events + self._active_model_requests = active_model_requests + self._request_duration = request_duration + self._log_records = log_records + self._log_handler = handler + self._initialized = True + + def _rebind(self, port: int) -> tuple[Any, threading.Thread | None]: + if self._start_http_server is None: + return None, None + server, thread = self._start_http_server(port=port, addr=_HOST, registry=self._registry) + old_server, old_thread = self._server, self._server_thread + self._server, self._server_thread, self._port = server, thread, port + return old_server, old_thread + + @staticmethod + def _event_run_id(correlation: object) -> str | None: + run_id = correlation.get("run_id") if isinstance(correlation, Mapping) else None + return run_id if isinstance(run_id, str) else None + + def _start_dataset_progress(self, run_id: str, scheduled: int) -> None: + with self._lock: + if run_id in self._active_dataset_runs: + return + self._active_dataset_runs[run_id] = (scheduled, 0) + self._dataset_progress_value = self._active_dataset_progress() + + def _advance_dataset_progress(self, run_id: str, processed: int) -> None: + with self._lock: + state = self._active_dataset_runs.get(run_id) + if state is None: + return + scheduled, current = state + self._active_dataset_runs[run_id] = (scheduled, min(scheduled, current + processed)) + self._dataset_progress_value = self._active_dataset_progress() + + def _finish_active_dataset(self, run_id: str) -> None: + with self._lock: + state = self._active_dataset_runs.pop(run_id, None) + if state is None: + return + if self._active_dataset_runs: + progress = self._active_dataset_progress() + else: + scheduled, processed = state + progress = processed / scheduled if scheduled else 1.0 + self._dataset_progress_value = progress + + def _active_dataset_progress(self) -> float: + scheduled = sum(state[0] for state in self._active_dataset_runs.values()) + processed = sum(state[1] for state in self._active_dataset_runs.values()) + return processed / scheduled if scheduled else 1.0 + + def _finish_active_requests(self, run_id: str) -> list[Exception]: + failures = [] + with self._lock: + active = [ + (key, attributes) for key, attributes in self._active_request_attempts.items() if key[0] == run_id + ] + for key, attributes in active: + try: + self._active_model_requests.add(-1, attributes) + except Exception as exc: + failures.append(exc) + finally: + self._active_request_attempts.pop(key, None) + return failures + + def _record_create_duration(self, duration: float, error_type: str | None = None) -> None: + attributes = {"error.type": _bounded_attribute(error_type)} if error_type is not None else None + try: + self._create_duration.record(duration, attributes) + except Exception: + logger.warning("Failed to record OpenTelemetry create duration.", exc_info=True) + + +def _stop_server(server: Any, thread: threading.Thread | None) -> None: + if server is None: + return + with contextlib.suppress(Exception): + server.shutdown() + with contextlib.suppress(Exception): + server.server_close() + if thread is not None: + with contextlib.suppress(Exception): + thread.join(timeout=5.0) + + +def _non_negative_int(value: object) -> int: + return max(0, value) if isinstance(value, int) else 0 + + +def _bounded_attribute(value: object) -> str: + normalized = "".join(character for character in str(value) if character.isalnum() or character in "._-") + return normalized[:_MAX_ATTRIBUTE_LENGTH] or "_OTHER" + + +def _bounded_identity_attribute(value: object) -> str: + identity = str(value) + if len(identity) <= _MAX_ATTRIBUTE_LENGTH: + return identity or "_OTHER" + digest = hashlib.sha256(identity.encode()).hexdigest()[:16] + prefix = identity[: _MAX_ATTRIBUTE_LENGTH - len(digest) - 1] + return f"{prefix}-{digest}" + + +def _log_severity(level: int) -> str: + if level >= logging.CRITICAL: + return "fatal" + if level >= logging.ERROR: + return "error" + if level >= logging.WARNING: + return "warn" + if level >= logging.INFO: + return "info" + return "debug" + + +_RUNTIME = OpenTelemetryRuntime() +atexit.register(_RUNTIME.shutdown) + + +def get_open_telemetry_runtime() -> OpenTelemetryRuntime: + """Return the process-owned DataDesigner OpenTelemetry runtime.""" + return _RUNTIME diff --git a/packages/data-designer/src/data_designer/interface/data_designer.py b/packages/data-designer/src/data_designer/interface/data_designer.py index 9b530655b..7dd2259a2 100644 --- a/packages/data-designer/src/data_designer/interface/data_designer.py +++ b/packages/data-designer/src/data_designer/interface/data_designer.py @@ -4,11 +4,12 @@ from __future__ import annotations import asyncio +import functools import logging import time from collections.abc import Callable from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Concatenate, ParamSpec, TypeVar import data_designer.lazy_heavy_imports as lazy from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults @@ -66,6 +67,7 @@ SecretResolver, ) from data_designer.engine.storage.artifact_storage import ArtifactStorage, ResumeMode +from data_designer.integrations.opentelemetry import OpenTelemetryRuntime, get_open_telemetry_runtime from data_designer.interface.composite_workflow import CompositeWorkflow from data_designer.interface.errors import ( DataDesignerEarlyShutdownError, @@ -86,6 +88,8 @@ _CHECK_MODELS_RETRYABLE_ERRORS = RETRYABLE_MODEL_ERRORS + (TimeoutError,) _interface_runtime_initialized = False +_P = ParamSpec("_P") +_R = TypeVar("_R") def _initialize_interface_runtime(*, auto_configure_logging: bool = True) -> None: @@ -99,6 +103,17 @@ def _initialize_interface_runtime(*, auto_configure_logging: bool = True) -> Non _interface_runtime_initialized = True +def _observe_create( + method: Callable[Concatenate[DataDesigner, _P], _R], +) -> Callable[Concatenate[DataDesigner, _P], _R]: + @functools.wraps(method) + def wrapper(self: DataDesigner, *args: _P.args, **kwargs: _P.kwargs) -> _R: + with self._open_telemetry.observe_create(self._run_config.otel_metrics_port): + return method(self, *args, **kwargs) + + return wrapper + + DEFAULT_SECRET_RESOLVER = CompositeResolver([EnvironmentResolver(), PlaintextResolver()]) DEFAULT_SEED_READERS = [ @@ -163,6 +178,7 @@ def __init__( auto_configure_logging: bool = True, ): _initialize_interface_runtime(auto_configure_logging=auto_configure_logging) + self._open_telemetry: OpenTelemetryRuntime = get_open_telemetry_runtime() self._secret_resolver = secret_resolver or DEFAULT_SECRET_RESOLVER self._artifact_path = Path(artifact_path) if artifact_path is not None else Path.cwd() / "artifacts" self._run_config = RunConfig() @@ -207,6 +223,7 @@ def list_mcp_tool_names(self, mcp_provider_name: str, *, timeout_sec: float = 10 configured = [p.name for p in self._mcp_providers] raise ValueError(f"No MCP provider named {mcp_provider_name!r}. Configured providers: {configured}") + @_observe_create def create( self, config_builder: DataDesignerConfigBuilder, @@ -733,12 +750,17 @@ def _create_resource_provider( tool_configs=config_builder.tool_configs, client_concurrency_mode=client_concurrency_mode, request_admission=self._request_admission, + request_event_sink=self._open_telemetry if self._run_config.otel_metrics_port is not None else None, + scheduler_event_sink=self._open_telemetry if self._run_config.otel_metrics_port is not None else None, ) def _create_request_admission_controller(self) -> AdaptiveRequestAdmissionController: from data_designer.engine.models.factory import create_request_admission_controller - return create_request_admission_controller(self._run_config) + return create_request_admission_controller( + self._run_config, + request_event_sink=self._open_telemetry if self._run_config.otel_metrics_port is not None else None, + ) def _get_interface_info(self, model_providers: list[ModelProvider]) -> InterfaceInfo: return InterfaceInfo(model_providers=model_providers) diff --git a/packages/data-designer/tests/integrations/test_opentelemetry.py b/packages/data-designer/tests/integrations/test_opentelemetry.py new file mode 100644 index 000000000..8e0a79436 --- /dev/null +++ b/packages/data-designer/tests/integrations/test_opentelemetry.py @@ -0,0 +1,996 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import contextlib +import logging +import re +import socket +import subprocess +import sys +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Barrier, Event +from unittest.mock import AsyncMock, MagicMock, patch +from urllib.error import URLError +from urllib.request import urlopen + +import pytest + +import data_designer.integrations.opentelemetry as opentelemetry +from data_designer.config.column_configs import LLMTextColumnConfig +from data_designer.config.config_builder import DataDesignerConfigBuilder +from data_designer.config.models import ModelConfig, ModelProvider +from data_designer.config.run_config import RunConfig +from data_designer.engine.models.clients.adapters.openai_compatible import OpenAICompatibleClient +from data_designer.engine.models.request_admission.resources import RequestDomain, RequestResourceKey +from data_designer.engine.observability import ( + RequestAdmissionEvent, + RuntimeCorrelation, + SchedulerAdmissionEvent, + runtime_correlation_provider, +) +from data_designer.engine.secret_resolver import PlaintextResolver +from data_designer.engine.testing import make_stub_completion_response +from data_designer.integrations.opentelemetry import OpenTelemetryRuntime +from data_designer.interface.data_designer import DataDesigner + + +@pytest.fixture +def runtime() -> Iterator[OpenTelemetryRuntime]: + value = OpenTelemetryRuntime() + yield value + value.shutdown() + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _scrape(runtime: OpenTelemetryRuntime) -> str: + assert runtime.metrics_url is not None + with urlopen(runtime.metrics_url, timeout=2) as response: + return response.read().decode() + + +def _correlation() -> RuntimeCorrelation: + correlation = runtime_correlation_provider.current() + assert correlation is not None + return correlation + + +def _active_request_total(runtime: OpenTelemetryRuntime) -> float: + return sum( + float(line.rsplit(" ", maxsplit=1)[1]) + for line in _scrape(runtime).splitlines() + if line.startswith("data_designer_model_request_active{") + ) + + +def test_runtime_only_accepts_exported_events_during_active_runs(runtime: OpenTelemetryRuntime) -> None: + assert not runtime.accepts_scheduler_event("scheduler_job_started") + assert not runtime.accepts_request_event("model_request_started") + + with runtime.observe_create(_free_port()): + assert runtime.accepts_scheduler_event("scheduler_job_started") + assert not runtime.accepts_scheduler_event("selected") + assert runtime.accepts_request_event("model_request_started") + assert runtime.accepts_request_event("model_request_completed") + assert not runtime.accepts_request_event("request_wait_started") + + assert not runtime.accepts_scheduler_event("scheduler_job_started") + assert not runtime.accepts_request_event("model_request_started") + + +def test_runtime_exports_bounded_metrics_without_log_bodies( + runtime: OpenTelemetryRuntime, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + caplog.set_level(logging.ERROR) + monkeypatch.setenv( + "OTEL_RESOURCE_ATTRIBUTES", + "api.token=SECRET-RESOURCE-SENTINEL,customer.id=SECRET-CUSTOMER-SENTINEL", + ) + port = _free_port() + with runtime.observe_create(port): + correlation = _correlation() + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + "scheduler_job_started", + sequence=1, + correlation=correlation, + diagnostics={"row_group_total_rows": 4}, + ) + ) + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + "row_group_checkpointed", + sequence=2, + correlation=correlation, + diagnostics={"surviving_rows": 3, "dropped_rows": 1}, + ) + ) + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + "scheduler_job_completed", + sequence=3, + correlation=correlation, + reason_or_result="success", + ) + ) + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_started", + sequence=1, + correlation=correlation, + request_resource_key={ + "provider_name": "openai", + "model_id": "model-1", + "domain": "chat", + }, + request_attempt_id="SECRET-REQUEST-SENTINEL", + ) + ) + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_completed", + sequence=2, + correlation=correlation, + request_resource_key={ + "provider_name": "openai", + "model_id": "model-1", + "domain": "chat", + }, + request_attempt_id="SECRET-REQUEST-SENTINEL", + diagnostics={"duration_seconds": 0.25, "outcome": "success"}, + ) + ) + logging.getLogger("data_designer.test").error("SECRET-LOG-SENTINEL") + + metrics = _scrape(runtime) + target_line = next(line for line in metrics.splitlines() if line.startswith("target_info{")) + assert set(re.findall(r'(\w+)="', target_line)) == {"service_name", "service_version"} + assert 'service_name="data-designer"' in metrics + assert 'otel_scope_name="data_designer"' in metrics + assert "data_designer_dataset_records_total{" in metrics + generated_line = next(line for line in metrics.splitlines() if 'record_result="generated"' in line) + dropped_line = next(line for line in metrics.splitlines() if 'record_result="dropped"' in line) + assert generated_line.endswith(" 3.0") + assert dropped_line.endswith(" 1.0") + progress_line = next(line for line in metrics.splitlines() if line.startswith("data_designer_dataset_progress{")) + assert progress_line.endswith(" 1.0") + assert "data_designer_scheduler_events_total{" in metrics + assert 'event_name="job.started"' in metrics + assert 'event_name="job.completed"' in metrics + assert 'outcome="success"' in metrics + assert "gen_ai_client_operation_duration_seconds_count{" in metrics + assert 'gen_ai_operation_name="chat"' in metrics + assert 'gen_ai_provider_name="openai"' in metrics + assert 'gen_ai_request_model="model-1"' in metrics + active_request_line = next( + line for line in metrics.splitlines() if line.startswith("data_designer_model_request_active{") + ) + assert active_request_line.endswith(" 0.0") + assert "data_designer_log_records_total{" in metrics + assert 'log_severity="error"' in metrics + assert "data_designer_create_duration_seconds_count{" in metrics + assert metrics.count("gen_ai_client_operation_duration_seconds_count{") == 1 + assert metrics.count("data_designer_scheduler_events_total{") == 2 + assert "SECRET-" not in metrics + assert sum(record.getMessage() == "SECRET-LOG-SENTINEL" for record in caplog.records) == 1 + + +def test_runtime_exports_current_active_model_requests(runtime: OpenTelemetryRuntime) -> None: + resource = RequestResourceKey("provider", "model-1", RequestDomain.CHAT) + with runtime.observe_create(_free_port()): + correlation = _correlation() + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_started", + sequence=1, + correlation=correlation, + request_resource_key=resource, + request_attempt_id="SECRET-REQUEST-SENTINEL", + ) + ) + active_metrics = _scrape(runtime) + active_line = next( + line for line in active_metrics.splitlines() if line.startswith("data_designer_model_request_active{") + ) + assert active_line.endswith(" 1.0") + assert 'gen_ai_operation_name="chat"' in active_line + assert 'gen_ai_provider_name="provider"' in active_line + assert 'gen_ai_request_model="model-1"' in active_line + assert "SECRET-" not in active_metrics + + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_completed", + sequence=2, + correlation=correlation, + request_resource_key=resource, + request_attempt_id="SECRET-REQUEST-SENTINEL", + diagnostics={"duration_seconds": 0.1, "outcome": "provider_timeout"}, + ) + ) + + completed_metrics = _scrape(runtime) + completed_line = next( + line for line in completed_metrics.splitlines() if line.startswith("data_designer_model_request_active{") + ) + assert completed_line.endswith(" 0.0") + assert "error_type=" not in completed_line + assert 'error_type="provider_timeout"' in completed_metrics + assert "SECRET-" not in completed_metrics + + +def test_runtime_reconciles_duplicate_unmatched_and_orphan_model_requests(runtime: OpenTelemetryRuntime) -> None: + resource = RequestResourceKey("provider", "model-1", RequestDomain.CHAT) + port = _free_port() + with runtime.observe_create(port): + correlation = _correlation() + start = RequestAdmissionEvent.capture( + "model_request_started", + sequence=1, + correlation=correlation, + request_resource_key=resource, + request_attempt_id="attempt-1", + ) + runtime.emit_request_event(start) + assert _active_request_total(runtime) == 1 + + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_started", + sequence=2, + correlation=correlation, + request_resource_key=resource, + request_attempt_id="attempt-1", + ) + ) + assert _active_request_total(runtime) == 1 + + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_completed", + sequence=3, + correlation=correlation, + request_resource_key=resource, + request_attempt_id="unmatched-attempt", + diagnostics={"duration_seconds": 0.1, "outcome": "success"}, + ) + ) + assert _active_request_total(runtime) == 1 + + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_completed", + sequence=4, + correlation=correlation, + request_resource_key=resource, + request_attempt_id="attempt-1", + diagnostics={"duration_seconds": 0.1, "outcome": "local_cancelled"}, + ) + ) + assert _active_request_total(runtime) == 0 + + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_completed", + sequence=5, + correlation=correlation, + request_resource_key=resource, + request_attempt_id="attempt-1", + diagnostics={"duration_seconds": 0.1, "outcome": "local_cancelled"}, + ) + ) + assert _active_request_total(runtime) == 0 + + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_started", + sequence=6, + correlation=correlation, + request_resource_key=resource, + request_attempt_id="orphan-attempt", + ) + ) + assert _active_request_total(runtime) == 1 + + assert _active_request_total(runtime) == 0 + + +def test_runtime_rejects_request_that_reaches_state_update_after_teardown(runtime: OpenTelemetryRuntime) -> None: + entered = Event() + release = Event() + + def pause_delayed_provider(value: object) -> str: + if value == "delayed-provider": + entered.set() + assert release.wait(timeout=5) + return str(value) + + resource = RequestResourceKey("provider", "model-1", RequestDomain.CHAT) + with ( + patch( + "data_designer.integrations.opentelemetry._bounded_identity_attribute", + side_effect=pause_delayed_provider, + ), + ThreadPoolExecutor(max_workers=1) as executor, + ): + future = None + try: + with runtime.observe_create(_free_port()): + correlation = _correlation() + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_started", + sequence=1, + correlation=correlation, + request_resource_key=resource, + request_attempt_id="baseline", + ) + ) + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_completed", + sequence=2, + correlation=correlation, + request_resource_key=resource, + request_attempt_id="baseline", + diagnostics={"duration_seconds": 0.1, "outcome": "success"}, + ) + ) + future = executor.submit( + runtime.emit_request_event, + RequestAdmissionEvent.capture( + "model_request_started", + sequence=3, + correlation=correlation, + request_resource_key=RequestResourceKey( + "delayed-provider", + "model-1", + RequestDomain.CHAT, + ), + request_attempt_id="delayed", + ), + ) + assert entered.wait(timeout=5) + finally: + release.set() + assert future is not None + future.result(timeout=5) + + assert _active_request_total(runtime) == 0 + + +def test_runtime_deactivation_holds_lock_through_state_cleanup( + runtime: OpenTelemetryRuntime, + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_finish_dataset = runtime._finish_active_dataset + original_finish_requests = runtime._finish_active_requests + + def lock_is_available_to_another_thread() -> bool: + acquired = runtime._lock.acquire(blocking=False) + if acquired: + runtime._lock.release() + return acquired + + with ThreadPoolExecutor(max_workers=1) as executor: + + def finish_dataset(run_id: str) -> None: + assert executor.submit(lock_is_available_to_another_thread).result(timeout=5) is False + original_finish_dataset(run_id) + + def finish_requests(run_id: str) -> list[Exception]: + assert executor.submit(lock_is_available_to_another_thread).result(timeout=5) is False + return original_finish_requests(run_id) + + monkeypatch.setattr(runtime, "_finish_active_dataset", finish_dataset) + monkeypatch.setattr(runtime, "_finish_active_requests", finish_requests) + with runtime.observe_create(_free_port()): + correlation = _correlation() + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + "scheduler_job_started", + sequence=1, + correlation=correlation, + diagnostics={"row_group_total_rows": 2}, + ) + ) + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_started", + sequence=2, + correlation=correlation, + request_resource_key=RequestResourceKey("provider", "model-1", RequestDomain.CHAT), + request_attempt_id="orphan", + ) + ) + + assert _active_request_total(runtime) == 0 + progress_line = next( + line for line in _scrape(runtime).splitlines() if line.startswith("data_designer_dataset_progress{") + ) + assert progress_line.endswith(" 0.0") + + +def test_runtime_progress_resets_for_sequential_runs(runtime: OpenTelemetryRuntime) -> None: + port = _free_port() + with runtime.observe_create(port): + correlation = _correlation() + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + "scheduler_job_started", + sequence=1, + correlation=correlation, + diagnostics={"row_group_total_rows": 4}, + ) + ) + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + "row_group_checkpointed", + sequence=2, + correlation=correlation, + diagnostics={"surviving_rows": 2, "dropped_rows": 0}, + ) + ) + progress_line = next( + line for line in _scrape(runtime).splitlines() if line.startswith("data_designer_dataset_progress{") + ) + assert progress_line.endswith(" 0.5") + + with runtime.observe_create(port): + correlation = _correlation() + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + "scheduler_job_started", + sequence=3, + correlation=correlation, + diagnostics={"row_group_total_rows": 8}, + ) + ) + progress_line = next( + line for line in _scrape(runtime).splitlines() if line.startswith("data_designer_dataset_progress{") + ) + assert progress_line.endswith(" 0.0") + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + "row_group_checkpointed", + sequence=4, + correlation=correlation, + diagnostics={"surviving_rows": 1, "dropped_rows": 1}, + ) + ) + progress_line = next( + line for line in _scrape(runtime).splitlines() if line.startswith("data_designer_dataset_progress{") + ) + assert progress_line.endswith(" 0.25") + + +def test_runtime_records_create_error_type_without_error_text(runtime: OpenTelemetryRuntime) -> None: + with pytest.raises(ValueError, match="SECRET-EXCEPTION-SENTINEL"): + with runtime.observe_create(_free_port()): + raise ValueError("SECRET-EXCEPTION-SENTINEL") + + metrics = _scrape(runtime) + assert 'error_type="ValueError"' in metrics + assert "SECRET-EXCEPTION-SENTINEL" not in metrics + + +@pytest.mark.parametrize( + ("event_kind", "event_name", "error_type"), + [ + ("non_retryable_dropped", "task.error", "ValueError"), + ("worker_spawn_failed", "worker.spawn_failed", "worker_spawn_failed"), + ("retry_deferred", "task.retry_deferred", None), + ("cancelled", "task.cancelled", None), + ], +) +def test_runtime_maps_selected_scheduler_events( + runtime: OpenTelemetryRuntime, + event_kind: str, + event_name: str, + error_type: str | None, +) -> None: + with runtime.observe_create(_free_port()): + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + event_kind, # type: ignore[arg-type] + sequence=1, + correlation=_correlation(), + diagnostics={"error_type": "ValueError"}, + ) + ) + + metrics = _scrape(runtime) + line = next( + line + for line in metrics.splitlines() + if line.startswith("data_designer_scheduler_events_total{") and f'event_name="{event_name}"' in line + ) + if error_type is None: + assert "error_type=" not in line + else: + assert f'error_type="{error_type}"' in line + + +@pytest.mark.parametrize( + ("domain", "operation"), + [ + (RequestDomain.CHAT, "chat"), + (RequestDomain.EMBEDDING, "embeddings"), + (RequestDomain.IMAGE, "generate_content"), + (RequestDomain.HEALTHCHECK, "healthcheck"), + ], +) +def test_runtime_maps_genai_operations( + runtime: OpenTelemetryRuntime, + domain: RequestDomain, + operation: str, +) -> None: + with runtime.observe_create(_free_port()): + correlation = _correlation() + resource = RequestResourceKey("provider", "model-1", domain) + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_started", + sequence=1, + correlation=correlation, + request_resource_key=resource, + ) + ) + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_completed", + sequence=2, + correlation=correlation, + request_resource_key=resource, + diagnostics={"duration_seconds": 0.1, "outcome": "provider_timeout"}, + ) + ) + + metrics = _scrape(runtime) + line = next( + line for line in metrics.splitlines() if line.startswith("gen_ai_client_operation_duration_seconds_count{") + ) + assert f'gen_ai_operation_name="{operation}"' in line + assert 'error_type="provider_timeout"' in line + assert 'gen_ai_request_model="model-1"' in line + + +def test_runtime_preserves_distinct_bounded_identity_attributes(runtime: OpenTelemetryRuntime) -> None: + common_prefix = "organization/" + ("m" * 150) + resources = [ + RequestResourceKey("provider:v1", "llama3:8b", RequestDomain.CHAT), + RequestResourceKey("providerv1", "llama38b", RequestDomain.CHAT), + RequestResourceKey("provider", common_prefix + "a", RequestDomain.CHAT), + RequestResourceKey("provider", common_prefix + "b", RequestDomain.CHAT), + ] + with runtime.observe_create(_free_port()): + correlation = _correlation() + for index, resource in enumerate(resources): + attempt_id = f"attempt-{index}" + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_started", + sequence=index * 2 + 1, + correlation=correlation, + request_resource_key=resource, + request_attempt_id=attempt_id, + ) + ) + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_completed", + sequence=index * 2 + 2, + correlation=correlation, + request_resource_key=resource, + request_attempt_id=attempt_id, + diagnostics={"duration_seconds": 0.1, "outcome": "success"}, + ) + ) + + metrics = _scrape(runtime) + count_lines = [ + line for line in metrics.splitlines() if line.startswith("gen_ai_client_operation_duration_seconds_count{") + ] + model_labels = {re.search(r'gen_ai_request_model="([^"]+)"', line).group(1) for line in count_lines} + provider_labels = {re.search(r'gen_ai_provider_name="([^"]+)"', line).group(1) for line in count_lines} + + assert {"llama3:8b", "llama38b"} <= model_labels + assert {"provider:v1", "providerv1"} <= provider_labels + long_labels = model_labels - {"llama3:8b", "llama38b"} + assert len(long_labels) == 2 + assert all(len(label) == 128 for label in long_labels) + assert all(re.fullmatch(r"organization/m+-[0-9a-f]{16}", label) for label in long_labels) + + +def test_runtime_collapses_unknown_genai_operation(runtime: OpenTelemetryRuntime) -> None: + with runtime.observe_create(_free_port()): + correlation = _correlation() + resource = {"provider_name": "provider", "domain": "SECRET-DOMAIN-SENTINEL"} + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_started", + sequence=1, + correlation=correlation, + request_resource_key=resource, + ) + ) + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_completed", + sequence=2, + correlation=correlation, + request_resource_key=resource, + diagnostics={"duration_seconds": 0.1, "outcome": "success"}, + ) + ) + + metrics = _scrape(runtime) + assert 'gen_ai_operation_name="_OTHER"' in metrics + assert "SECRET-DOMAIN-SENTINEL" not in metrics + + +def test_runtime_scopes_log_counts_to_active_data_designer_run(runtime: OpenTelemetryRuntime) -> None: + port = _free_port() + with runtime.observe_create(port): + logging.getLogger().error("root") + logging.getLogger("data_designer.test").warning("enabled") + with runtime.observe_create(None): + logging.getLogger("data_designer.test").warning("disabled") + logging.getLogger("data_designer.test").warning("enabled-again") + + metrics = _scrape(runtime) + warning_line = next( + line + for line in metrics.splitlines() + if line.startswith("data_designer_log_records_total{") and 'log_severity="warn"' in line + ) + assert warning_line.endswith(" 2.0") + assert 'log_severity="error"' not in metrics + + +def test_runtime_logs_only_after_releasing_runtime_lock(runtime: OpenTelemetryRuntime) -> None: + lock_available_when_logged = [] + + def record_lock_availability(*_args: object, **_kwargs: object) -> None: + def lock_is_available() -> bool: + acquired = runtime._lock.acquire(blocking=False) + if acquired: + runtime._lock.release() + return acquired + + with ThreadPoolExecutor(max_workers=1) as executor: + lock_available_when_logged.append(executor.submit(lock_is_available).result(timeout=5)) + + with ( + patch.object(opentelemetry.logger, "info", side_effect=record_lock_availability), + patch.object(opentelemetry.logger, "warning", side_effect=record_lock_availability), + ): + with runtime.observe_create(_free_port()): + correlation = _correlation() + runtime.emit_request_event( + RequestAdmissionEvent.capture( + "model_request_started", + sequence=1, + correlation=correlation, + request_resource_key=RequestResourceKey("provider", "model", RequestDomain.CHAT), + request_attempt_id="orphan", + ) + ) + runtime._active_model_requests = MagicMock() + runtime._active_model_requests.add.side_effect = RuntimeError("reconciliation failed") + with runtime.observe_create(_free_port()): + pass + + assert lock_available_when_logged == [True, True, True] + + +def test_runtime_rebinds_idle_listener_and_preserves_metrics(runtime: OpenTelemetryRuntime) -> None: + first_port = _free_port() + second_port = _free_port() + with runtime.observe_create(first_port): + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + "row_group_checkpointed", + sequence=1, + correlation=_correlation(), + diagnostics={"surviving_rows": 1, "dropped_rows": 0}, + ) + ) + first_url = runtime.metrics_url + + with runtime.observe_create(second_port): + runtime.emit_scheduler_event( + SchedulerAdmissionEvent.capture( + "row_group_checkpointed", + sequence=2, + correlation=_correlation(), + diagnostics={"surviving_rows": 2, "dropped_rows": 0}, + ) + ) + + assert first_url == f"http://127.0.0.1:{first_port}/metrics" + assert runtime.metrics_url == f"http://127.0.0.1:{second_port}/metrics" + assert 'record_result="generated"' in _scrape(runtime) + assert " 3.0" in _scrape(runtime) + with pytest.raises(URLError): + urlopen(first_url, timeout=0.2) + + +def test_runtime_stops_rebound_listener_outside_runtime_lock(runtime: OpenTelemetryRuntime) -> None: + first_port = _free_port() + with runtime.observe_create(first_port): + pass + + lock_was_available = False + original_stop_server = opentelemetry._stop_server + + def stop_server(server: object, thread: object) -> None: + nonlocal lock_was_available + + def lock_is_available() -> bool: + acquired = runtime._lock.acquire(blocking=False) + if acquired: + runtime._lock.release() + return acquired + + with ThreadPoolExecutor(max_workers=1) as executor: + lock_was_available = executor.submit(lock_is_available).result(timeout=5) + original_stop_server(server, thread) # type: ignore[arg-type] + + with patch.object(opentelemetry, "_stop_server", side_effect=stop_server) as stop_server_mock: + with runtime.observe_create(_free_port()): + pass + + stop_server_mock.assert_called_once() + assert lock_was_available is True + + +def test_runtime_reuses_active_listener_for_conflicting_port( + runtime: OpenTelemetryRuntime, + caplog: pytest.LogCaptureFixture, +) -> None: + caplog.set_level(logging.WARNING) + first_port = _free_port() + with runtime.observe_create(first_port): + with runtime.observe_create(_free_port()): + assert runtime.metrics_url == f"http://127.0.0.1:{first_port}/metrics" + + assert "reusing it instead of requested port" in caplog.text + + +def test_runtime_port_failure_does_not_escape(runtime: OpenTelemetryRuntime) -> None: + with socket.socket() as occupied: + occupied.bind(("127.0.0.1", 0)) + occupied.listen() + with runtime.observe_create(int(occupied.getsockname()[1])): + assert runtime.metrics_url is None + + +def test_runtime_initialization_failure_cleans_up_partial_provider( + runtime: OpenTelemetryRuntime, + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = MagicMock() + provider.get_meter.return_value.create_histogram.side_effect = RuntimeError("initialization failed") + monkeypatch.setattr("opentelemetry.sdk.metrics.MeterProvider", lambda **_kwargs: provider) + + with runtime.observe_create(_free_port()): + assert runtime.metrics_url is None + + provider.shutdown.assert_called_once_with() + assert runtime._initialized is False + assert runtime._meter_provider is None + assert runtime._log_handler is None + + +def test_metric_log_failure_does_not_echo_record( + runtime: OpenTelemetryRuntime, + capsys: pytest.CaptureFixture[str], +) -> None: + with runtime.observe_create(_free_port()): + assert runtime._log_handler is not None + runtime._log_records = MagicMock() + runtime._log_records.add.side_effect = RuntimeError("recording failed") + record = logging.LogRecord( + "data_designer.test", + logging.ERROR, + __file__, + 1, + "SECRET-LOG-HANDLER-SENTINEL", + (), + None, + ) + runtime._log_handler.emit(record) + + assert "SECRET-LOG-HANDLER-SENTINEL" not in capsys.readouterr().err + + +def test_runtime_shutdown_is_idempotent_and_releases_port(runtime: OpenTelemetryRuntime) -> None: + port = _free_port() + with runtime.observe_create(port): + pass + + runtime.shutdown() + runtime.shutdown() + assert runtime.metrics_url is None + with socket.socket() as rebound: + rebound.bind(("127.0.0.1", port)) + + +def test_disabled_runtime_does_not_import_otel_sdk_in_fresh_process() -> None: + completed = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; " + "from data_designer.integrations.opentelemetry import OpenTelemetryRuntime; " + "runtime = OpenTelemetryRuntime(); " + "context = runtime.observe_create(None); context.__enter__(); context.__exit__(None, None, None); " + "assert 'opentelemetry.sdk' not in sys.modules; " + "assert 'opentelemetry.exporter.prometheus' not in sys.modules" + ), + ], + check=False, + capture_output=True, + text=True, + ) + assert completed.returncode == 0, completed.stderr + + +def _data_designer( + artifact_path: Path, + managed_assets_path: Path, + model_providers: list[ModelProvider], + runtime: OpenTelemetryRuntime, + port: int, +) -> DataDesigner: + managed_assets_path.mkdir(parents=True, exist_ok=True) + designer = DataDesigner( + artifact_path=artifact_path, + model_providers=model_providers, + secret_resolver=PlaintextResolver(), + managed_assets_path=managed_assets_path, + ) + designer._open_telemetry = runtime + designer.set_run_config(RunConfig(otel_metrics_port=port, display_tui=False)) + return designer + + +def test_data_designer_create_exports_and_accumulates_end_to_end( + runtime: OpenTelemetryRuntime, + tmp_path: Path, + stub_model_providers: list[ModelProvider], + stub_sampler_only_config_builder: DataDesignerConfigBuilder, +) -> None: + port = _free_port() + designer = _data_designer( + tmp_path / "artifacts", + tmp_path / "managed-assets", + stub_model_providers, + runtime, + port, + ) + + designer.create(stub_sampler_only_config_builder, dataset_name="first", num_records=3) + designer.create(stub_sampler_only_config_builder, dataset_name="second", num_records=2) + + metrics = _scrape(runtime) + generated_line = next( + line + for line in metrics.splitlines() + if line.startswith("data_designer_dataset_records_total{") and 'record_result="generated"' in line + ) + completed_line = next( + line + for line in metrics.splitlines() + if line.startswith("data_designer_scheduler_events_total{") and 'event_name="job.completed"' in line + ) + create_line = next( + line for line in metrics.splitlines() if line.startswith("data_designer_create_duration_seconds_count") + ) + assert generated_line.endswith(" 5.0") + assert completed_line.endswith(" 2.0") + assert create_line.endswith(" 2.0") + + +def test_data_designer_model_and_log_sinks_export_once_end_to_end( + runtime: OpenTelemetryRuntime, + tmp_path: Path, + stub_model_providers: list[ModelProvider], + stub_model_configs: list[ModelConfig], +) -> None: + port = _free_port() + designer = _data_designer( + tmp_path / "artifacts", + tmp_path / "managed-assets", + stub_model_providers, + runtime, + port, + ) + model_config = stub_model_configs[0].model_copy(update={"skip_health_check": True}) + config_builder = DataDesignerConfigBuilder(model_configs=[model_config]) + config_builder.add_column(LLMTextColumnConfig(name="text", prompt="Generate text", model_alias="stub-model")) + completion = AsyncMock(return_value=make_stub_completion_response(content="generated")) + + with patch.object(OpenAICompatibleClient, "acompletion", completion): + designer.create( + config_builder, + dataset_name="model", + num_records=1, + on_batch_complete=lambda _path: logging.getLogger("data_designer.test").critical("counted once"), + ) + + metrics = _scrape(runtime) + request_line = next( + line for line in metrics.splitlines() if line.startswith("gen_ai_client_operation_duration_seconds_count{") + ) + log_line = next( + line + for line in metrics.splitlines() + if line.startswith("data_designer_log_records_total{") and 'log_severity="fatal"' in line + ) + assert 'gen_ai_operation_name="chat"' in request_line + assert 'gen_ai_provider_name="provider-1"' in request_line + assert request_line.endswith(" 1.0") + assert log_line.endswith(" 1.0") + completion.assert_awaited_once() + + +def test_concurrent_data_designer_creates_share_one_exporter( + runtime: OpenTelemetryRuntime, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + stub_model_providers: list[ModelProvider], + stub_sampler_only_config_builder: DataDesignerConfigBuilder, +) -> None: + port = _free_port() + managed_assets_path = tmp_path / "managed-assets" + designers = [ + _data_designer(tmp_path / f"artifacts-{index}", managed_assets_path, stub_model_providers, runtime, port) + for index in range(2) + ] + barrier = Barrier(2) + observe_create = runtime.observe_create + + @contextlib.contextmanager + def synchronized_observe_create(configured_port: int | None) -> Iterator[None]: + with observe_create(configured_port): + barrier.wait(timeout=10) + yield + + monkeypatch.setattr(runtime, "observe_create", synchronized_observe_create) + + def create(index: int) -> None: + designers[index].create( + stub_sampler_only_config_builder, + dataset_name=f"concurrent-{index}", + num_records=1, + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + list(executor.map(create, range(2))) + + metrics = _scrape(runtime) + generated_line = next( + line + for line in metrics.splitlines() + if line.startswith("data_designer_dataset_records_total{") and 'record_result="generated"' in line + ) + completed_line = next( + line + for line in metrics.splitlines() + if line.startswith("data_designer_scheduler_events_total{") and 'event_name="job.completed"' in line + ) + assert generated_line.endswith(" 2.0") + assert completed_line.endswith(" 2.0") diff --git a/packages/data-designer/tests/interface/test_data_designer.py b/packages/data-designer/tests/interface/test_data_designer.py index ac04b0b46..c261031be 100644 --- a/packages/data-designer/tests/interface/test_data_designer.py +++ b/packages/data-designer/tests/interface/test_data_designer.py @@ -521,6 +521,46 @@ def test_run_config_setting_persists(stub_artifact_path, stub_model_providers): assert data_designer.run_config.max_conversation_correction_steps == 1 +def test_create_observes_configured_otel_port(stub_artifact_path, stub_model_providers) -> None: + data_designer = DataDesigner(artifact_path=stub_artifact_path, model_providers=stub_model_providers) + telemetry = MagicMock() + telemetry.observe_create.return_value = contextlib.nullcontext() + data_designer._open_telemetry = telemetry + + with ( + patch.object(data_designer, "_create_resource_provider", side_effect=RuntimeError("stop")), + pytest.raises(RuntimeError, match="stop"), + ): + data_designer.create(MagicMock()) + + telemetry.observe_create.assert_called_once_with(9464) + + +def test_resource_provider_uses_otel_sink_only_when_enabled( + stub_artifact_path, + stub_model_providers, + stub_sampler_only_config_builder, +) -> None: + person_reader = MagicMock() + data_designer = DataDesigner( + artifact_path=stub_artifact_path, + model_providers=stub_model_providers, + person_reader=person_reader, + ) + telemetry = MagicMock() + data_designer._open_telemetry = telemetry + + with patch.object(dd_mod, "create_resource_provider", return_value=MagicMock()) as create_provider: + data_designer._create_resource_provider("enabled", stub_sampler_only_config_builder) + assert create_provider.call_args.kwargs["request_event_sink"] is telemetry + assert create_provider.call_args.kwargs["scheduler_event_sink"] is telemetry + + data_designer.set_run_config(RunConfig(otel_metrics_port=None)) + data_designer._create_resource_provider("disabled", stub_sampler_only_config_builder) + assert create_provider.call_args.kwargs["request_event_sink"] is None + assert create_provider.call_args.kwargs["scheduler_event_sink"] is None + + def test_run_config_normalizes_error_rate_when_disabled(stub_artifact_path, stub_model_providers): """Test that shutdown_error_rate is normalized to 1.0 when disabled.""" data_designer = DataDesigner(artifact_path=stub_artifact_path, model_providers=stub_model_providers) diff --git a/uv.lock b/uv.lock index 44bf00656..ea1e338f3 100644 --- a/uv.lock +++ b/uv.lock @@ -810,7 +810,11 @@ source = { editable = "packages/data-designer" } dependencies = [ { name = "data-designer-config" }, { name = "data-designer-engine" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-prometheus" }, + { name = "opentelemetry-sdk" }, { name = "packaging" }, + { name = "prometheus-client" }, { name = "prompt-toolkit" }, { name = "typer" }, ] @@ -819,7 +823,11 @@ dependencies = [ requires-dist = [ { name = "data-designer-config", editable = "packages/data-designer-config" }, { name = "data-designer-engine", editable = "packages/data-designer-engine" }, + { name = "opentelemetry-api", specifier = ">=1.43,<1.44" }, + { name = "opentelemetry-exporter-prometheus", specifier = ">=0.64b0,<0.65" }, + { name = "opentelemetry-sdk", specifier = ">=1.43,<1.44" }, { name = "packaging", specifier = ">=25,<27" }, + { name = "prometheus-client", specifier = ">=0.24,<1" }, { name = "prompt-toolkit", specifier = ">=3.0.0,<4" }, { name = "typer", specifier = ">=0.12.0,<1" }, ] @@ -2826,6 +2834,59 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/57/a7/b35835e278c18b85206834b3aa3abe68e77a98769c59233d1f6300284781/numpy-2.4.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:4b42639cdde6d24e732ff823a3fa5b701d8acad89c4142bc1d0bd6dc85200ba5", size = 12504685, upload-time = "2026-03-09T07:58:50.525Z" }, ] +[[package]] +name = "opentelemetry-api" +version = "1.43.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ae/cc/e4c9584181f86494df0f6bdec1a4f3280c50db44704dc2a407e994fc87bb/opentelemetry_api-1.43.0.tar.gz", hash = "sha256:107d0d03857ea8fc7c5fcbbbd83f800c281f0d560553d61c1d675fccfd1761c1", size = 73476, upload-time = "2026-06-24T15:19:55.323Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/17/83/6dba32b85f31868400440dc7ad2ca1eab94cbbf3a7b0459ed39f8311a9e2/opentelemetry_api-1.43.0-py3-none-any.whl", hash = "sha256:20acf45e9b21851926835292e4045d290acade1edd2ff3de86d2f069687ba1fd", size = 61912, upload-time = "2026-06-24T15:19:35.434Z" }, +] + +[[package]] +name = "opentelemetry-exporter-prometheus" +version = "0.64b0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-sdk" }, + { name = "prometheus-client" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/97/b3/f778f705289ea5f1c54589ae4494d318c9cede052f18ca33979fda0ef016/opentelemetry_exporter_prometheus-0.64b0.tar.gz", hash = "sha256:96fec79be9527cb9dc994d7e663051df35161eb936fe2d41954725e4595abbc1", size = 16405, upload-time = "2026-06-24T15:20:02.278Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/dd/d95db7ac5bec037b3dcc8ea6fe64cd5dc4bc87017d2e06e7410318e85dc0/opentelemetry_exporter_prometheus-0.64b0-py3-none-any.whl", hash = "sha256:9979a15f8d007d442bc7a6e16f4cbde5e8e0c5e99689887ceb5f33da251f0655", size = 13031, upload-time = "2026-06-24T15:19:44.159Z" }, +] + +[[package]] +name = "opentelemetry-sdk" +version = "1.43.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3e/eb/5041074274ac0956b03637cc039d434569112468e875eddfcc9a0674ce06/opentelemetry_sdk-1.43.0.tar.gz", hash = "sha256:d8187c81c162df9913e4003dd6485f7390d9a24fc17026ec7387b8b8218b08e9", size = 254744, upload-time = "2026-06-24T15:20:08.467Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/e3/b17be23af124201c9f52eececd4cc8ddfed1597d37b4ee771895d325805c/opentelemetry_sdk-1.43.0-py3-none-any.whl", hash = "sha256:d1323a547c1ce69d6a069a17a44b7da82bb8b332051ecb074041f87642c86823", size = 178852, upload-time = "2026-06-24T15:19:52.169Z" }, +] + +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.64b0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/30/5f26df29509eccd86b99b481ac9ffa39da49ba9577cc69071c552ae30447/opentelemetry_semantic_conventions-0.64b0.tar.gz", hash = "sha256:72f76fb2d1582d9d033dd1fcd84532e961e6ff3d90d24ba6fabc72975a83864c", size = 148340, upload-time = "2026-06-24T15:20:09.267Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/ca/23ba87a221b574a7c5a99d48849d80bfe8b047624681357e2b002e566187/opentelemetry_semantic_conventions-0.64b0-py3-none-any.whl", hash = "sha256:ea77e85e354b8f604ddbe5f3d9135216f982fa4d77e5859ac30f6d8a50505aa6", size = 203713, upload-time = "2026-06-24T15:19:53.339Z" }, +] + [[package]] name = "overrides" version = "7.7.0"