sleap_nn/training/model_trainer.py (1)

223-230: Verify that np_chunks_path exists when use_existing_np_chunks is True.

When use_existing_np_chunks is True but there's insufficient memory for in-memory caching, the code creates new paths for chunks without checking if they exist. This could lead to FileNotFoundError.

Add a check to verify that the chunks path exists when use_existing_np_chunks is True:

             if total_cache_memory > available_memory:
+                if self.use_existing_np_chunks:
+                    if not (Path("./train_chunks").exists() and Path("./val_chunks").exists()):
+                        raise FileNotFoundError(
+                            "Chunks directories not found at ./train_chunks and ./val_chunks. "
+                            "Set use_existing_np_chunks=False to generate new chunks."
+                        )
                 self.data_pipeline_fw = "torch_dataset_np_chunks"
                 self.np_chunks = True
                 self.train_np_chunks_path = Path("./train_chunks")
                 self.val_np_chunks_path = Path("./val_chunks")
                 print(
-                    f"Insufficient memory for in-memory caching. `npz` files will be created."
+                    f"Insufficient memory for in-memory caching. Using existing `npz` files."
+                    if self.use_existing_np_chunks
+                    else f"Insufficient memory for in-memory caching. `npz` files will be created."
                 )

🧰 Tools

🪛 Ruff (0.8.2)

229-229: f-string without any placeholders

Remove extraneous f prefix

(F541)

sleap_nn/data/custom_datasets.py (2)

50-50: Document the new attribute in the BaseDataset class docstring.

The new use_existing_chunks attribute is added but not documented in the class docstring.

Add the following line to the docstring:

        np_chunks_path: Path to save the `.npz` chunks. If `None`, current working dir is used.
+        use_existing_chunks: If `True`, use existing chunks in the `np_chunks_path` instead of generating new ones.

---

144-151: Add logging when reusing or generating chunks.

The code silently reuses or generates chunks without informing the user. Adding logging would improve user experience.

Add logging statements to inform the user about chunk reuse or generation:

         if self.np_chunks:
+            if self.use_existing_chunks:
+                print(f"Reusing existing chunk at {f_name}")
+            else:
+                print(f"Generating new chunk at {f_name}")
             sample["image"] = self.transform_to_pil(sample["image"].squeeze(dim=0))
             for k, v in sample.items():
                 if k != "image" and isinstance(v, torch.Tensor):
                     sample[k] = v.numpy()
             f_name = f"{self.np_chunks_path}/sample_{lf_idx}.npz"
             np.savez_compressed(f_name, **sample)

sleap_nn/training/model_trainer.py (2)

79-80: Document the new parameter in the ModelTrainer class docstring.

The new use_existing_np_chunks parameter is added but not documented in the class docstring.

Add the following line to the docstring:

        np_chunks_path: Path to save `.npz` chunks created with `torch_dataset_np_chunks` data pipeline framework.
+        use_existing_np_chunks: If `True`, use existing train and val chunks in the `np_chunks_path` instead of generating new ones.

---

757-768: Add warning when deleting chunks with use_existing_np_chunks=True.

The code deletes the chunks after training without warning the user, even when use_existing_np_chunks is True. This could be unexpected as the user might want to reuse these chunks later.

Add a warning when deleting chunks that were reused:

             if self.np_chunks and delete_np_chunks_after_training:
+                if self.use_existing_np_chunks:
+                    print("Warning: Deleting reused chunks after training. Set delete_np_chunks_after_training=False to keep them.")
                 if (self.train_np_chunks_path).exists():
                     shutil.rmtree(
                         (self.train_np_chunks_path).as_posix(),
                         ignore_errors=True,
                     )

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

tests/training/test_model_trainer.py (1)

330-350: 🛠️ Refactor suggestion

Enhance test coverage for reusing NumPy chunks.

The test cases can be improved in several ways:

1. Use specific exception types instead of bare Exception
2. Add test cases for both success and failure scenarios
3. Verify the exact error message
4. Add test cases for validation chunks

Apply this diff to improve the test cases:

