refactor: replace multiprocessing.Pool with ThreadPoolExecutor in process_single_position

src/iohub/ngff/utils.py

@@ -2,9 +2,10 @@
 
 import inspect
 import itertools
-import multiprocessing as mp
+import os
 from collections import defaultdict
 from collections.abc import Callable, Sequence
+from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from pathlib import Path
 from typing import Any, Literal
@@ -395,21 +396,19 @@ def process_single_position(
         output_position_path,
         **kwargs,
     )
-    num_processes = min(num_processes, len(flat_iterable), mp.cpu_count())
-    click.echo(f"\nStarting multiprocess pool with {num_processes} processes")
-    if num_processes <= 1:
-        # Run serially — Pool(1) with spawn unnecessarily forks a subprocess
+    cpu_count = os.cpu_count() or 1
+    num_workers = min(num_processes, len(flat_iterable), cpu_count)
+    click.echo(f"\nStarting thread pool with {num_workers} workers")
+    if num_workers <= 1:
         for args in flat_iterable:
             partial_apply_transform_to_czyx_and_save(*args)
     else:
-        # NOTE: use spawn to work around tensorstore#61
-        context = mp.get_context("spawn")
-        with context.Pool(num_processes) as p:
-            p.starmap(
-                partial_apply_transform_to_czyx_and_save,
+        with ThreadPoolExecutor(max_workers=num_workers) as executor:
+            list(executor.map(
+                lambda args: partial_apply_transform_to_czyx_and_save(*args),
                 flat_iterable,
-            )
-    click.echo("Shut down multiprocess pool")
+            ))
+    click.echo("Shut down thread pool")
 
 
 # -- Pure utility functions ------------------------------------------------

tests/ngff/test_ngff_utils.py

@@ -792,3 +792,70 @@ def test_process_single_position(setup, constant):
                     dummy_transform,
                     **kwargs,
                 )
+
+
+@given(
+    setup=process_single_position_setup(),
+    constant=st.integers(min_value=1, max_value=3),
+)
+@settings(max_examples=3, deadline=None)
+def test_process_single_position_threaded(setup, constant):
+    (
+        position_keys,
+        channel_names,
+        shape,
+        chunks,
+        shards_ratio,
+        scale,
+        dtype,
+        channel_indices,
+        time_indices,
+        version,
+    ) = setup
+
+    with _temp_ome_zarr_stores(
+        position_keys=position_keys,
+        channel_names=channel_names,
+        shape=shape,
+        chunks=chunks,
+        shards_ratio=shards_ratio,
+        scale=scale,
+        dtype=dtype,
+        version=version,
+    ) as (input_store_path, output_store_path):
+        populate_store(input_store_path, position_keys, shape, dtype)
+
+        for position_key_tuple in position_keys:
+            input_position_path = input_store_path / Path(*position_key_tuple)
+            output_position_path = output_store_path / Path(*position_key_tuple)
+            kwargs = {"constant": constant, "extra_metadata": {"temp": 10}}
+
+            process_single_position(
+                func=dummy_transform,
+                input_position_path=input_position_path,
+                output_position_path=output_position_path,
+                input_channel_indices=channel_indices,
+                output_channel_indices=channel_indices,
+                input_time_indices=time_indices,
+                output_time_indices=time_indices,
+                num_processes=2,
+                **kwargs,
+            )
+
+            if time_indices is None:
+                time_indices = list(range(shape[0]))
+            if channel_indices is None:
+                channel_indices = [[c] for c in range(shape[1])]
+
+            iterable = itertools.product(time_indices, channel_indices)
+            for t_idx, chan_idx in iterable:
+                verify_transformation(
+                    input_store_path,
+                    output_store_path,
+                    position_key_tuple,
+                    shape,
+                    t_idx,
+                    chan_idx,
+                    dummy_transform,
+                    **kwargs,
+                )