diff --git a/src/anemoi/registry/entry/dataset.py b/src/anemoi/registry/entry/dataset.py index 749a64b..0d929bc 100644 --- a/src/anemoi/registry/entry/dataset.py +++ b/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): diff --git a/tests/test_all.py b/tests/test_all.py index c3a6719..ab75dcd 100755 --- a/tests/test_all.py +++ b/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