`WorkChain` defined in `__main__` produce `ImportError` not allowing to save checkpoints and is not outputed properly in a verdi shell

[agoscinski]

Steps to reproduce

Using code snippets that are available Tutorial, the How-To Guides and Topics in the aiida-core readthedocs I end up with this file to create a simple WorkChain.

"""Implementation of the MultiplyAddWorkChain for testing and demonstration purposes."""
from aiida.engine import ToContext, WorkChain, calcfunction
from aiida.orm import AbstractCode, Int
from aiida.plugins.factories import CalculationFactory
from aiida.manage.configuration import get_config
from aiida.engine import submit, run
from aiida import orm, load_profile

ArithmeticAddCalculation = CalculationFactory('core.arithmetic.add')


@calcfunction
def multiply(x, y):
    return x * y


class MultiplyAddWorkChain(WorkChain):
    """WorkChain to multiply two numbers and add a third, for testing and demonstration purposes."""

    @classmethod
    def define(cls, spec):
        """Specify inputs and outputs."""
        super().define(spec)

        spec.input('x', valid_type=Int)
        spec.input('y', valid_type=Int)
        spec.input('z', valid_type=Int)
        spec.input('code', valid_type=AbstractCode)
        spec.outline(
            cls.multiply,
            cls.add,
            cls.validate_result,
            cls.result,
        )
        spec.output('result', valid_type=Int)
        spec.exit_code(400, 'ERROR_NEGATIVE_NUMBER', message='The result is a negative number.')

    def multiply(self):
        """Multiply two integers."""
        self.ctx.product = multiply(self.inputs.x, self.inputs.y)

    def add(self):
        """Add two numbers using the `ArithmeticAddCalculation` calculation job plugin."""
        inputs = {'x': self.ctx.product, 'y': self.inputs.z, 'code': self.inputs.code}
        future = self.submit(ArithmeticAddCalculation, **inputs)

        return ToContext(addition=future)

    def validate_result(self):
        """Make sure the result is not negative."""
        result = self.ctx.addition.outputs.sum

        if result.value < 0:
            return self.exit_codes.ERROR_NEGATIVE_NUMBER

    def result(self):
        """Add the result to the outputs."""
        self.out('result', self.ctx.addition.outputs.sum)


builder = MultiplyAddWorkChain.get_builder()
builder.code = orm.load_code(label='add')
builder.x = orm.Int(2)
builder.y = orm.Int(3)
builder.z = orm.Int(5)

result = run(builder)
print(result)

With verdi run it runs but outputs this error message

Error: Exception trying to save checkpoint, this means you will not be able to restart in case of a crash until the next successful checkpoint.
Traceback (most recent call last):
  File "/home/alexgo/micromamba/envs/aiida-dev/lib/python3.11/site-packages/aiida/engine/persistence.py", line 51, in load_object
    return getattr(module, name)
           ^^^^^^^^^^^^^^^^^^^^^
AttributeError: module '__main__' has no attribute 'MultiplyAddWorkChain'
[...]
Error: Exception trying to save checkpoint, this means you will not be able to restart in case of a crash until the next successful checkpoint.
[...]
  File "/home/alexgo/micromamba/envs/aiida-dev/lib/python3.11/site-packages/plumpy/loaders.py", line 63, in identify_object
    self.load_object(identifier)
  File "/home/alexgo/micromamba/envs/aiida-dev/lib/python3.11/site-packages/aiida/engine/persistence.py", line 53, in load_object
    raise ImportError(f"object '{name}' from identifier '{identifier}' could not be loaded")
ImportError: object 'MultiplyAddWorkChain' from identifier '__main__:MultiplyAddWorkChain' could not be loaded

{'result': <Int: uuid: 2e56ab90-f413-4dcc-a6f1-b35e510f5c2a (pk: 113) value: 11>}

Describe the bugs / Expected behavior

• Here the whole log verdi_run_out.txt. One solution is to create a dedicated file for the WorkChain class and import it, but this is not clear from the documentation that this can be done. Ideally, running everything from one file should be also supported with verdi run, or at least be understandable communicated to the user that it does not work suggesting a solution.

• Running the script within a verdi shell or jupyter notebook pipes the error to the process report which is also not good, since the user is not informed about the

Error: Exception trying to save checkpoint, this means you will not be able to restart in case of a crash until the next successful checkpoint.

• A smaller issue here is that the whole error log as well as the print output results from the actual result at the endhas been piped to stdout. There should be distinguishment between the error and output

Your environment

• Operating system: Ubuntu 22.04 LTS
• Python version: 3.11

$ verdi status
✔ version:     AiiDA v2.5.1
✔ config:      /home/alexgo/code/aiida-core/.aiida
✔ profile:     alexgo
✔ storage:     SqliteDosStorage[/home/alexgo/code/aiida-core/.aiida/repository/sqlite_dos_f275ff0f10174e8e84c13e82dd5ba452]: open,
✔ rabbitmq:    Connected to RabbitMQ v3.8.14 as amqp://guest:[email protected]:5672?heartbeat=600