diff --git a/alchemiscale/compute/api.py b/alchemiscale/compute/api.py index ada77df4..6e81c6f7 100644 --- a/alchemiscale/compute/api.py +++ b/alchemiscale/compute/api.py @@ -400,7 +400,14 @@ async def set_task_result( n4js.add_protocol_dag_result_ref_tracebacks( pdr.protocol_unit_failures, result_sk ) - n4js.set_task_error(tasks=[task_sk]) + # Format protocol unit failures into a reason string + failure_reasons = [] + for failure in pdr.protocol_unit_failures: + failure_reasons.append( + f"ProtocolUnit '{failure}' failed:\n{failure.exception}" + ) + reason = "\n\n".join(failure_reasons) if failure_reasons else None + n4js.set_task_error(tasks=[task_sk], reason=reason) # report that the compute service experienced a failure now = datetime.datetime.now(tz=datetime.UTC) @@ -410,6 +417,39 @@ async def set_task_result( return result_sk +@router.post("/tasks/{task_scoped_key}/error") +async def set_task_error( + task_scoped_key, + *, + request: Request, + n4js: Neo4jStore = Depends(get_n4js_depends), + token: TokenData = Depends(get_token_data_depends), +): + """Set a task to error status with an optional reason. + + This endpoint is used when a task fails before a ProtocolDAGResult can be created, + such as during ProtocolDAG creation from a Transformation. + """ + body = await request.body() + body_ = json.loads(body.decode("utf-8"), cls=JSON_HANDLER.decoder) + + reason = body_.get("reason") + compute_service_id = body_["compute_service_id"] + + task_sk = ScopedKey.from_str(task_scoped_key) + validate_scopes(task_sk.scope, token) + + # Set task to error status with the provided reason + n4js.set_task_error(tasks=[task_sk], reason=reason) + + # Report that the compute service experienced a failure + now = datetime.datetime.now(tz=datetime.UTC) + n4js.log_failure_compute_service(compute_service_id, now) + n4js.resolve_task_restarts(task_scoped_keys=[task_sk]) + + return task_sk + + def process_compute_manager_id_string( compute_manager_id_string: str, ) -> ComputeManagerID: diff --git a/alchemiscale/compute/client.py b/alchemiscale/compute/client.py index 7a685925..b7bf06cb 100644 --- a/alchemiscale/compute/client.py +++ b/alchemiscale/compute/client.py @@ -159,6 +159,26 @@ def set_task_result( return ScopedKey.from_dict(pdr_sk) + def set_task_error( + self, + task: ScopedKey, + reason: str | None = None, + compute_service_id: ComputeServiceID | None = None, + ) -> ScopedKey: + """Set a task to error status with an optional reason. + + This is used when a task fails before a ProtocolDAGResult can be created, + such as during ProtocolDAG creation from a Transformation. + """ + data = dict( + reason=reason, + compute_service_id=str(compute_service_id), + ) + + task_sk = self._post_resource(f"/tasks/{task}/error", data) + + return ScopedKey.from_dict(task_sk) + class AlchemiscaleComputeManagerClientError(AlchemiscaleBaseClientError): ... diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py index dc4ee30c..cbbd5cee 100644 --- a/alchemiscale/compute/service.py +++ b/alchemiscale/compute/service.py @@ -12,6 +12,7 @@ import threading from pathlib import Path import shutil +import traceback from gufe import Transformation from gufe.protocols.protocoldag import execute_DAG, ProtocolDAG, ProtocolDAGResult @@ -220,13 +221,30 @@ def execute(self, task: ScopedKey) -> ScopedKey: """ # obtain a ProtocolDAG from the task self.logger.info("Creating ProtocolDAG from '%s'...", task) - protocoldag, transformation, extends = self.task_to_protocoldag(task) - self.logger.info( - "Created '%s' from '%s' performing '%s'", - protocoldag, - task, - transformation.protocol, - ) + try: + protocoldag, transformation, extends = self.task_to_protocoldag(task) + self.logger.info( + "Created '%s' from '%s' performing '%s'", + protocoldag, + task, + transformation.protocol, + ) + except Exception as e: + # If ProtocolDAG creation fails, set task to error with traceback + error_traceback = traceback.format_exc() + self.logger.error( + "Failed to create ProtocolDAG from '%s': %s\n%s", + task, + str(e), + error_traceback, + ) + # Set task to error status with the exception traceback as reason + task_sk = self.client.set_task_error( + task=task, + reason=error_traceback, + compute_service_id=self.compute_service_id, + ) + return task_sk # execute the task; this looks the same whether the ProtocolDAG is a # success or failure diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py index a9aaf8bd..c44e096f 100644 --- a/alchemiscale/storage/models.py +++ b/alchemiscale/storage/models.py @@ -180,6 +180,14 @@ class Task(GufeTokenizable): claim Identifier of the compute service that has a claim on this task. datetime_created + Timestamp when the task was created. + creator + Identifier of who/what created the task. + extends + Reference to another task this task extends from. + reason + Optional reason field for task state changes, e.g., error tracebacks + or user-provided reasons for manual state changes to deleted/invalid. """ @@ -189,6 +197,7 @@ class Task(GufeTokenizable): datetime_created: datetime.datetime | None creator: str | None extends: str | None + reason: str | None def __init__( self, @@ -199,6 +208,7 @@ def __init__( creator: str | None = None, extends: str | None = None, claim: str | None = None, + reason: str | None = None, _key: str = None, ): if _key is not None: @@ -216,6 +226,7 @@ def __init__( self.creator = creator self.extends = extends self.claim = claim + self.reason = reason def _gufe_tokenize(self): # tokenize with uuid @@ -229,6 +240,7 @@ def _to_dict(self): "creator": self.creator, "extends": self.extends, "claim": self.claim, + "reason": self.reason, "_key": str(self.key), } diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index ce6de601..de920b0e 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -3610,11 +3610,11 @@ def get_task_status(self, tasks: list[ScopedKey]) -> list[TaskStatusEnum]: return statuses def _set_task_status( - self, tasks, q: str, err_msg_func, raise_error + self, tasks, q: str, err_msg_func, raise_error, **kwargs ) -> list[ScopedKey | None]: tasks_statused = [] with self.transaction() as tx: - res = tx.run(q, scoped_keys=[str(t) for t in tasks]) + res = tx.run(q, scoped_keys=[str(t) for t in tasks], **kwargs) for record in res: task_i = record["t"] @@ -3740,16 +3740,25 @@ def err_msg(t, status): return self._set_task_status(tasks, q, err_msg, raise_error=raise_error) def set_task_error( - self, tasks: list[ScopedKey], raise_error: bool = False + self, tasks: list[ScopedKey], reason: str | None = None, raise_error: bool = False ) -> list[ScopedKey | None]: """Set the status of a list of Tasks to `error`. Only `running` Tasks can be set to `error`. + Parameters + ---------- + tasks + List of task ScopedKeys to set to error status. + reason + Optional reason for the error (e.g., exception traceback). + raise_error + Whether to raise an error if the task cannot be set to error. + """ q = f""" - WITH $scoped_keys AS batch + WITH $scoped_keys AS batch, $reason AS reason UNWIND batch AS scoped_key OPTIONAL MATCH (t:Task {{_scoped_key: scoped_key}}) @@ -3757,6 +3766,7 @@ def set_task_error( OPTIONAL MATCH (t_:Task {{_scoped_key: scoped_key}}) WHERE t_.status IN ['{TaskStatusEnum.error.value}', '{TaskStatusEnum.running.value}'] SET t_.status = '{TaskStatusEnum.error.value}' + SET t_.reason = reason WITH scoped_key, t, t_ @@ -3771,23 +3781,32 @@ def set_task_error( def err_msg(t, status): return f"Cannot set task {t} with current status: {status} to `error` as it is not currently `running`." - return self._set_task_status(tasks, q, err_msg, raise_error=raise_error) + return self._set_task_status(tasks, q, err_msg, raise_error=raise_error, reason=reason) def set_task_invalid( - self, tasks: list[ScopedKey], raise_error: bool = False + self, tasks: list[ScopedKey], reason: str | None = None, raise_error: bool = False ) -> list[ScopedKey | None]: """Set the status of a list of Tasks to `invalid`. Any Task can be set to `invalid`; an `invalid` Task cannot change to any other status. + Parameters + ---------- + tasks + List of task ScopedKeys to set to invalid status. + reason + Optional reason for invalidating the task (e.g., user-provided explanation). + raise_error + Whether to raise an error if the task cannot be set to invalid. + """ # set the status and delete the ACTIONS relationship # make sure we follow the extends chain and set all tasks to invalid # and remove actions relationships q = f""" - WITH $scoped_keys AS batch + WITH $scoped_keys AS batch, $reason AS reason UNWIND batch AS scoped_key OPTIONAL MATCH (t:Task {{_scoped_key: scoped_key}}) @@ -3795,11 +3814,13 @@ def set_task_invalid( OPTIONAL MATCH (t_:Task {{_scoped_key: scoped_key}}) WHERE NOT t_.status IN ['{TaskStatusEnum.deleted.value}'] SET t_.status = '{TaskStatusEnum.invalid.value}' + SET t_.reason = reason WITH scoped_key, t, t_ OPTIONAL MATCH (t_)<-[er:EXTENDS*]-(extends_task:Task) SET extends_task.status = '{TaskStatusEnum.invalid.value}' + SET extends_task.reason = reason WITH scoped_key, t, t_, extends_task @@ -3825,23 +3846,32 @@ def set_task_invalid( def err_msg(t, status): return f"Cannot set task {t} with current status: {status} to `invalid` as it is `deleted`." - return self._set_task_status(tasks, q, err_msg, raise_error=raise_error) + return self._set_task_status(tasks, q, err_msg, raise_error=raise_error, reason=reason) def set_task_deleted( - self, tasks: list[ScopedKey], raise_error: bool = False + self, tasks: list[ScopedKey], reason: str | None = None, raise_error: bool = False ) -> list[ScopedKey | None]: """Set the status of a list of Tasks to `deleted`. Any Task can be set to `deleted`; a `deleted` Task cannot change to any other status. + Parameters + ---------- + tasks + List of task ScopedKeys to set to deleted status. + reason + Optional reason for deleting the task (e.g., user-provided explanation). + raise_error + Whether to raise an error if the task cannot be set to deleted. + """ # set the status and delete the ACTIONS relationship # make sure we follow the extends chain and set all tasks to deleted # and remove actions relationships q = f""" - WITH $scoped_keys AS batch + WITH $scoped_keys AS batch, $reason AS reason UNWIND batch AS scoped_key OPTIONAL MATCH (t:Task {{_scoped_key: scoped_key}}) @@ -3849,11 +3879,13 @@ def set_task_deleted( OPTIONAL MATCH (t_:Task {{_scoped_key: scoped_key}}) WHERE NOT t_.status IN ['{TaskStatusEnum.invalid.value}'] SET t_.status = '{TaskStatusEnum.deleted.value}' + SET t_.reason = reason WITH scoped_key, t, t_ OPTIONAL MATCH (t_)<-[er:EXTENDS*]-(extends_task:Task) SET extends_task.status = '{TaskStatusEnum.deleted.value}' + SET extends_task.reason = reason WITH scoped_key, t, t_, extends_task @@ -3879,7 +3911,7 @@ def set_task_deleted( def err_msg(t, status): return f"Cannot set task {t} with current status: {status} to `deleted` as it is `invalid`." - return self._set_task_status(tasks, q, err_msg, raise_error=raise_error) + return self._set_task_status(tasks, q, err_msg, raise_error=raise_error, reason=reason) ## task restart policies