Addressing performance burden from #6990

[edan-bainglass]

#6990 introduced a Pydantic model patching system to reduce the burden on plugin developers. However, it has been reported that this inccurs a hit on performance per interpreter (e.g., daemon worker). For example, adding a print statement to Entity.__init_subclass__ (from where all patches trigger) yields the following:

In [1]: from aiida import orm
Initializing subclass User
Initializing subclass Computer
Initializing subclass AuthInfo
Initializing subclass Comment
Initializing subclass Group
Initializing subclass AutoGroup
Initializing subclass ImportGroup
Initializing subclass UpfFamily
Initializing subclass Log
Initializing subclass Node
Initializing subclass Data
Initializing subclass BaseType
Initializing subclass ArrayData
Initializing subclass KpointsData
Initializing subclass BandsData
Initializing subclass OrbitalData
Initializing subclass ProjectionData
Initializing subclass TrajectoryData
Initializing subclass XyData
Initializing subclass Bool
Initializing subclass SinglefileData
Initializing subclass CifData
Initializing subclass AbstractCode
Initializing subclass Code
Initializing subclass InstalledCode
Initializing subclass ContainerizedCode
Initializing subclass PortableCode
Initializing subclass Dict
Initializing subclass EnumData
Initializing subclass NumericType
Initializing subclass Float
Initializing subclass FolderData
Initializing subclass Int
Initializing subclass JsonableData
Initializing subclass List
Initializing subclass RemoteData
Initializing subclass RemoteStashData
Initializing subclass RemoteStashCompressedData
Initializing subclass RemoteStashCustomData
Initializing subclass RemoteStashFolderData
Initializing subclass Str
Initializing subclass StructureData
Initializing subclass UpfData
Initializing subclass ProcessNode
Initializing subclass CalculationNode
Initializing subclass CalcFunctionNode
Initializing subclass CalcJobNode
Initializing subclass WorkflowNode
Initializing subclass WorkChainNode
Initializing subclass WorkFunctionNode

Not ideal to be patching the full ORM on import, as one might be interested in only a single node, or even no intention to use the model system.

A better solution would be to patch things (and only those that absolutely are worth patching - to reduce plugin developer burdens) on first use, i.e., JIT, i.e., lazy loading!

The planned ORM restructure (see #7556) may naturally resolve this, but unclear at the moment.

[agoscinski]

Very clear to me to move the pydantic models from the orm layer to the application layers that actually require json serialization. It will be always a performance hit if we include it by default in the orm classes. If we really need it, field validation can be unified differently. But even that, the current attempt to unify the validation for different application layers introduces more complexity than doing the same validation at different points in the code. Also pydantic validation seems slower than doing them "by hand" in the python object initialization (this was just a quick benchmark, I did not really look into this much why).

[edan-bainglass]

Thanks @agoscinski. At the moment, the presence of the models on the ORM layer allows for:

1. Node serialization via node.serialize() (and similarly, deserialization)
2. Deriving QueryBuilder fields

This is not to say that it must be on the ORM layer, but just clarifying that perhaps moving it away from the ORM complicated these use cases. But perhaps there are other ways to achieve this. For example, in #7559, I'm exploring a different approach to defining QbFields, one that is more lazy in nature. It still relies on a model to define the property types, but perhaps that model does not need to be on the ORM (currently assumed on the ORM).

Naively, but also perhaps sufficient, is to move the entire model system to a separate module. In practice, if a plugin developer extends Data with MyData(Data) (which we are reconsidering anyhow), they would simply extend DataReadModel with MyDataReadModel. So we shift a bit of burden to the developer, to gain the benefit of isolation - ORM is pure w.r.t models (i.e., not carrying them around).

[agoscinski]

To actually use the pydantic models functionalities, we need to generate the model. This initial model generation is costly. The advantage when paying this cost is we have a json (de)serialization conforming to a schema which is useful for applications that cross the process boundary like the rest api. For any other use case its not worth to pay this initial cost of the schema generation. For CLI this initial cost is even horrible because you cannot even cache the schema generation effectively since every command is starting a separate process. So for QueryBuilder we are effectively just using the Fields as data container. Therefore I would go back to use custom Fields that do not have the pydantic baggage. If needed, this custom field can be translated to a pydantic field in the rest api for code reusability. This would remove basically all models except the Read/WriteModel, and I don't see the problem to move it to the rest api module and put meta-programming patch subclass logic just there.

we shift a bit of burden to the developer, to gain the benefit of isolation

I don't see this point. As plugin developer one implements the logic for the rest api or one does not. If one does not, one get defaults which might be bad. That has nothing to do with where the logic resides.

[edan-bainglass]

Thanks @agoscinski. I'll present my solution when I'm back from holiday 🙂

[agoscinski]

I made a proof of concept describing to what structure I would like to move #7591 to solve performance and API concerns. Please check out the README.md there, and look at structure (what is in core API and what is not), and how are the ColumnFields and AttributeFields converted to pydantic.Fields to build a model. I iterated on this quite a lot to consider different use cases and caveats. When you back I am happy to discuss it. I meanwhile start a vertical slice to see if I am still missing any important points in the proof of concept.

[edan-bainglass]

Thanks @agoscinski. As I've stated above, and since #7510 is assigned to me 🙂, I'm already working on the ORM redesign (WIP, to be presented when I'm back before publishing the draft). Could you clarify if your draft is meant as a suggestion, or if you intend to take over this issue? Good to know early to avoid redundant effort 🙂🙏 Cheers

[edan-bainglass]

@agoscinski in any case, at a first glance, I do like the file hierarchy of your draft, with the core separation you had mentioned. Let's discuss this specifically when I'm back 🙂

[agoscinski]

I did not want to take over. It is just a PoC to discuss the design. The actual implementation of the PoC is still quite some work. It is more for me to be clearer about what the constraints are. If one is not aware about the actual constraints of the orm on the rest of the code, then a design discussion is not so meaningful. I hope by both of us digging into the code, we can speed up the process of deciding on a good design.