-    ##### test for reusing np chunks path
-    with pytest.raises(Exception):
-        model_trainer = ModelTrainer(
-            config,
-            data_pipeline_fw="torch_dataset_np_chunks",
-            np_chunks_path=tmp_path,
-            use_existing_np_chunks=True,
-        )
-
-    Path.mkdir(Path(tmp_path) / "train_chunks", parents=True)
-    file_path = Path(tmp_path) / "train_chunks" / "sample.npz"
-    np.savez_compressed(file_path, {1: 10})
-
-    with pytest.raises(Exception):
-        model_trainer = ModelTrainer(
-            config,
-            data_pipeline_fw="torch_dataset_np_chunks",
-            np_chunks_path=tmp_path,
-            use_existing_np_chunks=True,
-        )
+    ##### test for reusing np chunks path
+    # Test failure case: non-existent chunks
+    with pytest.raises(FileNotFoundError, match=r"There are no numpy chunks in the path:.*"):
+        model_trainer = ModelTrainer(
+            config,
+            data_pipeline_fw="torch_dataset_np_chunks",
+            np_chunks_path=tmp_path,
+            use_existing_np_chunks=True,
+        )
+
+    # Test failure case: missing validation chunks
+    train_chunks_path = Path(tmp_path) / "train_chunks"
+    train_chunks_path.mkdir(parents=True)
+    file_path = train_chunks_path / "sample.npz"
+    np.savez_compressed(file_path, {"data": np.zeros((10, 10))})
+
+    with pytest.raises(FileNotFoundError, match=r"There are no numpy chunks in the path:.*"):
+        model_trainer = ModelTrainer(
+            config,
+            data_pipeline_fw="torch_dataset_np_chunks",
+            np_chunks_path=tmp_path,
+            use_existing_np_chunks=True,
+        )
+
+    # Test success case: both train and validation chunks exist
+    val_chunks_path = Path(tmp_path) / "val_chunks"
+    val_chunks_path.mkdir(parents=True)
+    val_file_path = val_chunks_path / "sample.npz"
+    np.savez_compressed(val_file_path, {"data": np.zeros((10, 10))})
+
+    model_trainer = ModelTrainer(
+        config,
+        data_pipeline_fw="torch_dataset_np_chunks",
+        np_chunks_path=tmp_path,
+        use_existing_np_chunks=True,
+    )
+    assert model_trainer.use_existing_np_chunks is True

🧰 Tools

🪛 Ruff (0.8.2)

331-331: pytest.raises(Exception) should be considered evil

(B017)

---

343-343: pytest.raises(Exception) should be considered evil

(B017)

sleap_nn/training/model_trainer.py (1)

103-119: 🛠️ Refactor suggestion

Improve error handling when checking for existing chunks.

The current implementation can be improved by:

1. Using a more specific exception type
2. Extracting the validation logic into a helper method
3. Using a more specific file pattern check

Apply this diff to improve error handling:

-        if self.use_existing_np_chunks:
-            if not (
-                self.train_np_chunks_path.exists()
-                and self.train_np_chunks_path.is_dir()
-                and any(self.train_np_chunks_path.glob("*.npz"))
-            ):
-                raise Exception(
-                    f"There are no numpy chunks in the path: {self.train_np_chunks_path}"
-                )
-            if not (
-                self.val_np_chunks_path.exists()
-                and self.val_np_chunks_path.is_dir()
-                and any(self.val_np_chunks_path.glob("*.npz"))
-            ):
-                raise Exception(
-                    f"There are no numpy chunks in the path: {self.val_np_chunks_path}"
-                )
+        if self.use_existing_np_chunks:
+            self._validate_chunks_path(self.train_np_chunks_path, "train")
+            self._validate_chunks_path(self.val_np_chunks_path, "validation")
+
+    def _validate_chunks_path(self, path: Path, split: str) -> None:
+        """Validate that the chunks path exists and contains .npz files.
+        
+        Args:
+            path: Path to the chunks directory
+            split: Name of the split (train/validation) for error messages
+            
+        Raises:
+            FileNotFoundError: If the path doesn't exist or contain .npz files
+        """
+        if not path.exists():
+            raise FileNotFoundError(
+                f"The {split} chunks directory does not exist: {path}"
+            )
+        if not path.is_dir():
+            raise FileNotFoundError(
+                f"The {split} chunks path is not a directory: {path}"
+            )
+        if not any(p.suffix == '.npz' for p in path.iterdir()):
+            raise FileNotFoundError(
+                f"No .npz files found in the {split} chunks directory: {path}"
+            )

sleap_nn/training/model_trainer.py (1)

80-81: Enhance parameter documentation.

The docstring could provide more details about the parameter's behavior and requirements.

Apply this diff to improve the documentation:

-        use_existing_np_chunks: Use existing train and val chunks in the `np_chunks_path`.
+        use_existing_np_chunks: If True, use existing train and validation chunks from the
+            `np_chunks_path` instead of generating new ones. Both train_chunks and val_chunks
+            directories must exist and contain valid .npz files. Raises FileNotFoundError if
+            the directories don't exist or don't contain .npz files.

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro