diff --git a/waveorder/calib/calibration_workers.py b/waveorder/calib/calibration_workers.py index 43e11695..902b7356 100644 --- a/waveorder/calib/calibration_workers.py +++ b/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 diff --git a/waveorder/cli/apply_inverse_transfer_function.py b/waveorder/cli/apply_inverse_transfer_function.py index 81c443b1..a08940b3 100644 --- a/waveorder/cli/apply_inverse_transfer_function.py +++ b/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, ) diff --git a/waveorder/cli/parsing.py b/waveorder/cli/parsing.py index 6f6b2620..71059893 100644 --- a/waveorder/cli/parsing.py +++ b/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( diff --git a/waveorder/cli/reconstruct.py b/waveorder/cli/reconstruct.py index 9dfeaf22..67f0486b 100644 --- a/waveorder/cli/reconstruct.py +++ b/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, ) diff --git a/waveorder/cli/utils.py b/waveorder/cli/utils.py index 165dc465..279c62ef 100644 --- a/waveorder/cli/utils.py +++ b/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