`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:guest@127.0.0.1:5672?heartbeat=600

[sphuber]

I am not sure if this is really a bug as such. The code actually works and the error is correct. I think just that rather than it being an error, the message should be a warning. The message is actually correct: since the workchain class cannot be imported, it cannot be persisted, which normally would be a problem. In this case it is expected though. The fact that it wouldn't be able to be restarted is anyway not relevant since we are running it (and not submitting) where restarts are anyway not supported.

I am not sure if there is anything we can do to hide the warning in this case. We don't want to hide it by default because in certain cases it not being persistable could be a real problem. We could just add some code in the tutorial script to silence this warning, but that would also be quite "ugly" which will confuse the user.

So not quite sure what to do here...