|
| 1 | +"""Task orchestration utilities backed by Exorcist and a warehouse.""" |
| 2 | + |
| 3 | +from collections.abc import Iterable |
| 4 | +from dataclasses import dataclass |
| 5 | +from pathlib import Path |
| 6 | + |
| 7 | +from exorcist.taskdb import TaskStatusDB |
| 8 | +from gufe.protocols.protocoldag import _pu_to_pur |
| 9 | +from gufe.protocols.protocolunit import ( |
| 10 | + Context, |
| 11 | + ProtocolUnit, |
| 12 | + ProtocolUnitResult, |
| 13 | +) |
| 14 | +from gufe.storage.externalresource.base import ExternalStorage |
| 15 | +from gufe.storage.externalresource.filestorage import FileStorage |
| 16 | +from gufe.tokenization import GufeKey |
| 17 | + |
| 18 | +from openfe.storage.warehouse import FileSystemWarehouse |
| 19 | + |
| 20 | +from .exorcist_utils import ( |
| 21 | + alchemical_network_to_task_graph, |
| 22 | + build_task_db_from_alchemical_network, |
| 23 | +) |
| 24 | + |
| 25 | + |
| 26 | +@dataclass |
| 27 | +class Worker: |
| 28 | + """Execute protocol units from an Exorcist task database. |
| 29 | +
|
| 30 | + Parameters |
| 31 | + ---------- |
| 32 | + warehouse : FileSystemWarehouse |
| 33 | + Warehouse used to load queued tasks and store execution results. |
| 34 | + task_db_path : pathlib.Path, default=Path("./warehouse/tasks.db") |
| 35 | + Path to the Exorcist SQLite task database. |
| 36 | + """ |
| 37 | + |
| 38 | + warehouse: FileSystemWarehouse |
| 39 | + task_db_path: Path = Path("./warehouse/tasks.db") |
| 40 | + |
| 41 | + _RESULT_INDEX_PREFIX = "protocol_unit_results" |
| 42 | + _TASK_WORKDIR_PREFIX = "task_workdirs" |
| 43 | + |
| 44 | + @staticmethod |
| 45 | + def _collect_protocol_unit_keys(value: object) -> set[GufeKey]: |
| 46 | + """Collect `ProtocolUnit` keys from nested unit inputs.""" |
| 47 | + |
| 48 | + if isinstance(value, ProtocolUnit): |
| 49 | + return {value.key} |
| 50 | + |
| 51 | + found: set[GufeKey] = set() |
| 52 | + items: Iterable # TODO: update this to dict_values | list after python 3.13 min? |
| 53 | + if isinstance(value, dict): |
| 54 | + items = value.values() |
| 55 | + elif isinstance(value, list): |
| 56 | + items = value |
| 57 | + else: |
| 58 | + return found |
| 59 | + |
| 60 | + for item in items: |
| 61 | + found.update(Worker._collect_protocol_unit_keys(item)) |
| 62 | + return found |
| 63 | + |
| 64 | + @classmethod |
| 65 | + def _result_index_location(cls, source_key: GufeKey) -> str: |
| 66 | + return f"{cls._RESULT_INDEX_PREFIX}/{source_key}" |
| 67 | + |
| 68 | + @classmethod |
| 69 | + def _task_workdir_name(cls, taskid: str) -> str: |
| 70 | + return taskid.replace(":", "__") |
| 71 | + |
| 72 | + def _task_workspace_paths( |
| 73 | + self, taskid: str, scratch_root: Path, shared_root: Path |
| 74 | + ) -> tuple[Path, Path]: |
| 75 | + workdir_name = self._task_workdir_name(taskid) |
| 76 | + task_scratch = scratch_root / self._TASK_WORKDIR_PREFIX / workdir_name |
| 77 | + task_shared = shared_root / self._TASK_WORKDIR_PREFIX / workdir_name |
| 78 | + return task_scratch, task_shared |
| 79 | + |
| 80 | + def _store_result_index(self, result: ProtocolUnitResult) -> None: |
| 81 | + shared_store: ExternalStorage = self.warehouse.stores["shared"] |
| 82 | + location = self._result_index_location(result.source_key) |
| 83 | + shared_store.store_bytes(location, str(result.key).encode("utf-8")) |
| 84 | + |
| 85 | + def _load_result_from_index(self, source_key: GufeKey) -> ProtocolUnitResult | None: |
| 86 | + shared_store: ExternalStorage = self.warehouse.stores["shared"] |
| 87 | + location = self._result_index_location(source_key) |
| 88 | + |
| 89 | + if not shared_store.exists(location): |
| 90 | + return None |
| 91 | + |
| 92 | + with shared_store.load_stream(location) as stream: |
| 93 | + result_key = stream.read().decode("utf-8").strip() |
| 94 | + |
| 95 | + loaded = self.warehouse.load_result_tokenizable(GufeKey(result_key)) |
| 96 | + if isinstance(loaded, ProtocolUnitResult): |
| 97 | + return loaded |
| 98 | + |
| 99 | + return None |
| 100 | + |
| 101 | + def _scan_result_store_for_sources( |
| 102 | + self, source_keys: set[GufeKey] |
| 103 | + ) -> dict[GufeKey, ProtocolUnitResult]: |
| 104 | + found: dict[GufeKey, ProtocolUnitResult] = {} |
| 105 | + |
| 106 | + for location in self.warehouse.result_store.iter_contents(): |
| 107 | + if len(found) == len(source_keys): |
| 108 | + break |
| 109 | + |
| 110 | + loaded = self.warehouse.load_result_tokenizable(GufeKey(location)) |
| 111 | + if not isinstance(loaded, ProtocolUnitResult): |
| 112 | + continue |
| 113 | + |
| 114 | + source_key = loaded.source_key |
| 115 | + if source_key in source_keys and source_key not in found: |
| 116 | + found[source_key] = loaded |
| 117 | + |
| 118 | + return found |
| 119 | + |
| 120 | + def _build_input_result_mapping(self, unit: ProtocolUnit) -> dict[GufeKey, ProtocolUnitResult]: |
| 121 | + required_keys = self._collect_protocol_unit_keys(unit.inputs) |
| 122 | + if not required_keys: |
| 123 | + return {} |
| 124 | + |
| 125 | + results: dict[GufeKey, ProtocolUnitResult] = {} |
| 126 | + unresolved = set(required_keys) |
| 127 | + |
| 128 | + for source_key in required_keys: |
| 129 | + loaded = self._load_result_from_index(source_key) |
| 130 | + if loaded is not None: |
| 131 | + results[source_key] = loaded |
| 132 | + unresolved.discard(source_key) |
| 133 | + |
| 134 | + if unresolved: |
| 135 | + scanned = self._scan_result_store_for_sources(unresolved) |
| 136 | + for source_key, loaded in scanned.items(): |
| 137 | + results[source_key] = loaded |
| 138 | + self._store_result_index(loaded) |
| 139 | + unresolved.discard(source_key) |
| 140 | + |
| 141 | + if unresolved: |
| 142 | + missing_keys = ", ".join(sorted(str(k) for k in unresolved)) |
| 143 | + raise RuntimeError( |
| 144 | + "Missing ProtocolUnitResult(s) for dependency key(s): " |
| 145 | + f"{missing_keys}. Ensure upstream tasks completed successfully." |
| 146 | + ) |
| 147 | + |
| 148 | + return results |
| 149 | + |
| 150 | + def _checkout_task(self) -> tuple[TaskStatusDB, str, ProtocolUnit] | None: |
| 151 | + """Check out one available task and load its protocol unit. |
| 152 | +
|
| 153 | + Returns |
| 154 | + ------- |
| 155 | + tuple[TaskStatusDB, str, ProtocolUnit] or None |
| 156 | + The open database connection, checked-out task ID, and corresponding |
| 157 | + protocol unit, or ``None`` if no task is currently available. |
| 158 | + The caller is responsible for calling ``mark_task_completed`` on the |
| 159 | + returned database using the returned task ID. |
| 160 | + """ |
| 161 | + |
| 162 | + db: TaskStatusDB = TaskStatusDB.from_filename(self.task_db_path) |
| 163 | + # The format for the taskid is "Transformation-<HASH>:ProtocolUnit-<HASH>" |
| 164 | + taskid = db.check_out_task() |
| 165 | + if taskid is None: |
| 166 | + return None |
| 167 | + |
| 168 | + _, protocol_unit_key = taskid.split(":", maxsplit=1) |
| 169 | + unit = self.warehouse.load_task(GufeKey(protocol_unit_key)) |
| 170 | + return db, taskid, unit |
| 171 | + |
| 172 | + def _get_task(self) -> tuple[str, ProtocolUnit]: |
| 173 | + """Return the next available task ID and protocol unit. |
| 174 | +
|
| 175 | + Returns |
| 176 | + ------- |
| 177 | + tuple[str, ProtocolUnit] |
| 178 | + The checked-out task ID and corresponding protocol unit. |
| 179 | +
|
| 180 | + Raises |
| 181 | + ------ |
| 182 | + RuntimeError |
| 183 | + Raised when no task is available in the task database. |
| 184 | + """ |
| 185 | + |
| 186 | + task = self._checkout_task() |
| 187 | + if task is None: |
| 188 | + raise RuntimeError("No AVAILABLE tasks found in the task database.") |
| 189 | + db, taskid, unit = task |
| 190 | + return taskid, unit |
| 191 | + |
| 192 | + def execute_unit(self, scratch: Path) -> tuple[str, ProtocolUnitResult] | None: |
| 193 | + """Execute one checked-out protocol unit and persist its result. |
| 194 | +
|
| 195 | + Parameters |
| 196 | + ---------- |
| 197 | + scratch : pathlib.Path |
| 198 | + Scratch directory passed to the protocol execution context. |
| 199 | +
|
| 200 | + Returns |
| 201 | + ------- |
| 202 | + tuple[str, ProtocolUnitResult] or None |
| 203 | + The task ID and execution result for the processed task, or |
| 204 | + ``None`` if no task is currently available. |
| 205 | +
|
| 206 | + Raises |
| 207 | + ------ |
| 208 | + Exception |
| 209 | + Re-raises any exception thrown during protocol unit execution after |
| 210 | + marking the task as failed. |
| 211 | + """ |
| 212 | + |
| 213 | + # 1. Get task/unit |
| 214 | + task = self._checkout_task() |
| 215 | + if task is None: |
| 216 | + return None |
| 217 | + db, taskid, unit = task |
| 218 | + # 2. Construct the context |
| 219 | + # NOTE: On changes to context (gufe PR #753), this can easily be replaced with external storage objects |
| 220 | + # However, to satisfy the current work, we will use this implementation where we |
| 221 | + # force the use of a FileSystemWarehouse and in turn can assert that an object is FileStorage. |
| 222 | + shared_store = self.warehouse.stores["shared"] |
| 223 | + if not isinstance(shared_store, FileStorage): |
| 224 | + raise TypeError("Expected a FileStorage backend for the shared store") |
| 225 | + shared_root_dir = shared_store.root_dir |
| 226 | + task_scratch, task_shared = self._task_workspace_paths(taskid, scratch, shared_root_dir) |
| 227 | + task_scratch.mkdir(parents=True, exist_ok=True) |
| 228 | + task_shared.mkdir(parents=True, exist_ok=True) |
| 229 | + ctx = Context(task_scratch, shared=task_shared) |
| 230 | + # 3. Execute unit |
| 231 | + try: |
| 232 | + results = self._build_input_result_mapping(unit) |
| 233 | + inputs = _pu_to_pur(unit.inputs, results) |
| 234 | + result = unit.execute(context=ctx, **inputs) |
| 235 | + except Exception: |
| 236 | + db.mark_task_completed(taskid, success=False) |
| 237 | + raise |
| 238 | + |
| 239 | + db.mark_task_completed(taskid, success=result.ok()) |
| 240 | + # 4. output result to warehouse |
| 241 | + # TODO: we may need to end up handling namespacing on the warehouse side for tokenizables |
| 242 | + self.warehouse.store_result_tokenizable(result) |
| 243 | + self._store_result_index(result) |
| 244 | + return taskid, result |
0 commit comments