Threading-based parallelism in the peelers

Author: cwindolf

src/dartsort/config.py

@@ -142,7 +142,7 @@ class FeaturizationConfig:
     tpca_fit_radius: float = 75.0
     tpca_rank: int = 8
     tpca_centered: bool = False
-    # todo: use a WaveformConfig...
+    learn_cleaned_tpca_basis: bool = False
     input_tpca_waveform_config: WaveformConfig | None = WaveformConfig(
         ms_before=0.75, ms_after=1.25
     )
@@ -343,9 +343,9 @@ class ClusteringConfig:
 class ComputationConfig:
     n_jobs_cpu: int = 0
     n_jobs_gpu: int = 0
-    device: torch.device | None = argfield(default=None, arg_type=torch.device)
+    executor: str = "ThreadPoolExecutor"
+    device: str | None = argfield(default=None, arg_type=str)
 
-    @property
     def actual_device(self):
         if self.device is None:
             have_cuda = torch.cuda.is_available()
@@ -354,9 +354,8 @@ def actual_device(self):
             return torch.device("cpu")
         return torch.device(self.device)
 
-    @property
-    def actual_n_jobs_gpu(self):
-        if self.actual_device.type == "cuda":
+    def actual_n_jobs(self):
+        if self.actual_device().type == "cuda":
             return self.n_jobs_gpu
         return self.n_jobs_cpu
 

src/dartsort/main.py

@@ -16,6 +16,7 @@
     default_subtraction_config,
     default_template_config,
     default_waveform_config,
+    default_computation_config,
 )
 from dartsort.peel import ObjectiveUpdateTemplateMatchingPeeler, SubtractionPeeler
 from dartsort.templates import TemplateData
@@ -143,12 +144,11 @@ def subtract(
     waveform_config=default_waveform_config,
     featurization_config=default_featurization_config,
     subtraction_config=default_subtraction_config,
+    computation_config=default_computation_config,
     chunk_starts_samples=None,
-    n_jobs=0,
     overwrite=False,
     residual_filename=None,
     show_progress=True,
-    device=None,
     hdf5_filename="subtraction.h5",
     model_subdir="subtraction_models",
 ):
@@ -163,14 +163,13 @@ def subtract(
         subtraction_peeler,
         output_directory,
         hdf5_filename,
-        model_subdir,
-        featurization_config,
+        model_subdir=model_subdir,
+        featurization_config=featurization_config,
         chunk_starts_samples=chunk_starts_samples,
         overwrite=overwrite,
-        n_jobs=n_jobs,
+        computation_config=computation_config,
         residual_filename=residual_filename,
         show_progress=show_progress,
-        device=device,
     )
     return detections, output_hdf5_filename
 

src/dartsort/peel/subtract.py

@@ -289,15 +289,15 @@ def precompute_peeler_models(self):
         self.subtraction_denoising_pipeline.precompute()
 
     def fit_featurization_pipeline(
-        self, save_folder, tmp_dir=None, n_jobs=0, device=None
+        self, save_folder, tmp_dir=None, computation_config=None
     ):
         super().fit_featurization_pipeline(
-            save_folder, tmp_dir=tmp_dir, n_jobs=n_jobs, device=device
+            save_folder, tmp_dir=tmp_dir, computation_config=computation_config
         )
         if self.use_singlechan_templates:
             self.have_singlechan_templates = True
 
-    def fit_peeler_models(self, save_folder, tmp_dir=None, n_jobs=0, device=None):
+    def fit_peeler_models(self, save_folder, tmp_dir=None, computation_config=None):
         # when fitting peelers for subtraction, there are basically
         # two cases. fitting featurizers is easy -- they don't modify
         # the waveforms. fitting denoisers is hard -- they do. each
@@ -309,24 +309,23 @@ def fit_peeler_models(self, save_folder, tmp_dir=None, n_jobs=0, device=None):
         # so we will cheat for now:
         # just remove all the denoisers that need fitting, run peeling,
         # and fit everything
