Work on refactor progress code

Author: jeromekelleher

bio2zarr/core.py

@@ -1,10 +1,14 @@
 import dataclasses
 import contextlib
 import concurrent.futures as cf
+import multiprocessing
+import threading
 import logging
+import time
 
 import zarr
 import numpy as np
+import tqdm
 
 
 logger = logging.getLogger(__name__)
@@ -118,3 +122,73 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                 future.cancel()
         self.executor.shutdown()
         return False
+
+
+
+progress_counter = multiprocessing.Value("Q", 0)
+
+import os
+
+def update_progress(inc):
+    print("update progress", os.getpid(), inc)
+    with progress_counter.get_lock():
+        progress_counter.value += 1
+
+def progress_thread_worker(config):
+    pbar = tqdm.tqdm(
+        total=config.total, desc=config.title, unit_scale=True, unit=config.units,
+        smoothing=0.1
+    )
+
+    while (current := progress_counter.value) < config.total:
+        inc = current - pbar.n
+        pbar.update(inc)
+        time.sleep(0.1)
+    pbar.close()
+
+
+def init_workers(counter):
+    global progress_counter
+    progress_counter = counter
+
+
+@dataclasses.dataclass
+class ProgressConfig:
+    total: int
+    units: str
+    title: str
+
+
+class ParallelWorkManager(contextlib.AbstractContextManager):
+
+    def __init__(self, worker_processes=1, progress_config=None):
+        self.executor = cf.ProcessPoolExecutor(
+            max_workers=worker_processes,
+            initializer=init_workers,
+            initargs=(progress_counter,),
+        )
+
+        self.bar_thread = None
+        if progress_config is not None:
+            bar_thread = threading.Thread(
+                target=progress_thread_worker,
+                args=(progress_config,),
+                name="progress",
+                daemon=True,
+            )
+            bar_thread.start()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # if exc_type is None:
+        #     # Normal exit condition
+        #     self.next_row += 1
+        #     self.swap_buffers()
+        #     self.wait_on_futures()
+        # else:
+        #     for future in self.futures:
+        #         future.cancel()
+        self.executor.shutdown()
+        if self.bar_thread is not None:
+            self.bar_thread.join()
+        return False
+

bio2zarr/vcf.py

@@ -729,6 +729,7 @@ def convert(
         vcfs, out_path, *, column_chunk_size=16, worker_processes=1, show_progress=False
     ):
         out_path = pathlib.Path(out_path)
+        # TODO make scan work in parallel using general progress code too
         vcf_metadata = scan_vcfs(vcfs, show_progress=show_progress)
         pcvcf = PickleChunkedVcf(out_path, vcf_metadata)
         pcvcf.mkdirs()
@@ -741,29 +742,16 @@ def convert(
             f"Exploding {pcvcf.num_columns} columns {total_variants} variants "
             f"{pcvcf.num_samples} samples"
         )
-        global progress_counter
-        progress_counter = multiprocessing.Value("Q", 0)
-
-        # start update progress bar process
-        bar_thread = None
+        progress_config = None
         if show_progress:
-            bar_thread = threading.Thread(
-                target=update_bar,
-                args=(progress_counter, total_variants, "Explode", "vars"),
-                name="progress",
-                daemon=True,
-            )
-            bar_thread.start()
+            progress_config = core.ProgressConfig(
+                total=total_variants, units="vars", title="Explode")
+        with core.ParallelWorkManager(worker_processes, progress_config) as pwm:
 
-        with cf.ProcessPoolExecutor(
-            max_workers=worker_processes,
-            initializer=init_workers,
-            initargs=(progress_counter,),
-        ) as executor:
             futures = []
             for j, partition in enumerate(vcf_metadata.partitions):
                 futures.append(
-                    executor.submit(
+                    pwm.executor.submit(
                         PickleChunkedVcf.convert_partition,
                         vcf_metadata,
                         j,
@@ -775,9 +763,43 @@ def convert(
                 future.result() for future in cf.as_completed(futures)
             ]
 
-        assert progress_counter.value == total_variants
-        if bar_thread is not None:
-            bar_thread.join()
+#         global progress_counter
+#         progress_counter = multiprocessing.Value("Q", 0)
+
+#         # start update progress bar process
+#         bar_thread = None
+#         if show_progress:
+#             bar_thread = threading.Thread(
+#                 target=update_bar,
+#                 args=(progress_counter, total_variants, "Explode", "vars"),
+#                 name="progress",
+#                 daemon=True,
+#             )
+#             bar_thread.start()
+
+#         with cf.ProcessPoolExecutor(
+#             max_workers=worker_processes,
+#             initializer=init_workers,
+#             initargs=(progress_counter,),
+#         ) as executor:
+#             futures = []
+#             for j, partition in enumerate(vcf_metadata.partitions):
+#                 futures.append(
+#                     executor.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)
+#             ]
+
+#         assert progress_counter.value == total_variants
+#         if bar_thread is not None:
+#             bar_thread.join()
 
         for field in vcf_metadata.fields:
             for summary in partition_summaries:
@@ -862,11 +884,11 @@ def service_futures(max_waiting=2 * flush_threads):
 
                 service_futures()
 
+
                 # Note: an issue with updating the progress per variant here like this
                 # is that we get a significant pause at the end of the counter while
                 # all the "small" fields get flushed. Possibly not much to be done about it.
-                with progress_counter.get_lock():
-                    progress_counter.value += 1
+                core.update_progress(1)
 
             for col in columns.values():
                 col.flush()
@@ -876,21 +898,21 @@ def service_futures(max_waiting=2 * flush_threads):
             return summaries
 
 
-def update_bar(progress_counter, total, title, units):
-    pbar = tqdm.tqdm(
-        total=total, desc=title, unit_scale=True, unit=units, smoothing=0.1
-    )
+# def update_bar(progress_counter, total, title, units):
+#     pbar = tqdm.tqdm(
+#         total=total, desc=title, unit_scale=True, unit=units, smoothing=0.1
+#     )
 
-    while (current := progress_counter.value) < total:
-        inc = current - pbar.n
-        pbar.update(inc)
-        time.sleep(0.1)
-    pbar.close()
+#     while (current := progress_counter.value) < total:
+#         inc = current - pbar.n
+#         pbar.update(inc)
+#         time.sleep(0.1)
+#     pbar.close()
 
 
-def init_workers(counter):
-    global progress_counter
-    progress_counter = counter
+# def init_workers(counter):
+#     global progress_counter
+#     progress_counter = counter
 
 
 def explode(
@@ -1418,7 +1440,9 @@ def convert_vcf(
         )
 
 
-def encode_bed_partition_genotypes(bed_path, zarr_path, start_variant, end_variant, encoder_threads=8):
+def encode_bed_partition_genotypes(
+    bed_path, zarr_path, start_variant, end_variant, encoder_threads=8
+):
     bed = bed_reader.open_bed(bed_path, num_threads=1)
 
     store = zarr.DirectoryStore(zarr_path)
@@ -1432,7 +1456,6 @@ def encode_bed_partition_genotypes(bed_path, zarr_path, start_variant, end_varia
     buffered_arrays = [gt, gt_phased, gt_mask]
 
     with core.ThreadedZarrEncoder(buffered_arrays, encoder_threads) as te:
-
         start = start_variant
         while start < end_variant:
             stop = min(start + chunk_length, end_variant)