[feat] Kafka: event-time driven checkpointing from message timestamp

[tiankongdeguiji]

Motivation

For online/streaming training from Kafka, KafkaReader previously surfaced only the message offset (as the resume cursor); the Kafka message timestamp was never read, so there was no way to checkpoint based on the event-time of the consumed data. This PR surfaces the message timestamp and uses it to drive checkpoint saves on consumed event-time, which is what you want for time-aligned snapshots of a stream.

What it does

Surfaces the Kafka message timestamp on the batch and adds event-time–driven checkpoint triggers, alongside the existing step/epoch triggers:

• batch.data_timestamp — the max consumed event-time in the batch, as Unix-epoch seconds (float; -1.0 when the source has no timestamps). KafkaReader normalizes its ms timestamp to seconds; BaseDataset._build_batch just surfaces the per-batch max (unit-agnostic, so non-Kafka datasets are unaffected).
• New TrainConfig fields:
  • save_checkpoints_timestamp_interval — save every N seconds of consumed event-time, aligned to the Unix epoch (wall-clock), not training epochs. 0 disables.
  • save_checkpoints_timestamps — absolute event-time targets (Unix-epoch seconds); save once when consumed data crosses each. Empty disables.
  • save_checkpoints_timestamp_quorum — fraction of workers (0, 1], default 0.5, that must have consumed past a boundary/target before a save fires.

Design

• Centralized save policy. All save cadence (step / epoch / event-time) + the single per-step dedupe + the event-time watermark live in CheckpointManager.maybe_save(...) -> bool; the training loop calls it at the step / epoch / final points and runs eval when it returns True. This removes the four hand-guarded save sites in the loop (the source of an epoch-vs-timestamp double-save) and makes the decision logic unit-testable.
• Worker quorum (not min/max). Each rank passes its local event-time to maybe_save, which all-gathers across workers and reduces to the (1 - quorum) upper quantile — the event-time that ≥ quorum fraction of workers have reached. This generalizes min (quorum=1.0) / max (quorum→0), is robust to a single outlier/garbage timestamp at the default 0.5, and counts a worker without a timestamp (-1.0) as "not past" so the quorum is over all workers.
• Off the GPU. The per-step reconciliation gathers over a dedicated CPU (gloo) process group that CheckpointManager creates in set_save_policy and owns, so it never runs on the model device or forces a per-step device sync that would stall the TrainPipelineSparseDist prefetch.
• Resume-correct. The event-time watermark is persisted in dataloader_state (saved every checkpoint, restored on continue_train) and advanced on every save, so resume re-fires no already-saved boundary and misses none. It's a single global float, identical across ranks (the per-key max-merge is a no-op on it) and ignored by the existing dataloader_state consumers. dataloader_state is typed Dict[str, Any] since it now holds the float watermark alongside the int offsets.
• Distributed-safe. The reconciliation is gated only by config (identical on all ranks) and the save decision is deterministic + identical across ranks, so the collective save is entered in lockstep; nothing collective runs when not saving. It relies on the same equal-step-count invariant the model's own gradient collectives already require.
• No behavior change when unset — with the new config fields unset, training is identical to before.

Tests

• Pure helpers quorum_event_time (quantile incl. outlier/sentinel cases) and should_save_on_timestamp (interval/target crossings, no-refire, first-batch init) — parameterized.
• CheckpointManager.maybe_save — step/epoch/timestamp triggers, the dedupe (incl. the epoch-after-same-step regression), watermark stamping, restore seeding, and the -1.0-sentinel path; plus _reconcile_event_time single-process behavior (save mocked, no GPU/dist needed).
• _build_batch data_timestamp extraction; Kafka roundtrip asserts a positive seconds value.
• Docs: the three new TrainConfig fields documented in train.md.

🤖 Generated with Claude Code

[github-actions]

When timestamp-checkpointing is active, this runs a synchronous all_reduce + .item() every training step. The .item() forces a host/GPU sync that isn't overlapped with compute, and unlike the gradient all-reduce it can't piggyback on the backward pass — so it partially defeats TrainPipelineSparseDist's prefetch/overlap. Checkpoint decisions only need coarse event-time resolution; consider reducing every N steps (with N derived identically on all ranks so the collective stays in lockstep) or using async_op=True and reading the result one step later.

Separately, this per-step collective assumes all ranks run the same number of steps. That invariant is already required by the model's fwd/bwd collectives, but it's worth noting that with Kafka's uneven partition data and check_all_workers_data_status=False (i.e. no batch_cost_size), a rank whose partitions drain first hits StopIteration and breaks before this all_reduce, hanging the others. The "every rank participates identically" comment holds only under that equal-step-count invariant — worth documenting that this feature needs batch_cost_size/even data, or gating the all-reduce on check_all_workers_data_status.

[github-actions]

