refactor

src/anemoi/registry/entry/dataset.py

@@ -59,11 +59,8 @@ def upload(self, source, target, platform="unknown", resume=True):
         )
         LOG.info(f"Task: {kwargs}")
 
-        from anemoi.utils.s3 import upload
-
         from anemoi.registry.tasks import TaskCatalogueEntry
         from anemoi.registry.tasks import TaskCatalogueEntryList
-        from anemoi.registry.workers.transfer_dataset import Progress
 
         def find_or_create_task(**kwargs):
             lst = TaskCatalogueEntryList(**kwargs)
@@ -84,16 +81,21 @@ def find_or_create_task(**kwargs):
             return task
 
         task = find_or_create_task(**kwargs)
-        task.set_status("running")
+        self.transfer(task, source_path, target, resume=True, threads=2)
+
+    def transfer(self, task, source_path, target, resume, threads):
+        from anemoi.utils.s3 import upload
+
+        from anemoi.registry.workers.transfer_dataset import Progress
 
         progress = Progress(task, frequency=10)
         LOG.info(f"Upload('{source_path}','{target}', resume=True, threads=2)")
+        task.set_status("running")
         try:
-            upload(source_path, target, resume=True, threads=2, progress=progress)
+            upload(source_path, target, resume=resume, threads=threads, progress=progress)
         except:
             task.set_status("stopped")
             raise
-
         task.unregister()
 
     def set_recipe(self, file):

tests/test_all.py

@@ -84,6 +84,9 @@ def test_datasets():
     )
     run("anemoi-registry", "datasets", TMP_DATASET, "--add-location", "ewc")
 
+    # do not upload the dataset to avoid polluting the s3 bucket, until we have a way to clean it up automatically
+    # run("anemoi-registry", "datasets", TMP_DATASET_PATH, "--add-location", "ewc", "--upload")
+
 
 def test_weights():
     # assert run("anemoi-registry", "weights", "a5275e04-0000-0000-a0f6-be19591b09fe") == 1