Merge branch 'main' into codex/docs-nemotron-ocr-pipeline

Author: lbliii

.github/workflows/claude-review.yml

@@ -13,8 +13,9 @@ permissions:
 jobs:
   claude-review:
     if: github.event.issue.pull_request != null && contains(fromJson('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.comment.author_association)
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_claude_review.yml@v0.88.0
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_claude_review.yml@v1.7.0
     with:
+      model: ${{ vars.CLAUDE_MODEL }}
       prompt: |
         You are doing a light code review. Keep it concise and actionable.
 
@@ -36,4 +37,5 @@ jobs:
         It's perfectly acceptable to not have anything to comment on.
         If you do not have anything to comment on, post "LGTM".
     secrets:
-      ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+      NVIDIA_INFERENCE_URL: ${{ secrets.NVIDIA_INFERENCE_URL }}
+      NVIDIA_INFERENCE_KEY: ${{ secrets.NVIDIA_INFERENCE_KEY }}

nemo_curator/backends/base.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,14 +17,32 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
+from loguru import logger
+
+from nemo_curator.backends.failed_task_markers import record_failed_tasks
+from nemo_curator.backends.slurm_array import (
+    filter_slurm_array_source_tasks,
+    resolve_slurm_array_config,
+)
 from nemo_curator.core.utils import ignore_ray_head_node
 from nemo_curator.tasks import Task
+from nemo_curator.tasks.sentinels import FailedTask, NoneTask
 from nemo_curator.utils.performance_utils import StageTimer
+from nemo_curator.utils.resumability_client import (
+    completed_resumability_sources,
+    flush_resumability_deltas,
+    is_resumability_actor_active,
+)
 
 if TYPE_CHECKING:
     from nemo_curator.stages.base import ProcessingStage
 
 
+def _is_sentinel(task: Task) -> bool:
+    """A payload-less marker (NoneTask/FailedTask), stripped before the next stage."""
+    return isinstance(task, (NoneTask, FailedTask))
+
+
 @dataclass
 class NodeInfo:
     """Generic node information for setup_on_node calls across backends.
@@ -85,9 +103,47 @@ def process_batch(self, tasks: list[Task]) -> list[Task]:
             # Use the batch processing logic
             results = self.stage.process_batch(tasks)
 
+        # A returned ``None`` ("filter this slot") becomes a NoneTask so every
+        # output is a real Task that gets a task_id. Sentinels (NoneTask /
+        # FailedTask) carry no identity and are stripped again before this
+        # method returns.
+        results = [NoneTask() if r is None else r for r in results]
+
         # Guarantee every emitted task has a task_id (derived id, or uuid fallback).
         results = self._post_process_task_ids(tasks, results)
 
+        # Failed tasks on the source stage are not supported.
+        is_source_stage = getattr(self.stage, "is_source_stage", False)
+        failed_tasks = [r for r in results if isinstance(r, FailedTask)]
+        if failed_tasks and is_source_stage:
+            msg = (
+                f"Source stage {self.stage.name} emitted FailedTask, which is not supported."
+            )
+            raise ValueError(msg)
+
+        # Record failed tasks for later inspection or retry bookkeeping.
+        if failed_tasks:
+            record_failed_tasks()
+
+        # Source-stage sentinels (NoneTask only; FailedTask already raised above) are
+        # not real partitions and must not influence shard assignment or resumability
+        # counters. Non-source stages keep sentinels here so _apply_resumability_counters
+        # can fire the correct -1 delta for filtered (NoneTask) slots in the 1:1 path.
+        if is_source_stage:
+            results = [r for r in results if not _is_sentinel(r)]
+
+        # Filter tasks based on the Slurm array configuration.
+        slurm_array = resolve_slurm_array_config(is_source_stage=is_source_stage)
+        if slurm_array is not None and is_source_stage:
+            results = filter_slurm_array_source_tasks(results, slurm_array, self.stage.name)
+
+        # Opt-in resumability: fire per-source deltas (no-op when no actor registered).
+        if is_resumability_actor_active():
+            results = self._apply_resumability_counters(tasks, results)
+
+        # Sentinels never propagate to the next stage.
+        results = [r for r in results if not _is_sentinel(r)]
+
         # Log performance stats and add to result tasks
         _, stage_perf_stats = self._timer.log_stats()
         # Consume and attach any custom metrics recorded by the stage during this call
@@ -99,62 +155,41 @@ def process_batch(self, tasks: list[Task]) -> list[Task]:
 
         return results
 
-    def _post_process_task_ids(self, input_tasks: list[Task], output_tasks: list[Task | None]) -> list[Task]:
-        """Assign a deterministic ``task_id`` to every emitted task.
-
-        This is the single place task ids are assigned — it runs for every
-        stage on every backend (all backend adapters subclass this), so it
-        makes no difference whether a stage defines ``process`` or overrides
-        ``process_batch``. ``task_id`` is the task's id path (parents + own segment); ids are
-        re-derived at each stage boundary so the same object passing through
-        N stages gets N ids.
-
-        The input→output mapping decides each output's PARENT; whether the
-        stage is a source decides each output's SEGMENT (content id vs index)
-        — the two are independent. ``None`` outputs (Curator's "return None to
-        filter") are NOT removed before the length check — keeping them in
-        place preserves positional alignment for filter stages — and are then
-        dropped from the returned list.
-
-        - single input → every output is its child (fan-out): ``parent_<seg>``
-        - ``len(output) == len(input)`` → positional 1:1: each ``parent_i_<seg>``;
-          a ``None`` slot just means input ``i`` was filtered.
-        - any other (ambiguous) cardinality across a batch → a random ``uuid``
-          prefixed with ``"r"`` (e.g. ``"r3f9a…"``), so ``task_id`` is never
-          empty even when a derived id is not possible. The ``"r"`` prefix flags
-          the id as non-deterministic / ancestry-not-tracked (see
-          ``Task.task_id`` docstring).
-
-        ``seg`` is the output's content id (``Task.get_deterministic_id()``)
-        for a source stage when available, else the positional index — so a
-        source partition keeps a stable id across reorderings regardless of
-        whether the source is 1→N or N→N.
-
-        Note: a stage that BOTH filters and fans out within a single batch
-        (returning a flat list rather than a per-input slot) cannot be mapped
-        positionally; if its length happens to equal the input length the 1:1
-        assumption may misattribute parents. That combination is unsupported
-        until per-slot sentinels (NoneTask/FailedTask) land in a later PR.
+    def _post_process_task_ids(self, input_tasks: list[Task], output_tasks: list[Task]) -> list[Task]:
+        """Assign a deterministic ``task_id`` (parent id + own segment) to every
+        emitted task. Runs once per stage on every backend, so ``process`` vs
+        ``process_batch`` makes no difference; ids are re-derived at each stage
+        boundary, so one object passing through N stages gets N ids.
+
+        - single input → fan-out: each output is ``parent_<seg>``
+        - ``len(output) == len(input)`` → positional 1:1: ``parent_i_<seg>``; a
+          ``NoneTask`` slot means input ``i`` was filtered (kept for alignment, then
+          dropped from the result)
+        - any other cardinality → a random ``"r"``-prefixed uuid (non-deterministic,
+          ancestry-not-tracked; see ``Task.task_id``)
+
+        ``seg`` is the content id (``get_deterministic_id()``) for a source stage,
+        else the positional index. A stage that both filters and fans out in one
+        batch can't be mapped positionally and falls to the ``"r"`` case — return
+        one value (or ``None``) per input to stay positional.
         """
         is_source = getattr(self.stage, "is_source_stage", False)
 
         if len(input_tasks) == 1:
-            # Fan-out (incl. a source reading from EmptyTask): every non-None
+            # Fan-out (including a source reading from EmptyTask): every
             # output is a child of the single input.
             parent_id = input_tasks[0].task_id
-            out: list[Task] = [t for t in output_tasks if t is not None]
+            out = list(output_tasks)
             for i, task in enumerate(out):
                 suffix = (task.get_deterministic_id() or i) if is_source else i
                 task._set_task_id(parent_id, suffix)
             return out
 
         if len(output_tasks) == len(input_tasks):
-            # Positional 1:1. None is kept above so a filtered slot still lines
-            # up with its own parent; drop the None slots from the result.
+            # Positional 1:1. A NoneTask sentinel remains aligned with the
+            # parent whose output was filtered.
             out = []
             for parent, task in zip(input_tasks, output_tasks, strict=True):
-                if task is None:
-                    continue
                 suffix = (task.get_deterministic_id() or 0) if is_source else 0
                 task._set_task_id(parent.task_id, suffix)
                 out.append(task)
@@ -163,11 +198,91 @@ def _post_process_task_ids(self, input_tasks: list[Task], output_tasks: list[Tas
         # Ambiguous cardinality across a batch: a derived id is not possible. Use a
         # random "r"-prefixed uuid so task_id is non-empty but clearly flagged
         # non-deterministic.
-        out = [t for t in output_tasks if t is not None]
+        out = list(output_tasks)
         for task in out:
             task.task_id = "r" + uuid.uuid4().hex
         return out
 
+    # Resumability (opt-in): stamp _source_id, fire per-source deltas, drop
+    # completed sources. task_ids are already assigned; sentinels stripped by caller.
+    def _apply_resumability_counters(self, input_tasks: list[Task], output_tasks: list[Task]) -> list[Task]:  # noqa: C901
+        # Dedup key is always an OUTPUT task_id, never the input's: the source
+        # already keyed its +1 on that id, and an output id is one level deeper,
+        # so it's unique to the (task, stage) that produced it.
+        stage = self.stage
+        if getattr(stage, "is_source_stage", False):
+            return self._source_counters(output_tasks)
+
+        # No outputs (e.g. a batch entirely filtered, or an end-of-pipeline
+        # no-op): nothing to attribute a delta to, so skip.
+        if not output_tasks:
+            return output_tasks
+
+        # Pre-source: inputs have no _source_id yet; nothing to track.
+        if all(not t._source_id for t in input_tasks):
+            return output_tasks
+
+        is_sink = stage.is_sink_stage
+        per_task: list[tuple[str, str, int]] = []
+
+        if len(input_tasks) == 1 and len(output_tasks) > 1:
+            # Fan-out (1->N): parent consumed (-1); each real child continues
+            # (+1, or 0 at a sink); each FailedTask keeps the source open (+1);
+            # NoneTask contributes 0.
+            parent = input_tasks[0]
+            real = [t for t in output_tasks if not _is_sentinel(t)]
+            n_failed = sum(1 for t in output_tasks if isinstance(t, FailedTask))
+            continuing = 0 if is_sink else len(real)
+            delta = continuing + n_failed - 1
+            # Key on output[0].task_id (not parent.task_id, which collides with the
+            # source's +1). Non-source children are indexed positionally, so
+            # output[0] is always "<parent>_0".
+            per_task.append((output_tasks[0].task_id, parent._source_id, delta))
+            for c in real:
+                if not c._source_id:
+                    c._source_id = parent._source_id
+        elif len(output_tasks) == len(input_tasks):
+            # Positional 1:1; each delta keys on the output id (r.task_id).
+            for parent, r in zip(input_tasks, output_tasks, strict=True):
+                sid = parent._source_id
+                if isinstance(r, NoneTask):  # filtered -> consumed
+                    per_task.append((r.task_id, sid, -1))
+                    continue
+                if isinstance(r, FailedTask):  # failed -> source stays open (no sink test)
+                    per_task.append((r.task_id, sid, 0))
+                    continue
+                per_task.append((r.task_id, sid, -1 if is_sink else 0))  # real: sink -1, else 0
+                if not r._source_id:
+                    r._source_id = sid
+        else:
+            # M->K (M!=K): can't attribute parents; skip (source stays pending -> reprocessed).
+            logger.warning(
+                f"resumability: {type(stage).__name__} produced {len(output_tasks)} outputs "
+                f"for {len(input_tasks)} inputs; can't attribute sources, skipping counter "
+                f"update for this batch."
+            )
+            return output_tasks
+
+        flush_resumability_deltas(per_task)
+        return output_tasks
+
+    def _source_counters(self, output_tasks: list[Task]) -> list[Task]:
+        """Source stage: each output is a source partition; its ``_source_id`` is
+        ``Task.get_source_id()``. Drop already-completed sources; each survivor fires ``+1``."""
+        sources = [t for t in output_tasks if not _is_sentinel(t)]
+        for t in sources:
+            t._source_id = t.get_source_id()
+        completed = completed_resumability_sources([t._source_id for t in sources])
+        per_task: list[tuple[str, str, int]] = []
+        survivors: list[Task] = []
+        for t in sources:
+            if t._source_id in completed:
+                continue
+            per_task.append((t.task_id, t._source_id, +1))
+            survivors.append(t)
+        flush_resumability_deltas(per_task)
+        return survivors
+
     def setup_on_node(self, node_info: NodeInfo | None = None, worker_metadata: WorkerMetadata | None = None) -> None:
         """Setup the stage on a node.
 

nemo_curator/backends/failed_task_markers.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import uuid
+from pathlib import Path
+
+from nemo_curator.utils.retry_manifest import METADATA_DIRNAME
+
+FAILED_TASKS_DIR_ENV_VAR = "NEMO_CURATOR_FAILED_TASKS_DIR"
+FAILED_TASK_MANIFEST_FILENAME = "failed_tasks.json"
+
+
+def _configure_failed_task_manifest_dir(default_dir: Path) -> Path:
+    existing = os.environ.get(FAILED_TASKS_DIR_ENV_VAR)
+    if existing:
+        return Path(existing)
+
+    manifest_dir = default_dir.absolute()
+    os.environ[FAILED_TASKS_DIR_ENV_VAR] = str(manifest_dir)
+    return manifest_dir
+
+
+def configure_failed_task_manifest_dir(checkpoint_path: str | Path) -> Path:
+    """Configure a local attempt-scoped FailedTask manifest directory unless overridden."""
+    manifest_dir = Path(
+        checkpoint_path,
+        METADATA_DIRNAME,
+        ".failed_tasks",
+        f"local_attempt_{uuid.uuid4().hex}",
+    )
+    return _configure_failed_task_manifest_dir(manifest_dir)
+
+
+def configure_slurm_array_failed_task_manifest_dir(checkpoint_path: str | Path, shard_index: int) -> Path:
+    """Configure an attempt-scoped FailedTask manifest directory unless overridden."""
+    job_id = os.environ.get("SLURM_JOB_ID", f"local_{os.getpid()}")
+    array_task_id = os.environ.get("SLURM_ARRAY_TASK_ID", "local")
+    restart_count = os.environ.get("SLURM_RESTART_COUNT", "0")
+    manifest_dir = Path(
+        checkpoint_path,
+        METADATA_DIRNAME,
+        ".failed_tasks",
+        f"slurm_job_{job_id}",
+        f"array_task_{array_task_id}",
+        f"restart_{restart_count}",
+        f"shard_{shard_index}",
+    )
+    return _configure_failed_task_manifest_dir(manifest_dir)
+
+
+def record_failed_tasks() -> None:
+    """Write one attempt-scoped manifest after any FailedTask is detected."""
+    manifest_dir = os.environ.get(FAILED_TASKS_DIR_ENV_VAR)
+    if not manifest_dir:
+        return
+
+    manifest_path = Path(manifest_dir, FAILED_TASK_MANIFEST_FILENAME)
+    if manifest_path.is_file():
+        return
+
+    manifest_path.parent.mkdir(parents=True, exist_ok=True)
+    manifest_path.touch(exist_ok=True)
+
+
+def failed_task_manifest_exists(manifest_dir: str | Path | None = None) -> bool:
+    """Return whether the current attempt has recorded any FailedTask."""
+    resolved_manifest_dir = manifest_dir if manifest_dir is not None else os.environ.get(FAILED_TASKS_DIR_ENV_VAR)
+    if not resolved_manifest_dir:
+        return False
+    return Path(resolved_manifest_dir, FAILED_TASK_MANIFEST_FILENAME).is_file()