Resumable Processing

Pass checkpoint_path to Pipeline.run() to skip source partitions that completed in an earlier run. This is the primary resumability mechanism for NeMo Curator pipelines.

1
pipeline.run(checkpoint_path="./checkpoints/my-pipeline")

Without checkpoint_path, execution is unchanged and no checkpoint metadata is created.

Completion Model

NeMo Curator automatically identifies one source and one sink when it builds the pipeline:

• If no stage is explicitly marked, the first execution stage is the source and the last is the sink.
• Set is_source_stage=True or is_sink_stage=True on a stage to override a default.
• A pipeline can have at most one explicit source and one explicit sink.

The framework tracks the descendants of each source output with an in-memory counter:

1. A new source partition opens with one pending descendant.
2. A one-to-many stage replaces its parent with its continuing children.
3. A task that reaches the sink or is intentionally filtered consumes one pending descendant.
4. The source is persisted as completed when its pending count reaches zero.
5. A later run drops completed partitions immediately after the source stage.

Only completed source IDs are persisted. Pending counters disappear when a run exits, so any incomplete source is deliberately replayed.

Output Results and Retry Behavior

NoneTask and FailedTask are framework marker tasks derived from SentinelTask. They carry no payload and are removed before the next stage. A custom stage normally returns None for a successful filter and returns FailedTask() only when it deliberately converts a recoverable per-task failure into a retry marker.

Interrupt and Resume Example

Create several JSONL input files so each source partition can complete independently. Save this script as resume_demo.py:

1
import time
2
from dataclasses import dataclass
3
4
from nemo_curator.pipeline import Pipeline
5
from nemo_curator.stages.base import ProcessingStage
6
from nemo_curator.stages.text.io.reader import JsonlReader
7
from nemo_curator.stages.text.io.writer import JsonlWriter
8
from nemo_curator.tasks import DocumentBatch
9
10
11
@dataclass
12
class SlowPassThrough(ProcessingStage[DocumentBatch, DocumentBatch]):
13
    """Delay each partition so an interrupted run is easy to observe."""
14
15
    name: str = "slow_pass_through"
16
17
    def process(self, task: DocumentBatch) -> DocumentBatch:
18
        time.sleep(2)
19
        return task
20
21
22
pipeline = Pipeline(name="resume-demo")
23
pipeline.add_stage(
24
    JsonlReader(
25
        file_paths="./input/*.jsonl",
26
        files_per_partition=1,
27
        fields=["text"],
28
    )
29
)
30
pipeline.add_stage(SlowPassThrough())
31
pipeline.add_stage(JsonlWriter(path="./output"))
32
33
pipeline.run(checkpoint_path="./checkpoints/resume-demo")

Run the script, wait for several output files, and interrupt it with Ctrl+C:

$
python resume_demo.py

Run the same command again. The file-partitioning source derives stable IDs from each partition’s sorted paths, reads the completion set, and emits only partitions that did not finish the first run.

Keep all resume-sensitive inputs stable between launches:

• Use the same checkpoint and output paths.
• Keep the pipeline topology, source/sink roles, partitioning, and filtering behavior compatible.
• Do not change the contents behind an unchanged source path and expect automatic invalidation. FileGroupTask hashes sorted paths, not file contents.
• Use a new checkpoint directory after a semantic input change or incompatible pipeline change.

Checkpoint Storage

Given checkpoint_path="/shared/checkpoints/run-42", state is stored as:

/shared/checkpoints/run-42/
└── .nemo_curator_metadata/
    ├── host-a-4128.mdb
    ├── host-b-7319.mdb
    └── ...

Each run creates or reuses one sparse LMDB file named from its hostname and process ID. A writer modifies only its own file. At startup it reads the union of completed source IDs from every *.mdb file in .nemo_curator_metadata.

This design avoids cross-host LMDB writer locks and allows jobs such as a SLURM array to use a checkpoint directory on a shared filesystem. It has these boundaries:

• Every host must resolve checkpoint_path to the same durable shared directory.
• Completed IDs are unioned when each actor starts; the files are not a live work queue.
• Concurrent jobs should process disjoint input shards. Jobs that start together with overlapping inputs can both begin the same source before either records completion.
• Run only one checkpoint-enabled pipeline at a time in a Ray cluster. The resumability actor has a fixed cluster-wide name; overlapping pipelines would share the first actor rather than create independent checkpoint owners.
• An unreadable writer file is skipped with a warning. Preserve the other files and investigate the damaged file rather than assuming its sources completed.
• The default LMDB map is 1 GiB but sparse on Linux; physical usage grows with stored completion keys.

