Pipeline resumability via source-level counter checkpointing

[abhinavg4]

Discussion (Design Doc)

#2034

Summary

Adds opt-in resumability to NeMo Curator pipelines. When the user passes Pipeline.run(checkpoint_path="..."), completed source partitions are tracked in an LMDB file via a single Ray actor. On rerun against the same checkpoint, the source stage skips sources already marked complete and processes only the remaining work.

Design discussion: the full design rationale, principles, and algorithm details will go into a GitHub Discussion (to be linked here once posted). The top-level invariants:

• Workers never block on the resumability layer (fire-and-forget event flow; backpressure via Ray's _max_pending_calls).
• One LMDB write per source, at end-of-life — counter math runs in memory on the actor; we touch the disk only when a source's pending counter reaches zero.
• Single integration point in BaseStageAdapter.process_batch + one method in Pipeline.run; per-backend executors (Xenna, RayData, RayActorPool) are not touched.
• No code change for existing pipelines that don't pass checkpoint_path — the resumability hook is a no-op when no actor is registered.

What's in this PR

• Task field additions (nemo_curator/tasks/tasks.py)
  • _lineage_path: index path through the pipeline DAG.
  • _deterministic_lineage_path_hash: sha256(lineage_path)[:32] — stable across reruns.
  • _source_id: explicitly stamped by the source-stage adapter and inherited downstream. Empty for pre-source tasks.
  • _set_lineage(parent_paths, child_index) method.
• NoneTask / FailedTask sentinels (nemo_curator/tasks/sentinels.py) — Task subclasses with a msg field. Returned from user code in process_batch to mark specific slots as filtered or failed.
• ProcessingStage flags (nemo_curator/stages/base.py)
  • is_source_stage: bool = False — user-overridable; defaults to the first stage at build time.
  • _is_terminal_stage: bool = False — same; defaults to the last stage.
  • assign_root_lineage and assign_child_lineage helpers for lineage propagation.
• ResumabilityActor (nemo_curator/utils/resumability_actor.py) — singleton Ray actor with apply_deltas, are_completed, errors, close methods. Maintains in-memory _pending (per-source counter) and _applied (per-task-hash dedup + non-determinism detection). Writes to LMDB only when a source's counter hits zero.
• Client helpers (nemo_curator/utils/resumability_client.py) — _is_active, _flush_deltas, _skip_completed_sources.
• Stage adapter hook (nemo_curator/backends/base.py) — _resumability_postprocess on BaseStageAdapter that auto-wraps None → NoneTask, validates batch fan-out rules, computes counter deltas, and fires them fire-and-forget.
• Pipeline.run lifecycle (nemo_curator/pipeline/pipeline.py)
  • New checkpoint_path argument.
  • Build-time validation: at most one stage may be explicitly marked is_source_stage=True or _is_terminal_stage=True; otherwise defaults to first/last.
  • _run_with_resumability(...) method: spawns the actor, starts the watchdog thread, runs the pipeline, does a final error check, closes the actor.

Test plan

Unit tests cover:

• ResumabilityActor.apply_deltas counter math for all delta combinations.
• NonDeterministicTaskError fires on conflicting delta for the same task hash.
• NegativePendingCountError fires when a counter would go below zero.
• Same delta re-fired for the same task is silently idempotent.
• LMDB write happens exactly when pending[sid] hits 0; LMDB state survives close + reopen.
• are_completed returns parallel bool list consistent with _completed set.
• errors() returns a snapshot; mutation does not affect actor state.
• NoneTask / FailedTask.from_input constructs with inherited identity.
• Task._set_lineage sets lineage path and hash, idempotent on re-call, filters empty parent paths.
• assign_root_lineage sets lineage on user-provided tasks, skips EmptyTask, does not set _source_id.
• assign_child_lineage propagates lineage and inherits _source_id from parents.
• _resumability_postprocess auto-wraps None as NoneTask, no-ops for pre-source tasks (empty _source_id), raises on batched fan-out with len(input)>1 and len(output)!=len(input).
• Source-stage handling: stamps _source_id, filters already-completed sources, fires +1 per surviving source.
• Pipeline.build raises on multiple explicit source stages and on multiple explicit terminal stages; defaults to first/last when unmarked.

Functional / end-to-end (SIGKILL-resume integration, autoscale mid-run) is noted as follow-up — see "Future work" below.

Future work (will be updated as this PR evolves)

• Functional / SIGINT integration test — drive a real pipeline, kill mid-run, restart, verify only remaining sources are reprocessed.
• Actor-death survival — periodically snapshot _pending and _applied to LMDB so the actor can recover its state after restart.
• Bounded _applied dict — drop entries for sources that moved to _completed (current implementation lets _applied grow unbounded for the duration of the run).
• 64-bit hashes — extend get_deterministic_hash to take a length parameter; use 16 hex chars (64 bits) for resumability source_ids to avoid birthday collisions at 100M+ scale.
• Configurable watchdog poll interval — currently hard-coded at 60 seconds; expose as Pipeline.run(..., resumability_poll_interval=...).
• Explicit is_source_stage requirement — currently defaults to the first stage; consider requiring an explicit marker to avoid surprising "filter-as-source" behavior in pipelines without a reader.
• Reader ordering invariant — formalize the requirement that source stages emit sources in a deterministic order across runs (currently informally enforced by FilePartitioningStage.sort_by_size / sorted file paths).
• Composite stages — verify is_source_stage / _is_terminal_stage defaults work correctly when the user's pipeline contains CompositeStage instances that decompose into multiple primitive stages.
• Pipeline.describe() integration — include resumability config in the pipeline description output.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

Adds opt-in pipeline resumability via a singleton Ray actor (ResumabilityActor) that tracks per-source pending counters in an LMDB file. When Pipeline.run(checkpoint_path=...) is set, completed source partitions are persisted across runs and skipped on rerun; existing pipelines without checkpoint_path pay zero cost.

• Core mechanism: source stage fires +1 per new source; each downstream stage fires 0 (pass-through), −1 (filter/sink), or +(n−1) (fan-out); when a counter reaches zero the source_id is written to LMDB and skipped on rerun.
• Sentinel types: NoneTask (filter, decrements counter) and FailedTask (error, counter held so source reruns) are new public Task subclasses returned from process_batch.
• Design trade-offs: fire-and-forget delta delivery, "never raises" actor, and detached actor lifetime are deliberate choices with recovery paths for most anomaly cases; several edge cases (stale actor reuse, actor-death recovery) are called out as future work.

Confidence Score: 3/5

Two defects in the changed code path need fixes before merging: a raw-None crash in the pre-source fan-out branch, and uncertainty about silent delta loss at high throughput when the actor queue is full.

The pre-source sentinel filter omits r is not None, so any pre-source fan-out stage returning bare None slots crashes the perf-stats loop immediately. The _max_pending_calls=100 + fire-and-forget combination could silently discard deltas if Ray rejects calls when the queue is full rather than blocking — at high batch concurrency this leaves sources permanently pending in the checkpoint, silently defeating the resumability guarantee. Both issues are on the new resumability code path. The rest of the implementation is well-structured and unit-tested.

nemo_curator/backends/base.py (pre-source None filter) and nemo_curator/pipeline/pipeline.py (_max_pending_calls + fire-and-forget interaction) are the two files that need attention before this is merged.

Important Files Changed

Filename	Overview
nemo_curator/backends/base.py	Adds _resumability_postprocess and _handle_source_stage to BaseStageAdapter. The pre-source fan-out path fails to exclude raw None values from its output, which crashes the add_stage_perf loop.
nemo_curator/utils/resumability_actor.py	New singleton Ray actor; "never raises" design is a sound simplification. The newly_done mid-batch zero-crossing anomaly is low-risk with correct delta accounting but worth a guard for robustness.
nemo_curator/utils/resumability_client.py	Module-level helpers for worker→actor communication. Each call independently invokes _actor() with a GCS round-trip; caching the handle would reduce per-batch overhead (flagged in previous review).
nemo_curator/pipeline/pipeline.py	Adds checkpoint_path to Pipeline.run and the _run_with_resumability lifecycle. The _max_pending_calls=100 + fire-and-forget combination may silently lose deltas at high throughput.
nemo_curator/stages/base.py	Adds assign_root_lineage and assign_child_lineage helpers; logic is clean and None entries are filtered correctly.
nemo_curator/tasks/tasks.py	Adds _source_id field and _set_lineage / get_deterministic_id methods to Task. Implementation is straightforward and consistent with the design.
nemo_curator/tasks/sentinels.py	New NoneTask and FailedTask sentinel classes; well-structured with from_input constructors that inherit task identity correctly.
pyproject.toml	Adds lmdb>=1.4 dependency; appropriate minimum version for a stable API.

Sequence Diagram

sequenceDiagram
    participant PR as Pipeline.run(checkpoint_path)
    participant RA as ResumabilityActor (detached)
    participant LMDB as LMDB file
    participant SA as BaseStageAdapter (source stage)
    participant DS as BaseStageAdapter (downstream)
    participant SK as BaseStageAdapter (sink)

    PR->>RA: "spawn (get_if_exists=True, max_pending_calls=100)"
    PR->>SA: executor.execute(stages, initial_tasks)

    SA->>RA: are_completed([sid0, sid1, ...]) [ray.get blocking]
    RA-->>SA: [False, True, ...]
    SA->>SA: filter completed sources
    SA->>RA: apply_deltas([(task_id, sid, +1), ...]) [fire-and-forget]

    loop each non-sink stage batch
        DS->>RA: apply_deltas([(task_id, sid, 0 or -1), ...]) [fire-and-forget]
    end

    SK->>RA: apply_deltas([(task_id, sid, -1), ...]) [fire-and-forget]

    RA->>RA: _pending[sid] reaches 0
    RA->>LMDB: write sid as complete

    PR->>RA: "close() [ray.get, timeout=10s]"
    PR->>RA: ray.kill(actor)

    Note over PR,LMDB: On rerun: actor loads _completed from LMDB, are_completed() skips finished sources

Loading

_{Reviews (4): Last reviewed commit: "Adapt resumability layer to task_id coll..." | Re-trigger Greptile}

[greptile-apps]

Final error check silently swallowed by except Exception — the raise final_errs[0] on line 352 lives inside the try block that owns the broad except Exception as e handler. When the actor records a NonDeterministicTaskError or NegativePendingCountError that only surfaces at teardown (e.g., from the last in-flight batch), the raise is immediately re-caught and demoted to a logger.warning(...). The pipeline then exits cleanly with return results even though data-integrity errors were detected.

Suggested change

	try:
	ray.init(ignore_reinit_error=True)
	actor_handle = ray.get_actor(ACTOR_NAME)
	final_errs = ray.get(actor_handle.errors.remote(), timeout=30) # type: ignore[attr-defined]
	if final_errs and not collected_error:
	raise final_errs[0]
	except ray.exceptions.RayActorError: # type: ignore[attr-defined]
	pass
	except Exception as e: # noqa: BLE001
	logger.warning(f"final resumability error check failed: {e}")
	_deferred_error: Exception | None = None
	try:
	ray.init(ignore_reinit_error=True)
	actor_handle = ray.get_actor(ACTOR_NAME)
	final_errs = ray.get(actor_handle.errors.remote(), timeout=30) # type: ignore[attr-defined]
	if final_errs and not collected_error:
	_deferred_error = final_errs[0]
	except ray.exceptions.RayActorError: # type: ignore[attr-defined]
	pass
	except Exception as e: # noqa: BLE001
	logger.warning(f"final resumability error check failed: {e}")
	if _deferred_error is not None:
	raise _deferred_error

[greptile-apps]

Inconsistent _pending state after mid-batch error

_apply_impl accumulates newly_done entries while iterating the batch, then calls _persist_completed only if the loop finishes cleanly. If the loop raises NegativePendingCountError (e.g., two -1 deltas for the same source in one batch), any source that had already reached zero earlier in that same batch is never persisted and never added to _completed, but _pending[sid] for the failing source is left at a negative value. Subsequent apply_deltas calls on the same actor see neither the correct _completed entry nor a sane counter, which can lead to a false "source completed" if a later delta happens to bring the negative value back to zero.

[greptile-apps]

ray.get_actor() GCS round-trip on every processed batch

_is_active() delegates to _actor(), which calls ray.get_actor(ACTOR_NAME) unconditionally whenever Ray is initialized — even for pipelines that never set checkpoint_path. For an executor that runs hundreds of thousands of small batches this adds a GCS lookup per batch. Caching the handle (or the False sentinel) at module level after the first lookup would eliminate the per-batch overhead.

[greptile-apps]

get_if_exists=True can silently reuse a stale detached actor

When get_if_exists=True, Ray returns the existing actor handle without calling the constructor if an actor named nemo_curator_resumability is already alive — for instance from a previous run that crashed before the cleanup block ran. That stale actor owns a different LMDB file (from its own checkpoint_path), so all new deltas accumulate in the wrong place and are_completed queries reflect the wrong state. The caller has no indication that the path argument was ignored. Consider checking whether the returned actor is pointing at the expected path (e.g., via a path() accessor), or kill any pre-existing actor and spawn fresh.

[abhinavg4]

Improvements we should probably apply to v0.

Priority:

1. Rename terminal stage to sink. Source and sink is better than source and terminal
2. Here, instead of erroring out, can we just rewrite the new value? and removing the whole demon logic
3. Unify task_id, _uuid and _deterministic_lineage_path_hash.
4. support a dynamic dataset

[abhinavg4]

Do we really need this? We have task_id, _uuid and _deterministic_lineage_path_hash. They can be combined into a single variable. But before that we need to ensure the usage of task_id and uuid is not affected. Eg:- Writer stage might use task_id right now, we should change to this.

[abhinavg4]

I do not like this, the reason this is being done is probably coz we call ray init inside the executor and not here, so we need to differentiate that. Not ideal for sure. Maybe we can make a function for _start_resumability_actor and call it inside each executor. Design choice.

[greptile-apps]

The pre-source sentinel filter must also exclude raw None values; otherwise a fan-out stage that returns None in its output list will pass those values through to the add_stage_perf loop, causing an AttributeError at runtime.

Suggested change

	if all(not t._source_id for t in tasks):
	return [r for r in results if not isinstance(r, (NoneTask, FailedTask))]
	if all(not t._source_id for t in tasks):
	return [r for r in results if r is not None and not isinstance(r, (NoneTask, FailedTask))]

[greptile-apps]

Fire-and-forget delta loss when _max_pending_calls=100 is exceeded

_max_pending_calls=100 on the actor causes Ray to embed a RayActorError in the returned ObjectRef when the queue is full. Because _flush_deltas never calls ray.get() on that ref, the rejection is silently swallowed and the delta is permanently lost. For a pipeline processing more than 100 batches before the actor can drain its queue (trivially possible at high throughput), any source whose counter delta was lost will have _pending stuck above 0 and will never be written to LMDB. The PR description claims "_max_pending_calls provides backpressure", but the actual behavior is silent delta loss, not blocking — the two have opposite throughput implications. A limit of 100 is also very low relative to a typical pipeline's batch concurrency. Consider either raising the limit substantially, adding a ray.get() health check in the watchdog thread to surface errors, or verifying the actual Ray behavior for the fire-and-forget case.