Abstract progress and multiprocessing

Author: jeromekelleher

bio2zarr/core.py

@@ -4,6 +4,7 @@
 import multiprocessing
 import threading
 import logging
+import functools
 import time
 
 import zarr
@@ -14,6 +15,25 @@
 logger = logging.getLogger(__name__)
 
 
+class SynchronousExecutor(cf.Executor):
+    def submit(self, fn, /, *args, **kwargs):
+        future = cf.Future()
+        future.set_result(fn(*args, **kwargs))
+        return future
+
+
+def wait_on_futures(futures):
+    for future in cf.as_completed(futures):
+        exception = future.exception()
+        if exception is not None:
+            raise exception
+
+
+def cancel_futures(futures):
+    for future in futures:
+        future.cancel()
+
+
 @dataclasses.dataclass
 class BufferedArray:
     array: zarr.Array
@@ -98,14 +118,8 @@ def next_buffer_row(self):
             self.next_row = 0
         return self.next_row
 
-    def wait_on_futures(self):
-        for future in cf.as_completed(self.futures):
-            exception = future.exception()
-            if exception is not None:
-                raise exception
-
     def swap_buffers(self):
-        self.wait_on_futures()
+        wait_on_futures(self.futures)
         self.futures = []
         for ba in self.buffered_arrays:
             self.futures.extend(
@@ -118,19 +132,20 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             # Normal exit condition
             self.next_row += 1
             self.swap_buffers()
-            self.wait_on_futures()
+            wait_on_futures(self.futures)
         else:
-            for future in self.futures:
-                future.cancel()
+            cancel_futures(self.futures)
         self.executor.shutdown()
         return False
 
 
 @dataclasses.dataclass
 class ProgressConfig:
-    total: int
-    units: str
-    title: str
+    total: int = 0
+    units: str = ""
+    title: str = ""
+    show: bool = False
+    poll_interval: float = 0.001
 
 
 _progress_counter = multiprocessing.Value("Q", 0)
@@ -141,44 +156,71 @@ def update_progress(inc):
         _progress_counter.value += inc
 
 
+def get_progress():
+    with _progress_counter.get_lock():
+        val = _progress_counter.value
+    return val
+
+
+def set_progress(value):
+    with _progress_counter.get_lock():
+        _progress_counter.value = value
+
+
 def progress_thread_worker(config):
     pbar = tqdm.tqdm(
         total=config.total,
         desc=config.title,
         unit_scale=True,
         unit=config.units,
         smoothing=0.1,
+        disable=not config.show,
     )
 
-    while (current := _progress_counter.value) < config.total:
+    while (current := get_progress()) < config.total:
         inc = current - pbar.n
         pbar.update(inc)
-        time.sleep(0.1)
+        time.sleep(config.poll_interval)
     pbar.close()
 
 
 class ParallelWorkManager(contextlib.AbstractContextManager):
     def __init__(self, worker_processes=1, progress_config=None):
-        self.executor = cf.ProcessPoolExecutor(
-            max_workers=worker_processes,
+        if worker_processes <= 0:
+            # NOTE: this is only for testing, not for production use!
+            self.executor = SynchronousExecutor()
+        else:
+            self.executor = cf.ProcessPoolExecutor(
+                max_workers=worker_processes,
+            )
+        set_progress(0)
+        if progress_config is None:
+            progress_config = ProgressConfig()
+        self.bar_thread = threading.Thread(
+            target=progress_thread_worker,
+            args=(progress_config,),
+            name="progress",
+            daemon=True,
         )
+        self.bar_thread.start()
+        self.progress_config = progress_config
+        self.futures = []
 
-        self.bar_thread = None
-        if progress_config is not None:
-            self.bar_thread = threading.Thread(
-                target=progress_thread_worker,
-                args=(progress_config,),
-                name="progress",
-                daemon=True,
-            )
-            self.bar_thread.start()
+    def submit(self, *args, **kwargs):
+        self.futures.append(self.executor.submit(*args, **kwargs))
+
+    def results_as_completed(self):
+        for future in cf.as_completed(self.futures):
+            yield future.result()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        # if exc_type is None:
-        #     print("normal exit")
-        # else:
-        #     print("Error occured")
-        if self.bar_thread is not None:
-            self.bar_thread.join(timeout=0)
+        if exc_type is None:
+            wait_on_futures(self.futures)
+            set_progress(self.progress_config.total)
+            timeout = None
+        else:
+            cancel_futures(self.futures)
+            timeout = 0
+        self.bar_thread.join(timeout)
         self.executor.shutdown()
         return False

bio2zarr/vcf.py

@@ -1,12 +1,9 @@
 import concurrent.futures as cf
 import dataclasses
-import multiprocessing
 import functools
 import logging
 import os
-import threading
 import pathlib
-import time
 import pickle
 import sys
 import shutil
@@ -111,15 +108,6 @@ def assert_prefix_float_equal_2d(vcf_val, zarr_val):
     # nt.assert_array_equal(v, z[:k])
 
 
-# TODO rename to wait_and_check_futures
-def flush_futures(futures):
-    # Make sure previous futures have completed
-    for future in cf.as_completed(futures):
-        exception = future.exception()
-        if exception is not None:
-            raise exception
-
-
 @dataclasses.dataclass
 class VcfFieldSummary:
     num_chunks: int = 0
@@ -742,26 +730,19 @@ def convert(
             f"Exploding {pcvcf.num_columns} columns {total_variants} variants "
             f"{pcvcf.num_samples} samples"
         )
-        progress_config = None
-        if show_progress:
-            progress_config = core.ProgressConfig(
-                total=total_variants, units="vars", title="Explode"
-            )
+        progress_config = core.ProgressConfig(
+            total=total_variants, units="vars", title="Explode", show=show_progress
+        )
         with core.ParallelWorkManager(worker_processes, progress_config) as pwm:
-            futures = []
             for j, partition in enumerate(vcf_metadata.partitions):
-                futures.append(
-                    pwm.executor.submit(
-                        PickleChunkedVcf.convert_partition,
-                        vcf_metadata,
-                        j,
-                        out_path,
-                        column_chunk_size=column_chunk_size,
-                    )
+                pwm.submit(
+                    PickleChunkedVcf.convert_partition,
+                    vcf_metadata,
+                    j,
+                    out_path,
+                    column_chunk_size=column_chunk_size,
                 )
-            partition_summaries = [
-                future.result() for future in cf.as_completed(futures)
-            ]
+            partition_summaries = list(pwm.results_as_completed())
 
         for field in vcf_metadata.fields:
             for summary in partition_summaries:
@@ -1258,31 +1239,28 @@ def convert(
         for variable in conversion_spec.variables[:]:
             sgvcf.create_array(variable)
 
-        progress_config = None
-        if show_progress:
-            progress_config = core.ProgressConfig(
-                total=pcvcf.total_uncompressed_bytes, title="Encode", units="b"
-            )
+        progress_config = core.ProgressConfig(
+            total=pcvcf.total_uncompressed_bytes,
+            title="Encode",
+            units="b",
+            show=show_progress,
+        )
         with core.ParallelWorkManager(worker_processes, progress_config) as pwm:
-            futures = [
-                pwm.executor.submit(
-                    sgvcf.encode_samples,
-                    pcvcf,
-                    conversion_spec.sample_id,
-                    conversion_spec.chunk_width,
-                ),
-                pwm.executor.submit(sgvcf.encode_alleles, pcvcf),
-                pwm.executor.submit(sgvcf.encode_id, pcvcf),
-                pwm.executor.submit(
-                    sgvcf.encode_contig,
-                    pcvcf,
-                    conversion_spec.contig_id,
-                    conversion_spec.contig_length,
-                ),
-                pwm.executor.submit(
-                    sgvcf.encode_filters, pcvcf, conversion_spec.filter_id
-                ),
-            ]
+            pwm.submit(
+                sgvcf.encode_samples,
+                pcvcf,
+                conversion_spec.sample_id,
+                conversion_spec.chunk_width,
+            )
+            pwm.submit(sgvcf.encode_alleles, pcvcf)
+            pwm.submit(sgvcf.encode_id, pcvcf)
+            pwm.submit(
+                sgvcf.encode_contig,
+                pcvcf,
+                conversion_spec.contig_id,
+                conversion_spec.contig_length,
+            )
+            pwm.submit(sgvcf.encode_filters, pcvcf, conversion_spec.filter_id)
             has_gt = False
             for variable in conversion_spec.variables[:]:
                 if variable.vcf_field is not None:
@@ -1292,21 +1270,14 @@ def convert(
                     # long wait for the largest GT columns to finish.
                     # Straightforward to do because we can chunk-align the work
                     # packages.
-                    future = pwm.executor.submit(sgvcf.encode_column, pcvcf, variable)
-                    futures.append(future)
+                    pwm.submit(sgvcf.encode_column, pcvcf, variable)
                 else:
                     if variable.name == "call_genotype":
                         has_gt = True
             if has_gt:
                 # TODO add mixed ploidy
-                futures.append(pwm.executor.submit(sgvcf.encode_genotypes, pcvcf))
-
-            flush_futures(futures)
+                pwm.executor.submit(sgvcf.encode_genotypes, pcvcf)
 
-        # FIXME can't join the bar_thread because we never get to the correct
-        # number of bytes
-        # if bar_thread is not None:
-        #     bar_thread.join()
         zarr.consolidate_metadata(write_path)
         # Atomic swap, now we've completely finished.
         logger.info(f"Moving to final path {path}")
@@ -1617,14 +1588,9 @@ def convert_plink(
             partitions.append((last_stop, m))
     # print(partitions)
 
-    progress_config = None
-    if show_progress:
-        progress_config = core.ProgressConfig(total=m, title="Convert", units="vars")
+    progress_config = core.ProgressConfig(
+        total=m, title="Convert", units="vars", show=show_progress
+    )
     with core.ParallelWorkManager(worker_processes, progress_config) as pwm:
-        futures = [
-            pwm.executor.submit(
-                encode_bed_partition_genotypes, bed_path, zarr_path, start, end
-            )
-            for start, end in partitions
-        ]
-        flush_futures(futures)
+        for start, end in partitions:
+            pwm.submit(encode_bed_partition_genotypes, bed_path, zarr_path, start, end)

tests/test_core.py

@@ -93,3 +93,52 @@ def test_error_in_encode(self):
                 # an error in the futures. In reality these will happen
                 # when we run out of disk space, but this is hard to simulate
                 ba.buff = np.array(["not an integer"])
+
+
+class TestParallelWorkManager:
+    @pytest.mark.parametrize("total", [1, 10, 2**63])
+    @pytest.mark.parametrize("workers", [0, 1])
+    def test_one_future_progress(self, total, workers):
+        progress_config = core.ProgressConfig(total=total)
+        with core.ParallelWorkManager(workers, progress_config) as pwm:
+            pwm.submit(core.update_progress, total)
+        assert core.get_progress() == total
+
+    @pytest.mark.parametrize("total", [1, 10, 1000])
+    @pytest.mark.parametrize("workers", [0, 1, 2, 3])
+    def test_n_futures_progress(self, total, workers):
+        progress_config = core.ProgressConfig(total=total)
+        with core.ParallelWorkManager(workers, progress_config) as pwm:
+            for _ in range(total):
+                pwm.submit(core.update_progress, 1)
+        assert core.get_progress() == total
+
+    @pytest.mark.parametrize("total", [1, 10, 20])
+    @pytest.mark.parametrize("workers", [0, 1, 2, 3])
+    def test_results_as_completed(self, total, workers):
+        with core.ParallelWorkManager(workers) as pwm:
+            for j in range(total):
+                pwm.submit(frozenset, range(j))
+            results = set(pwm.results_as_completed())
+            assert results == set(frozenset(range(j)) for j in range(total))
+
+    @pytest.mark.parametrize("total", [1, 10, 20])
+    @pytest.mark.parametrize("workers", [1, 2, 3])
+    def test_error_in_workers_as_completed(self, total, workers):
+        with pytest.raises(TypeError):
+            with core.ParallelWorkManager(workers) as pwm:
+                for j in range(total):
+                    pwm.submit(frozenset, range(j))
+                # Raises a TypeError:
+                pwm.submit(frozenset, j)
+                set(pwm.results_as_completed())
+
+    @pytest.mark.parametrize("total", [1, 10, 20])
+    @pytest.mark.parametrize("workers", [1, 2, 3])
+    def test_error_in_workers_on_exit(self, total, workers):
+        with pytest.raises(TypeError):
+            with core.ParallelWorkManager(workers) as pwm:
+                for j in range(total):
+                    pwm.submit(frozenset, range(j))
+                # Raises a TypeError:
+                pwm.submit(frozenset, j)