From 2a6d93221fe31c86e11f9db730b6a2abe276ec3d Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Tue, 7 Jul 2026 10:40:54 -0400 Subject: [PATCH 1/3] feat: persist scheduler events to JSONL Add an opt-in RunConfig flag that writes create-run scheduler events beside dataset artifacts. Keep preview and default runs disk-free, and make file I/O failures non-fatal. Signed-off-by: Eric W. Tramel --- .../src/data_designer/config/run_config.py | 5 ++ .../tests/config/test_run_config.py | 8 +++ .../dataset_builders/dataset_builder.py | 67 ++++++++++-------- .../src/data_designer/engine/observability.py | 50 ++++++++++++- .../dataset_builders/test_dataset_builder.py | 66 ++++++++++++++++- .../tests/engine/test_observability.py | 70 +++++++++++++++++++ 6 files changed, 235 insertions(+), 31 deletions(-) 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 c56d000dd..0e85b82d7 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 @@ -147,6 +147,10 @@ class RunConfig(ConfigBase): single conversation when generation tasks call `ModelFacade.generate(...)`. Must be >= 0. Default is 0. async_trace: If True, collect per-task tracing data. Default is False. + write_scheduler_events: If True, create runs write structured scheduler diagnostics to + ``scheduler_events.jsonl`` in the dataset directory. The file is JSONL, not direct + Perfetto input, and may contain sensitive column, provider, model, task, and resource + labels. Preview runs do not write it. Default is False. display_tui: If True, display the terminal throughput TUI instead of periodic log lines during generation. Requires a TTY; falls back to log lines in non-TTY environments. Default is True. @@ -185,6 +189,7 @@ class RunConfig(ConfigBase): max_conversation_restarts: int = Field(default=5, ge=0) max_conversation_correction_steps: int = Field(default=0, ge=0) async_trace: bool = False + write_scheduler_events: bool = False display_tui: bool = True progress_interval: float = Field(default=5.0, gt=0.0) preserve_dropped_columns: bool = Field( 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 e52705d5f..ec795215e 100644 --- a/packages/data-designer-config/tests/config/test_run_config.py +++ b/packages/data-designer-config/tests/config/test_run_config.py @@ -32,6 +32,14 @@ def test_run_config_accepts_display_tui() -> None: assert RunConfig(display_tui=False).display_tui is False +def test_run_config_does_not_write_scheduler_events_by_default() -> None: + assert RunConfig().write_scheduler_events is False + + +def test_run_config_accepts_scheduler_event_writes() -> None: + assert RunConfig(write_scheduler_events=True).write_scheduler_events is True + + 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/dataset_builder.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py index c29dee3b7..a7b22b57a 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 @@ -5,6 +5,7 @@ import asyncio import concurrent.futures +import contextlib import functools import json import logging @@ -49,6 +50,7 @@ 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.processing.processors.base import Processor from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor from data_designer.engine.readiness import run_readiness_check @@ -808,38 +810,45 @@ def on_complete(final_path: Path | str | None) -> None: buffer_size=buffer_size, ) - scheduler, buffer_manager = self._prepare_async_run( - generators, - num_records, - buffer_size, - on_finalize_row_group=finalize_row_group, - shutdown_error_rate=settings.shutdown_error_rate, - shutdown_error_window=settings.shutdown_error_window, - disable_early_shutdown=settings.disable_early_shutdown, - trace=trace_enabled, - precomputed_row_groups=precomputed_row_groups, - initial_actual_num_records=initial_actual_num_records, - initial_total_num_batches=initial_total_num_batches, - ) - # Telemetry snapshot group_id = uuid.uuid4().hex pre_batch_snapshot = self._resource_provider.model_registry.get_model_usage_snapshot() - # Run on background event loop. Capture scheduler state in `finally` - # so the structured signal is preserved even if `scheduler.run()` - # raises during the salvage path - otherwise callers see a generic - # error and lose the early-shutdown context. - loop = ensure_async_engine_loop() - future = asyncio.run_coroutine_threadsafe(scheduler.run(), loop) - try: - _await_async_scheduler_result(future, scheduler) - finally: - self._task_traces = scheduler.traces - self._early_shutdown = scheduler.early_shutdown - self._partial_row_groups = scheduler.partial_row_groups - self._actual_num_records = buffer_manager.actual_num_records - self._first_non_retryable_error = scheduler.first_non_retryable_error + event_sink_context = ( + JsonlSchedulerEventSink(self.artifact_storage.base_dataset_path / "scheduler_events.jsonl") + if settings.write_scheduler_events + else contextlib.nullcontext() + ) + with event_sink_context as scheduler_event_sink: + scheduler, buffer_manager = self._prepare_async_run( + generators, + num_records, + buffer_size, + on_finalize_row_group=finalize_row_group, + shutdown_error_rate=settings.shutdown_error_rate, + shutdown_error_window=settings.shutdown_error_window, + disable_early_shutdown=settings.disable_early_shutdown, + trace=trace_enabled, + precomputed_row_groups=precomputed_row_groups, + initial_actual_num_records=initial_actual_num_records, + initial_total_num_batches=initial_total_num_batches, + scheduler_event_sink=scheduler_event_sink, + ) + + # Run on background event loop. Capture scheduler state in `finally` + # so the structured signal is preserved even if `scheduler.run()` + # raises during the salvage path - otherwise callers see a generic + # error and lose the early-shutdown context. + loop = ensure_async_engine_loop() + future = asyncio.run_coroutine_threadsafe(scheduler.run(), loop) + try: + _await_async_scheduler_result(future, scheduler) + finally: + self._task_traces = scheduler.traces + self._early_shutdown = scheduler.early_shutdown + self._partial_row_groups = scheduler.partial_row_groups + self._actual_num_records = buffer_manager.actual_num_records + self._first_non_retryable_error = scheduler.first_non_retryable_error # Emit telemetry try: @@ -889,6 +898,7 @@ def _prepare_async_run( precomputed_row_groups: RowGroupInput | None = None, initial_actual_num_records: int = 0, initial_total_num_batches: int = 0, + scheduler_event_sink: SchedulerAdmissionEventSink | None = None, ) -> tuple[AsyncTaskScheduler, RowGroupBufferManager]: """Build a fully-wired scheduler and buffer manager for async generation. @@ -973,6 +983,7 @@ 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, 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/observability.py b/packages/data-designer-engine/src/data_designer/engine/observability.py index b43923557..353205b8e 100644 --- a/packages/data-designer-engine/src/data_designer/engine/observability.py +++ b/packages/data-designer-engine/src/data_designer/engine/observability.py @@ -5,14 +5,16 @@ import contextvars import itertools +import json import logging import math import time from collections.abc import Mapping -from dataclasses import dataclass, field, fields, is_dataclass +from dataclasses import asdict, dataclass, field, fields, is_dataclass from enum import Enum +from os import PathLike from threading import Lock -from typing import Callable, Literal, Protocol +from typing import Callable, Literal, Protocol, TextIO logger = logging.getLogger(__name__) @@ -216,6 +218,50 @@ class SchedulerAdmissionEventSink(Protocol): def emit_scheduler_event(self, event: SchedulerAdmissionEvent) -> None: ... +class JsonlSchedulerEventSink: + """Append scheduler events as newline-delimited JSON without affecting generation on file errors.""" + + def __init__(self, path: str | PathLike[str]) -> None: + self._path = path + self._file: TextIO | None = None + + def __enter__(self) -> JsonlSchedulerEventSink | None: + try: + self._file = open(self._path, "a", encoding="utf-8") + except Exception: + logger.warning("Failed to open scheduler event file %s", self._path, exc_info=True) + return None + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def emit_scheduler_event(self, event: SchedulerAdmissionEvent) -> None: + file = self._file + if file is None: + return + try: + file.write(json.dumps(asdict(event), ensure_ascii=False) + "\n") + file.flush() + except Exception: + self._file = None + try: + file.close() + except Exception: + pass + raise + + def close(self) -> None: + file = self._file + self._file = None + if file is None: + return + try: + file.close() + except Exception: + logger.warning("Failed to close scheduler event file %s", self._path, exc_info=True) + + class RequestAdmissionEventSink(Protocol): def emit_request_event(self, event: RequestAdmissionEvent) -> 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 3d855fd5a..4d625243f 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 @@ -1047,11 +1047,15 @@ def _make_sampler_only_builder( tmp_path: Path, *, resume: ResumeMode = ResumeMode.IF_POSSIBLE, + write_scheduler_events: bool = False, ) -> tuple[DatasetBuilder, ArtifactStorage]: """Create a builder that can run end-to-end without model or MCP stubs.""" storage = ArtifactStorage(artifact_path=tmp_path, resume=resume) stub_resource_provider.artifact_storage = storage - stub_resource_provider.run_config = RunConfig(buffer_size=2) + stub_resource_provider.run_config = RunConfig( + buffer_size=2, + write_scheduler_events=write_scheduler_events, + ) config_builder = DataDesignerConfigBuilder() config_builder.add_column( @@ -1066,6 +1070,66 @@ def _make_sampler_only_builder( ) +@pytest.mark.parametrize("write_scheduler_events", [False, True]) +def test_build_writes_scheduler_events_when_enabled( + stub_resource_provider: Mock, + tmp_path: Path, + write_scheduler_events: bool, +) -> None: + builder, _storage = _make_sampler_only_builder( + stub_resource_provider, + tmp_path, + resume=ResumeMode.NEVER, + write_scheduler_events=write_scheduler_events, + ) + + final_path = builder.build(num_records=2, resume=ResumeMode.NEVER) + event_path = final_path.parent / "scheduler_events.jsonl" + + assert event_path.exists() is write_scheduler_events + if write_scheduler_events: + events = [json.loads(line) for line in event_path.read_text(encoding="utf-8").splitlines()] + assert events[0]["event_kind"] == "scheduler_job_started" + assert events[-1]["event_kind"] == "scheduler_job_completed" + + +def test_preview_does_not_write_scheduler_events(stub_resource_provider: Mock, tmp_path: Path) -> None: + builder, _storage = _make_sampler_only_builder( + stub_resource_provider, + tmp_path, + resume=ResumeMode.NEVER, + write_scheduler_events=True, + ) + + builder.build_preview(num_records=1) + + assert list(tmp_path.rglob("scheduler_events.jsonl")) == [] + + +def test_resumed_build_appends_scheduler_event_segment(stub_resource_provider: Mock, tmp_path: Path) -> None: + builder, _storage = _make_sampler_only_builder( + stub_resource_provider, + tmp_path, + resume=ResumeMode.NEVER, + write_scheduler_events=True, + ) + final_path = builder.build(num_records=1, resume=ResumeMode.NEVER) + event_path = final_path.parent / "scheduler_events.jsonl" + + resumed_builder, _storage = _make_sampler_only_builder( + stub_resource_provider, + tmp_path, + resume=ResumeMode.ALWAYS, + write_scheduler_events=True, + ) + resumed_builder.build(num_records=3, resume=ResumeMode.ALWAYS) + + events = [json.loads(line) for line in event_path.read_text(encoding="utf-8").splitlines()] + starts = [event for event in events if event["event_kind"] == "scheduler_job_started"] + assert len(starts) == 2 + assert starts[0]["diagnostics"]["run_id"] != starts[1]["diagnostics"]["run_id"] + + def test_build_resume_ordered_seed_dataset_continues_from_next_planned_row(stub_resource_provider, tmp_path): """Regression for issue #709: resume must not replay ordered seed rows.""" diff --git a/packages/data-designer-engine/tests/engine/test_observability.py b/packages/data-designer-engine/tests/engine/test_observability.py index 78ec66424..5e67d3935 100644 --- a/packages/data-designer-engine/tests/engine/test_observability.py +++ b/packages/data-designer-engine/tests/engine/test_observability.py @@ -4,10 +4,16 @@ from __future__ import annotations import json +import logging from dataclasses import asdict, dataclass from enum import Enum +from pathlib import Path +from unittest.mock import Mock + +import pytest from data_designer.engine.observability import ( + JsonlSchedulerEventSink, RequestAdmissionEvent, RuntimeCorrelation, RuntimeCorrelationProvider, @@ -117,6 +123,70 @@ def test_in_memory_admission_event_sink_collects_scheduler_and_request_events() assert sink.request_events == [request_event] +def test_jsonl_scheduler_event_sink_flushes_closes_and_appends(tmp_path: Path) -> None: + path = tmp_path / "scheduler_events.jsonl" + events = [ + SchedulerAdmissionEvent.capture("selected", sequence=1, diagnostics={"label": "café"}), + SchedulerAdmissionEvent.capture("task_completed", sequence=2), + ] + + for event in events: + with JsonlSchedulerEventSink(path) as sink: + assert sink is not None + sink.emit_scheduler_event(event) + contents = path.read_text(encoding="utf-8") + assert json.loads(contents.splitlines()[-1]) == asdict(event) + + assert "café" in contents + assert [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()] == [ + asdict(event) for event in events + ] + + +def test_jsonl_scheduler_event_sink_open_failure_is_inactive( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr("builtins.open", Mock(side_effect=OSError("open failed"))) + + with caplog.at_level(logging.WARNING): + with JsonlSchedulerEventSink("events.jsonl") as sink: + assert sink is None + + assert [record.message for record in caplog.records] == ["Failed to open scheduler event file events.jsonl"] + + +def test_jsonl_scheduler_event_sink_close_failure_warns_once( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + file = Mock() + file.close.side_effect = OSError("close failed") + monkeypatch.setattr("builtins.open", Mock(return_value=file)) + + with caplog.at_level(logging.WARNING): + with JsonlSchedulerEventSink("events.jsonl") as sink: + assert sink is not None + sink.close() + + file.close.assert_called_once_with() + assert [record.message for record in caplog.records] == ["Failed to close scheduler event file events.jsonl"] + + +def test_jsonl_scheduler_event_sink_disables_after_write_failure(monkeypatch: pytest.MonkeyPatch) -> None: + file = Mock() + file.write.side_effect = OSError("write failed") + monkeypatch.setattr("builtins.open", Mock(return_value=file)) + event = SchedulerAdmissionEvent.capture("selected", sequence=1) + + with JsonlSchedulerEventSink("events.jsonl") as sink: + assert sink is not None + with pytest.raises(OSError, match="write failed"): + sink.emit_scheduler_event(event) + sink.emit_scheduler_event(event) + + file.write.assert_called_once() + file.close.assert_called_once_with() + + 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( From 382da97e77ce157762414b56aea74361dc5c5be4 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Tue, 7 Jul 2026 11:16:43 -0400 Subject: [PATCH 2/3] docs: note scheduler event I/O overhead --- .../src/data_designer/config/run_config.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 0e85b82d7..4f81df05f 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 @@ -150,7 +150,8 @@ class RunConfig(ConfigBase): write_scheduler_events: If True, create runs write structured scheduler diagnostics to ``scheduler_events.jsonl`` in the dataset directory. The file is JSONL, not direct Perfetto input, and may contain sensitive column, provider, model, task, and resource - labels. Preview runs do not write it. Default is False. + labels. Each event is flushed to disk, so enabling this option adds file I/O overhead. + Preview runs do not write it. Default is False. display_tui: If True, display the terminal throughput TUI instead of periodic log lines during generation. Requires a TTY; falls back to log lines in non-TTY environments. Default is True. From 254490fd405f25d5e84223d3aea10d090697b064 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Tue, 7 Jul 2026 15:17:24 -0400 Subject: [PATCH 3/3] docs: clarify scheduler event persistence --- architecture/dataset-builders.md | 2 ++ .../src/data_designer/engine/observability.py | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/architecture/dataset-builders.md b/architecture/dataset-builders.md index a218e08e9..280d96838 100644 --- a/architecture/dataset-builders.md +++ b/architecture/dataset-builders.md @@ -88,6 +88,8 @@ Metadata writes are atomic (`tmp` file + `fsync` + `os.replace`) because `metada `DatasetCreationResults` from a resume invocation reflects the full on-disk dataset for anything that reads the artifact directory (`load_dataset`, `count_records`, `load_analysis`, `export`, `push_to_hub`), but per-run observability (`task_traces`, model-usage logs, telemetry events) is scoped to the current invocation — the original run's in-memory state is not persisted across process boundaries. +When `RunConfig.write_scheduler_events` is enabled, each invocation also appends its scheduler event segment to `scheduler_events.jsonl` in the dataset directory. + ## Data Flow ``` 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 353205b8e..3d422e994 100644 --- a/packages/data-designer-engine/src/data_designer/engine/observability.py +++ b/packages/data-designer-engine/src/data_designer/engine/observability.py @@ -219,7 +219,10 @@ def emit_scheduler_event(self, event: SchedulerAdmissionEvent) -> None: ... class JsonlSchedulerEventSink: - """Append scheduler events as newline-delimited JSON without affecting generation on file errors.""" + """Append scheduler events as newline-delimited JSON from a single writer. + + A write or flush failure disables the sink and is re-raised once for the caller to handle. + """ def __init__(self, path: str | PathLike[str]) -> None: self._path = path