aiidateam / aiida-workgraph

Efficiently design and manage flexible workflows with AiiDA, featuring an interactive GUI, checkpoints, provenance tracking, and remote execution capabilities.
https://aiida-workgraph.readthedocs.io/en/latest/
MIT License
9 stars 5 forks source link

Nested while task stuck #218

Closed superstar54 closed 2 months ago

superstar54 commented 2 months ago

AiiDA-WorkGraph version 0.3.16. Here is an example with nested while task,

from aiida import load_profile
from aiida_workgraph import task, WorkGraph

load_profile()

# we need a compare task for `n<100`,
# it's a normal function instead of a calcfunction
@task()
def compare(x, y):
    return x < y

# define multiply task for n*2
@task.calcfunction()
def multiply(x, y):
    return x*y

# define add task for n+3
@task.calcfunction()
def add(x, y):
    return x + y

n = 1
wg = WorkGraph("test_while")
# set a context variable before running.
wg.context = {"n": 1,
              "should_run": True}
add1 = wg.add_task(add, name="add1", x="{{n}}", y=1)
add1.set_context({"result": "n"})
#---------------------------------------------------------------------
# Create a WorkGraph will repeat itself based on the conditions
# then we output the result of from the context (context)
add2 = wg.add_task(add, name="add2", x="{{n}}", y=1)
add2.wait.append("add1")
multiply1 = wg.add_task(multiply, name="multiply1",
                        x=add2.outputs["result"],
                        y=2)
# update the context variable
multiply1.set_context({"result": "n"})
compare1 = wg.add_task(compare, name="compare1", x=multiply1.outputs["result"], y=30)
compare1.set_context({"result": "should_run"})
while1 = wg.add_task("While", name = "while1", max_iterations=100,
                        conditions=["should_run"],
                        tasks=["add2", "multiply1", "compare1"],)
# the `result` of compare1 taskis used as condition
#---------------------------------------------------------------------
add3 = wg.add_task(add, name="add3", x=1, y=1)
add3.set_context({"result": "n"})
wg.add_link(multiply1.outputs["result"], add3.inputs["x"])
#---------------------------------------------------------------------
add4 = wg.add_task(add, name="add4", x=1, y=1)
add4.set_context({"result": "n"})
compare2 = wg.add_task(compare, name="compare2", x=add3.outputs["result"], y=100)
compare2.set_context({"result": "should_run"})
while2 = wg.add_task("While", name="while2",
                        conditions=["should_run"],
                        tasks=["add1", "while1", "add3", "compare2"],)
while2.wait.append("add4")
add5 = wg.add_task(add, name="add5", x=1, y=add3.outputs["result"])
add5.set_context({"result": "n"})

# the `result` of compare1 taskis used as condition
wg.run()

The WorkGraph stucks at this state:

Screenshot from 2024-08-11 22-48-04