Merge pull request #286 from OpenFreeEnergy/feature/iss-277-restart-p…

…olicy_resolve_restarts Restart policy: resolve restarts
OpenFreeEnergy · Sep 19, 2024 · cf0e961 · cf0e961
2 parents 2310fd5 + 7a4b114
commit cf0e961
Show file tree

Hide file tree

Showing 7 changed files with 665 additions and 52 deletions.
diff --git a/alchemiscale/compute/api.py b/alchemiscale/compute/api.py
@@ -12,6 +12,7 @@
 from fastapi import FastAPI, APIRouter, Body, Depends
 from fastapi.middleware.gzip import GZipMiddleware
 from gufe.tokenization import GufeTokenizable, JSON_HANDLER
+from gufe.protocols import ProtocolDAGResult
 
 from ..base.api import (
     QueryGUFEHandler,
@@ -248,7 +249,7 @@ def set_task_result(
     validate_scopes(task_sk.scope, token)
 
     pdr = json.loads(protocoldagresult, cls=JSON_HANDLER.decoder)
-    pdr = GufeTokenizable.from_dict(pdr)
+    pdr: ProtocolDAGResult = GufeTokenizable.from_dict(pdr)
 
     tf_sk, _ = n4js.get_task_transformation(
         task=task_scoped_key,
@@ -270,7 +271,11 @@ def set_task_result(
     if protocoldagresultref.ok:
         n4js.set_task_complete(tasks=[task_sk])
     else:
+        n4js.add_protocol_dag_result_ref_tracebacks(
+            pdr.protocol_unit_failures, result_sk
+        )
         n4js.set_task_error(tasks=[task_sk])
+        n4js.resolve_task_restarts(tasks=[task_sk])
 
     return result_sk
 

diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py
@@ -202,9 +202,11 @@ def __eq__(self, other):
         return self.pattern == other.pattern
 
 
-class Traceback(GufeTokenizable):
+class Tracebacks(GufeTokenizable):
 
-    def __init__(self, tracebacks: List[str]):
+    def __init__(
+        self, tracebacks: List[str], source_keys: List[str], failure_keys: List[str]
+    ):
         value_error = ValueError(
             "`tracebacks` must be a non-empty list of string values"
         )
@@ -216,21 +218,25 @@ def __init__(self, tracebacks: List[str]):
             if not all_string_values or "" in tracebacks:
                 raise value_error
 
+        # TODO: validate
         self.tracebacks = tracebacks
-
-    def _gufe_tokenize(self):
-        return hashlib.md5(str(self.tracebacks).encode()).hexdigest()
+        self.source_keys = source_keys
+        self.failure_keys = failure_keys
 
     @classmethod
     def _defaults(cls):
-        raise NotImplementedError
+        return super()._defaults()
 
     @classmethod
     def _from_dict(cls, dct):
-        return Traceback(**dct)
+        return cls(**dct)
 
     def _to_dict(self):
-        return {"tracebacks": self.tracebacks}
+        return {
+            "tracebacks": self.tracebacks,
+            "source_keys": self.source_keys,
+            "failure_keys": self.failure_keys,
+        }
 
 
 class TaskHub(GufeTokenizable):