Improve workchain/workgraph _context

[agoscinski]

The issue is description is kept very general, more specific issues will we be here referenced as subissues. We have several problems with the context and my description might be a bit vague because they are not isolated problems touching a lot of code, and I don't fully understand how to tackle them.

Origin of assignments to context

In aiida we use the context to overcome the constraint that the provenance graph is acyclic. In the workgraph the usage of the context prevents a correct representation in the GUI, the origin of the context variable is not correctly tracked. In workgraph the information is available in the provenance graph but not in the GUI, so it must me somewhere accessible in aiida-core and could be used to make it available in the GUI.

_context is used for different use cases

Currently workgraph creates an entry in context for user context variables (self.ctx._tasks[name]["context_mapping"]) to separate it from internal usage. In aiida-core this separation does not exist. I think here workgraph does the right thing to separate. Also at least the term "context" is sometimes also used for the usage of PortNamespace. I am not sure if there is a reason for it or it is just poor naming. The context I am referring to is a AttributeDict.

Inheritance problem in WorkGraph

This could be in aiida-workgraph but I just put it here for now. I think also WorkGraph is not correctly inheriting from aiida.engine.WorkChain or plumpy.mixins.ContextMixin, it just assumes a _context variable and property ctx.

Current representation of loop and conditionals in aiida-core/plumpy is not sufficient for a representation in the provenance graph

Workgraph is using context to create loops and conditionals which is very confusing as it stands in contrast all other usages. In plumpy there is a class for representing loops and conditional (the subclasses of the Stepper), but they are not good enough integrated in aiida to be used in the workgraph. The leading question for me is what prevents workgraph from using these Stepper subclasses for its implementation of loop and conditional logic.

[unkcpz]

Thanks for open the issue, btw I like the tag "design-issue" ;)

I change the name title process -> workchain/workgraph, hope it can be more clearly express the issue behind.

In workgraph the information is available in the provenance graph but not in the GUI, so it must me somewhere accessible in aiida-core and could be used to make it available in the GUI.

I don't understand the description here, because for things showing in the provenance graph , it needs to be a "node". The self.ctx is for hold the temporary states between steps.

Currently workgraph creates an entry in context for user context variables (self.ctx._tasks[name]["context_mapping"]) to separate it from internal usage. In aiida-core this separation does not exist.

Since self.ctx is not for storing node, it is not true that in aiida-core this separation not exist, it is because the self.ctx is not design for storing the node information. The self.ctx will be serialized and stored in the checkpoint so the workchain states can be restored.

But I agree that the steppers' information should be stored. If I understand correctly, the GUI you mentioned is not specificly to provenance graph but it is the workgraph gui that can read workgraph info and draw running logic? If so, then storing the stepper info from dynamically running workchain can be helpful for posterior usage.

I don't know how workgraph implement the storing the process state, but I think self.ctx is more for to be used in the user space, for managing the step transition information inside a workchain, it is better to have an internal datastructure to hold those.

For the aiida/plumpy basic process, it has no context concept since it is a single process running and handle its own states transitions. The context concept appears in the workchain since inside a workchain, different processes need to record and passing the variables in between. Thus it is straightforward to come with the idea of context that can be used to store all information when running workchain by just push the tmp data into the context.
It then requires the lifetime of the context is the same as the workchain, and the context contains all the temporary data that produced when workchain is running. So performance and storage space wise it is not ideal to have such design.

I was thinking that the context can be bind to each step, that every step has a context input from previous step and a mutated context that can pass to the next step.
The pros by having the local scope of context are 1. it is much easy to manage the ctx when writing steps. It was always a problem for me that when I write workchain I need to remember the execute logic in the outline so that the variables in the self.ctx is created first and then can be access by another steps. 2. it can then support defining the step as a independent function so it can be reused.
The cons is also obvious, it is not always guaranteed that the step (either it is a class function or a independent function) get the variable it needed from the context.

If we are really thinking about a "compiler" that can parsing the logic, we can think about provide syntax that for each step it declare what context in required and what context it consumed, and the "compiler" will validating it first and then parsing it to AiiDA/Plumpy process.