-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
8ee92a6
commit 1cf9f18
Showing
7 changed files
with
191 additions
and
98 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,162 @@ | ||
| # (C) Copyright 2023 European Centre for Medium-Range Weather Forecasts. | ||
| # This software is licensed under the terms of the Apache Licence Version 2.0 | ||
| # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. | ||
| # In applying this licence, ECMWF does not waive the privileges and immunities | ||
| # granted to it by virtue of its status as an intergovernmental organisation | ||
| # nor does it submit to any jurisdiction. | ||
|
|
||
| import datetime | ||
| import logging | ||
| import os | ||
|
|
||
| from anemoi.registry import config | ||
| from anemoi.registry.entry import CatalogueEntry | ||
| from anemoi.registry.rest import RestItemList | ||
| from anemoi.registry.rest import trace_info | ||
|
|
||
| # from anemoi.utils.provenance import trace_info | ||
|
|
||
| LOG = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| class Actor: | ||
| def __init__(self, **kwargs): | ||
| for k, v in kwargs.items(): | ||
| if k not in ["action", "status", "created", "updated", "uuid"]: | ||
| LOG.error(f"Unknown attribute {k}={v} in actor {kwargs}") | ||
|
|
||
| def check(self): | ||
| LOG.info(f"Checking {self}") | ||
|
|
||
|
|
||
| class TransferDataset(Actor): | ||
| def __init__(self, *, status, source, destination, dataset, target, **kwargs): | ||
| super().__init__(**kwargs) | ||
| self.source = source | ||
| self.destination = destination | ||
| self.target = target | ||
| self.dataset = dataset | ||
|
|
||
| def __repr__(self): | ||
| return f"{self.__class__.__name__}(source={self.source}, target={self.target}, dataset={self.dataset})" | ||
|
|
||
| def run(self): | ||
| LOG.info(f"Transferring {self.dataset} from {self.source} to {self.target}") | ||
| LOG.info(f"anemoi-datasets copy {self.source}/{self.dataset} {self.target}/{self.dataset}") | ||
| c = config() | ||
| print(c) | ||
|
|
||
| def check(self): | ||
| super().check() | ||
| if "/" in self.destination: | ||
| raise ValueError(f"Destination {self.destination} must not contain '/', this is a platform name") | ||
|
|
||
| if "/" in self.source: | ||
| raise ValueError(f"Source {self.source} must not contain '/', this is a platform name") | ||
|
|
||
| if not os.path.exists(self.target) or not os.path.isdir(self.target): | ||
| raise ValueError(f"Target {self.target} must exist and be a directory") | ||
|
|
||
| if "." in self.dataset: | ||
| raise ValueError(f"The dataset {self.dataset} must not contain a '.', this is the name of the dataset.") | ||
|
|
||
|
|
||
| class Worker: | ||
| def __init__(self, *args, timeout=None): | ||
| if timeout: | ||
| import signal | ||
|
|
||
| signal.alarm(args.timeout) | ||
|
|
||
| self.args = TaskCatalogueEntry.list_to_dict(args) | ||
|
|
||
| def run(self): | ||
| while self.run_one_task(): | ||
| pass | ||
|
|
||
| def run_one_task(self): | ||
| request = self.args.copy() | ||
| request["status"] = "queued" | ||
| data = RestItemList(TaskCatalogueEntry.collection).get(params=request) | ||
|
|
||
| if not data: | ||
| LOG.info("No tasks found") | ||
| return False | ||
|
|
||
| uuid = data[-1]["uuid"] | ||
| entry = TaskCatalogueEntry(key=uuid) | ||
| LOG.info(f"Processing task {uuid}: {entry}") | ||
| entry.to_actor().check() | ||
|
|
||
| entry.take_ownership() | ||
| actor = entry.to_actor() | ||
| actor.run() | ||
| LOG.info(f"Task {uuid} completed.") | ||
| entry.unregister() | ||
| LOG.info(f"Task {uuid} deleted.") | ||
| return True | ||
|
|
||
|
|
||
| def actor_factory(**record): | ||
| LOG.info(f"Converting task {record} to actor") | ||
| record = record.copy() | ||
| action = record.pop("action").replace("-", "_").lower() | ||
|
|
||
| ACTIONS = dict( | ||
| transfer_dataset=TransferDataset, | ||
| # delete_dataset=DeleteDataset, | ||
| ) | ||
| if action not in ACTIONS: | ||
| raise ValueError(f"Unknown action {action}") | ||
|
|
||
| return ACTIONS[action](**record) | ||
|
|
||
|
|
||
| class TaskCatalogueEntry(CatalogueEntry): | ||
| collection = "queues" | ||
| main_key = "uuid" | ||
|
|
||
| def to_actor(self): | ||
| return actor_factory(**self.record) | ||
|
|
||
| @classmethod | ||
| def new(cls, *args, **kwargs): | ||
| if args: | ||
| return cls.new(**cls.list_to_dict(args)) | ||
| rest_collection = RestItemList(cls.collection) | ||
| actor_factory(status=None, **kwargs).check() | ||
| return rest_collection.post(kwargs) | ||
|
|
||
| def set_status(self, status): | ||
| patch = [{"op": "add", "path": "/status", "value": status}] | ||
| self.rest_item.patch(patch) | ||
|
|
||
| def unregister(self): | ||
| return self.rest_item.delete() | ||
|
|
||
| def take_ownership(self): | ||
| trace = trace_info() | ||
| trace["timestamp"] = datetime.datetime.now().isoformat() | ||
| return self.rest_item.patch( | ||
| [ | ||
| {"op": "test", "path": "/status", "value": "queued"}, | ||
| {"op": "replace", "path": "/status", "value": "running"}, | ||
| {"op": "add", "path": "/worker", "value": trace}, | ||
| ] | ||
| ) | ||
|
|
||
| def release_ownership(self): | ||
| self.rest_item.patch( | ||
| [ | ||
| {"op": "test", "path": "/status", "value": "running"}, | ||
| {"op": "replace", "path": "/status", "value": "queued"}, | ||
| {"op": "remove", "path": "/worker"}, | ||
| ] | ||
| ) | ||
|
|
||
| def set_progress(self, progress): | ||
| assert isinstance(progress, int), progress | ||
| if not (0 <= progress <= 100): | ||
| raise ValueError("Progress must be between 0 and 100") | ||
| patch = [{"op": "add", "path": "/progress", "value": progress}] | ||
| self.rest_item.patch(patch) |