Merge pull request #12 from The-Obstacle-Is-The-Way/fix/pipeline-resource-leak

src/stroke_deepisles_demo/pipeline.py

@@ -25,11 +25,15 @@ logger = get_logger(__name__)
 
 @dataclass(frozen=True)
 class PipelineResult:
-    """Complete result of running the pipeline on a case."""
+    """Complete result of running the pipeline on a case.
+
+    All file paths in this result point to valid, accessible files in results_dir.
+    Callers are responsible for cleaning up results_dir when done (if desired).
+    """
 
     case_id: str
-    input_files: CaseFiles
-    staged_dir: Path
+    input_files: CaseFiles  # Copied to results_dir; always valid paths
+    results_dir: Path  # Directory containing all result files (for cleanup)
     prediction_mask: Path
     ground_truth: Path | None
     dice_score: float | None  # None if ground truth unavailable or not computed
@@ -81,38 +85,64 @@ def run_pipeline_on_case(
 
     start_time = time.time()
 
-    # 1. Load Dataset
-    dataset = load_isles_dataset()  # Uses default local path for now
-
-    # Resolve ID if integer
-    if isinstance(case_id, int):
-        all_ids = dataset.list_case_ids()
-        if case_id < 0 or case_id >= len(all_ids):
-            raise IndexError(f"Case index {case_id} out of range (0-{len(all_ids) - 1})")
-        resolved_case_id = all_ids[case_id]
-    else:
-        resolved_case_id = case_id
-
-    # Get case files
-    case_files = dataset.get_case(resolved_case_id)
-
-    # 2. Stage Files
-    # Use a temp dir for staging if output_dir not provided, or a subdir of output_dir
-    if output_dir:
-        output_dir = Path(output_dir)
-        output_dir.mkdir(parents=True, exist_ok=True)
-        staging_root = output_dir / "staging" / resolved_case_id
-        results_dir = output_dir / resolved_case_id
-    else:
-        # If no output dir, we create a temp dir that persists (unless cleanup requested)
-        # But wait, the user wants paths. If we use tempfile.TemporaryDirectory context,
-        # it disappears. We should use mkdtemp or let stage_case handle it.
-        # Let's use a temp dir for staging.
-        base_temp = Path(tempfile.mkdtemp(prefix="deepisles_pipeline_"))
-        staging_root = base_temp / "staging"
-        results_dir = base_temp / "results"
-
-    staged = stage_case_for_deepisles(case_files, staging_root)
+    # Use context manager to ensure HuggingFace temp files are cleaned up
+    # This prevents unbounded disk usage from accumulating temp NIfTI files
+    with load_isles_dataset() as dataset:
+        # Resolve ID if integer
+        if isinstance(case_id, int):
+            all_ids = dataset.list_case_ids()
+            if case_id < 0 or case_id >= len(all_ids):
+                raise IndexError(f"Case index {case_id} out of range (0-{len(all_ids) - 1})")
+            resolved_case_id = all_ids[case_id]
+        else:
+            resolved_case_id = case_id
+
+        # Set up output directories (now that we have resolved_case_id)
+        if output_dir:
+            output_dir = Path(output_dir)
+            output_dir.mkdir(parents=True, exist_ok=True)
+            staging_root = output_dir / "staging" / resolved_case_id
+            results_dir = output_dir / resolved_case_id
+        else:
+            base_temp = Path(tempfile.mkdtemp(prefix="deepisles_pipeline_"))
+            staging_root = base_temp / "staging"
+            results_dir = base_temp / "results"
+
+        # Get case files
+        case_files = dataset.get_case(resolved_case_id)
+
+        # Stage files (copies DWI/ADC to staging directory)
+        staged = stage_case_for_deepisles(case_files, staging_root)
+
+        # Copy input files to results_dir before dataset cleanup
+        # (HuggingFace mode stores files in temp dirs that get cleaned up)
+        # This ensures all paths in PipelineResult remain valid after function returns
+        results_dir.mkdir(parents=True, exist_ok=True)
+
+        # Copy DWI (required for UI visualization)
+        dwi_dest = results_dir / f"{resolved_case_id}_dwi.nii.gz"
+        shutil.copy2(case_files["dwi"], dwi_dest)
+
+        # Copy ADC
+        adc_dest = results_dir / f"{resolved_case_id}_adc.nii.gz"
+        shutil.copy2(case_files["adc"], adc_dest)
+
+        # Copy ground truth if available
+        ground_truth: Path | None = None
+        original_ground_truth = case_files.get("ground_truth")
+        if original_ground_truth and original_ground_truth.exists():
+            ground_truth = results_dir / f"{resolved_case_id}_ground_truth.nii.gz"
+            shutil.copy2(original_ground_truth, ground_truth)
+
+        # Build input_files with copied paths (always valid after function returns)
+        preserved_input_files: CaseFiles = {
+            "dwi": dwi_dest,
+            "adc": adc_dest,
+        }
+        if ground_truth:
+            preserved_input_files["ground_truth"] = ground_truth
+
+    # Dataset temp files cleaned up here (context manager __exit__)
 
     # 3. Run Inference
     inference_result = run_deepisles_on_folder(
@@ -122,10 +152,8 @@ def run_pipeline_on_case(
         gpu=gpu,
     )
 
-    # 4. Compute Metrics
+    # 4. Compute Metrics (using copied ground truth)
     dice_score: float | None = None
-    ground_truth = case_files.get("ground_truth")
-
     if compute_dice and ground_truth and ground_truth.exists():
         try:
             dice_score = metrics.compute_dice(inference_result.prediction_path, ground_truth)
@@ -140,8 +168,8 @@ def run_pipeline_on_case(
 
     return PipelineResult(
         case_id=resolved_case_id,
-        input_files=case_files,
-        staged_dir=staged.input_dir,
+        input_files=preserved_input_files,
+        results_dir=results_dir,
         prediction_mask=inference_result.prediction_path,
         ground_truth=ground_truth,
         dice_score=dice_score,

src/stroke_deepisles_demo/ui/app.py

@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from typing import Any
+import shutil
+from typing import TYPE_CHECKING, Any
 
 import gradio as gr
 from matplotlib.figure import Figure  # noqa: TC002
@@ -20,8 +21,14 @@ from stroke_deepisles_demo.ui.viewer import (
     render_slice_comparison,
 )
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
 logger = get_logger(__name__)
 
+# Shared output directory for UI results (cleaned up between runs to prevent disk accumulation)
+_previous_results_dir: Path | None = None
+
 
 def run_segmentation(
     case_id: str, fast_mode: bool, show_ground_truth: bool
@@ -47,6 +54,13 @@ def run_segmentation(
         )
 
     try:
+        global _previous_results_dir
+
+        # Clean up previous results to prevent disk accumulation on HF Spaces
+        if _previous_results_dir and _previous_results_dir.exists():
+            shutil.rmtree(_previous_results_dir, ignore_errors=True)
+            logger.debug("Cleaned up previous results: %s", _previous_results_dir)
+
         logger.info("Running segmentation for %s", case_id)
         result = run_pipeline_on_case(
             case_id,
@@ -55,6 +69,9 @@ def run_segmentation(
             cleanup_staging=True,
         )
 
+        # Track results_dir for cleanup on next run
+        _previous_results_dir = result.results_dir
+
         # 1. NiiVue Visualization
         # We need data URLs for the browser
         # Note: This reads the file content into memory (base64)

tests/test_cli.py

@@ -26,7 +26,7 @@ class TestCli:
         result = PipelineResult(
             case_id="sub-001",
             input_files=MagicMock(),
-            staged_dir=MagicMock(),
+            results_dir=MagicMock(),
             prediction_mask=MagicMock(),
             ground_truth=None,
             dice_score=None,
@@ -50,7 +50,7 @@ class TestCli:
         result = PipelineResult(
             case_id="sub-001",
             input_files=MagicMock(),
-            staged_dir=MagicMock(),
+            results_dir=MagicMock(),
             prediction_mask=MagicMock(),
             ground_truth=None,
             dice_score=None,

tests/test_pipeline.py

@@ -35,21 +35,23 @@ class TestRunPipelineOnCase:
             # Configure mocks
             mock_dataset = MagicMock()
 
-            # Mock paths that "exist"
-            dwi_path = MagicMock(spec=Path)
-            dwi_path.exists.return_value = True
-            adc_path = MagicMock(spec=Path)
-            adc_path.exists.return_value = True
-            gt_path = MagicMock(spec=Path)
-            gt_path.exists.return_value = True
+            # Create real temp files (pipeline copies these to results_dir)
+            dwi_file = temp_dir / "dwi_mock.nii.gz"
+            dwi_file.write_bytes(b"fake dwi nifti")
+            adc_file = temp_dir / "adc_mock.nii.gz"
+            adc_file.write_bytes(b"fake adc nifti")
+            gt_file = temp_dir / "gt_mock.nii.gz"
+            gt_file.write_bytes(b"fake gt nifti")
 
             mock_dataset.get_case.return_value = CaseFiles(
-                dwi=dwi_path,
-                adc=adc_path,
-                ground_truth=gt_path,
+                dwi=dwi_file,
+                adc=adc_file,
+                ground_truth=gt_file,
                 # flair omitted
             )
-            mock_load.return_value = mock_dataset
+            # Support context manager protocol: with load_isles_dataset() as dataset:
+            mock_load.return_value.__enter__ = MagicMock(return_value=mock_dataset)
+            mock_load.return_value.__exit__ = MagicMock(return_value=None)
 
             mock_stage.return_value = MagicMock(
                 input_dir=temp_dir / "staged",
@@ -142,15 +144,18 @@ class TestRunPipelineOnCase:
     def test_handles_missing_ground_truth(
         self,
         mock_dependencies: dict[str, MagicMock],
-        temp_dir: Path,  # noqa: ARG002
+        temp_dir: Path,
     ) -> None:
         """Handles cases without ground truth gracefully."""
-        # Modify mock to return no ground truth
-        dwi = MagicMock(spec=Path)
-        adc = MagicMock(spec=Path)
+        # Create real files for DWI/ADC (pipeline copies these)
+        dwi_file = temp_dir / "dwi_no_gt.nii.gz"
+        dwi_file.write_bytes(b"fake dwi")
+        adc_file = temp_dir / "adc_no_gt.nii.gz"
+        adc_file.write_bytes(b"fake adc")
+
         mock_dependencies["dataset"].get_case.return_value = CaseFiles(
-            dwi=dwi,
-            adc=adc,
+            dwi=dwi_file,
+            adc=adc_file,
             # ground_truth omitted
         )
 
@@ -231,7 +236,7 @@ class TestRunPipelineOnBatch:
                 PipelineResult(
                     case_id="sub-001",
                     input_files=MagicMock(),
-                    staged_dir=MagicMock(),
+                    results_dir=MagicMock(),
                     prediction_mask=MagicMock(),
                     ground_truth=None,
                     dice_score=0.8,
@@ -240,7 +245,7 @@ class TestRunPipelineOnBatch:
                 PipelineResult(
                     case_id="sub-002",
                     input_files=MagicMock(),
-                    staged_dir=MagicMock(),
+                    results_dir=MagicMock(),
                     prediction_mask=MagicMock(),
                     ground_truth=None,
                     dice_score=0.9,
@@ -261,7 +266,7 @@ class TestRunPipelineOnBatch:
             mock_run.return_value = PipelineResult(
                 case_id="sub-001",
                 input_files=MagicMock(),
-                staged_dir=MagicMock(),
+                results_dir=MagicMock(),
                 prediction_mask=MagicMock(),
                 ground_truth=None,
                 dice_score=0.8,

tests/test_pipeline_cleanup.py

@@ -4,8 +4,13 @@ from unittest.mock import MagicMock, patch
 from stroke_deepisles_demo.pipeline import run_pipeline_on_case
 
 
-def test_pipeline_cleanup_default() -> None:
+def test_pipeline_cleanup_default(temp_dir: Path) -> None:
     """Test that pipeline cleans up staging directory by default."""
+    # Create real files (pipeline now copies input files to results_dir)
+    dwi_file = temp_dir / "dwi.nii.gz"
+    dwi_file.write_bytes(b"fake dwi")
+    adc_file = temp_dir / "adc.nii.gz"
+    adc_file.write_bytes(b"fake adc")
 
     # Mock everything to avoid running actual heavy inference
     with (
@@ -13,13 +18,17 @@ def test_pipeline_cleanup_default() -> None:
         patch("stroke_deepisles_demo.pipeline.stage_case_for_deepisles") as mock_stage,
         patch("stroke_deepisles_demo.pipeline.run_deepisles_on_folder") as mock_run,
         patch("stroke_deepisles_demo.pipeline.metrics.compute_dice"),
-        patch("shutil.rmtree") as mock_rmtree,
+        patch("stroke_deepisles_demo.pipeline.shutil.rmtree") as mock_rmtree,
     ):
         # Setup mocks
         mock_dataset = MagicMock()
-        mock_load.return_value = mock_dataset
         mock_dataset.list_case_ids.return_value = ["case1"]
-        mock_dataset.get_case.return_value = {"dwi": Path("dwi.nii.gz")}
+        # Return dict with real files (no ground_truth)
+        mock_dataset.get_case.return_value = {"dwi": dwi_file, "adc": adc_file}
+
+        # Support context manager protocol: with load_isles_dataset() as dataset:
+        mock_load.return_value.__enter__ = MagicMock(return_value=mock_dataset)
+        mock_load.return_value.__exit__ = MagicMock(return_value=None)
 
         mock_staged = MagicMock()
         mock_staged.input_dir = Path("/tmp/mock_staging")
@@ -32,7 +41,7 @@ def test_pipeline_cleanup_default() -> None:
         # Run pipeline with defaults (cleanup_staging=True is the default)
         run_pipeline_on_case("case1")
 
-        # Verify that rmtree was called
+        # Verify that rmtree was called (for staging cleanup)
         assert mock_rmtree.called
 
         # Get the path passed to stage_case_for_deepisles

tests/ui/test_app.py

@@ -51,7 +51,7 @@ def test_run_segmentation_logic() -> None:
     mock_result = PipelineResult(
         case_id="sub-001",
         input_files={"dwi": MagicMock(), "adc": MagicMock()},
-        staged_dir=MagicMock(),
+        results_dir=MagicMock(),
         prediction_mask=MagicMock(),
         ground_truth=MagicMock(),
         dice_score=0.85,