wip

src/anemoi/registry/commands/worker.py

@@ -45,6 +45,7 @@ def add_arguments(self, command_parser):
         command_parser.add_argument("--target-dir", help="Target directory", default=".")
         command_parser.add_argument("--destination", help="Platform destination (e.g. leonardo, lumi)")
         command_parser.add_argument("--request", help="Filter tasks to process (key=value list)", nargs="*", default=[])
+        command_parser.add_argument("--threads", help="Number of threads to use", type=int, default=1)
 
     def run(self, args):
         kwargs = vars(args)

src/anemoi/registry/workers.py

@@ -28,6 +28,7 @@ def __init__(
         stop_if_finished=True,
         target_dir=".",
         auto_register=True,
+        threads=1,
         request={},
     ):
         """Run a worker that will process tasks in the queue.
@@ -44,6 +45,7 @@ def __init__(
         self.destination = destination
         self.target_dir = target_dir
         self.request = request
+        self.threads = threads
 
         self.wait = wait
         self.stop_if_finished = stop_if_finished
@@ -92,7 +94,7 @@ def process_entry(self, entry):
 
         dataset_entry = DatasetCatalogueEntry(key=dataset)
 
-        LOG.info(f"Transferring {dataset} from {source} to {destination}")
+        LOG.info(f"Transferring {dataset} from '{source}' to '{destination}'")
 
         def get_source_path():
             e = dataset_entry.record
@@ -129,12 +131,12 @@ def get_source_path():
 
         if target_path.startswith("s3://"):
             # untested
-            download(source_path, target_path, resume=True)
+            download(source_path, target_path, resume=True, threads=self.threads)
             return
         else:
             target_tmp_path = os.path.join(self.target_dir + "-downloading", basename)
             os.makedirs(os.path.dirname(target_tmp_path), exist_ok=True)
-            download(source_path, target_tmp_path, resume=True)
+            download(source_path, target_tmp_path, resume=True, threads=self.threads)
             os.rename(target_tmp_path, target_path)
 
         if self.auto_register: