Fix multiprocessing deadlock and parallelize directory walk in 01_scanDicom.py

code/preprocessing/01_scanDicom.py

@@ -208,8 +208,9 @@ def _has_dcm_magic(path: str) -> bool:
 # Pipeline functions
 # ---------------------------------------------------------------------------
 
-def _extractDicom_impl(f: str, logger: logging.Logger, slice_counts: Dict[str, int] = None) -> Optional[Dict[str, Any]]:
+def _extractDicom_impl(f: str, slice_counts: Dict[str, int] = None) -> Optional[Dict[str, Any]]:
     """Extract DICOM information from a specific file path."""
+    logger = logging.getLogger('01_scanDicom')
     try:
         logger.debug(f'Extracting information for file: {f}')
         directory = os.path.dirname(f)
@@ -391,29 +392,53 @@ def _scan_subdir_worker(subdir: str, sample_pct: float, sample_seed: Optional[in
     """Worker for multiprocessing directory scanning.
 
     Calls `_find_dicom_worker` on its assigned single top-level subdirectory.
-    Creates its own logger inside the child process to avoid pickle-lock deadlock.
     Returns list of (dicom_files, slice_counts).
     """
-    wlogger = logging.getLogger(__name__ + '.worker')
-    return [_find_dicom_worker(subdir, sample_pct, sample_seed, wlogger)]
+    return [_find_dicom_worker(subdir, sample_pct, sample_seed)]
 
 
-def _scan_subdir(topdir: str):
-    """Return a list of subdirectories containing .dcm files."""
-    dirs_with_dcm = []
-    for root, _, files in os.walk(topdir, followlinks=False):
-        if any(f.lower().endswith('.dcm') for f in files):
-            dirs_with_dcm.append(root)
-    return dirs_with_dcm
+def _scan_subdir(topdir: str, min_targets: int = 16):
+    """Return a list of disjoint subdirectories that cover the entire tree.
+    We gather directories using BFS until we have enough targets. If a directory contains
+    files directly, we stop expanding it to keep subtrees disjoint for the os.walk workers."""
+    dirs_to_scan = []
+    queue = [topdir]
+
+    while queue and (len(dirs_to_scan) + len(queue)) < min_targets:
+        curr = queue.pop(0)
+        try:
+            with os.scandir(curr) as it:
+                subdirs = []
+                has_files = False
+                for entry in it:
+                    if entry.is_dir(follow_symlinks=False):
+                        subdirs.append(entry.path)
+                    elif entry.is_file():
+                        if entry.name.lower().endswith('.dcm'):
+                            has_files = True
+
+                if has_files:
+                    dirs_to_scan.append(curr)
+                else:
+                    queue.extend(subdirs)
+        except Exception:
+            pass
+
+    dirs_to_scan.extend(queue)
+
+    if not dirs_to_scan:
+        dirs_to_scan = [topdir]
+
+    return dirs_to_scan
 
 
-def _find_dicom_worker(directory: str, sample_pct: float, sample_seed: Optional[int],
-                       logger: logging.Logger) -> tuple:
+def _find_dicom_worker(directory: str, sample_pct: float, sample_seed: Optional[int]) -> tuple:
     """Worker for findDicom — called per directory, accepts only plain args.
 
     Returns:
         (dicom_files, slice_counts)
     """
+    logger = logging.getLogger('01_scanDicom')
     dicom_files = []
     slice_counts = {}
 
@@ -524,20 +549,20 @@ def main(cfg: ScanConfig, logger: logging.Logger, out_name: str = 'Data_table.cs
 
     if combined_result is None:
         if cfg.parallel:
-            # Parallel scan: walk the tree once to find subdirectories with .dcm files,
-            # then dispatch multiprocessing workers across them.
-            dirs_with_dcm = _scan_subdir(scan_dir)
-            logger.info(f'Found {len(dirs_with_dcm)} directories with DICOM files to scan')
+            # Parallel scan: we parallelize the walk by getting immediate subdirectories
+            # and dispatching multiprocessing workers to walk those subtrees independently.
+            target_dirs = _scan_subdir(scan_dir, min_targets=cfg.n_cpus * 4)
+            logger.info(f'Found {len(target_dirs)} branch directories to scan in parallel')
 
-            if len(dirs_with_dcm) > 1:
+            if len(target_dirs) > 1:
                 worker_results = run_function(
-                    logger, _scan_subdir_worker, dirs_with_dcm,
+                    logger, _scan_subdir_worker, target_dirs,
                     Parallel=True, P_type='hybrid', N_CPUS=cfg.n_cpus,
                     sample_pct=cfg.sample_pct, sample_seed=cfg.sample_seed,
                 )
             else:
                 worker_results = [run_function(
-                    logger, _scan_subdir_worker, dirs_with_dcm,
+                    logger, _scan_subdir_worker, target_dirs,
                     Parallel=False, P_type='thread', N_CPUS=1,
                     sample_pct=cfg.sample_pct, sample_seed=cfg.sample_seed,
                 )]
@@ -580,7 +605,7 @@ def main(cfg: ScanConfig, logger: logging.Logger, out_name: str = 'Data_table.cs
         worker_results = run_function(
             logger, _find_dicom_worker, dicom_dirs,
             Parallel=cfg.parallel, P_type='hybrid', N_CPUS=cfg.n_cpus,
-            sample_pct=cfg.sample_pct, sample_seed=cfg.sample_seed, logger=logger,
+            sample_pct=cfg.sample_pct, sample_seed=cfg.sample_seed,
         )
         dicom_files = [f for files, _ in worker_results for f in files]
         slice_counts = {}
@@ -605,7 +630,7 @@ def main(cfg: ScanConfig, logger: logging.Logger, out_name: str = 'Data_table.cs
             info_list = None
 
     if info_list is None:
-        extract_partial = partial(_extractDicom_impl, logger=logger, slice_counts=slice_counts)
+        extract_partial = partial(_extractDicom_impl, slice_counts=slice_counts)
         info_list = run_function(
             logger, extract_partial, dicom_files,
             Parallel=cfg.parallel, P_type='hybrid', N_CPUS=cfg.n_cpus,

test/test_scanDicom_full.py

@@ -232,7 +232,7 @@ def test_A4_missing_series_number_no_crash(tmp_path):
     d.mkdir()
     make_realistic_mr_dcm(str(d / "ns.dcm"), modality='MR', series_number=1)
     logger = _scan_logger()
-    found_files, _ = scan._find_dicom_worker(str(d), sample_pct=0.0, sample_seed=None, logger=logger)
+    found_files, _ = scan._find_dicom_worker(str(d), sample_pct=0.0, sample_seed=None)
     assert isinstance(found_files, list)
 
 
@@ -243,7 +243,7 @@ def test_A5_duplicate_series_returns_one(tmp_path):
     for i in range(5):
         make_minimal_dcm(str(root / f"dup_{i}.dcm"), modality='MR', series_number=42)
     logger = _scan_logger()
-    found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None, logger=logger)
+    found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None)
     assert len(found_files) == 1
 
 
@@ -256,7 +256,7 @@ def test_A6_corrupt_files(tmp_path):
     (d / "bad2.dcm").write_bytes(b'\xff' * 512)
     (d / "bad3.dcm").write_bytes(b'\0' * 100)
     logger = _scan_logger()
-    found_files, _ = scan._find_dicom_worker(str(d), sample_pct=0.0, sample_seed=None, logger=logger)
+    found_files, _ = scan._find_dicom_worker(str(d), sample_pct=0.0, sample_seed=None)
     assert len(found_files) == 1
     assert "good.dcm" in found_files[0]
 
@@ -278,8 +278,8 @@ def test_A8_sampling_deterministic(tmp_path):
     for i in range(20):
         make_minimal_dcm(str(root / f"f_{i:02d}.dcm"), modality='MR', series_number=(i % 5) + 1)
     logger = _scan_logger()
-    first = scan._find_dicom_worker(str(root), sample_pct=15.0, sample_seed=99, logger=logger)
-    second = scan._find_dicom_worker(str(root), sample_pct=15.0, sample_seed=99, logger=logger)
+    first = scan._find_dicom_worker(str(root), sample_pct=15.0, sample_seed=99)
+    second = scan._find_dicom_worker(str(root), sample_pct=15.0, sample_seed=99)
     assert first == second
 
 
@@ -320,7 +320,7 @@ def test_B1_extractDicom_has_all_keys(tmp_path):
     f = tmp_path / "extract_test.dcm"
     make_realistic_mr_dcm(str(f), repetition_time=500.0)
     logger = _scan_logger()
-    result = scan._extractDicom_impl(str(f), logger)
+    result = scan._extractDicom_impl(str(f))
     assert result is not None
     assert isinstance(result, dict)
     assert EXPECTED_KEYS.issubset(result.keys()), f"Missing keys: {EXPECTED_KEYS - result.keys()}"
@@ -331,21 +331,21 @@ def test_B2_T1_vs_T2_modality(tmp_path):
     logger = _scan_logger()
     t1_path = tmp_path / "t1.dcm"
     make_realistic_mr_dcm(str(t1_path), repetition_time=779.0)
-    t1_result = scan._extractDicom_impl(str(t1_path), logger)
+    t1_result = scan._extractDicom_impl(str(t1_path))
     assert t1_result['Modality'] == 'T1', f"Expected T1, got {t1_result['Modality']}"
 
     t2_path = tmp_path / "t2.dcm"
     make_realistic_mr_dcm(str(t2_path), repetition_time=780.0)
-    t2_result = scan._extractDicom_impl(str(t2_path), logger)
+    t2_result = scan._extractDicom_impl(str(t2_path))
     assert t2_result['Modality'] == 'T2', f"Expected T2, got {t2_result['Modality']}"
 
     t1_edge = tmp_path / "t1_edge.dcm"
     make_realistic_mr_dcm(str(t1_edge), repetition_time=779.999)
-    assert scan._extractDicom_impl(str(t1_edge), logger)['Modality'] == 'T1'
+    assert scan._extractDicom_impl(str(t1_edge))['Modality'] == 'T1'
 
     t2_edge = tmp_path / "t2_edge.dcm"
     make_realistic_mr_dcm(str(t2_edge), repetition_time=780.001)
-    assert scan._extractDicom_impl(str(t2_edge), logger)['Modality'] == 'T2'
+    assert scan._extractDicom_impl(str(t2_edge))['Modality'] == 'T2'
 
 
 # B3 — Unknown fields for missing tags
@@ -354,7 +354,7 @@ def test_B3_unknown_fields_missing_tags(tmp_path):
     d.mkdir()
     make_minimal_dcm(str(d / "sparse.dcm"), modality='MR', series_number=1)
     logger = _scan_logger()
-    result = scan._extractDicom_impl(str(d / "sparse.dcm"), logger)
+    result = scan._extractDicom_impl(str(d / "sparse.dcm"))
     assert result is not None
     for key in ['Accession', 'DOB', 'Lat']:
         assert result[key] == 'Unknown', f"{key} should be 'Unknown' but is '{result[key]}'"

test/test_scanDicom_unit.py

@@ -95,15 +95,15 @@ def test_findDicom_series(tmp_path):
     make_minimal_dcm(str(root / "b.dcm"), modality='MR', series_number=2)
     make_minimal_dcm(str(root / "c.dcm"), modality='CT', series_number=3)
     logger = _make_logger()
-    found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None, logger=logger)
+    found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None)
     assert any("a.dcm" in f or "b.dcm" in f for f in found_files)
 
 
 def test_extractDicom_basic(tmp_path):
     f = tmp_path / "x.dcm"
     make_minimal_dcm(str(f), modality='MR', series_number=5, patient_id='P1')
     logger = _make_logger()
-    out = scan._extractDicom_impl(str(f), logger)
+    out = scan._extractDicom_impl(str(f))
     assert isinstance(out, dict)
     assert isinstance(out['Modality'], str)
 
@@ -125,7 +125,7 @@ def test_findDicom_handles_unreadable_and_returns_mr_only(tmp_path):
     make_minimal_dcm(str(root / "mri.dcm"), modality='MR', series_number=10)
     (root / "garbage.dcm").write_text("corrupt")
     logger = _make_logger()
-    found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None, logger=logger)
+    found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None)
     assert any("mri.dcm" in f for f in found_files)
 
 
@@ -137,6 +137,6 @@ def test_findDicom_sampling_is_deterministic_with_seed(tmp_path):
         make_minimal_dcm(str(root / f"img_{i}.dcm"), modality='MR', series_number=series)
 
     logger = _make_logger()
-    first = scan._find_dicom_worker(str(root), sample_pct=20.0, sample_seed=123, logger=logger)
-    second = scan._find_dicom_worker(str(root), sample_pct=20.0, sample_seed=123, logger=logger)
+    first = scan._find_dicom_worker(str(root), sample_pct=20.0, sample_seed=123)
+    second = scan._find_dicom_worker(str(root), sample_pct=20.0, sample_seed=123)
     assert first == second