+
         while self._fit_subtraction_transformers(
             save_folder,
             tmp_dir=tmp_dir,
-            n_jobs=n_jobs,
-            device=device,
+            computation_config=computation_config,
             which="denoisers",
         ):
             pass
         self._fit_subtraction_transformers(
             save_folder,
             tmp_dir=tmp_dir,
-            n_jobs=n_jobs,
-            device=device,
+            computation_config=computation_config,
             which="featurizers",
         )
 
     def _fit_subtraction_transformers(
-        self, save_folder, tmp_dir=None, n_jobs=0, device=None, which="denoisers"
+        self, save_folder, tmp_dir=None, computation_config=None, which="denoisers"
     ):
         """Fit models which are run during the subtraction step
 
@@ -352,9 +351,11 @@ def _fit_subtraction_transformers(
         if not needs_fit:
             return False
 
-        if device is None:
+        if computation_config is None:
             device = "cuda" if torch.cuda.is_available() else "cpu"
-        device = torch.device(device)
+            device = torch.device(device)
+        else:
+            device = computation_config.actual_device()
 
         orig_denoise = self.subtraction_denoising_pipeline
         init_voltage_feature = Voltage(
@@ -396,8 +397,7 @@ def _fit_subtraction_transformers(
             try:
                 self.run_subsampled_peeling(
                     temp_hdf5_filename,
-                    n_jobs=n_jobs,
-                    device=device,
+                    computation_config=computation_config,
                     task_name=f"Load examples for {which[:-1]} fitting",
                 )
 
@@ -563,6 +563,7 @@ def subtract_chunk(
             times_samples >= trough_offset_samples, times_samples < max_trough_time
         )
         (keep,) = keep.nonzero(as_tuple=True)
+
         if not keep.numel():
             break
         times_samples = times_samples[keep]

src/dartsort/transform/enforce_decrease.py

@@ -21,7 +21,7 @@ def __init__(
         self,
         channel_index,
         geom,
-        batch_size=32,
+        batch_size=1024,
         name=None,
         name_prefix="",
     ):

src/dartsort/transform/single_channel_denoiser.py

@@ -25,8 +25,8 @@ def __init__(
         channel_index,
         geom=None,
         denoiser=None,
-        batch_size=32,
-        in_place=True,
+        batch_size=1024,
+        in_place=False,
         pretrained_path=default_pretrained_path,
         name=None,
         name_prefix="",
@@ -41,8 +41,6 @@ def __init__(
 
         if denoiser is None:
             denoiser = dnclss[clsname]().load(pretrained_path)
-            denoiser.eval()
-            denoiser.requires_grad_(False)
         self.denoiser = denoiser
 
     def forward(self, waveforms, max_channels=None):
@@ -97,6 +95,8 @@ def forward(self, x):
     def load(self, pretrained_path=default_pretrained_path):
         checkpoint = torch.load(pretrained_path, map_location="cpu", weights_only=True)
         self.load_state_dict(checkpoint)
+        self.eval()
+        self.requires_grad_(False)
         return self
 
 

src/dartsort/util/multiprocessing_util.py

@@ -3,6 +3,7 @@
 from concurrent.futures import ThreadPoolExecutor as _ThreadPoolExecutor
 from multiprocessing import get_context
 
+import torch
 import torch.multiprocessing as torchmp
 
 # TODO: torch.multiprocessing?
@@ -125,6 +126,25 @@ def rank_init(queue):
     print(f"rank init got {rank_init.rank=}")
 
 
+def pool_from_cfg(computation_config=None, with_rank_queue=False, check_local=False):
+    if computation_config is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = torch.device(device)
+        n_jobs = 0
+        cls = MockPoolExecutor
+    else:
+        device = computation_config.actual_device()
+        n_jobs = computation_config.actual_n_jobs()
+        cls = computation_config.executor
+
+    return get_pool(
+        n_jobs,
+        cls=cls,
+        with_rank_queue=with_rank_queue,
+        check_local=check_local,
+    )
+
+
 def get_pool(
     n_jobs,
     context="spawn",
@@ -133,16 +153,30 @@ def get_pool(
     rank_queue_empty=False,
     n_tasks=None,
     max_tasks_per_child=None,
+    check_local=False,
 ):
     if n_jobs == -1:
         n_jobs = multiprocessing.cpu_count()
     do_parallel = n_jobs >= 1
     n_jobs = max(1, n_jobs)
 
+    if isinstance(cls, str):
+        if cls == "CloudpicklePoolExecutor":
+            cls = CloudpicklePoolExecutor
+        elif cls == "ThreadPoolExecutor":
+            cls = ThreadPoolExecutor
+        elif cls == "ProcessPoolExecutor":
+            cls = ProcessPoolExecutor
+        elif cls == "MockPoolExecutor":
+            cls = MockPoolExecutor
+        else:
+            assert False
+
     if cls == CloudpicklePoolExecutor and not have_cloudpickle:
         cls = ProcessPoolExecutor
 
     Executor = cls if do_parallel else MockPoolExecutor
+    is_local = cls in (MockPoolExecutor, ThreadPoolExecutor)
     if context == "torchspawn":
         context = torchmp.get_context("spawn")
     else:
@@ -163,6 +197,10 @@ def get_pool(
                 for rank in range(n_jobs):
                     rank_queue.put(rank)
 
+        if check_local:
+            return n_jobs, Executor, context, rank_queue, is_local
         return n_jobs, Executor, context, rank_queue
 
+    if check_local:
+        return n_jobs, Executor, context, is_local
     return n_jobs, Executor, context

src/dartsort/util/peel_util.py

@@ -14,12 +14,11 @@ def run_peeler(
     hdf5_filename,
     model_subdir,
     featurization_config,
+    computation_config,
     chunk_starts_samples=None,
     overwrite=False,
-    n_jobs=0,
     residual_filename=None,
     show_progress=True,
-    device=None,
     localization_dataset_name="point_source_localizations",
 ):
     output_directory = Path(output_directory)
@@ -49,20 +48,21 @@ def run_peeler(
 
     # fit models if needed
     peeler.load_or_fit_and_save_models(
-        model_dir, overwrite=overwrite, n_jobs=n_jobs, device=device
+        model_dir, overwrite=overwrite, computation_config=computation_config
     )
 
     # run main
     peeler.peel(
         output_hdf5_filename,
         chunk_starts_samples=chunk_starts_samples,
-        n_jobs=n_jobs,
         overwrite=overwrite,
         residual_filename=residual_filename,
         show_progress=show_progress,
-        device=device,
+        computation_config=computation_config,
+    )
+    _gc(
+        computation_config.actual_n_jobs(), computation_config.actual_device()
     )
-    _gc(n_jobs, device)
 
     # do localization
     if do_localization:
@@ -74,21 +74,22 @@ def run_peeler(
             amplitude_vectors_dataset_name=f"{wf_name}_{loc_amp_type}_amplitude_vectors",
             output_dataset_name=localization_dataset_name,
             show_progress=show_progress,
-            n_jobs=n_jobs,
-            device=device,
+            n_jobs=computation_config.actual_n_jobs(),
+            device=computation_config.actual_device(),
             localization_model=featurization_config.localization_model,
         )
-        _gc(n_jobs, device)
+        _gc(
+            computation_config.actual_n_jobs(), computation_config.actual_device()
+        )
 
     if featurization_config.n_residual_snips:
         peeler.run_subsampled_peeling(
             output_hdf5_filename,
-            n_jobs=n_jobs,
             chunk_length_samples=peeler.spike_length_samples,
             residual_to_h5=True,
             skip_features=True,
             ignore_resuming=True,
-            device=device,
+            computation_config=computation_config,
             n_chunks=featurization_config.n_residual_snips,
             task_name="Residual snips",
             overwrite=False,