👌 Defer ORM Pydantic model rebuilding (#7507)

Avoid rebuilding generated ORM Pydantic models during class setup. Leave
validator and schema construction to Pydantic's deferred build machinery
until models are validated, serialized, or used for schema generation.

Add focused regression coverage to ensure generated node model setup
does not force model rebuilds while explicit model use still works.

Author: agoscinski

src/aiida/orm/entities.py

@@ -415,7 +415,6 @@ def optionalize(annotation: Any) -> Any:
         )
         model.__qualname__ = f'{cast(Any, cls).__name__}.Model'
         model.model_config = deepcopy(cls.ReadModel.model_config)
-        model.model_rebuild(force=True)
 
         cls._COMPAT_MODEL = model
 
@@ -602,8 +601,6 @@ def copy_model_field(field: pdt.fields.FieldInfo) -> tuple[Any, pdt.fields.Field
             WriteModel.__pydantic_decorators__.field_validators = validators
             WriteModel.model_config = deepcopy(model_cls.model_config)
 
-            WriteModel.model_rebuild(force=True)
-
             return WriteModel
 
         if 'WriteModel' not in cls.__dict__:

src/aiida/orm/nodes/data/code/abstract.py

@@ -227,7 +227,6 @@ def _patch_cli_model(cls):
         )
         CliModel.__qualname__ = f'{cls.__name__}.CliModel'
         CliModel.model_config['arbitrary_types_allowed'] = True
-        CliModel.model_rebuild(force=True)
         cls._CliModel = CliModel
 
     @abc.abstractmethod

src/aiida/orm/nodes/node.py

@@ -455,7 +455,6 @@ def merge_field(
         )
         model.__qualname__ = f'{cast(Any, cls).__name__}.Model'
         model.model_config = deepcopy(cls.ReadModel.model_config)
-        model.model_rebuild(force=True)
 
         cls._COMPAT_MODEL = model
 
@@ -1177,7 +1176,6 @@ def _patch_attributes_model(cls):
             )
             AttributesModel.__qualname__ = f'{cls.__name__}.AttributesModel'
             cls.AttributesModel = AttributesModel  # type: ignore[misc]
-        cls.AttributesModel.model_rebuild(force=True)
 
     @classmethod
     def _get_patched_node_type_field(cls):
@@ -1236,7 +1234,6 @@ def _patch_read_model(cls):
             ),
         )
         ReadModel.__qualname__ = f'{cls.__name__}.ReadModel'
-        ReadModel.model_rebuild(force=True)
 
         cls.ReadModel = ReadModel  # type: ignore[misc]
 
@@ -1265,7 +1262,6 @@ def _patch_constructor_model(cls):
             ),
         )
         ConstructorModel.__qualname__ = f'{cls.__name__}.ConstructorModel'
-        ConstructorModel.model_rebuild(force=True)
 
         cls._ConstructorModel = ConstructorModel
 

src/aiida/orm/pydantic.py

@@ -86,8 +86,6 @@ def _as_minimal_model(cls: type[OrmModel]) -> type[OrmModel]:
                     field.default_factory = None
             MinimalModel.model_fields[key] = field
 
-        MinimalModel.model_rebuild(force=True)
-
         # Make subsequent calls idempotent for this specific class and the derived model
         cls._AIIDA_MINIMAL_MODEL = MinimalModel  # type: ignore[attr-defined]
         MinimalModel._AIIDA_MINIMAL_MODEL = MinimalModel  # type: ignore[attr-defined]

tests/orm/models/test_models.py

@@ -13,6 +13,7 @@
 from aiida import orm
 from aiida.common.datastructures import StashMode
 from aiida.common.exceptions import UnsupportedSchemaError
+from aiida.orm.pydantic import OrmModel
 
 orm_to_test = (
     orm.AuthInfo,
@@ -517,6 +518,31 @@ def test_minimal_model_idempotency():
     assert RepeatedDynamicModel is DynamicModel
 
 
+def test_generated_orm_model_setup_defers_pydantic_rebuild(monkeypatch):
+    """Test generated ORM models are not rebuilt eagerly during class setup."""
+    rebuilt: list[type[OrmModel]] = []
+
+    def model_rebuild(cls, *args, **kwargs):
+        rebuilt.append(cls)
+        return True
+
+    with monkeypatch.context() as context:
+        context.setattr(OrmModel, 'model_rebuild', classmethod(model_rebuild))
+
+        class TestData(orm.Data):
+            class AttributesModel(orm.Data.AttributesModel):
+                value: int
+
+            class ConstructorArgsModel(OrmModel):
+                value: int
+
+        assert rebuilt == []
+
+    model = TestData.WriteModel(node_type=TestData.class_node_type, attributes={'value': '1'})
+    assert model.attributes.value == 1
+    assert TestData.ReadModel.model_json_schema()['title'] == 'TestDataReadModel'
+
+
 @pytest.mark.parametrize(
     'required_arguments',
     orm_to_test,
@@ -527,10 +553,10 @@ def test_model_overrides(required_arguments: RequiredEntityArguments):
     name = cls.__name__
 
     assert cls.ReadModel.__qualname__ == f'{name}.ReadModel'
-    assert cls.ReadModel.model_config.get('title') == f'{name}ReadModel'
+    assert cls.ReadModel.model_json_schema()['title'] == f'{name}ReadModel'
 
     assert cls.WriteModel.__qualname__ == f'{name}.WriteModel'
-    assert cls.WriteModel.model_config.get('title') == f'{name}WriteModel'
+    assert cls.WriteModel.model_json_schema()['title'] == f'{name}WriteModel'
 
 
 def _clean_and_sort(dictionary: dict) -> dict:
@@ -616,11 +642,11 @@ def test_node_attributes_model_overrides(required_arguments: RequiredNodeArgumen
     AttributesModel = cls.ReadModel.model_fields['attributes'].annotation  # noqa: N806
     assert AttributesModel is cls.AttributesModel
     assert AttributesModel.__qualname__ == f'{name}.AttributesModel'
-    assert AttributesModel.model_config.get('title') == f'{name}AttributesModel'
+    assert AttributesModel.model_json_schema()['title'] == f'{name}AttributesModel'
 
     AttributesWriteModel = cls.WriteModel.model_fields['attributes'].annotation  # noqa: N806
     assert AttributesWriteModel.__qualname__ == f'{name}.AttributesWriteModel'
-    assert AttributesWriteModel.model_config.get('title') == f'{name}AttributesWriteModel'
+    assert AttributesWriteModel.model_json_schema()['title'] == f'{name}AttributesWriteModel'
 
 
 def _validate_value(value):