Add get_task_tracebacks method to expose stored tracebacks for Tasks

alchemiscale/interface/api.py

@@ -1176,6 +1176,24 @@ def get_task_failures(
     return [str(sk) for sk in n4js.get_task_failures(sk)]
 
 
+@router.get("/tasks/{task_scoped_key}/tracebacks")
+def get_task_tracebacks(
+    task_scoped_key: str,
+    *,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+    token: TokenData = Depends(get_token_data_depends),
+) -> list[dict[str, str]]:
+    """Get all stored tracebacks associated with a given Task.
+
+    Returns a list of dictionaries, one per failed ProtocolDAGResultRef,
+    each mapping ProtocolUnitFailure GufeKey to traceback string.
+    """
+    sk = ScopedKey.from_str(task_scoped_key)
+    validate_scopes(sk.scope, token)
+
+    return n4js.get_task_tracebacks(sk)
+
+
 ### strategies
 
 

alchemiscale/interface/client.py

@@ -1973,6 +1973,38 @@ def get_task_failures(
 
         return pdrs
 
+    def get_task_tracebacks(self, task: ScopedKey) -> list[dict[str, str]]:
+        """Get stored tracebacks associated with a given `Task`.
+
+        This method retrieves traceback information from failed
+        `ProtocolDAGResult`s associated with the `Task`. Each traceback is
+        returned as a dictionary mapping the `ProtocolUnitFailure` `GufeKey` to
+        its corresponding traceback string.
+
+        Parameters
+        ----------
+        task
+            The `ScopedKey` of the `Task` to retrieve tracebacks for.
+
+        Returns
+        -------
+        list[dict[str, str]]
+            A list of dictionaries, one per failed `ProtocolDAGResultRef`
+            associated with the `Task`. Each dictionary maps
+            `ProtocolUnitFailure` `GufeKey` strings to their corresponding
+            traceback strings. Returns an empty list if no tracebacks are found,
+            including for Tasks that have no failures or do not exist.
+
+        Examples
+        --------
+        >>> tracebacks = client.get_task_tracebacks(task_sk)
+        >>> for tb_dict in tracebacks:
+        ...     for failure_key, traceback_str in tb_dict.items():
+        ...         print(f"Failure {failure_key}: {traceback_str[:100]}...")
+
+        """
+        return self._get_resource(f"/tasks/{task}/tracebacks")
+
     def add_task_restart_patterns(
         self,
         network_scoped_key: ScopedKey,

alchemiscale/storage/statestore.py

@@ -3547,6 +3547,52 @@ def add_protocol_dag_result_ref_tracebacks(
 
             merge_subgraph(tx, subgraph, "GufeTokenizable", "_scoped_key")
 
+    def get_task_tracebacks(self, task: ScopedKey) -> list[dict[str, str]]:
+        """Get all stored tracebacks associated with a given Task.
+
+        This method retrieves traceback information from failed ProtocolDAGResults
+        associated with the Task. Each traceback is returned as a dictionary mapping
+        the ProtocolUnitFailure GufeKey to its corresponding traceback string.
+
+        Parameters
+        ----------
+        task
+            The ScopedKey of the Task to retrieve tracebacks for.
+
+        Returns
+        -------
+        list[dict[str, str]]
+            A list of dictionaries, one per failed ProtocolDAGResultRef associated
+            with the Task. Each dictionary maps ProtocolUnitFailure GufeKey strings
+            to their corresponding traceback strings. Returns an empty list if no
+            tracebacks are found, including for Tasks that have no failures or do
+            not exist.
+
+        """
+        q = """
+        MATCH (task:Task {_scoped_key: $scoped_key})-[:RESULTS_IN]->(pdrr:ProtocolDAGResultRef)<-[:DETAILS]-(tracebacks:Tracebacks)
+        WHERE pdrr.ok = false
+        RETURN tracebacks.tracebacks AS tracebacks, tracebacks.failure_keys AS failure_keys
+        ORDER BY pdrr.datetime_created DESC
+        """
+
+        results = []
+        with self.transaction() as tx:
+            res = tx.run(q, scoped_key=str(task))
+
+            for record in res:
+                tracebacks = record["tracebacks"]
+                failure_keys = record["failure_keys"]
+
+                # Create a mapping of failure_key -> traceback
+                traceback_mapping = {}
+                for failure_key, traceback in zip(failure_keys, tracebacks):
+                    traceback_mapping[failure_key] = traceback
+
+                results.append(traceback_mapping)
+
+        return results
+
     def set_task_status(
         self, tasks: list[ScopedKey], status: TaskStatusEnum, raise_error: bool = False
     ) -> list[ScopedKey | None]:

alchemiscale/tests/integration/interface/client/test_client.py

@@ -2633,6 +2633,137 @@ def test_get_task_failures(
         # TODO: can we mix in a success in here somewhere?
         # not possible with current BrokenProtocol, unfortunately
 
+    def test_get_task_tracebacks(
+        self,
+        scope_test,
+        n4js_preloaded,
+        s3os_server,
+        user_client: client.AlchemiscaleClient,
+        network_tyk2_failure,
+        tmpdir,
+    ):
+        """Test that get_task_tracebacks returns traceback information for failed tasks."""
+        n4js = n4js_preloaded
+
+        # select the transformation we want to compute
+        an = network_tyk2_failure
+        network_sk = user_client.create_network(an, scope_test)
+
+        # ensure full AlchemicalNetwork is present before we proceed
+        while True:
+            try:
+                an_ = user_client.get_network(network_sk)
+
+                if an_ != an:
+                    raise Exception("Network out doesn't exactly match network in yet")
+                else:
+                    break
+            except Exception:
+                sleep(0.1)
+
+        tf_sks = user_client.get_network_transformations(network_sk)
+
+        # select the transformation we want to compute
+        for tf_sk in tf_sks:
+            tf = user_client.get_transformation(tf_sk)
+            if tf.name == "broken":
+                transformation_sk = tf_sk
+                break
+
+        # user client : create tasks for the transformation
+        tasks = user_client.create_tasks(transformation_sk, count=1)
+
+        # action the tasks
+        actioned_tasks = user_client.action_tasks(tasks, network_sk)
+
+        # execute the task and push results directly
+        with tmpdir.as_cwd():
+            try:
+                protocoldagresults = self._execute_tasks(
+                    actioned_tasks, n4js, s3os_server
+                )
+            except TypeError as e:
+                if (
+                    str(e)
+                    == "Transformation.__init__() missing 3 required positional arguments: 'stateA', 'stateB', and 'protocol'"
+                ):
+                    pytest.xfail()
+                else:
+                    raise e
+
+            # push results and add tracebacks for failures
+            for task_sk, pdr in zip(actioned_tasks, protocoldagresults):
+                transformation_sk_task, _ = n4js.get_task_transformation(
+                    task_sk, return_gufe=False
+                )
+                protocoldagresultref = s3os_server.push_protocoldagresult(
+                    compress_gufe_zstd(pdr),
+                    pdr.ok(),
+                    pdr.key,
+                    transformation=transformation_sk_task,
+                )
+                result_sk = n4js.set_task_result(
+                    task=task_sk, protocoldagresultref=protocoldagresultref
+                )
+
+                # add tracebacks for failures (this is what we're testing)
+                if not pdr.ok():
+                    n4js.add_protocol_dag_result_ref_tracebacks(
+                        pdr.protocol_unit_failures, result_sk
+                    )
+                    n4js.set_task_error(tasks=[task_sk])
+
+        # now test get_task_tracebacks
+        for task in tasks:
+            tracebacks = user_client.get_task_tracebacks(task)
+
+            # should have at least one set of tracebacks for the failed task
+            assert len(tracebacks) >= 1
+
+            # each entry should be a dict mapping failure key to traceback
+            for tb_dict in tracebacks:
+                assert isinstance(tb_dict, dict)
+                for failure_key, traceback_str in tb_dict.items():
+                    assert isinstance(failure_key, str)
+                    assert isinstance(traceback_str, str)
+                    # traceback should not be empty
+                    assert len(traceback_str) > 0
+
+    def test_get_task_tracebacks_no_failures(
+        self,
+        scope_test,
+        n4js_preloaded,
+        user_client: client.AlchemiscaleClient,
+        network_tyk2,
+    ):
+        """Test that get_task_tracebacks returns empty list for task with no failures."""
+        n4js = n4js_preloaded
+
+        # create a simple network and task but don't execute it
+        an = network_tyk2
+        network_sk = user_client.create_network(an, scope_test)
+
+        # ensure full AlchemicalNetwork is present before we proceed
+        while True:
+            try:
+                an_ = user_client.get_network(network_sk)
+                if an_ != an:
+                    raise Exception("Network out doesn't exactly match network in yet")
+                else:
+                    break
+            except Exception:
+                sleep(0.1)
+
+        tf_sks = user_client.get_network_transformations(network_sk)
+        transformation_sk = tf_sks[0]
+
+        # create a task but don't execute it
+        tasks = user_client.create_tasks(transformation_sk, count=1)
+
+        # get_task_tracebacks should return empty list
+        tracebacks = user_client.get_task_tracebacks(tasks[0])
+        assert tracebacks == []
+
 
 class TestTaskRestartPolicy:
     default_max_retries = 3

alchemiscale/tests/integration/storage/test_statestore.py

@@ -2264,6 +2264,125 @@ def test_add_protocol_dag_result_ref_traceback(
 
         assert returned_tracebacks == [puf.traceback for puf in protocol_unit_failures]
 
+    @pytest.mark.parametrize("failure_count", (1, 2, 3))
+    def test_get_task_tracebacks(
+        self,
+        network_tyk2_failure,
+        n4js,
+        scope_test,
+        transformation_failure,
+        protocoldagresults_failure,
+        failure_count: int,
+    ):
+        """Test retrieving tracebacks associated with a Task."""
+        an = network_tyk2_failure.copy_with_replacements(
+            name=network_tyk2_failure.name + "_test_get_task_tracebacks"
+        )
+        n4js.assemble_network(an, scope_test)
+        transformation_scoped_key = n4js.get_scoped_key(
+            transformation_failure, scope_test
+        )
+
+        # create a task; pretend we computed it, submit reference for pre-baked result
+        task_scoped_key = n4js.create_task(transformation_scoped_key)
+
+        protocol_unit_failure = protocoldagresults_failure[0].protocol_unit_failures[0]
+
+        pdrr = ProtocolDAGResultRef(
+            scope=task_scoped_key.scope,
+            obj_key=protocoldagresults_failure[0].key,
+            ok=protocoldagresults_failure[0].ok(),
+        )
+
+        # push the result
+        pdrr_scoped_key = n4js.set_task_result(task_scoped_key, pdrr)
+
+        # simulating many failures
+        protocol_unit_failures = []
+        for failure_index in range(failure_count):
+            protocol_unit_failures.append(
+                protocol_unit_failure.copy_with_replacements(
+                    traceback=protocol_unit_failure.traceback + "_" + str(failure_index)
+                )
+            )
+
+        n4js.add_protocol_dag_result_ref_tracebacks(
+            protocol_unit_failures, pdrr_scoped_key
+        )
+
+        # now test the get_task_tracebacks method
+        tracebacks = n4js.get_task_tracebacks(task_scoped_key)
+
+        assert len(tracebacks) == 1  # one ProtocolDAGResultRef
+
+        # check that the traceback mapping is correct
+        traceback_dict = tracebacks[0]
+        assert len(traceback_dict) == failure_count
+
+        for puf in protocol_unit_failures:
+            assert puf.key in traceback_dict
+            assert traceback_dict[puf.key] == puf.traceback
+
+    def test_get_task_tracebacks_no_failures(
+        self,
+        network_tyk2,
+        n4js,
+        scope_test,
+    ):
+        """Test that get_task_tracebacks returns empty list for task with no failures."""
+        an = network_tyk2.copy_with_replacements(
+            name=network_tyk2.name + "_test_get_task_tracebacks_no_failures"
+        )
+        n4js.assemble_network(an, scope_test)
+        transformation = list(an.edges)[0]
+        transformation_scoped_key = n4js.get_scoped_key(transformation, scope_test)
+
+        # create a task without any results
+        task_scoped_key = n4js.create_task(transformation_scoped_key)
+
+        # should return empty list
+        tracebacks = n4js.get_task_tracebacks(task_scoped_key)
+        assert tracebacks == []
+
+    def test_get_task_tracebacks_multiple_failures(
+        self,
+        network_tyk2_failure,
+        n4js,
+        scope_test,
+        transformation_failure,
+        protocoldagresults_failure,
+    ):
+        """Test retrieving tracebacks when a task has multiple failed results."""
+        an = network_tyk2_failure.copy_with_replacements(
+            name=network_tyk2_failure.name + "_test_get_task_tracebacks_multiple"
+        )
+        n4js.assemble_network(an, scope_test)
+        transformation_scoped_key = n4js.get_scoped_key(
+            transformation_failure, scope_test
+        )
+
+        task_scoped_key = n4js.create_task(transformation_scoped_key)
+
+        protocol_unit_failure = protocoldagresults_failure[0].protocol_unit_failures[0]
+
+        # Create two failed ProtocolDAGResultRefs
+        for i in range(2):
+            pdrr = ProtocolDAGResultRef(
+                scope=task_scoped_key.scope,
+                obj_key=protocoldagresults_failure[i].key,
+                ok=protocoldagresults_failure[i].ok(),
+            )
+            pdrr_scoped_key = n4js.set_task_result(task_scoped_key, pdrr)
+
+            puf = protocol_unit_failure.copy_with_replacements(
+                traceback=f"Error traceback for failure {i}"
+            )
+            n4js.add_protocol_dag_result_ref_tracebacks([puf], pdrr_scoped_key)
+
+        # get tracebacks - should have 2 entries
+        tracebacks = n4js.get_task_tracebacks(task_scoped_key)
+        assert len(tracebacks) == 2
+
     ### task restart policies
 
     class TestTaskRestartPolicy: