Processing calcjob outputs in BaseRestartWorkChain

General Usage
1

Hi all,

I’ve started converting my workchains to BaseRestartWorkChains and I’m generally confused about how to grab the outputs from the calcjobs I define in the _process_class.

I’m currently using the orca plugin in a work chain to relax a molecule and extract the homo/lumo energy values:

OrcaCalculation = CalculationFactory("orca.orca")

class HomoLumoWorkChain(BaseRestartWorkChain):

    _process_class = OrcaCalculation

    @classmethod
    def define(cls, spec):
        """ Homo Lumo Calculation WorkChain """
        super().define(spec)
        spec.input('structure', valid_type=StructureData)
        spec.input('code', valid_type=AbstractCode)
        spec.input('parameters', valid_type=Dict)
        spec.input('resources', valid_type=Dict) # keep this a regular dictionary
        spec.outline(
            cls.setup,
            while_(cls.should_run_process)(
                cls.run_process,
                cls.inspect_process,
            ),
            cls.extract_homo_lumo,
            cls.results,
        )
        spec.expose_outputs(OrcaCalculation)
        spec.output('homo', valid_type=Float)
        spec.output('lumo', valid_type=Float)

    def setup(self):
        super().setup()
        wallclock_seconds = 60 * 60 * 8 # 8 hours
        for param in self.inputs.parameters.get_dict()["input_keywords"]:
            if param.startswith("PAL"):
                nprocs = int(param[-1])
        
        metadata = {
            "options": {
                "resources": self.inputs.resources.get_dict(),
                "max_wallclock_seconds": wallclock_seconds,
                "withmpi": False,
                "max_memory_kb": int(3.8 * 1e6 * nprocs),
                "account": "account_xxxx"
            }
        }
        self.ctx.inputs = {'structure': self.inputs.structure, 'code': self.inputs.code, 'parameters': 
        self.inputs.parameters, "metadata": metadata}

    def extract_homo_lumo(self):
        outputs = self.exposed_outputs(OrcaCalculation)
        outputdata = outputs["output_parameters"]
        E_homo = get_homo(outputdata)
        E_lumo = get_lumo(outputdata)
        self.out('homo', E_homo)
        self.out('lumo', E_lumo)

I just ran this workchain and was met with this (contracted) error:

  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/plumpy/workchains.py", line 246, in step
    return True, self._fn(self._workchain)
  File "/home/coopy/onedrive/Research/Batteries/aiida/WorkChains/HomoLumoWorkChain.py", line 104, in extract_homo_lumo
    # outputs = self.get_outputs(OrcaCalculation)
  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/aiida/engine/processes/workchains/restart.py", line 316, in get_outputs
    return self.exposed_outputs(node, self.process_class)
  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/aiida/engine/processes/process.py", line 965, in exposed_outputs
    process_outputs_dict = node.base.links.get_outgoing(link_type=link_types).nested()
AttributeError: type object 'OrcaCalculation' has no attribute 'base'

I’ve also tried taking the outputs in the extract_homo_lumo() method using:

outputs = self.get_outputs(OrcaCalculation)

which also failed with a similar error message.

How should I be extracting the outputs and processing them in other steps? I’ve searched pretty thoroughly through the docs and think this may be a hole in the current documentation. Let me know if I missed something.

Thanks!

2

There seems to be something funky with your workchain. The exception is referencing a commented out line:

That should never happen so I think there is a weird state of your daemon. Did you remember to restart your daemon after editing the workchain file?

Another weird thing is the exception itself claiming that the node doesn’t have the base attribute. What version of AiiDA is installed? Do you have multiple virtual environments perhaps and are you running in the correct one?

Either way, there a few ways you can retrieve outputs of a completed process node.

  1. Directly: through node.ouputs.some_link_label
    2: If you have called expose_outputs in the define method, you can use self.exposed_outputs(ProcessClass, namespace=namespace).

You are doing this in the version of the workchain you pasted:

outputs = self.exposed_outputs(OrcaCalculation)

and that should be correct.

So I think there is nothing wrong with that code, just that your daemon is probably not using the correct version of the file.

3

Hi @sphuber

I generally restart my daemon after changing code. I may not have remembered in this instance, although I have reproduced this exact issue many times after restarting the daemon. That base attribute issue has been persistent.

I only have one venv and it has aiida-core 2.5.1.

I just restarted the daemon with the --reset flag and reran the workflow with this code in the problematic section

def extract_homo_lumo(self):
        outputs = self.exposed_outputs(OrcaCalculation)

and got a new error message:

  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/plumpy/process_states.py", line 228, in execute
    result = self.run_fn(*self.args, **self.kwargs)
  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/aiida/engine/processes/workchains/workchain.py", line 313, in _do_step
    finished, stepper_result = self._stepper.step()
  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/plumpy/workchains.py", line 295, in step
    finished, result = self._child_stepper.step()
  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/plumpy/workchains.py", line 246, in step
    return True, self._fn(self._workchain)
  File "/home/coopy/onedrive/Research/Batteries/aiida/WorkChains/HomoLumoWorkChain.py", line 103, in extract_homo_lumo
    outputs = self.exposed_outputs(OrcaCalculation)
TypeError: Process.exposed_outputs() missing 1 required positional argument: 'process_class'
4

And using the self.get_outputs(OrcaCalculation) version after a daemon restart with the --reset flag gives:

  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/plumpy/process_states.py", line 228, in execute
    result = self.run_fn(*self.args, **self.kwargs)
  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/aiida/engine/processes/workchains/workchain.py", line 313, in _do_step
    finished, stepper_result = self._stepper.step()
  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/plumpy/workchains.py", line 295, in step
    finished, result = self._child_stepper.step()
  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/plumpy/workchains.py", line 246, in step
    return True, self._fn(self._workchain)
  File "/home/coopy/onedrive/Research/Batteries/aiida/WorkChains/HomoLumoWorkChain.py", line 102, in extract_homo_lumo
    outputs = self.get_outputs(OrcaCalculation)
  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/aiida/engine/processes/workchains/restart.py", line 316, in get_outputs
    return self.exposed_outputs(node, self.process_class)
  File "/home/coopy/envs/aiida/lib/python3.10/site-packages/aiida/engine/processes/process.py", line 965, in exposed_outputs
    process_outputs_dict = node.base.links.get_outgoing(link_type=link_types).nested()
AttributeError: type object 'OrcaCalculation' has no attribute 'base'

Is this possibly an issue with the aiida-orca plugin?

5

Hi @cote3804!

This is not a bug of the aiida-orca plugin. You only pass the OrcaCalculation class to the get_outputs or self.exposed_outputs methods. To extract the outputs, you would need to pass a node of an actual calculation. See the corresponding implementation of the BaseRestartWorkChain for reference (aiida-core/src/aiida/engine/processes/workchains/restart.py at main · aiidateam/aiida-core · GitHub):

    def get_outputs(self, node) -> Mapping[str, orm.Node]:
        # Comments and doc string removed
        return self.exposed_outputs(node, self.process_class)

    def results(self) -> Optional['ExitCode']:
        # Comments and doc string removed
        node = self.ctx.children[self.ctx.iteration - 1]
        max_iterations = self.inputs.max_iterations.value
        if not self.ctx.is_finished and self.ctx.iteration >= max_iterations:
            self.report(
                f'reached the maximum number of iterations {max_iterations}: '
                f'last ran {self.ctx.process_name}<{node.pk}>'
            )
            return self.exit_codes.ERROR_MAXIMUM_ITERATIONS_EXCEEDED

        self.report(f'work chain completed after {self.ctx.iteration} iterations')
        self._attach_outputs(node)
        return None

    def _attach_outputs(self, node) -> Mapping[str, orm.Node]:
        # Comments and doc string removed
        outputs = self.get_outputs(node)
        existing_outputs = self.node.base.links.get_outgoing(link_type=LinkType.RETURN).all_link_labels()

There you can see that the CalcJob of the last iteration node = self.ctx.children[self.ctx.iteration - 1] is passed to those methods (instead, you are basically passing self._process_class at the moment). Note that get_output takes only the node as an input, whereas exposed_outputs requires the actual node and the process_class as an input.

This being said, changing your code in the following way should fix the issue:

def extract_homo_lumo(self):
        last_calc = self.ctx.children[self.ctx.iteration - 1]
        outputs = self.exposed_outputs(last_calc, OrcaCalculation)
        outputdata = outputs["output_parameters"]
        E_homo = get_homo(outputdata)
        E_lumo = get_lumo(outputdata)
        self.out('homo', E_homo)
        self.out('lumo', E_lumo)

Hope that helps.

2 Likes
6

My bad, I completely read over the fact that the actual node was missing from the exposed_outputs call. @t-reents is right and his suggestion should indeed solve it

1 Like
7

Sorry for the late response. My cluster went down and I didn’t have the chance to test this solution.

This did completely resolve my issues. I think it may still be worth adding an example like this, where the outputs from the process class are processed in the work chain, to the “How to write error resistant workflows” section of the docs.

Thanks for the help!

8

This topic was automatically closed 5 days after the last reply. New replies are no longer allowed.