minor cleaning

Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>

Author: sarahyurick

nemo_curator/backends/base.py

@@ -110,7 +110,7 @@ def process_batch(self, tasks: list[Task]) -> list[Task]:
             raise ValueError(msg)
 
         # Record failed tasks for later inspection or retry bookkeeping.
-        record_failed_tasks(self.stage.name, failed_tasks)
+        record_failed_tasks(failed_tasks)
 
         # Sentinels never propagate to the next stage.
         results = [r for r in results if not isinstance(r, (NoneTask, FailedTask))]

nemo_curator/backends/failed_task_markers.py

@@ -62,7 +62,7 @@ def configure_slurm_array_failed_task_manifest_dir(checkpoint_path: str | Path,
     return _configure_failed_task_manifest_dir(manifest_dir)
 
 
-def record_failed_tasks(_stage_name: str, failed_tasks: list[FailedTask]) -> None:
+def record_failed_tasks(failed_tasks: list[FailedTask]) -> None:
     """Write one attempt-scoped manifest after any FailedTask is detected."""
     if not failed_tasks:
         return

nemo_curator/backends/slurm_array.py

@@ -39,23 +39,26 @@
 
 def _get_int_env_var(env_var: str, fallback_name: str | None = None, default: int | None = None) -> int:
     """Read an integer env var, with optional fallback/default."""
+    resolved_var = env_var
     env_value = os.environ.get(env_var)
-    if env_value is None:
-        if fallback_name is not None:
-            env_var = fallback_name
-            env_value = os.environ.get(env_var)
+    if env_value is None and fallback_name is not None:
+        resolved_var = fallback_name
+        env_value = os.environ.get(fallback_name)
 
-        if env_value is None:
-            if default is not None:
-                return default
+    if env_value is None:
+        if default is not None:
+            return default
 
+        if fallback_name is not None:
+            msg = f"Environment variable {env_var} (or {fallback_name}) is not set"
+        else:
             msg = f"Environment variable {env_var} is not set"
-            raise ValueError(msg)
+        raise ValueError(msg)
 
     try:
         return int(env_value)
     except ValueError as e:
-        msg = f"Environment variable {env_var} must contain an integer, got {env_value!r}"
+        msg = f"Environment variable {resolved_var} must contain an integer, got {env_value!r}"
         raise ValueError(msg) from e
 
 
@@ -195,7 +198,13 @@ def filter_slurm_array_source_tasks(
 
 
 def is_slurm_array_driver_process(use_slurm: bool) -> bool:
-    """Return true for the process that owns retry metadata."""
+    """Return true for the process that owns retry metadata.
+
+    When ``use_slurm`` is False (local or single-node) every process is the
+    driver. When ``use_slurm`` is True, only the Slurm head node
+    (``SLURM_NODEID == 0``) is the driver; if the variable is absent (e.g.
+    bare ``srun`` without an array) the process is treated as the head.
+    """
     return not use_slurm or os.environ.get("SLURM_NODEID", "0") == "0"
 
 

nemo_curator/utils/atomic_io.py

@@ -32,6 +32,14 @@ def fsync_directory(path: Path) -> None:
         os.close(dir_fd)
 
 
+def _unlink_best_effort(path: Path) -> None:
+    """Remove a temporary file without masking the primary result."""
+    try:
+        path.unlink(missing_ok=True)
+    except OSError:
+        pass
+
+
 def _write_json_temp_file(
     path: Path,
     payload: Any,  # noqa: ANN401
@@ -57,10 +65,10 @@ def _write_json_temp_file(
             tmp.write("\n")
             tmp.flush()
             os.fsync(tmp.fileno())
-        return Path(tmp.name)
+        return tmp_path
     except Exception:
         if tmp_path is not None:
-            tmp_path.unlink(missing_ok=True)
+            _unlink_best_effort(tmp_path)
         raise
 
 
@@ -97,7 +105,7 @@ def write_json_atomically(
         os.replace(tmp_path, path)
         _fsync_directory_best_effort(path.parent)
     except Exception:
-        tmp_path.unlink(missing_ok=True)
+        _unlink_best_effort(tmp_path)
         raise
 
 
@@ -120,9 +128,12 @@ def write_json_atomically_if_absent(
     try:
         os.link(tmp_path, path)
     except FileExistsError:
+        _unlink_best_effort(tmp_path)
         return False
-    finally:
-        tmp_path.unlink(missing_ok=True)
+    except Exception:
+        _unlink_best_effort(tmp_path)
+        raise
 
+    _unlink_best_effort(tmp_path)
     _fsync_directory_best_effort(path.parent)
     return True

tests/backends/test_base_stage_adapter.py

@@ -104,7 +104,7 @@ def filter_tasks(
             calls["filter_tasks"] += 1
             return tasks
 
-        def record_failed_tasks(_stage_name: str, _failed_tasks: list[FailedTask]) -> None:
+        def record_failed_tasks(_failed_tasks: list[FailedTask]) -> None:
             calls["record_failed_tasks"] += 1
 
         monkeypatch.setattr(base_module, "resolve_slurm_array_config", resolve_config)

tests/backends/test_failed_task_markers.py

@@ -83,7 +83,7 @@ def test_record_failed_tasks_writes_single_manifest(
         manifest_dir = tmp_path / "failed-tasks"
         monkeypatch.setenv(FAILED_TASKS_DIR_ENV_VAR, str(manifest_dir))
 
-        record_failed_tasks("failed", [_failed_task("0_7_0"), _failed_task("0_8_0")])
+        record_failed_tasks([_failed_task("0_7_0"), _failed_task("0_8_0")])
 
         manifest_files = list(manifest_dir.glob("*.json"))
         assert manifest_files == [manifest_dir / FAILED_TASK_MANIFEST_FILENAME]
@@ -95,11 +95,11 @@ def test_additional_failed_tasks_leave_existing_manifest_unchanged(
     ) -> None:
         manifest_dir = tmp_path / "failed-tasks"
         monkeypatch.setenv(FAILED_TASKS_DIR_ENV_VAR, str(manifest_dir))
-        record_failed_tasks("stage-a", [_failed_task("0_7_0")])
+        record_failed_tasks([_failed_task("0_7_0")])
         manifest_file = manifest_dir / FAILED_TASK_MANIFEST_FILENAME
         original_manifest = manifest_file.read_text()
 
-        record_failed_tasks("stage-b", [_failed_task("0_8_0")])
+        record_failed_tasks([_failed_task("0_8_0")])
 
         assert list(manifest_dir.glob("*.json")) == [manifest_file]
         assert manifest_file.read_text() == original_manifest
@@ -110,7 +110,7 @@ def test_record_failed_tasks_without_configured_attempt_is_noop(
         monkeypatch.delenv(FAILED_TASKS_DIR_ENV_VAR, raising=False)
         monkeypatch.chdir(tmp_path)
 
-        record_failed_tasks("failed", [_failed_task()])
+        record_failed_tasks([_failed_task()])
 
         assert not (tmp_path / ".nemo_curator_metadata").exists()
 
@@ -120,7 +120,7 @@ def test_record_failed_tasks_does_not_write_manifest_for_empty_list(
         manifest_dir = tmp_path / "failed-tasks"
         monkeypatch.setenv(FAILED_TASKS_DIR_ENV_VAR, str(manifest_dir))
 
-        record_failed_tasks("failed", [])
+        record_failed_tasks([])
 
         assert not manifest_dir.exists()
         assert not failed_task_manifest_exists()
@@ -137,7 +137,7 @@ def fail_write(*_args: object, **_kwargs: object) -> None:
         monkeypatch.setattr(failed_task_markers_module, "write_json_atomically", fail_write)
 
         with pytest.raises(OSError, match="storage unavailable"):
-            record_failed_tasks("failed", [_failed_task()])
+            record_failed_tasks([_failed_task()])
 
     def test_failed_task_manifest_exists_accepts_explicit_directory(self, tmp_path: Path) -> None:
         manifest_dir = tmp_path / "failed-tasks"

tests/backends/test_slurm_array.py

@@ -21,6 +21,7 @@
 
 import nemo_curator.backends.slurm_array as slurm_array_module
 from nemo_curator.backends.slurm_array import (
+    SLURM_ARRAY_COMPLETION_MANIFEST_NAMESPACE,
     SLURM_ARRAY_ENABLED_ENV_VAR,
     SLURM_ARRAY_MINIMUM_SHARD_INDEX_ENV_VAR,
     SLURM_ARRAY_SHARD_INDEX_ENV_VAR,
@@ -205,6 +206,11 @@ def test_filtering_is_disabled_without_config(self) -> None:
 
         assert filter_slurm_array_source_tasks(tasks, None, "source") == tasks
 
+    def test_filtering_returns_empty_for_empty_input(self) -> None:
+        slurm_array = SlurmArrayConfig(shard_index=0, total_shards=3)
+
+        assert filter_slurm_array_source_tasks([], slurm_array, "source") == []
+
     def test_assigns_each_source_task_to_one_shard(self) -> None:
         tasks = [_task(f"0_{i}") for i in range(8)]
         assigned_task_ids = []
@@ -256,6 +262,7 @@ def test_is_driver_process_for_local_and_slurm_head(self, monkeypatch: MonkeyPat
         monkeypatch.setenv("SLURM_NODEID", "1")
 
         assert is_slurm_array_driver_process(use_slurm=True) is False
+        assert is_slurm_array_driver_process(use_slurm=False) is True
 
     def test_build_completion_manifest_writes_run_config_and_shard_identity(self, tmp_path: Path) -> None:
         manifest = build_slurm_array_completion_manifest(
@@ -335,6 +342,27 @@ def test_find_retries_returns_empty_plan_when_all_shards_completed(self, tmp_pat
     def test_find_retries_returns_none_without_run_config(self, tmp_path: Path) -> None:
         assert find_slurm_array_retries(tmp_path) is None
 
+    def test_find_retries_rejects_zero_total_shards_in_run_config(self, tmp_path: Path) -> None:
+        manifest_dir = tmp_path / METADATA_DIRNAME / ".slurm_array_completion"
+        manifest_dir.mkdir(parents=True)
+        (manifest_dir / "run.json").write_text('{"minimum_shard_index":0,"total_shards":0}\n')
+
+        with pytest.raises(ValueError, match="total_shards greater than 0"):
+            find_slurm_array_retries(tmp_path)
+
+    def test_find_retries_rejects_out_of_range_completed_shard(self, tmp_path: Path) -> None:
+        manifest_dir = tmp_path / METADATA_DIRNAME / ".slurm_array_completion"
+        manifest_dir.mkdir(parents=True)
+        (manifest_dir / "run.json").write_text('{"minimum_shard_index":1,"total_shards":3}\n')
+        # shard_index=5 is outside the valid range [1, 3]
+        filename = f"completed_{SLURM_ARRAY_COMPLETION_MANIFEST_NAMESPACE}_bad.json"
+        (manifest_dir / filename).write_text(
+            '{"minimum_shard_index":1,"shard_index":5,"status":"completed","total_shards":3}'
+        )
+
+        with pytest.raises(ValueError, match="outside the original shard range"):
+            find_slurm_array_retries(tmp_path)
+
     def test_find_retries_rejects_negative_minimum_in_run_config(self, tmp_path: Path) -> None:
         manifest_dir = tmp_path / METADATA_DIRNAME / ".slurm_array_completion"
         manifest_dir.mkdir(parents=True)

tests/utils/test_atomic_io.py

@@ -61,3 +61,19 @@ def test_write_json_atomically_if_absent_does_not_replace_existing_file(self, tm
 
         assert json.loads(output_path.read_text()) == {"writer": 1}
         assert not list(output_path.parent.glob(f".{output_path.name}.*.tmp"))
+
+    def test_write_json_atomically_if_absent_does_not_fail_after_commit_when_cleanup_fails(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        output_path = tmp_path / "payload.json"
+        original_unlink = Path.unlink
+
+        def fail_temp_cleanup(path: Path, *, missing_ok: bool = False) -> None:
+            if path.suffix == ".tmp":
+                raise OSError("cleanup unavailable")
+            original_unlink(path, missing_ok=missing_ok)
+
+        monkeypatch.setattr(Path, "unlink", fail_temp_cleanup)
+
+        assert write_json_atomically_if_absent(output_path, {"writer": 1}) is True
+        assert json.loads(output_path.read_text()) == {"writer": 1}

tutorials/slurm/README.md

@@ -342,7 +342,7 @@ python tutorials/slurm/retry_array.py \
     --format fields
 ```
 
-For example, this output means that shards `1`, `2`, `5` through `10`, and `99` should be retried with offset `0`, using the original logical shard range of `0` through `99`:
+Without `--max-array-size`, the output is always at most one line (all missing logical shards are emitted as a single submission). For example, this output means that shards `1`, `2`, `5` through `10`, and `99` should be retried with offset `0`, using the original logical shard range of `0` through `99`:
 
 ```text
 1-2,5-10,99 0 0 100

tutorials/slurm/array_pipeline.py

@@ -157,6 +157,9 @@ def main() -> None:
 
         pipeline.run()
 
+        # Check unconditionally: the user may have set NEMO_CURATOR_FAILED_TASKS_DIR
+        # manually even without --checkpoint-path, and the function safely returns
+        # False when no directory is configured.
         if failed_task_manifest_exists():
             logger.warning(
                 "Pipeline completed without raising, but a FailedTask manifest exists. "