nemo_curator.utils.resumability_actor

Per-writer LMDB owner tracking per-source completion for resumability.

LMDB can’t be safely shared by writers across hosts (its lock lives in an mmap’d file not shared on a networked FS), so each actor writes ONLY its own <dir>/<host>-<pid>.mdb and on startup reads the UNION of completed sources across every *.mdb in the dir. A rerun thus skips everything any prior writer finished — letting the tasks of a SLURM array share one checkpoint dir.

apply_deltas is fire-and-forget and never raises; see its docstring for the dedup/rewrite/anomaly rules.

Module Contents

Classes

Name	Description
`ResumabilityActor`	Per-writer counter + LMDB owner. Spawned by `Pipeline.run` with

Functions

Name	Description
`create_resumability_actor`	Spawn the detached resumability actor and block until it has scanned the
`shutdown_resumability_actor`	Flush and kill the detached actor. `ray.kill` always runs even if

Data

METADATA_DIRNAME

_COMPLETED_DB

_DEFAULT_MAP_SIZE

API

class nemo_curator.utils.resumability_actor.ResumabilityActor(
    base_dir: str,
    map_size: int = _DEFAULT_MAP_SIZE,
    writer_id: str | None = None
)

Per-writer counter + LMDB owner. Spawned by Pipeline.run with lifetime="detached" and closed at end-of-run; apply_deltas is fire-and-forget and never raises.

_applied

dict[str, int] = {}

_completed

set[str] = self._load_completed()

_db

= self._env.open_db(_COMPLETED_DB)

_dir

= Path(base_dir).absolute() / METADATA_DIRNAME

_env

_path

= str(self._dir / f'{wid}.mdb')

_pending

dict[str, int] = {}

nemo_curator.utils.resumability_actor.ResumabilityActor._load_completed() -> set[str]

Union of completed sources across all writer files; unreadable files (mid-write, or open in-process during tests) are skipped with a warning.

nemo_curator.utils.resumability_actor.ResumabilityActor._persist_completed(
    sids: collections.abc.Iterable[str]
) -> None

nemo_curator.utils.resumability_actor.ResumabilityActor._read_completed_from(
    env: lmdb.Environment
) -> set[str]

Completed-source ids from an open LMDB env (empty if it has no completed-sources db yet).

nemo_curator.utils.resumability_actor.ResumabilityActor._remove_from_completed(
    sid: str
) -> None

Un-complete sid (in-memory + our LMDB file) so it reruns. If a different writer completed it, that entry can’t be removed and may reappear from the union next startup — acceptable for this rare path.

nemo_curator.utils.resumability_actor.ResumabilityActor.apply_deltas(
    per_task: list[tuple[str, str, int]]
) -> None

Apply per-task counter deltas (fire-and-forget; no ray.get).

Each tuple is (task_id, source_id, delta):

seen task_id, same delta → skip (Ray-retry idempotency).
seen task_id, different delta → rewrite _pending by -old+new.
any delta for an already-completed source → warn and un-complete it (in-memory + LMDB) so it reprocesses next run (indicates a bug).
else → apply; persist the source when its counter hits 0.

Never raises.

nemo_curator.utils.resumability_actor.ResumabilityActor.are_completed(
    source_ids: list[str]
) -> list[bool]

Parallel bool list: which source_ids are complete (skip on rerun).

nemo_curator.utils.resumability_actor.ResumabilityActor.close() -> None

nemo_curator.utils.resumability_actor.ResumabilityActor.wait() -> None

No-op the caller ray.gets after spawning the actor: it blocks until __init__ (the checkpoint scan) has finished and surfaces any startup error (e.g. an LMDB open failure) before the pipeline begins.

nemo_curator.utils.resumability_actor.create_resumability_actor(
    checkpoint_path: str
) -> None

Spawn the detached resumability actor and block until it has scanned the checkpoint dir (so the first apply_deltas/are_completed works, and any LMDB startup error surfaces here). Must be called with an active Ray connection — the pipeline wraps it in with ray.init().

nemo_curator.utils.resumability_actor.shutdown_resumability_actor() -> None

Flush and kill the detached actor. ray.kill always runs even if close fails/times out, so a stale actor can’t leak into the next run. A no-op if Ray is already down (the actor then dies with the cluster; its LMDB rows are durable, written sync per delta).

nemo_curator.utils.resumability_actor.METADATA_DIRNAME = '.nemo_curator_metadata'

nemo_curator.utils.resumability_actor._COMPLETED_DB = b'completed_sources'

nemo_curator.utils.resumability_actor._DEFAULT_MAP_SIZE = 1 << 30