fix: replace ProcessPoolExecutor with ThreadPoolExecutor (#552)

waveorder/calib/calibration_workers.py

@@ -334,7 +334,7 @@ def work(self):
             transfer_function_dirpath=transfer_function_path,
             config_filepath=reconstruction_config_path,
             output_dirpath=reconstruction_path,
-            num_processes=1,
+            num_threads=1,
         )
 
         # Load reconstructions from file for layers

waveorder/cli/apply_inverse_transfer_function.py

@@ -9,7 +9,7 @@
     config_filepath,
     input_position_dirpaths,
     output_dirpath,
-    processes_option,
+    threads_option,
     transfer_function_dirpath,
     write_config_scale_to_output,
 )
@@ -203,17 +203,16 @@ def apply_inverse_transfer_function_single_position(
     transfer_function_dirpath: Path,
     config_filepath: Path,
     output_position_dirpath: Path,
-    num_processes,
+    num_threads,
     output_channel_names: list[str],
     verbose: bool = True,
 ) -> None:
 
     # Deferred imports for fast CLI help
-    from concurrent.futures import ProcessPoolExecutor, as_completed
+    from concurrent.futures import ThreadPoolExecutor, as_completed
     from functools import partial
 
     import numpy as np
-    import torch.multiprocessing as mp
     from iohub import open_ome_zarr
 
     from waveorder.api import (
@@ -356,18 +355,11 @@ def apply_inverse_transfer_function_single_position(
         **apply_inverse_args,
     )
 
-    # Multiprocessing logic
-    if num_processes > 1:
+    # Threading logic
+    if num_threads > 1:
         if verbose:
-            click.echo(f"\nStarting multiprocess pool with {num_processes} processes")
-        # NOTE: spawn (not fork) — tensorstore runs internal C++ threads
-        # that are not fork-safe, so a forked worker can deadlock or
-        # segfault before our code runs. See google/tensorstore#61.
-        # NOTE: ProcessPoolExecutor (not mp.Pool) so silent worker death
-        # (e.g. cgroup OOM-kill) surfaces as BrokenProcessPool instead
-        # of hanging indefinitely on pool.starmap.
-        context = mp.get_context("spawn")
-        with ProcessPoolExecutor(max_workers=num_processes, mp_context=context) as p:
+            click.echo(f"\nStarting thread pool with {num_threads} threads")
+        with ThreadPoolExecutor(max_workers=num_threads) as p:
             futures = [p.submit(partial_apply_inverse_to_zyx_and_save, t_idx) for t_idx in time_indices]
             for fut in as_completed(futures):
                 fut.result()
@@ -393,7 +385,7 @@ def apply_inverse_transfer_function_cli(
     transfer_function_dirpath: Path,
     config_filepath: Path,
     output_dirpath: Path,
-    num_processes,
+    num_threads,
     write_config_scale_to_output: bool = False,
 ) -> None:
     # Deferred imports for fast CLI help
@@ -437,10 +429,6 @@ def apply_inverse_transfer_function_cli(
         with open_ome_zarr(str(output_dirpath), mode="r+") as output_plate:
             output_plate.zattrs.update(plate_metadata)
 
-    # Initialize torch threads
-    if num_processes > 1:
-        torch.set_num_threads(1)
-        torch.set_num_interop_threads(1)
 
     # Loop through positions
     for i, input_position_dirpath in enumerate(input_position_dirpaths):
@@ -457,7 +445,7 @@ def apply_inverse_transfer_function_cli(
             transfer_function_dirpath,
             config_filepath,
             output_position_path,
-            num_processes,
+            num_threads,
             output_metadata["channel_names"],
         )
 
@@ -467,14 +455,14 @@ def apply_inverse_transfer_function_cli(
 @transfer_function_dirpath()
 @config_filepath()
 @output_dirpath()
-@processes_option(default=1)
+@threads_option(default=1)
 @write_config_scale_to_output()
 def _apply_inverse_transfer_function_cli(
     input_position_dirpaths: list[Path],
     transfer_function_dirpath: Path,
     config_filepath: Path,
     output_dirpath: Path,
-    num_processes,
+    num_threads,
     write_config_scale_to_output: bool,
 ) -> None:
     """Apply an inverse transfer function to a dataset.
@@ -491,6 +479,6 @@ def _apply_inverse_transfer_function_cli(
         transfer_function_dirpath,
         config_filepath,
         output_dirpath,
-        num_processes,
+        num_threads,
         write_config_scale_to_output,
     )

waveorder/cli/parsing.py

@@ -87,30 +87,61 @@ def decorator(f: Callable) -> Callable:
     return decorator
 
 
-# TODO: this setting will have to be collected from SLURM?
-def processes_option(default: int = None) -> Callable:
-    def check_processes_option(ctx, param, value):
-        # Deferred: torch.multiprocessing pulls in all of torch.
-        import torch.multiprocessing as mp
-
-        max_processes = mp.cpu_count()
-        if value > max_processes:
-            raise click.BadParameter(f"Maximum number of processes is {max_processes}")
+def threads_option(default: int = None) -> Callable:
+    """CLI option for number of threads to run in parallel.
+
+    Accepts --num-threads (canonical) or --num-processes / --num_processes
+    (deprecated aliases, kept for backward compatibility with existing scripts).
+    """
+    import os
+
+    def check_threads_option(ctx, param, value):
+        if value is None:
+            return default or 1
+        max_threads = os.cpu_count() or 1
+        if value > max_threads:
+            raise click.BadParameter(f"Maximum number of threads is {max_threads}")
+        return value
+
+    def deprecated_processes_callback(ctx, param, value):
+        if value is not None:
+            click.echo(
+                "Warning: --num-processes / --num_processes is deprecated. "
+                "Use --num-threads instead.",
+                err=True,
+            )
+            ctx.params["num_threads"] = check_threads_option(ctx, param, value)
         return value
 
     def decorator(f: Callable) -> Callable:
-        return click.option(
-            "--num_processes",
+        f = click.option(
+            "--num-threads",
             "-j",
             default=default or 1,
             type=int,
-            help="Number of processes to run in parallel.",
-            callback=check_processes_option,
+            help="Number of threads to run in parallel.",
+            callback=check_threads_option,
         )(f)
+        f = click.option(
+            "--num-processes",
+            "--num_processes",
+            default=None,
+            type=int,
+            is_eager=True,
+            expose_value=False,
+            hidden=True,
+            help="Deprecated: use --num-threads instead.",
+            callback=deprecated_processes_callback,
+        )(f)
+        return f
 
     return decorator
 
 
+# Keep old name as alias for backward compatibility with any direct Python imports
+processes_option = threads_option
+
+
 def write_config_scale_to_output() -> Callable:
     def decorator(f: Callable) -> Callable:
         return click.option(

waveorder/cli/reconstruct.py

@@ -6,7 +6,7 @@
     config_filepath,
     input_position_dirpaths,
     output_dirpath,
-    processes_option,
+    threads_option,
     unique_id,
     write_config_scale_to_output,
 )
@@ -16,14 +16,14 @@
 @input_position_dirpaths()
 @config_filepath()
 @output_dirpath()
-@processes_option(default=1)
+@threads_option(default=1)
 @unique_id()
 @write_config_scale_to_output()
 def _reconstruct_cli(
     input_position_dirpaths,
     config_filepath,
     output_dirpath,
-    num_processes,
+    num_threads,
     unique_id,
     write_config_scale_to_output,
 ):
@@ -78,7 +78,7 @@ def _reconstruct_cli(
         transfer_function_path,
         config_filepath,
         output_dirpath,
-        num_processes,
+        num_threads,
         write_config_scale_to_output,
     )
 

waveorder/cli/utils.py

@@ -105,7 +105,7 @@ def apply_inverse_to_zyx_and_save(
         click.echo(f"Finished writing t={t_idx}")
 
 
-def estimate_resources(shape, settings, num_processes):
+def estimate_resources(shape, settings, num_threads):
     T, C, Z, Y, X = shape
 
     gb_ram_per_cpu = 0
@@ -122,7 +122,7 @@ def estimate_resources(shape, settings, num_processes):
         gb_ram_per_cpu += input_memory * fourier_resource_multiplier
     ram_multiplier = 1
     gb_ram_per_cpu = np.ceil(np.max([1, ram_multiplier * gb_ram_per_cpu])).astype(int)
-    num_cpus = np.min([32, num_processes])
+    num_cpus = np.min([32, num_threads])
 
     return num_cpus, gb_ram_per_cpu