For deterministic array sharding, shard-completion manifests, scheduler limits, and incomplete-shard resubmission, see SLURM Job Arrays.

Reset or Clean Up

To restart every source, stop all jobs using the checkpoint and remove its .nemo_curator_metadata directory, or select a fresh checkpoint_path. Removing only an output directory does not reset completion state; a subsequent run would still skip sources recorded as complete.

Keep checkpoint metadata until you have validated the output. After a successful run is accepted, archive or remove the entire checkpoint directory according to your retention policy.

Deterministic Task Identity

task_id is framework-owned and must not be passed to task constructors or assigned by stage code. It starts empty on a user-created task and is re-derived at every stage boundary as an underscore-separated lineage path.

0                                      EmptyTask root
└── 0_<source-id>                      source partition
    └── 0_<source-id>_0                next one-to-one stage
        ├── 0_<source-id>_0_0          first fan-out child
        └── 0_<source-id>_0_1          second fan-out child

• EmptyTask() is a class instance with root ID "0". It replaces the former shared EmptyTask singleton pattern.
• Source outputs use Task.get_deterministic_id() when the task type implements it. FileGroupTask hashes its sorted paths, which keeps unchanged file groups stable if other groups are added or removed.
• A custom source task’s deterministic ID must be stable, unique within the source, and must not contain underscores. Resumability uses the final underscore-separated lineage segment as the source ID, so underscores can make distinct sources collide during checkpoint lookup.
• Otherwise, a source output uses its position. Reordering positional sources can make a checkpoint refer to different work.
• One-to-one results inherit their positional parent and append 0.
• A single input can fan out deterministically; child positions are appended.
• For a custom batch mapping with multiple inputs and a different number of outputs (M→K where M is greater than one and M does not equal K), parentage is ambiguous. Outputs receive random IDs prefixed with r, and resumability logs a warning and leaves those sources pending.

You can inspect task.task_id for logs or deterministic filenames after a stage boundary. Treat it as read-only; returning the same Python task object still causes the adapter to append a new lineage segment.

Custom Stage Responsibilities

The execution adapter assigns IDs and updates completion counters for both process() and process_batch(). Custom stages should follow these mapping rules:

Do not drop filtered entries from an overridden process_batch(). Keep a None placeholder in each filtered input’s position so the adapter can attribute the filter to the correct source.

When returning a new task object, construct it from payload and dataset fields without setting task_id:

1
def process(self, task: DocumentBatch) -> DocumentBatch:
2
    data = task.to_pandas().copy()
3
    data["normalized_text"] = data["text"].str.strip()
4
    return DocumentBatch(
5
        dataset_name=task.dataset_name,
6
        data=data,
7
        _metadata=task._metadata,
8
        _stage_perf=task._stage_perf,
9
    )

For a recoverable per-task failure:

1
from nemo_curator.tasks import FailedTask
2
3
def process(self, task: DocumentBatch) -> DocumentBatch | FailedTask:
4
    try:
5
        # Replace with the external call used by your stage.
6
        return call_external_service(task)
7
    # Replace TemporaryError with the recoverable exception raised by your service.
8
    except TemporaryError:
9
        return FailedTask()

An unhandled exception can still abort execution. Because the affected source never reaches completed state, it is eligible to run again with the same checkpoint.

Operational Checklist

Before relying on a resume:

• Use source tasks with stable, content-meaningful get_deterministic_id() values.
• Keep source data, pipeline shape, stage settings, software version, and model behavior compatible.
• Use deterministic, replay-safe sinks.
• Confirm the log reports the expected number of pending partitions after restart.
• Treat random r-prefixed task IDs and M→K ambiguity warnings as non-resumable paths to fix.
• Give concurrent hosts disjoint source shards even when they share the checkpoint directory.
• Avoid running multiple checkpoint-enabled pipelines concurrently in one Ray cluster.
• Retain checkpoint metadata until output validation is complete.

Related References

• SLURM Job Arrays
• ProcessingStage API
• Pipeline Execution Backends
• Memory Management