perf: short-circuit skip-existing reuse

Author: yousefkotp

atlas_patch/cli.py

@@ -449,46 +449,84 @@ def _run_pipeline_from_config(
 
     configure_logging(verbose)
 
-    from atlas_patch.orchestration.runner import ProcessingRunner
-    from atlas_patch.services.extraction import PatchExtractionService
-    from atlas_patch.services.feature_embedding import PatchFeatureEmbeddingService
-    from atlas_patch.services.mpp import CSVMPPResolver
-    from atlas_patch.services.segmentation import SAM2SegmentationService
-    from atlas_patch.services.visualization import DefaultVisualizationService
-    from atlas_patch.services.wsi_loader import DefaultWSILoader
+    from atlas_patch.orchestration.runner import ProcessingRunner, classify_existing_slide_output
+
+    if slides is None:
+        slides = [
+            Slide(path=Path(path))
+            for path in get_wsi_files(
+                str(app_cfg.processing.input_path),
+                recursive=app_cfg.processing.recursive,
+            )
+        ]
+    else:
+        slides = list(slides)
 
-    segmentation_service = SAM2SegmentationService(app_cfg.segmentation)
-    extractor_service = PatchExtractionService(app_cfg.extraction, app_cfg.output)
-    visualizer_service = None
-    if (
-        app_cfg.output.visualize_grids
-        or app_cfg.output.visualize_mask
-        or app_cfg.output.visualize_contours
-    ):
-        visualizer_service = DefaultVisualizationService(
-            app_cfg.output,
-            app_cfg.extraction,
-            app_cfg.visualization,
+    if not slides:
+        logging.getLogger("atlas_patch.runner").warning("No slides found to process.")
+        return [], []
+
+    preflight_results: list = []
+    slides_to_process: list[Slide] = []
+    for slide in slides:
+        decision, existing_result = classify_existing_slide_output(app_cfg, slide)
+        if decision == "skip":
+            continue
+        if decision == "reuse" and existing_result is not None:
+            preflight_results.append(existing_result)
+            continue
+        slides_to_process.append(slide)
+
+    results = list(preflight_results)
+    failures: list = []
+    wsi_loader = None
+
+    if slides_to_process:
+        from atlas_patch.services.extraction import PatchExtractionService
+        from atlas_patch.services.mpp import CSVMPPResolver
+        from atlas_patch.services.segmentation import SAM2SegmentationService
+        from atlas_patch.services.visualization import DefaultVisualizationService
+        from atlas_patch.services.wsi_loader import DefaultWSILoader
+
+        wsi_loader = DefaultWSILoader()
+        segmentation_service = SAM2SegmentationService(app_cfg.segmentation)
+        extractor_service = PatchExtractionService(app_cfg.extraction, app_cfg.output)
+        visualizer_service = None
+        if (
+            app_cfg.output.visualize_grids
+            or app_cfg.output.visualize_mask
+            or app_cfg.output.visualize_contours
+        ):
+            visualizer_service = DefaultVisualizationService(
+                app_cfg.output,
+                app_cfg.extraction,
+                app_cfg.visualization,
+            )
+
+        mpp_resolver = CSVMPPResolver(app_cfg.processing.mpp_csv)
+        runner = ProcessingRunner(
+            config=app_cfg,
+            segmentation=segmentation_service,
+            extractor=extractor_service,
+            visualizer=visualizer_service,
+            mpp_resolver=mpp_resolver,
+            wsi_loader=wsi_loader,
+            show_progress=not verbose,
         )
 
-    mpp_resolver = CSVMPPResolver(app_cfg.processing.mpp_csv)
-    wsi_loader = DefaultWSILoader()
-    runner = ProcessingRunner(
-        config=app_cfg,
-        segmentation=segmentation_service,
-        extractor=extractor_service,
-        visualizer=visualizer_service,
-        mpp_resolver=mpp_resolver,
-        wsi_loader=wsi_loader,
-        show_progress=not verbose,
-    )
+        try:
+            run_results, run_failures = runner.run(slides=slides_to_process)
+            results.extend(run_results)
+            failures.extend(run_failures)
+        finally:
+            segmentation_service.close()
 
-    try:
-        results, failures = runner.run(slides=slides)
-    finally:
-        segmentation_service.close()
+    if app_cfg.features is not None and results:
+        from atlas_patch.services.feature_embedding import PatchFeatureEmbeddingService
+        from atlas_patch.services.wsi_loader import DefaultWSILoader
 
-    if app_cfg.features is not None:
+        if wsi_loader is None:
+            wsi_loader = DefaultWSILoader()
         feature_service = PatchFeatureEmbeddingService(
             app_cfg.extraction,
             app_cfg.output,

atlas_patch/orchestration/runner.py

@@ -36,6 +36,78 @@ def _chunked(items: Sequence[Slide], size: int) -> Iterable[Sequence[Slide]]:
         yield items[i : i + size]
 
 
+def build_existing_extraction_result(slide: Slide, h5_path: Path) -> ExtractionResult | None:
+    """Create a lightweight ExtractionResult from an existing H5."""
+    metadata: dict[str, Any] = {}
+    num_patches: int | None = None
+    patch_size_level0: int | None = None
+    try:
+        with h5py.File(h5_path, "r") as f:
+            num_attr = f.attrs.get("num_patches")
+            if num_attr is not None:
+                num_patches = int(num_attr)
+            elif "coords" in f:
+                num_patches = int(f["coords"].shape[0])
+
+            ps_level0_attr = f.attrs.get("patch_size_level0")
+            if ps_level0_attr is not None:
+                patch_size_level0 = int(ps_level0_attr)
+    except Exception as e:  # noqa: BLE001
+        logger.warning(
+            "Failed to read existing output for %s; will reprocess. Error: %s",
+            slide.path.name,
+            e,
+        )
+        return None
+
+    if num_patches is None or num_patches <= 0:
+        return None
+
+    return ExtractionResult(
+        slide=slide,
+        h5_path=h5_path,
+        num_patches=int(num_patches),
+        patch_size_level0=patch_size_level0,
+        metadata=metadata,
+    )
+
+
+def classify_existing_slide_output(
+    config: AppConfig,
+    slide: Slide,
+) -> tuple[str | None, ExtractionResult | None]:
+    """Return how an existing patch H5 should be treated for this run.
+
+    Returns one of:
+    - (None, None): no reusable output; the slide needs full processing
+    - ("skip", None): output is fully complete for this run
+    - ("reuse", ExtractionResult): reuse patches/H5 and continue with downstream features
+    """
+    if not config.output.skip_existing:
+        return None, None
+
+    existing_path = find_existing_patch(slide, config.output, config.extraction)
+    if existing_path is None:
+        return None, None
+
+    feat_cfg = config.features
+    if feat_cfg is None or not feat_cfg.extractors:
+        return "skip", None
+
+    existing_result = build_existing_extraction_result(slide, existing_path)
+    if existing_result is None:
+        return None, None
+
+    missing = missing_features(
+        existing_path,
+        feat_cfg.extractors,
+        expected_total=existing_result.num_patches,
+    )
+    if not missing:
+        return "skip", None
+    return "reuse", existing_result
+
+
 class ProcessingRunner:
     """High-level orchestration of WSI segmentation, patch extraction, and visualization."""
 
@@ -68,41 +140,6 @@ def discover_slides(self) -> list[Slide]:
             slides.append(slide)
         return slides
 
-    def _build_existing_result(self, slide: Slide, h5_path: Path) -> ExtractionResult | None:
-        """Create a lightweight ExtractionResult from an existing H5 (no re-segmentation)."""
-        metadata: dict[str, Any] = {}
-        num_patches: int | None = None
-        patch_size_level0: int | None = None
-        try:
-            with h5py.File(h5_path, "r") as f:
-                num_attr = f.attrs.get("num_patches")
-                if num_attr is not None:
-                    num_patches = int(num_attr)
-                elif "coords" in f:
-                    num_patches = int(f["coords"].shape[0])
-
-                ps_level0_attr = f.attrs.get("patch_size_level0")
-                if ps_level0_attr is not None:
-                    patch_size_level0 = int(ps_level0_attr)
-        except Exception as e:  # noqa: BLE001
-            logger.warning(
-                "Failed to read existing output for %s; will reprocess. Error: %s",
-                slide.path.name,
-                e,
-            )
-            return None
-
-        if num_patches is None or num_patches <= 0:
-            return None
-
-        return ExtractionResult(
-            slide=slide,
-            h5_path=h5_path,
-            num_patches=int(num_patches),
-            patch_size_level0=patch_size_level0,
-            metadata=metadata,
-        )
-
     def _handle_existing_slide(
         self,
         slide: Slide,
@@ -113,35 +150,22 @@ def _handle_existing_slide(
 
         Returns True when the slide is fully handled (skip or reuse), False to continue processing.
         """
-        if not self.config.output.skip_existing:
-            return False
-
-        existing_path = find_existing_patch(slide, self.config.output, self.config.extraction)
-        if existing_path is None:
+        decision, existing_result = classify_existing_slide_output(self.config, slide)
+        if decision is None:
             return False
-
-        feat_cfg = self.config.features
-        if feat_cfg is None or not feat_cfg.extractors:
+        if decision == "skip":
             logger.info("Skipping %s (already processed).", slide.path.name)
             if progress:
                 progress.update(1)
             return True
-
-        existing_result = self._build_existing_result(slide, existing_path)
         if existing_result is None:
-            logger.info("Existing output invalid for %s; reprocessing.", slide.path.name)
             return False
-
+        results.append(existing_result)
         missing = missing_features(
-            existing_path, feat_cfg.extractors, expected_total=existing_result.num_patches
+            existing_result.h5_path,
+            self.config.features.extractors,
+            expected_total=existing_result.num_patches,
         )
-        if not missing:
-            logger.info("Skipping %s (features complete).", slide.path.name)
-            if progress:
-                progress.update(1)
-            return True
-
-        results.append(existing_result)
         logger.info(
             "Reusing existing patches for %s; missing features: %s",
             slide.path.name,

atlas_patch/services/feature_embedding.py

@@ -275,6 +275,8 @@ def embed_all(
             progress.update(completed_units)
 
         for name in self.extractor_names:
+            if not any(name in missing_for_slide for missing_for_slide in pending.values()):
+                continue
             try:
                 extractor = self.registry.create(name)
             except Exception as e:  # noqa: BLE001