feat: keep buit specs in registry

tests/unit/test_builder.py

@@ -141,8 +141,7 @@ def test_register(register_dummy_a):
 def test_build_unversioned(spec: dict, expected: Any, register_dummy_a, register_dummy_b):
     result = builder.build(spec)
     assert result == expected
-    if hasattr(result, "__dict__"):
-        assert result.__built_with_spec == spec
+    assert builder.get_initial_builder_spec(result) == spec
 
 
 @pytest.mark.parametrize(
@@ -166,8 +165,7 @@ def test_build_unversioned(spec: dict, expected: Any, register_dummy_a, register
 def test_build_versioned(spec: dict, expected: Any, register_dummy_a_v0, register_dummy_a_v2):
     result = builder.build(spec)
     assert result == expected
-    if hasattr(result, "__dict__"):
-        assert result.__built_with_spec == spec
+    assert builder.get_initial_builder_spec(result) == spec
 
 
 def test_build_partial(register_dummy_a, register_dummy_c):

zetta_utils/builder/__init__.py

@@ -1,5 +1,5 @@
 """Building objects from specs"""
 from . import constants
 from .registry import REGISTRY, register, get_matching_entry, unregister
-from .build import SPECIAL_KEYS, build, BuilderPartial
+from .build import SPECIAL_KEYS, build, BuilderPartial, get_initial_builder_spec
 from . import built_in_registrations

zetta_utils/builder/build.py

@@ -9,6 +9,7 @@
 from zetta_utils import parsing
 from zetta_utils.common import ctx_managers
 from zetta_utils.parsing import json
+from zetta_utils.typing import JsonDict
 
 from . import constants
 from .registry import get_matching_entry
@@ -19,6 +20,19 @@
     "version": "@version",
 }
 
+BUILT_OBJECT_ID_REGISTRY: dict[int, JsonDict] = {}
+
+
+def get_initial_builder_spec(obj: Any) -> JsonDict | None:
+    """Returns the builder spec that the object was initially built with.
+    Note that mutations to the object after it was built will not be
+    reflected in the spec. Returns `None` if the object was not built with
+    builder
+    """
+    # breakpoint()
+    result = BUILT_OBJECT_ID_REGISTRY.get(id(obj), None)
+    return result
+
 
 @typechecked
 def build(
@@ -41,12 +55,18 @@ def build(
 
     # error check the spec
     _traverse_spec(
-        final_spec, _check_type_value, name_prefix="spec", version=constants.DEFAULT_VERSION
+        final_spec,
+        _check_type_value,
+        name_prefix="spec",
+        version=constants.DEFAULT_VERSION,
     )
 
     # build the spec
     result = _traverse_spec(
-        final_spec, _build_dict_spec, name_prefix="spec", version=constants.DEFAULT_VERSION
+        final_spec,
+        _build_dict_spec,
+        name_prefix="spec",
+        version=constants.DEFAULT_VERSION,
     )
 
     return result
@@ -64,7 +84,10 @@ def _traverse_spec(spec: Any, apply_fn: Callable, name_prefix: str, version: str
         elif isinstance(spec, list):
             result = [
                 _traverse_spec(
-                    spec=e, apply_fn=apply_fn, name_prefix=f"{name_prefix}[{i}]", version=version
+                    spec=e,
+                    apply_fn=apply_fn,
+                    name_prefix=f"{name_prefix}[{i}]",
+                    version=version,
                 )
                 for i, e in enumerate(spec)
             ]
@@ -93,7 +116,7 @@ def _traverse_spec(spec: Any, apply_fn: Callable, name_prefix: str, version: str
                 }
         else:
             result = spec
-
+    BUILT_OBJECT_ID_REGISTRY[id(result)] = spec
     return result
 
 
@@ -106,7 +129,10 @@ def _check_type_value(spec: dict[str, Any], name_prefix: str, version: str) -> A
     get_matching_entry(this_type, version=version)
     for k, v in spec.items():
         _traverse_spec(
-            v, apply_fn=_check_type_value, name_prefix=f"{name_prefix}.{k}", version=version
+            v,
+            apply_fn=_check_type_value,
+            name_prefix=f"{name_prefix}.{k}",
+            version=version,
         )
 
 
@@ -153,11 +179,6 @@ def _build_dict_spec(spec: dict[str, Any], name_prefix: str, version: str) -> An
     else:
         raise ValueError(f"Unsupported mode: {this_mode}")
 
-    # save the spec that was used to create the object if possible
-    # slotted classes won't allow adding new attributes
-    if hasattr(result, "__dict__"):
-        object.__setattr__(result, "__built_with_spec", spec)
-
     return result
 
 

zetta_utils/training/lightning/trainers/default.py

@@ -14,7 +14,9 @@
 from pytorch_lightning.strategies import ddp
 
 from zetta_utils import builder, log
+from zetta_utils.builder.build import get_initial_builder_spec
 from zetta_utils.parsing import json
+from zetta_utils.typing import JsonSerializableValue
 
 logger = log.get_logger("zetta_utils")
 ONNX_OPSET_VERSION = 17
@@ -133,15 +135,20 @@ def save_checkpoint(
 
         regime = self.lightning_module
         for k, v in regime._modules.items():  # pylint: disable=protected-access
-            if hasattr(v, "__built_with_spec"):
-                model_spec = getattr(v, "__built_with_spec")  # pylint: disable=protected-access
-                while "@type" in model_spec and model_spec["@type"] == "load_weights_file":
-                    model_spec = model_spec["model"]
+            model_spec: JsonSerializableValue = get_initial_builder_spec(v)
+            if model_spec is not None:
+                unrolled_spec: JsonSerializableValue = model_spec
+                while (
+                    isinstance(unrolled_spec, dict)
+                    and "@type" in unrolled_spec
+                    and unrolled_spec["@type"] == "load_weights_file"
+                ):
+                    unrolled_spec = unrolled_spec["model"]
 
                 spec = {
                     "@type": "load_weights_file",
                     "@version": importlib.metadata.version("zetta_utils"),
-                    "model": model_spec,
+                    "model": unrolled_spec,
                     "ckpt_path": filepath,
                     "component_names": [k],
                     "remove_component_prefix": True,

zetta_utils/typing.py

@@ -121,3 +121,16 @@ def ensure_seq_of_seq(x, length):
         result = [x] * length
 
     return result
+
+
+JsonSerializableValue = Union[
+    str,
+    int,
+    float,
+    bool,
+    None,
+    list["JsonSerializableValue"],
+    dict[str, "JsonSerializableValue"],
+]
+
+JsonDict = dict[str, JsonSerializableValue]