feat(lightning): add MultiloaderLifecycle callback

docs/reference/lightning/multiloader_lifecycle.md

@@ -0,0 +1,3 @@
+# ratiopath.lightning.callbacks.MultiloaderLifecycle
+
+::: ratiopath.lightning.callbacks.MultiloaderLifecycle

mkdocs.yml

@@ -57,6 +57,8 @@ nav:
           - read_slides: reference/ray/read_slides.md
           - VipsTiffDatasink: reference/ray/vips_tiff_datasink.md
       - OpenSlide: reference/openslide.md
+      - Lightning:
+          - MultiloaderLifecycle: reference/lightning/multiloader_lifecycle.md
 
 theme:
   name: material

pyproject.toml

@@ -31,6 +31,9 @@ dependencies = [
     "albumentationsx>=2.3.0",
 ]
 
+[project.optional-dependencies]
+lightning = ["lightning>=2.4.0"]
+
 [dependency-groups]
 dev = [
     "mypy",
@@ -43,7 +46,7 @@ docs = [
     "mkdocs-material>=9.6.18",
     "mkdocstrings[python]>=0.30.0",
 ]
-tests = ["openslide-bin>=4.0.0.8", "pytest>=8.4.1"]
+tests = ["lightning>=2.4.0", "openslide-bin>=4.0.0.8", "pytest>=8.4.1"]
 
 
 [project.urls]

ratiopath/lightning/callbacks/__init__.py

@@ -0,0 +1,4 @@
+from ratiopath.lightning.callbacks.multiloader_lifecycle import MultiloaderLifecycle
+
+
+__all__ = ["MultiloaderLifecycle"]

ratiopath/lightning/callbacks/multiloader_lifecycle.py

@@ -0,0 +1,143 @@
+from abc import ABC
+from typing import Any
+
+from lightning import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
+
+
+class MultiloaderLifecycle(Callback, ABC):
+    """Callback base class with hooks for the start and end of each dataloader.
+
+    With several validation, test or predict dataloaders (e.g. one per slide), Lightning
+    reports only batches and epochs. This class adds `on_{stage}_dataloader_start` and
+    `on_{stage}_dataloader_end` hooks, where `stage` is `validation`, `test` or
+    `predict`. They are called whenever the dataloader index changes and once more for
+    the last dataloader at the end of the epoch.
+
+    Requires the `lightning` extra: `pip install ratiopath[lightning]`.
+
+    Note:
+        Subclasses that override `on_{stage}_batch_start` or `on_{stage}_epoch_end`
+        must call the parent implementation, otherwise the dataloader hooks are not
+        called.
+
+    Examples:
+        ```python
+        from lightning import LightningModule, Trainer
+
+        from ratiopath.lightning.callbacks import MultiloaderLifecycle
+
+
+        class SlideHeatmaps(MultiloaderLifecycle):
+            def on_predict_dataloader_start(
+                self, trainer: Trainer, pl_module: LightningModule, dataloader_idx: int
+            ) -> None: ...  # e.g. allocate a heatmap for the slide
+
+            def on_predict_dataloader_end(
+                self, trainer: Trainer, pl_module: LightningModule, dataloader_idx: int
+            ) -> None: ...  # e.g. save the finished heatmap
+        ```
+    """
+
+    def __init__(self) -> None:
+        self._dataloader_idxs = {"validation": -1, "test": -1, "predict": -1}
+
+    def on_validation_dataloader_start(
+        self, trainer: Trainer, pl_module: LightningModule, dataloader_idx: int
+    ) -> None:
+        """Called before the first validation batch of a dataloader."""
+
+    def on_validation_dataloader_end(
+        self, trainer: Trainer, pl_module: LightningModule, dataloader_idx: int
+    ) -> None:
+        """Called after the last validation batch of a dataloader."""
+
+    def on_test_dataloader_start(
+        self, trainer: Trainer, pl_module: LightningModule, dataloader_idx: int
+    ) -> None:
+        """Called before the first test batch of a dataloader."""
+
+    def on_test_dataloader_end(
+        self, trainer: Trainer, pl_module: LightningModule, dataloader_idx: int
+    ) -> None:
+        """Called after the last test batch of a dataloader."""
+
+    def on_predict_dataloader_start(
+        self, trainer: Trainer, pl_module: LightningModule, dataloader_idx: int
+    ) -> None:
+        """Called before the first predict batch of a dataloader."""
+
+    def on_predict_dataloader_end(
+        self, trainer: Trainer, pl_module: LightningModule, dataloader_idx: int
+    ) -> None:
+        """Called after the last predict batch of a dataloader."""
+
+    def on_batch_start(
+        self,
+        stage: str,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        dataloader_idx: int,
+    ) -> None:
+        if dataloader_idx != self._dataloader_idxs[stage]:
+            if self._dataloader_idxs[stage] != -1:
+                getattr(self, f"on_{stage}_dataloader_end")(
+                    trainer, pl_module, self._dataloader_idxs[stage]
+                )
+
+            self._dataloader_idxs[stage] = dataloader_idx
+            getattr(self, f"on_{stage}_dataloader_start")(
+                trainer, pl_module, dataloader_idx
+            )
+
+    def on_epoch_end(
+        self, stage: str, trainer: Trainer, pl_module: LightningModule
+    ) -> None:
+        if self._dataloader_idxs[stage] != -1:
+            getattr(self, f"on_{stage}_dataloader_end")(
+                trainer, pl_module, self._dataloader_idxs[stage]
+            )
+            self._dataloader_idxs[stage] = -1
+
+    def on_validation_batch_start(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        self.on_batch_start("validation", trainer, pl_module, dataloader_idx)
+
+    def on_validation_epoch_end(
+        self, trainer: Trainer, pl_module: LightningModule
+    ) -> None:
+        self.on_epoch_end("validation", trainer, pl_module)
+
+    def on_test_batch_start(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        self.on_batch_start("test", trainer, pl_module, dataloader_idx)
+
+    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        self.on_epoch_end("test", trainer, pl_module)
+
+    def on_predict_batch_start(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        self.on_batch_start("predict", trainer, pl_module, dataloader_idx)
+
+    def on_predict_epoch_end(
+        self, trainer: Trainer, pl_module: LightningModule
+    ) -> None:
+        self.on_epoch_end("predict", trainer, pl_module)

tests/test_multiloader_lifecycle.py

@@ -0,0 +1,72 @@
+import pytest
+
+from ratiopath.lightning.callbacks import MultiloaderLifecycle
+
+
+class Recorder(MultiloaderLifecycle):
+    def __init__(self) -> None:
+        super().__init__()
+        self.events: list[tuple[str, str, int]] = []
+
+    def on_validation_dataloader_start(self, trainer, pl_module, dataloader_idx):
+        self.events.append(("validation", "start", dataloader_idx))
+
+    def on_validation_dataloader_end(self, trainer, pl_module, dataloader_idx):
+        self.events.append(("validation", "end", dataloader_idx))
+
+    def on_test_dataloader_start(self, trainer, pl_module, dataloader_idx):
+        self.events.append(("test", "start", dataloader_idx))
+
+    def on_test_dataloader_end(self, trainer, pl_module, dataloader_idx):
+        self.events.append(("test", "end", dataloader_idx))
+
+    def on_predict_dataloader_start(self, trainer, pl_module, dataloader_idx):
+        self.events.append(("predict", "start", dataloader_idx))
+
+    def on_predict_dataloader_end(self, trainer, pl_module, dataloader_idx):
+        self.events.append(("predict", "end", dataloader_idx))
+
+
+@pytest.mark.parametrize("stage", ["validation", "test", "predict"])
+def test_dataloader_hooks_follow_index_changes(stage):
+    callback = Recorder()
+    batch_start = getattr(callback, f"on_{stage}_batch_start")
+    epoch_end = getattr(callback, f"on_{stage}_epoch_end")
+
+    for batch_idx, dataloader_idx in enumerate([0, 0, 1, 1, 1, 2]):
+        batch_start(None, None, None, batch_idx, dataloader_idx=dataloader_idx)
+    epoch_end(None, None)
+
+    assert callback.events == [
+        (stage, "start", 0),
+        (stage, "end", 0),
+        (stage, "start", 1),
+        (stage, "end", 1),
+        (stage, "start", 2),
+        (stage, "end", 2),
+    ]
+
+
+def test_state_resets_between_epochs():
+    callback = Recorder()
+
+    for _ in range(2):
+        callback.on_validation_batch_start(None, None, None, 0, dataloader_idx=0)
+        callback.on_validation_epoch_end(None, None)
+
+    assert (
+        callback.events
+        == [
+            ("validation", "start", 0),
+            ("validation", "end", 0),
+        ]
+        * 2
+    )
+
+
+def test_epoch_end_without_batches_calls_nothing():
+    callback = Recorder()
+
+    callback.on_test_epoch_end(None, None)
+
+    assert callback.events == []