The step-trigger and timestamp-trigger both guard against re-saving a step already saved (i_step != last_ckpt_step), and so does the final save below — but the epoch-trigger doesn't. If a timestamp-triggered save fired on the last step of the epoch (so last_ckpt_step == i_step), this re-saves the same step: a redundant full checkpoint write and a redundant eval.

Suggested change

	if save_checkpoints_epochs > 0 and i_step > 0:
	if (i_epoch + 1) % save_checkpoints_epochs == 0:
	last_ckpt_step = i_step
	ckpt_manager.save(i_step, model, optimizer, dataloader_state)
	if eval_dataloader is not None:
	_evaluate(
	model,
	eval_dataloader,
	eval_config,
	eval_result_filename=eval_result_filename,
	global_step=i_step,
	eval_summary_writer=eval_summary_writer,
	global_epoch=i_epoch,
	check_all_workers_data_status=check_all_workers_data_status,
	)
	model.train()
	do_checkpoint(i_step)
	if save_checkpoints_epochs > 0 and i_step > 0 and i_step != last_ckpt_step:
	if (i_epoch + 1) % save_checkpoints_epochs == 0:
	do_checkpoint(i_step)

Alternatively, centralize the guard inside do_checkpoint (early-return when step == last_ckpt_step) so no call site can forget it.

[github-actions]

"epoch-aligned" is misleading here — alongside save_checkpoints_epochs/num_epochs it reads as "aligned to a training epoch", but the code is floor(ts / interval_s), i.e. boundaries are multiples of interval_s since the Unix epoch (wall-clock), unrelated to training epochs. Same wording appears in checkpoint_util.py (should_save_on_timestamp docstring/comments) and train.md ("按epoch对齐") — worth fixing in all three.

Suggested change

	// event-time interval (seconds) for saving checkpoint, based on the consumed
	// data timestamp (e.g. kafka message timestamp). saves once per epoch-aligned
	// interval boundary that the consumed event-time crosses. 0 disables.
	// event-time interval (seconds) for saving checkpoint, based on the consumed
	// data timestamp (e.g. kafka message timestamp). saves once per interval
	// boundary (aligned to the Unix epoch, not training epochs) crossed. 0 disables.

[github-actions]

Review summary

Clean, well-documented feature. The do_checkpoint(step) de-duplication of the four save sites is a real readability win, should_save_on_timestamp is pure and well unit-tested, and the transient __data_timestamp__ column is correctly popped before feature parsing so it can never leak into features or checkpoint state. When the feature is unset, the training path is unchanged. Nice work.

I left inline comments on three points; the substantive ones:

• Per-step all_reduce + .item() (main.py): when active, a synchronous collective + host sync runs every step and isn't overlapped with compute. Consider reducing every N steps (N identical across ranks) or async_op. Also worth documenting the equal-step-count invariant for uneven Kafka partitions (batch_cost_size/check_all_workers_data_status).
• Epoch-trigger double-save (main.py): the epoch save lacks the i_step != last_ckpt_step guard the other sites have, so it can re-save (checkpoint + eval) a step a timestamp-trigger already saved.
• "epoch-aligned" wording (train.proto / checkpoint_util.py / train.md): means Unix-epoch-aligned, not training-epoch-aligned — easy to misread.

A couple of additional notes not worth inlining:

• Test gaps (no Kafka needed): the _build_batch data_timestamp extraction (None when column absent, the >= 0 filter that drops -1 "no timestamp" batches, max selection) and Batch.to()/pin_memory() propagation of data_timestamp are both untested and both fail silently (event-time checkpointing becomes a no-op) rather than loudly. Both are unit-testable via the existing _TestReader/_TestDataset harness. The checkpoint_timestamp_reduce validation (ValueError on bad value) is also a cheap test to add.
• Untrusted timestamps: dataset.py only rejects negative values. A single far-future/garbage msg.timestamp() advances last_ckpt_data_ts_s and can then silently suppress all future saves (mitigated under the default min reduce, exposed under max). A sanity bound + warning would harden the streaming path.

[tiankongdeguiji]

Thanks for the review — addressed in 2dd3dcc. Summary of how each point was handled:

Epoch-trigger double-save (real bug) — fixed at the root. Rather than adding the missing guard to one more site, the save-cadence logic (step + epoch + event-time) is now centralized in CheckpointManager.maybe_save(...), which owns last_ckpt_step and does a single dedupe. No call site can re-save a step another already saved. This also makes the decision logic unit-testable (it wasn't, buried in the loop) — added a regression test test_maybe_save_dedupe_epoch_after_same_step.

Per-step all_reduce + .item(). The cross-rank reconciliation is now an all_gather driving a worker quorum (replacing the min/max knob): save_checkpoints_timestamp_quorum (float, default 0.5) — save once ≥ quorum fraction of data-carrying workers have consumed past the boundary/target, computed as the (1-quorum) upper quantile. It generalizes the old knobs (1.0≡min, →0≡max) and is robust to a single far-future/garbage timestamp at the default, so the "untrusted timestamp" concern needs no separate wall-clock guard. (Kept per-step per design discussion; the loop does the gather, the manager does the quorum + decision with no collective when not saving.)

"epoch-aligned" wording. Fixed in train.proto, checkpoint_util.py, and train.md to clarify it means the Unix epoch (wall-clock), not training epochs.

Test gaps. Added: quorum_event_time (quantile edge cases incl. outlier robustness), maybe_save (step/epoch/timestamp/final + dedupe + watermark + restore seeding), _build_batch data_timestamp extraction (max valid / -1 filter / absent→None), and Batch.to()/pin_memory() propagation (the silent-failure guard).

"Uneven partitions → deadlock". Not a new deadlock: a rank that drains first hangs at the model's own gradient all-reduce inside pipeline.progress, which precedes our collective — the equal-step-count invariant is already required by the existing collectives. Corrected the overclaiming comment to say so.

New: resume correctness. The event-time watermark is now persisted in dataloader_state and seeded on continue_train, so resume re-fires no already-saved boundary and misses none (advanced on every save). It's a single global float, identical across ranks, and ignored by every existing dataloader_state consumer.

[github-actions]

When timestamp checkpointing is enabled, this runs every step: gather_float_scalar builds a CUDA float64 tensor, does an all_gather_into_tensor, then .tolist() — a blocking D2H sync. The baseline loop has no per-step host sync (train metrics only .item() every log_step_count_steps), so this new sync caps how far TrainPipelineSparseDist can run ahead to prefetch/input_dist the next batch and can cost throughput — exactly in the streaming-Kafka scenario where this feature is on.

Since saves are rare relative to steps, consider gathering/deciding only every N steps (keep it a deterministic function of i_step so it stays in lockstep across ranks), or running the gather on a gloo/CPU PG so the .tolist() doesn't synchronize the compute stream. No impact when the feature is disabled.

[github-actions]

This quorum-range validation (and its ValueError) has no test, and gather_float_scalar — the new distributed primitive that feeds the trigger — is entirely untested (there's no dist_util_test.py). Both are cheap and high-value:

• assertRaises(ValueError) for quorum<=0 / quorum>1 with a trigger active, plus a no-raise case when both triggers are off (the conjunction here is exactly what silently breaks on refactor).
• At least the non-distributed gather_float_scalar path (return [value]); the gloo mp.spawn pattern already in checkpoint_util_test.py can cover the distributed rank-order path.

[github-actions]

Minor: dataloader_state is annotated Dict[str, int] throughout (save, maybe_save, save_dataloader_state, update_dataloder_state), but this stores a float watermark. Functionally harmless — the cross-rank max-merge in save_dataloader_state is a no-op since the value is rank-identical — but in this pyre-checked codebase the annotation is now inaccurate. Consider widening to Dict[str, Union[int, float]].

[github-actions]

The save_checkpoints_steps/epochs bullets above note 保存模型后会做一次评估, but these three timestamp bullets don't. Timestamp-triggered saves also run an eval (run_eval fires whenever maybe_save returns True), so worth adding the same note for parity.

[github-actions]

Review summary

Reviewed across code quality, performance, tests, docs, and distributed safety. This is a well-built PR — the centralization of all save cadence into CheckpointManager.maybe_save(...) is a real improvement (it removes the three near-identical save+eval blocks and the epoch-vs-timestamp double-save), the quorum quantile is a nice robustness choice over min/max, and the resume path (persisted watermark seeded in restore_dataloader_state, untouched by set_save_policy) is correct. I verified the key distributed-safety claim: with timestamp checkpointing on, the per-step gather + the maybe_save decision are a deterministic function of identical all-gathered inputs on every rank, and TrainPipelineSparseDist._next_batch drops the remainder on all workers together, so the collective save stays in lockstep — no deadlock from uneven/exhausted data. Defaults disable every trigger, so existing runs are unaffected.

Noteworthy items, posted inline:

• Perf (main only when enabled): the per-step gather_float_scalar adds a per-step .tolist() D2H sync the baseline loop doesn't have — can throttle the prefetch pipeline in the streaming case this targets. Suggest gathering/deciding every N steps or on a CPU PG.
• Tests: the set_save_policy quorum ValueError path and gather_float_scalar (both paths) are untested — cheap, valuable additions.
• Types: dataloader_state is Dict[str, int] but now carries a float watermark (harmless, but inaccurate under pyre).
• Docs: the timestamp bullets in train.md should mention that a save also runs an eval, for parity with the step/epoch bullets.

One optional hardening (not blocking): the watermark read from dataloader_state.json on resume is fed straight into int(... // interval_s) without a type/finiteness check. The framework only ever writes valid floats, so this only bites on a hand-edited/corrupted state file (a loud crash at the first resumed step), but a float(...) + math.isfinite guard in restore_dataloader_state would make resume robust to it.

Nothing here blocks merge.