WorkChain: return exit_code in while_ will not terminate the workchain.

[unkcpz]

Describe the bug

I have a work chain with the following outline:

    @classmethod
    def define(cls, spec):
        """Define the process specification."""
        # yapf: disable
        super().define(spec)
        ....
        spec.outline(
            ....<setup parameters>

            cls.run_bands,
            while_(cls.not_enough_bands)(
                cls.increase_nbands,
                cls.run_bands,
            ),
            ....<more>

.....
    def run_bands(self):
        """run bands calculation"""
        inputs = self._get_base_bands_inputs()
        inputs["nbands_factor"] = self.ctx.nbands_factor
        inputs["bands_kpoints"] = self.ctx.bands_kpoints

        running = self.submit(PwBandsWorkChain, **inputs)
        self.report(f"Running pw bands calculation pk={running.pk}")
        return ToContext(workchain_bands=running)

    def not_enough_bands(self):
        """inspect and check if the number of bands enough for fermi shift (_FERMI_SHIFT)"""
        workchain = self.ctx.workchain_bands

        if not workchain.is_finished_ok:
            self.report(
                f"PwBandsWorkChain for bands evaluation failed with exit status {workchain.exit_status}"
            )
            return self.exit_codes.ERROR_SUB_PROCESS_FAILED_BANDS

        fermi_energy = workchain.outputs.band_parameters["fermi_energy"]
        bands = workchain.outputs.band_structure.get_array("bands")

In the not_enough_bands step, I do a check to make sure the previous process by cls.run_bands is finished okay. If not, I return self.exit_codes.ERROR_SUB_PROCESS_FAILED_BANDS and expect the whole work chain terminated.
However, I run into the infinite loop that the loop does not break by this return.
I think this is a bug since if I move the return of exit_code out of while_ works all good to me.

Expected behavior

return exit_code from step in while_ and terminate the workflow.

solution proposed

We can introduce a break_ to explicitly terminate the loop in work chain.

Your environment

• aiida-core version 1.6.5:

[sphuber]

I am not sure if this can be considered a bug per se. In a sense, it only makes sense for an expression that is evaluated in the while() to return a boolean. I think that the original implementation by @muhrin was also made with this concept in mind, but he should confirm this.

I think it would be possible to have the while() construct to check the type of the return value of the step that is called and make an exception in the case of an ExitCode being returned.

However, I am not sure if this is the best approach. Personally, I would have made the outline as follows:

while_(cls.not_enough_bands)(
    cls.increase_nbands,
    cls.run_bands,
    cls.inspect_bands,
),

And have the check of the workchain having finished correctly in inspect_bands. Logically this makes more sense to me and seems cleaner.

If you agree, we could of course improve while() to except or warn when anything but a boolean is returned from the expression it evaluates.

What do you think @unkcpz @muhrin

[sphuber]

@unkcpz This needs to be fixed in plumpy if we want this incorrect usage to raise an exception. I have opened a PR there: aiidateam/plumpy#259