feat: persist scheduler events to JSONL

[eric-tramel]

📋 Summary

Adds an opt-in RunConfig.write_scheduler_events setting that persists the existing structured scheduler event stream beside dataset artifacts. This gives operators a durable JSONL timeline for create runs while preserving the current zero-I/O default and keeping preview memory-only.

The sink appends and flushes each event, closes with the scheduler lifecycle, and treats open, write, flush, and close failures as non-fatal diagnostics failures.

Perfetto visualization derived from the opt-in scheduler JSONL for an 8-row gpt-5.5 workflow with 24 successful model calls. The scheduler saturates its eight-task concurrency cap, briefly queues work, and cleanly drains all workers and resources.

🔗 Related Issue

Related to #727. This PR exposes the raw opt-in scheduler event stream; it does not implement the broader automatic coarse diagnostics contract.

🔄 Changes

• Add RunConfig.write_scheduler_events=False.
• Write enabled create-run events to <dataset root>/scheduler_events.jsonl.
• Append a new run segment on resume instead of truncating prior diagnostics.
• Keep preview runs disk-free.
• Disable only the JSONL sink after its first write failure so generation continues without warning storms.
• Add config, sink lifecycle, create, preview, and resume coverage.

Usage

data_designer.set_run_config(dd.RunConfig(write_scheduler_events=True))
results = data_designer.create(...)

The file is raw JSON Lines and may include column, provider, model, task, and resource labels. It is not direct Perfetto input.

Changelog

• Adds opt-in scheduler event JSONL artifacts for dataset create runs.

🧪 Testing

Ran make check-all-fix successfully. Ran the full make test suite with 3,797 tests passing and 1 skipped. After review-driven failure-handling changes, reran all 2,173 engine tests successfully.

• make test passes
• Unit tests added/updated
• E2E tests added/updated (not applicable; covered by engine integration tests)

✅ Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Architecture docs updated (not applicable; the public setting is documented in RunConfig)

[github-actions]

Fern preview: https://nvidia-preview-pr-800.docs.buildwithfern.com/nemo/datadesigner

Fern previews include the docs-website version archive with PR changes synced into latest. Notebook tutorials are rendered without execution outputs in previews.

[greptile-apps]

Greptile Summary

This PR adds an opt-in RunConfig.write_scheduler_events flag (defaulting to False) that, when enabled, appends structured scheduler diagnostics to scheduler_events.jsonl alongside dataset artifacts during create runs.

• Introduces JsonlSchedulerEventSink as a context-manager-backed JSONL writer; open/write/flush/close failures are treated as non-fatal diagnostics (the sink self-disables on first write error to avoid warning storms; the scheduler's own try/except wrapper around emit_scheduler_event prevents any failure from aborting generation).
• Plumbs the optional sink through _prepare_async_run into AsyncTaskScheduler; preview runs are explicitly excluded by not entering the sink context; resumed runs append a new segment via "a" mode, preserving prior run history.
• Adds comprehensive unit and integration tests covering the default-off behavior, create vs. preview distinction, resume appending, and all failure modes (open, write, flush, close).

Confidence Score: 5/5

Safe to merge — the change is strictly opt-in, the default is off, and all failure modes leave generation unaffected.

The feature is entirely additive behind a false-defaulted flag. The sink self-disables after the first write error, the scheduler's own exception wrapper ensures no I/O failure can propagate to abort a run, and the preview path is untouched. Tests cover the happy path, resume appending, and every documented failure mode.

No files require special attention.

Important Files Changed

Filename	Overview
packages/data-designer-engine/src/data_designer/engine/observability.py	Adds JsonlSchedulerEventSink: context manager that opens the file in append mode, writes/flushes each event, self-disables on first write failure (re-raises once for caller logging), and closes cleanly.
packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py	Wraps the async-run block in the event-sink context when write_scheduler_events is enabled; preview path unchanged (no sink); nullcontext fallback cleanly passes None to _prepare_async_run.
packages/data-designer-config/src/data_designer/config/run_config.py	Adds write_scheduler_events: bool = False field with docstring; clean addition with no side effects on existing config.
packages/data-designer-engine/tests/engine/test_observability.py	Four new test cases covering: append and flush, open failure, close failure (idempotent), and write-failure disabling.
packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py	Three integration tests covering create (enabled/disabled), preview exclusion, and resume append with distinct run_ids.
packages/data-designer-config/tests/config/test_run_config.py	Two simple unit tests confirming default-off and explicit-on behavior.
architecture/dataset-builders.md	One-line documentation addition for the scheduler_events.jsonl artifact.

_{Reviews (7): Last reviewed commit: "Merge remote-tracking branch 'origin/mai..." | Re-trigger Greptile}

[github-actions]

Code Review: PR #800 — feat: persist scheduler events to JSONL

Author: @eric-tramel · Base: main · Files: 6 (+235 / −31) · Related to #727

Summary

Adds an opt-in RunConfig.write_scheduler_events flag (default False). When enabled, create runs persist the already-existing structured scheduler event stream to <dataset root>/scheduler_events.jsonl via a new JsonlSchedulerEventSink in observability.py. The sink appends+flushes each event, is disabled on its first write failure, and is opened/closed with the scheduler lifecycle through a context manager. Preview runs remain disk-free, and resume appends a new run segment rather than truncating. The change is well-scoped and the design (opt-in, non-fatal on I/O failure, zero-I/O default) matches the stated intent.

Findings

Correctness — solid, one subtlety worth confirming

• File-open ordering is safe. The sink opens base_dataset_path / "scheduler_events.jsonl" at dataset_builder.py:818 (context __enter__). The directory is created earlier in build() via _write_builder_config() → mkdir_if_needed(base_dataset_path) (dataset_builder.py:1019, called at :327), which runs before _build_async. So the parent dir exists by the time the file is opened. Good. Note this relies on the build() call path; if _build_async were ever reached without _write_builder_config() first, the open would fail — but that failure is non-fatal by design, so it degrades gracefully rather than crashing.

• Double-layer failure handling is correct and non-redundant. JsonlSchedulerEventSink.emit_scheduler_event re-raises after disabling itself (sets _file = None, closes), and the scheduler's _emit_scheduler_event (async_scheduler.py:479) catches the exception and logs one warning. On the next event the sink sees _file is None and returns silently. Net effect: exactly one warning, then quiet — the "no warning storm" claim holds end-to-end. The re-raise is load-bearing (it's how the scheduler learns to stop), so it should not be "simplified" to a swallow later.

• Import direction preserved. run_config.py (config package) adds only a plain bool field with no engine import (verified). The new sink lives in the engine package, consumed by the engine builder — correct per interface → engine → config. The structural analyzer's "config → engine" violations are false positives from .set()/emit name collisions, not real reverse imports (see Structural Impact note below).

Concurrency / Performance

• Synchronous flush() on every event runs on the async loop thread. Events are emitted from within scheduler.run() on the background event loop (async_scheduler.py:1067, _emit_scheduler_event called from ~30 sites). Each emit does file.write(...) + file.flush() synchronously. For large runs this is many blocking flush() calls on the loop thread, which can add scheduler latency. Because the feature is opt-in and off by default, this is acceptable, but the per-event flush() is the most likely thing an operator enabling this on a big run would feel. Consider documenting that this flag adds I/O overhead, or (follow-up) buffering flushes. Not blocking.

Robustness

• asdict(event) assumes fully-serializable dataclass contents. emit_scheduler_event does json.dumps(asdict(event), ensure_ascii=False). SchedulerAdmissionEvent diagnostics carry task_group_key, resource_request amounts, snapshots, etc. (async_scheduler.py:448-476). If any diagnostic value is ever a non-JSON-serializable object (e.g. an enum without a value, a set, a custom key type), json.dumps raises TypeError — which the sink treats as a write failure and re-raises, permanently disabling the sink after the first event. That's a safe failure mode (generation continues), but it would silently drop all diagnostics for the run. A default=str fallback in json.dumps would make the artifact resilient to future diagnostic additions. Worth considering given the diagnostics dict is open-ended. The existing test only exercises string/int diagnostics, so this path is untested.

Style / Conventions

• Missing from __future__ import annotations — observability.py uses JsonlSchedulerEventSink | None as the __enter__ return annotation. Since the file already imported from typing import ... TextIO and the return type references the class by forward name inside its own body, this works, but AGENTS.md requires from __future__ import annotations in every source file. Confirm the file already has it at the top (the diff doesn't show the header); if not, it should be added. (The engine test file test_observability.py correctly has it.)
• __exit__(self, *args: object) is a reasonable terse signature; matches the "match neighboring code" guidance if other context managers here do the same. Minor: the conventional typed form is (exc_type, exc, tb), but *args: object is acceptable.

Test Coverage — strong

• Config: default-false and accepts-true (test_run_config.py).
• Sink unit tests cover the four important paths: flush/close/append + UTF-8 (café), open-failure → inactive + one warning, close-failure → one warning, write-failure → disabled after first write (file.write.assert_called_once). This directly validates the anti-warning-storm behavior.
• Builder integration: parametrized enabled/disabled asserts file presence and first/last event_kind; preview writes nothing; resume produces two scheduler_job_started segments with distinct run_ids. The resume test is a nice touch — it locks in the append-don't-truncate contract.
• Gap: no test with a non-string/non-primitive diagnostic value to confirm serialization behavior (ties to the asdict/json.dumps note above).

Suggestions (priority order)

1. (nice-to-have) Add default=str to json.dumps in emit_scheduler_event so future non-serializable diagnostics degrade to a string rather than disabling the whole artifact; add a test with an enum/complex diagnostic value.
2. (verify) Confirm observability.py has from __future__ import annotations at the top (required by AGENTS.md).
3. (doc) Mention the per-event flush I/O cost in the RunConfig.write_scheduler_events docstring, alongside the existing sensitive-labels warning.

None of these are blocking; the PR is safe to merge as-is given the opt-in default.

Structural Impact (graphify, 2.3s)

Risk: HIGH (18 import direction violation(s))

• 6 Python files, 102 AST entities, 7/83 clusters

Import Direction Violations (18)

Legal direction: interface -> engine -> config

• ._is_missing_value() (config) --calls--> .set() (engine)
• column_types() (config) --calls--> .set() (engine)
• generate_analysis_report() (config) --calls--> .set() (engine)
• _validate_skip_scope() (config) --calls--> .set() (engine)
• required_columns() (config) --calls--> .set() (engine)
• +13 more

High-Connectivity Changes

• .set() (114 deps) in packages/data-designer-engine/src/data_designer/engine/observability.py
• DatasetBuilder (75 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py
• JinjaRenderingEngine (66 deps) in packages/data-designer-config/src/data_designer/config/run_config.py
• SchedulerAdmissionEventSink (65 deps) in packages/data-designer-engine/src/data_designer/engine/observability.py
• RunConfig (59 deps) in packages/data-designer-config/src/data_designer/config/run_config.py
• SchedulerAdmissionEvent (45 deps) in packages/data-designer-engine/src/data_designer/engine/observability.py
• RuntimeCorrelation (43 deps) in packages/data-designer-engine/src/data_designer/engine/observability.py
• RowGroupResumePlan (36 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py
• +52 more

Cross-Package Dependencies

• .run_create() (interface) --calls--> .model_copy() (config)
• .__init__() (interface) --calls--> .set() (engine)
• _validate_install_metadata() (interface) --calls--> .set() (engine)
• _required_catalog_object() (interface) --calls--> .set() (engine)
• _validate_catalog_object_keys() (interface) --calls--> .set() (engine)
• select_multiple_with_arrows() (interface) --calls--> .set() (engine)
• +288 more

Reviewer note on structural impact: The 18 flagged "config → engine" violations are almost certainly false positives from method-name collisions (.set(), emit_*) rather than real reverse imports. I verified run_config.py adds only a scalar bool field with no engine import, and the new sink lives in and is consumed within the engine package. The genuinely useful signal here is the high-connectivity of observability.py (.set() at 114 deps, SchedulerAdmissionEventSink at 65 deps): the new JsonlSchedulerEventSink implements the widely-used SchedulerAdmissionEventSink Protocol, so its emit_scheduler_event contract (including the re-raise-on-failure semantics) is on a hot path touched by ~30 call sites in the scheduler. The blast radius is contained because the sink is only instantiated behind the opt-in flag, but the Protocol conformance and failure semantics deserved the extra scrutiny applied above.

Verdict

Approve with minor follow-ups. The change is well-designed, correctly layered, non-fatal on I/O errors, and thoroughly tested for its stated behavior. The failure-handling flow (disable-and-warn-once) is correct end-to-end across the sink and scheduler. Recommended non-blocking follow-ups: default=str in json.dumps + a matching test for non-primitive diagnostics, and confirming the from __future__ import annotations header in observability.py.

[nabinchha]

Thanks for putting this together, @eric-tramel — clean, well-scoped feature with genuinely thorough tests.

Summary

Adds an opt-in RunConfig.write_scheduler_events flag that persists the existing structured scheduler event stream to <dataset root>/scheduler_events.jsonl for create runs, appending a new segment on resume while keeping the default and preview paths disk-free. The implementation matches the stated intent in the PR description, and the failure handling (open/write/flush/close treated as non-fatal) is well covered by tests.

Findings

Warnings — Worth addressing

packages/data-designer-engine/src/data_designer/engine/observability.py:239 — Docstring promises more than the sink enforces

• What: The class docstring says events are appended "without affecting generation on file errors," but emit_scheduler_event actually re-raises on write/flush failure. The non-fatal guarantee is really provided by AsyncTaskScheduler._emit_scheduler_event, which wraps the call in try/except and logs a single warning.
• Why: The error-handling contract is inconsistent between the docstring and the implementation, and the "non-fatal" guarantee leaks to the caller. The sink implements the general SchedulerAdmissionEventSink protocol, so a future caller using it directly — reasonably trusting the docstring — would get an exception on the first write error and could crash generation.
• Suggestion: Either self-contain the guarantee (log a single warning + disable internally and return instead of raise), or update the docstring/emit_scheduler_event to state explicitly that it disables itself and re-raises, so callers know they must handle it. The current split works for today's single caller; it's the contract clarity that's worth tightening.

architecture/dataset-builders.md:89 — Doc now understates observability persistence

• What: That line states 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." This PR makes a category of per-run scheduler observability persistable to disk and appended across resume segments.
• Why: Contributors rely on that architecture note to reason about what does and doesn't survive a run. Leaving it as-is means the recorded architecture is now incomplete relative to the new behavior.
• Suggestion: Add a sentence noting the opt-in scheduler_events.jsonl artifact (in this PR, or an explicit follow-up issue) so the doc stays aligned with the code.

Suggestions — Take it or leave it

packages/data-designer-engine/src/data_designer/engine/observability.py:239 — Single-writer assumption is implicit

• What: write, flush, and the self._file = None reset aren't synchronized. This is safe today because scheduler events are only emitted from the single background asyncio loop thread, so writes are serialized.
• Why: The class is otherwise a general-purpose sink; if it's ever wired to a multi-threaded emitter later, interleaved writes/flush could corrupt lines or race on the disable path.
• Suggestion: A one-line comment documenting the single-writer expectation would save a future reader (or refactorer) from having to re-derive it.

What Looks Good

• Clean opt-in surface — contextlib.nullcontext() for the disabled path keeps the hot path allocation-free and the default zero-I/O; preview stays memory-only by construction.
• Genuinely thorough tests — config defaults, unicode round-trip, append semantics, and every failure mode (open/close/write) are covered, plus create/preview/resume at the builder level with the nice run_id segment-differentiation assertion on resume.
• Failure handling avoids warning storms — disabling the sink after the first write failure (and the idempotent, warn-once close()) is a thoughtful touch, and the test that asserts write/close are each called exactly once nails it down.
• Docstring calls out the real trade-offs — flagging potential sensitivity of the labels and the per-event flush overhead in the RunConfig docs is exactly the right place for operators to see it.

Verdict

Needs changes — two light Warnings worth resolving before merge:

• Reconcile the sink's docstring with its re-raising behavior (or make the non-fatal guarantee self-contained).
• Update architecture/dataset-builders.md (or file a follow-up) to reflect the new opt-in on-disk scheduler-event artifact.

Both are quick; the core implementation is solid and ruff check / ruff format --check are clean on the changed files.

---

This review was generated by an AI assistant.

[eric-tramel]

@nabinchha Thanks for the review — addressed all three points in 254490fd:

• Clarified that JsonlSchedulerEventSink disables itself and re-raises the first write/flush failure for the caller to handle.
• Documented the sink's single-writer expectation.
• Updated the dataset-builder architecture guide to describe the opt-in scheduler_events.jsonl artifact and its appended per-invocation segments.

Focused validation: 9 observability tests passed, plus Ruff check and format verification. I’ve re-requested your review.

[nabinchha]

Thanks @eric-tramel — both warnings are addressed: the docstring now documents the single-writer + re-raise-once contract, and the architecture doc notes the opt-in scheduler_events.jsonl artifact. ruff check/format are clean on the changed files. Ship it.