Continue a finished workgraph

Introduction

One of the key features of WorkGraph is its ability to continue previous workflows. When a workgraph finishes its execution, it saves its state in the AiiDA process node. This allows you to rebuild the workgraph from the process and add new tasks to continue the workflow.

In the following sections, we will demonstrate how you can load a finished workgraph, amend it with new tasks, and continue the workflow.

from aiida_workgraph import WorkGraph
from aiida import load_profile

load_profile()

Profile<uuid='fa705ae5c8fb4460b0d168fea18da213' name='presto'>

Submit a workgraph

We first create and submit a simply workgraph, so that we can reload and continue it later.

with WorkGraph("AddMultiplyToBeContinued") as wg1:
    wg1.inputs = dict.fromkeys(["x", "y"])
    the_sum = wg1.inputs.x + wg1.inputs.y
    the_product = the_sum * 3
    wg1.outputs.result = {
        "sum": the_sum,
        "product": the_product,
    }

wg1.submit(
    inputs={
        "x": 1,
        "y": 2,
    },
    wait=True,
)

WorkGraph process created, PK: 488
Process 488 finished with state: FINISHED

<WorkGraphNode: uuid: 3cb23da0-21d5-41a1-be92-caf71c182d0e (pk: 488) (aiida.workflows:workgraph.engine)>

print("Results:")
print(f"  Sum: {wg1.outputs.result.sum.value}")
print(f"  Product: {wg1.outputs.result.product.value}")

Results:
  Sum: uuid: 56074b84-d04d-46f0-8ed9-401d9efb793b (pk: 492) value: 3
  Product: uuid: fbeae64f-b085-47ae-b010-28abf8a7d711 (pk: 496) value: 9

Modify the workgraph and resubmit

It is possible to reload the workgraph, modify an input, and resubmit. We use the load method to load the previous workgraph by its pk, mark it for restart, and modify the input of the multiplication task. This allows us to change the behavior of the workgraph without having to redefine it from scratch.

with WorkGraph.load(wg1.pk) as wg2:
    wg2.name = "AddMultiplyModified"
    wg2.restart()
    wg2.tasks.op_mul.inputs.y = 4

wg2.to_html()

wg2.submit(wait=True)

Reset tasks: ['op_mul']
WorkGraph process created, PK: 498
Process 498 finished with state: EXCEPTED

<WorkGraphNode: uuid: 3cf1cfa5-826d-4043-914f-cea5b523f156 (pk: 498) (aiida.workflows:workgraph.engine)>

print("Results:")
print(f"  Sum: {wg2.outputs.result.sum.value}")
print(f"  Product: {wg2.outputs.result.product.value}")

Results:
  Sum: uuid: 56074b84-d04d-46f0-8ed9-401d9efb793b (pk: 492) value: 3
  Product: uuid: fbeae64f-b085-47ae-b010-28abf8a7d711 (pk: 496) value: 9

Note that the sum has not changed (the value, but more importantly, the pk, as it is the same node). The product, however, is the result of the calculation repeating with the new input, hence a brand new node.

Continue a workgraph

Let’s now pick up the previous workgraph and extend it by a second addition, leveraging the results of the previous work.

with WorkGraph.load(wg2.pk) as wg3:
    wg3.name = "AddMultiplyContinued"
    wg3.inputs = dict.fromkeys(["z"])
    wg3.restart()
    wg3.outputs.result.new_sum = wg3.tasks.op_mul.outputs.result + wg3.inputs.z

wg3.to_html()

print(f"State of WorkGraph : {wg3.state}")
print(f"State of add       : {wg3.tasks.op_add.state}")
print(f"State of multiply  : {wg3.tasks.op_mul.state}")
print(f"State of new add   : {wg3.tasks.op_add1.state}")

State of WorkGraph : PLANNED
State of add       : FINISHED
State of multiply  : PLANNED
State of new add   : PLANNED

Note the PLANNED new addition task. Let’s run it.

wg3.submit(
    inputs={
        "z": 5,
    },
    wait=True,
)

Reset tasks: ['graph_inputs', 'graph_outputs']
WorkGraph process created, PK: 500
Process 500 finished with state: EXCEPTED

<WorkGraphNode: uuid: c7ad78e5-c845-4d81-a750-a3ecf3a0e6b1 (pk: 500) (aiida.workflows:workgraph.engine)>

print("Results:")
print(f"  Sum: {wg3.outputs.result.sum.value}")
print(f"  Product: {wg3.outputs.result.product.value}")
print(f"  New sum: {wg3.outputs.result.new_sum.value}")

Results:
  Sum: uuid: 56074b84-d04d-46f0-8ed9-401d9efb793b (pk: 492) value: 3
  Product: uuid: fbeae64f-b085-47ae-b010-28abf8a7d711 (pk: 496) value: 9
  New sum: None

Again, note that the previous data nodes are the same. Only the new addition task ran and created a new data node.

The extended provenance

The full provenance remains intact, including the original workflow.

from aiida_workgraph.utils import generate_node_graph

generate_node_graph(wg3.pk)

Summary

In this section, you learned how to load an existing workflow and rerun it with modified input, or continue it with new tasks. These features allow you to extend workflows without losing the original provenance, enabling a flexible and iterative approach to workflow management in AiiDA.