cli argument for thread-count

Author: jonatantreijs

h5pyd/_apps/hsload.py

@@ -124,6 +124,7 @@ def main():
     cfg.setitem("ignorefilters", False, flags=["--ignore-filters"], help="ignore any filters used by source dataset")
     cfg.setitem("retries", 3, flags=["--retries",], choices=["N",], help="Set number of server retry attempts")
     cfg.setitem("no_checks", False, flags=["--no-checks"], help="do not check for existence before creating resources")
+    cfg.setitem("thread_count", 30, flags=["--thread-count"], choices=["N",] ,help="The number of threads to allocate when making requests in parallel, defaults to 30")
     cfg.setitem("help", False, flags=["-h", "--help"], help="this message")
 
     try:
@@ -259,6 +260,7 @@ def main():
                     "bucket": cfg["hs_bucket"],
                     "mode": mode,
                     "retries": int(cfg["retries"]),
+                    "thread_count": int(cfg["thread_count"]),
                 }
 
                 fout = h5pyd.File(tgt, **kwargs)
@@ -325,6 +327,7 @@ def main():
                     "ignore_error": cfg["ignore_error"],
                     "no_clobber": no_clobber,
                     "no_checks": cfg["no_checks"],
+                    "thread_count": int(cfg["thread_count"]),
                 }
                 load_file(fin, fout, **kwargs)
 

h5pyd/_apps/utillib.py

@@ -1170,7 +1170,7 @@ def create_dataset(dobj, ctx):
         print(msg)
     fout = ctx["fout"]
 
-    if not ctx["no-checks"] and dobj.name in fout:
+    if not ctx["no_checks"] and dobj.name in fout:
         dset = fout[dobj.name]
         logging.debug(f"{dobj.name} already exists")
         if ctx["no_clobber"]:
@@ -1688,7 +1688,7 @@ def create_group(gobj, ctx):
 
     grp = None
 
-    if not ctx["no-checks"] and gobj.name in fout:
+    if not ctx["no_checks"] and gobj.name in fout:
         grp = fout[gobj.name]
         logging.debug(f"{gobj.name} already exists")
         if ctx["no_clobber"]:
@@ -1709,7 +1709,7 @@ def create_group(gobj, ctx):
             if not ctx["ignore_error"]:
                 raise IOError(msg)
     else:
-        if not ctx["no-checks"] and ctx["verbose"]:
+        if not ctx["no_checks"] and ctx["verbose"]:
             print(f"{gobj.name} not found")
 
         grp = fout.create_group(gobj.name)
@@ -1800,7 +1800,8 @@ def load_file(
     extend_dim=None,
     extend_offset=0,
     ignore_error=False,
-    no_checks=False
+    no_checks=False,
+    thread_count=30,
 ):
 
     logging.info(f"input file: {fin.filename}")
@@ -1836,7 +1837,8 @@ def load_file(
     ctx["extend_offset"] = extend_offset
     ctx["srcid_desobj_map"] = {}
     ctx["ignore_error"] = ignore_error
-    ctx["no-checks"] = no_checks
+    ctx["no_checks"] = no_checks
+    ctx["thread_count"] = thread_count
 
     def copy_attribute_helper(name, obj):
         logging.info(f"copy attribute - name: {name}  obj: {obj.name}")
@@ -1885,7 +1887,7 @@ def _visit_in_parallell(func):
         def _add_to_jobs(name, obj):
             jobs.append((name, obj))
 
-        with concurrent.futures.ThreadPoolExecutor(max_workers=30) as executor:
+        with concurrent.futures.ThreadPoolExecutor(max_workers=thread_count) as executor:
             fin.visititems(_add_to_jobs)
             futures = [executor.submit(func, item[0], item[1]) for item in jobs]
 

h5pyd/_hl/files.py

@@ -279,6 +279,7 @@ def __init__(
         track_order=None,
         retries=10,
         timeout=180,
+        thread_count=30,
         **kwds,
     ):
         """Create a new file object.
@@ -419,6 +420,7 @@ def __init__(
                 logger=logger,
                 retries=retries,
                 timeout=timeout,
+                thread_count=thread_count,
             )
 
             root_json = None

h5pyd/_hl/folders.py

@@ -83,6 +83,7 @@ def __init__(
         owner=None,
         batch_size=1000,
         retries=3,
+        thread_count=30,
         **kwds,
     ):
         """Create a new Folders object.
@@ -179,6 +180,7 @@ def __init__(
             mode=mode,
             logger=logger,
             retries=retries,
+            thread_count=thread_count
         )
         self.log = self._http_conn.logging
 

h5pyd/_hl/httpconn.py

@@ -166,6 +166,7 @@ def __init__(
         logger=None,
         retries=3,
         timeout=DEFAULT_TIMEOUT,
+        thread_count=30,
         **kwds,
     ):
         self._domain = domain_name
@@ -179,6 +180,8 @@ def __init__(
         self._api_key = api_key
         self._s = None  # Sessions
         self._server_info = None
+        self._thread_count = thread_count
+        
         if use_cache:
             self._cache = {}
             self._objdb = {}
@@ -739,11 +742,11 @@ def session(self):
 
                 s.mount(
                     "http://",
-                    HTTPAdapter(max_retries=retry, pool_connections=30, pool_maxsize=30),
+                    HTTPAdapter(max_retries=retry, pool_connections=self._thread_count, pool_maxsize=self._thread_count),
                 )
                 s.mount(
                     "https://",
-                    HTTPAdapter(max_retries=retry, pool_connections=30, pool_maxsize=30),
+                    HTTPAdapter(max_retries=retry, pool_connections=self._thread_count, pool_maxsize=self._thread_count),
                 )
                 self._s = s
             else: