feat: persist scheduler events to JSONL

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
 
 ```

packages/data-designer-config/src/data_designer/config/run_config.py

@@ -147,6 +147,11 @@ 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. 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 False.
@@ -185,6 +190,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 = False
     progress_interval: float = Field(default=5.0, gt=0.0)
     preserve_dropped_columns: bool = Field(

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)

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

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,53 @@ class SchedulerAdmissionEventSink(Protocol):
     def emit_scheduler_event(self, event: SchedulerAdmissionEvent) -> None: ...
 
 
+class JsonlSchedulerEventSink:
+    """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
+        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: ...
 

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."""
 

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(