get num cpus at function runtime (#471)

Author: tasansal

src/mdio/segy/blocked_io.py

@@ -38,7 +38,6 @@
     zfpy = None
 
 default_cpus = cpu_count(logical=True)
-NUM_CPUS = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus))
 
 
 def to_zarr(
@@ -119,7 +118,8 @@ def to_zarr(
     # For Unix async writes with s3fs/fsspec & multiprocessing,
     # use 'spawn' instead of default 'fork' to avoid deadlocks
     # on cloud stores. Slower but necessary. Default on Windows.
-    num_workers = min(num_chunks, NUM_CPUS)
+    num_cpus = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus))
+    num_workers = min(num_chunks, num_cpus)
     context = mp.get_context("spawn")
     executor = ProcessPoolExecutor(max_workers=num_workers, mp_context=context)
 

src/mdio/segy/parsers.py

@@ -21,7 +21,6 @@
     from segy import SegyFile
 
 default_cpus = cpu_count(logical=True)
-NUM_CPUS = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus))
 
 
 def parse_index_headers(
@@ -54,7 +53,8 @@ def parse_index_headers(
     # For Unix async reads with s3fs/fsspec & multiprocessing,
     # use 'spawn' instead of default 'fork' to avoid deadlocks
     # on cloud stores. Slower but necessary. Default on Windows.
-    num_workers = min(n_blocks, NUM_CPUS)
+    num_cpus = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus))
+    num_workers = min(n_blocks, num_cpus)
     context = mp.get_context("spawn")
 
     tqdm_kw = dict(unit="block", dynamic_ncols=True)