From 5784ab30880ad048f855ffd0f9e31d196de5d587 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Mon, 19 Jan 2026 16:25:04 -0700 Subject: [PATCH 1/3] Add get_task_tracebacks method to expose stored tracebacks for Tasks (#347) Adds the ability to retrieve traceback information from failed Tasks: - Add `Neo4jStore.get_task_tracebacks()` method to statestore - Add `/tasks/{task_scoped_key}/tracebacks` API endpoint - Add `AlchemiscaleClient.get_task_tracebacks()` client method - Add integration tests for all layers The method returns a list of dictionaries, one per failed ProtocolDAGResultRef, mapping ProtocolUnitFailure GufeKey to traceback string. Closes #347 Co-Authored-By: Claude Opus 4.5 --- alchemiscale/interface/api.py | 18 +++ alchemiscale/interface/client.py | 31 +++++ alchemiscale/storage/statestore.py | 47 +++++++ .../interface/client/test_client.py | 131 ++++++++++++++++++ .../integration/storage/test_statestore.py | 119 ++++++++++++++++ 5 files changed, 346 insertions(+) diff --git a/alchemiscale/interface/api.py b/alchemiscale/interface/api.py index e41a1c87..147af5e4 100644 --- a/alchemiscale/interface/api.py +++ b/alchemiscale/interface/api.py @@ -1177,6 +1177,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, + *, + 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 diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index 9a321a39..726e17af 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -1920,6 +1920,37 @@ 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. + + 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, diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index ce6de601..3dedb32e 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -3551,6 +3551,53 @@ 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. + + """ + 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]: diff --git a/alchemiscale/tests/integration/interface/client/test_client.py b/alchemiscale/tests/integration/interface/client/test_client.py index 974574e0..8839f1a0 100644 --- a/alchemiscale/tests/integration/interface/client/test_client.py +++ b/alchemiscale/tests/integration/interface/client/test_client.py @@ -2559,6 +2559,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 diff --git a/alchemiscale/tests/integration/storage/test_statestore.py b/alchemiscale/tests/integration/storage/test_statestore.py index 621644a3..d18b4d77 100644 --- a/alchemiscale/tests/integration/storage/test_statestore.py +++ b/alchemiscale/tests/integration/storage/test_statestore.py @@ -2270,6 +2270,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: From 24a63b66392c7d80d29dc367fda2f6a323863245 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Mon, 19 Jan 2026 16:54:31 -0700 Subject: [PATCH 2/3] Address review feedback: add type annotation and improve docstrings - Add `str` type annotation to task_scoped_key parameter in API endpoint - Clarify in docstrings that empty list is returned for non-existent tasks Co-Authored-By: Claude Opus 4.5 --- alchemiscale/interface/api.py | 2 +- alchemiscale/interface/client.py | 3 ++- alchemiscale/storage/statestore.py | 3 ++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/alchemiscale/interface/api.py b/alchemiscale/interface/api.py index 147af5e4..0b2d3a0e 100644 --- a/alchemiscale/interface/api.py +++ b/alchemiscale/interface/api.py @@ -1179,7 +1179,7 @@ def get_task_failures( @router.get("/tasks/{task_scoped_key}/tracebacks") def get_task_tracebacks( - task_scoped_key, + task_scoped_key: str, *, n4js: Neo4jStore = Depends(get_n4js_depends), token: TokenData = Depends(get_token_data_depends), diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index 726e17af..521fb232 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -1939,7 +1939,8 @@ def get_task_tracebacks(self, task: ScopedKey) -> 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. + traceback strings. Returns an empty list if no tracebacks are found, + including for Tasks that have no failures or do not exist. Examples -------- diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index 3dedb32e..eb184267 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -3571,7 +3571,8 @@ def get_task_tracebacks( 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. + tracebacks are found, including for Tasks that have no failures or do + not exist. """ q = """ From 75ec00d62e57a2cf210af8e67cc0d4705222877b Mon Sep 17 00:00:00 2001 From: David Dotson Date: Wed, 10 Jun 2026 17:47:08 -0600 Subject: [PATCH 3/3] Black! --- alchemiscale/storage/statestore.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index 19b2f1f9..5403d310 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -3547,9 +3547,7 @@ 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]]: + 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