mlbricks-core 0.1.1__tar.gz → 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/PKG-INFO +5 -2
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/README.md +4 -1
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/pyproject.toml +1 -1
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/pyproject.toml.orig +1 -1
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/src/mlbricks/__init__.py +12 -0
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/src/mlbricks/configurable.py +64 -1
- mlbricks_core-0.2.0/src/mlbricks/magic_registry.py +204 -0
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/src/mlbricks/registry.py +43 -31
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/src/mlbricks/resolve_config.py +55 -9
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/src/mlbricks/cache.py +0 -0
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/src/mlbricks/featurization.py +0 -0
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/src/mlbricks/py.typed +0 -0
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/src/mlbricks/storage.py +0 -0
- {mlbricks_core-0.1.1 → mlbricks_core-0.2.0}/src/mlbricks/utils.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: mlbricks-core
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
4
4
|
Summary: Core ML tooling, including config resolution, registry, and storage.
|
|
5
5
|
Author: Shivam Verma
|
|
6
6
|
Author-email: Shivam Verma <shivam.vvm.verma@gmail.com>
|
|
@@ -65,7 +65,7 @@ seed: 0
|
|
|
65
65
|
import torch.nn as nn
|
|
66
66
|
from pydantic import PositiveInt
|
|
67
67
|
import yaml
|
|
68
|
-
from mlbricks import Configurable, Registry, resolve
|
|
68
|
+
from mlbricks import Configurable, Registry, resolve, to_dict
|
|
69
69
|
|
|
70
70
|
REGISTRY = Registry()
|
|
71
71
|
MODEL_REGISTRY = REGISTRY.subgroup("models")
|
|
@@ -80,6 +80,7 @@ class MLP(nn.Module, Configurable):
|
|
|
80
80
|
|
|
81
81
|
def __init__(self, cfg: "MLP.Config") -> None:
|
|
82
82
|
super().__init__()
|
|
83
|
+
self.hparams = to_dict(cfg, registry=REGISTRY)
|
|
83
84
|
self.layers = nn.Sequential(
|
|
84
85
|
nn.Linear(cfg.in_dim, cfg.hidden_dim),
|
|
85
86
|
nn.ReLU(),
|
|
@@ -93,6 +94,8 @@ class MLP(nn.Module, Configurable):
|
|
|
93
94
|
raw = yaml.safe_load(open("config.yaml")) # or OmegaConf.to_object(OmegaConf.load(...))
|
|
94
95
|
cfg = resolve(raw, registry=REGISTRY)
|
|
95
96
|
model = cfg["model"].build()
|
|
97
|
+
print(model.hparams)
|
|
98
|
+
# {'_registry_': 'models.mlp', 'in_dim': 784, 'hidden_dim': 128, 'out_dim': 10}
|
|
96
99
|
```
|
|
97
100
|
|
|
98
101
|
## Installation
|
|
@@ -52,7 +52,7 @@ seed: 0
|
|
|
52
52
|
import torch.nn as nn
|
|
53
53
|
from pydantic import PositiveInt
|
|
54
54
|
import yaml
|
|
55
|
-
from mlbricks import Configurable, Registry, resolve
|
|
55
|
+
from mlbricks import Configurable, Registry, resolve, to_dict
|
|
56
56
|
|
|
57
57
|
REGISTRY = Registry()
|
|
58
58
|
MODEL_REGISTRY = REGISTRY.subgroup("models")
|
|
@@ -67,6 +67,7 @@ class MLP(nn.Module, Configurable):
|
|
|
67
67
|
|
|
68
68
|
def __init__(self, cfg: "MLP.Config") -> None:
|
|
69
69
|
super().__init__()
|
|
70
|
+
self.hparams = to_dict(cfg, registry=REGISTRY)
|
|
70
71
|
self.layers = nn.Sequential(
|
|
71
72
|
nn.Linear(cfg.in_dim, cfg.hidden_dim),
|
|
72
73
|
nn.ReLU(),
|
|
@@ -80,6 +81,8 @@ class MLP(nn.Module, Configurable):
|
|
|
80
81
|
raw = yaml.safe_load(open("config.yaml")) # or OmegaConf.to_object(OmegaConf.load(...))
|
|
81
82
|
cfg = resolve(raw, registry=REGISTRY)
|
|
82
83
|
model = cfg["model"].build()
|
|
84
|
+
print(model.hparams)
|
|
85
|
+
# {'_registry_': 'models.mlp', 'in_dim': 784, 'hidden_dim': 128, 'out_dim': 10}
|
|
83
86
|
```
|
|
84
87
|
|
|
85
88
|
## Installation
|
|
@@ -1,5 +1,7 @@
|
|
|
1
1
|
from mlbricks.cache import StorageCache
|
|
2
2
|
from mlbricks.configurable import (
|
|
3
|
+
UNSET,
|
|
4
|
+
BuildTimeField,
|
|
3
5
|
ConfigFieldNameClashError,
|
|
4
6
|
Configurable,
|
|
5
7
|
ConfigurableBuildError,
|
|
@@ -19,6 +21,11 @@ from mlbricks.featurization import (
|
|
|
19
21
|
TransformHooks,
|
|
20
22
|
TransformPipeline,
|
|
21
23
|
)
|
|
24
|
+
from mlbricks.magic_registry import (
|
|
25
|
+
MagicRegistryError,
|
|
26
|
+
MagicRegistryImportError,
|
|
27
|
+
MagicRegistryUnknownFieldError,
|
|
28
|
+
)
|
|
22
29
|
from mlbricks.registry import (
|
|
23
30
|
DuplicateRegistrationError,
|
|
24
31
|
InvalidRegistryNameError,
|
|
@@ -36,6 +43,8 @@ from mlbricks.storage import StorageDriver
|
|
|
36
43
|
from mlbricks.utils import path_slug, uri_slug
|
|
37
44
|
|
|
38
45
|
__all__ = [
|
|
46
|
+
"UNSET",
|
|
47
|
+
"BuildTimeField",
|
|
39
48
|
"ConfigFieldNameClashError",
|
|
40
49
|
"ConfigInstantiationError",
|
|
41
50
|
"ConfigSerializationError",
|
|
@@ -47,6 +56,9 @@ __all__ = [
|
|
|
47
56
|
"Context",
|
|
48
57
|
"DuplicateRegistrationError",
|
|
49
58
|
"InvalidRegistryNameError",
|
|
59
|
+
"MagicRegistryError",
|
|
60
|
+
"MagicRegistryImportError",
|
|
61
|
+
"MagicRegistryUnknownFieldError",
|
|
50
62
|
"MissingCfgParameterError",
|
|
51
63
|
"MissingConfigError",
|
|
52
64
|
"NamedTransformStep",
|
|
@@ -1,11 +1,14 @@
|
|
|
1
1
|
import inspect
|
|
2
2
|
import types
|
|
3
|
-
from typing import Any, ClassVar
|
|
3
|
+
from typing import Any, ClassVar, Self
|
|
4
4
|
|
|
5
5
|
import pydantic
|
|
6
6
|
from pydantic import BaseModel, ConfigDict, ValidationInfo, model_validator
|
|
7
|
+
from pydantic.fields import FieldInfo
|
|
7
8
|
|
|
8
9
|
__all__ = [
|
|
10
|
+
"UNSET",
|
|
11
|
+
"BuildTimeField",
|
|
9
12
|
"ConfigFieldNameClashError",
|
|
10
13
|
"ConfigValidationError",
|
|
11
14
|
"Configurable",
|
|
@@ -86,6 +89,55 @@ def build_validator(func: types.FunctionType) -> Any:
|
|
|
86
89
|
return model_validator(mode="after")(wrapper)
|
|
87
90
|
|
|
88
91
|
|
|
92
|
+
class _Unset:
|
|
93
|
+
"""Sentinel for a BuildTimeField that has not been resolved yet.
|
|
94
|
+
|
|
95
|
+
Distinct from None so a BuildTimeField[T] can hold a T that is itself
|
|
96
|
+
None. Never constructed outside this module.
|
|
97
|
+
"""
|
|
98
|
+
|
|
99
|
+
__slots__ = ()
|
|
100
|
+
|
|
101
|
+
def __repr__(self) -> str:
|
|
102
|
+
return "UNSET"
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
UNSET = _Unset()
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
# Config field annotation for build-time-only, non-serializable values.
|
|
109
|
+
#
|
|
110
|
+
# Never required at Config construction, never emitted by
|
|
111
|
+
# resolve_config.to_dict(), and required at build() time unless the field is
|
|
112
|
+
# given its own default (same precedence as any pydantic field). Declare the
|
|
113
|
+
# "= UNSET" default explicitly -- it is not implicit.
|
|
114
|
+
#
|
|
115
|
+
# Example:
|
|
116
|
+
# class Trainer(Configurable):
|
|
117
|
+
# class Config(Configurable.Config["Trainer"]):
|
|
118
|
+
# dataloader: BuildTimeField[DataLoader] = UNSET
|
|
119
|
+
#
|
|
120
|
+
# def __init__(self, cfg: "Trainer.Config") -> None:
|
|
121
|
+
# self.loader = cfg.dataloader # plain DataLoader, no unwrap
|
|
122
|
+
#
|
|
123
|
+
# cfg = Trainer.Config()
|
|
124
|
+
# trainer = cfg.build(dataloader=my_loader)
|
|
125
|
+
type BuildTimeField[T] = T | _Unset
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def is_build_time_field(field: FieldInfo) -> bool:
|
|
129
|
+
"""Return True if field's class-declared default is the BuildTimeField sentinel.
|
|
130
|
+
|
|
131
|
+
Checks the declared default (``FieldInfo.default``), not the current
|
|
132
|
+
instance value, so this stays correct after build() resolves a real value
|
|
133
|
+
on a validated Config instance.
|
|
134
|
+
|
|
135
|
+
Used by build()'s automatic unresolved-field check and by
|
|
136
|
+
resolve_config.to_dict() to skip these fields during serialization.
|
|
137
|
+
"""
|
|
138
|
+
return field.default is UNSET
|
|
139
|
+
|
|
140
|
+
|
|
89
141
|
class Configurable:
|
|
90
142
|
"""Mixin for classes that are constructed from a nested Config pydantic model.
|
|
91
143
|
|
|
@@ -159,6 +211,17 @@ class Configurable:
|
|
|
159
211
|
except pydantic.ValidationError as exc:
|
|
160
212
|
raise ConfigValidationError(exc) from exc
|
|
161
213
|
|
|
214
|
+
@build_validator
|
|
215
|
+
def _check_build_time_fields(self) -> Self:
|
|
216
|
+
unresolved = [
|
|
217
|
+
name
|
|
218
|
+
for name, field in type(self).model_fields.items()
|
|
219
|
+
if is_build_time_field(field) and getattr(self, name) is UNSET
|
|
220
|
+
]
|
|
221
|
+
if unresolved:
|
|
222
|
+
raise ValueError(f"BuildTimeField(s) required at build(): {unresolved}")
|
|
223
|
+
return self
|
|
224
|
+
|
|
162
225
|
def build(self, **kwargs: Any) -> T:
|
|
163
226
|
"""Construct the owning class from this config.
|
|
164
227
|
|
|
@@ -0,0 +1,204 @@
|
|
|
1
|
+
import importlib
|
|
2
|
+
import inspect
|
|
3
|
+
import typing
|
|
4
|
+
from collections.abc import Iterable
|
|
5
|
+
from typing import Any, cast
|
|
6
|
+
|
|
7
|
+
from pydantic import ConfigDict, create_model
|
|
8
|
+
|
|
9
|
+
from mlbricks.configurable import Configurable
|
|
10
|
+
|
|
11
|
+
__all__ = [
|
|
12
|
+
"MagicRegistryError",
|
|
13
|
+
"MagicRegistryImportError",
|
|
14
|
+
"MagicRegistryUnknownFieldError",
|
|
15
|
+
"magic_wrap",
|
|
16
|
+
]
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class MagicRegistryError(ValueError):
|
|
20
|
+
"""Base for magic-registry resolution errors."""
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class MagicRegistryImportError(MagicRegistryError):
|
|
24
|
+
"""Raised when a dotted path cannot be imported or resolved to a class."""
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
class MagicRegistryUnknownFieldError(MagicRegistryError):
|
|
28
|
+
"""Raised when a field key has no matching __init__ parameter on the
|
|
29
|
+
target class and the target has no **kwargs catch-all."""
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _import_dotted(path: str) -> type:
|
|
33
|
+
"""Import a dotted path to a class, trying the longest module prefix first.
|
|
34
|
+
|
|
35
|
+
Args:
|
|
36
|
+
path: A dotted path, e.g. "torch.optim.Adam".
|
|
37
|
+
|
|
38
|
+
Returns:
|
|
39
|
+
The resolved class.
|
|
40
|
+
|
|
41
|
+
Raises:
|
|
42
|
+
MagicRegistryImportError: If no prefix of `path` imports as a module,
|
|
43
|
+
an attribute in the remaining chain is missing, or the resolved
|
|
44
|
+
object is not a class.
|
|
45
|
+
"""
|
|
46
|
+
parts = path.split(".")
|
|
47
|
+
|
|
48
|
+
for split in range(len(parts) - 1, 0, -1):
|
|
49
|
+
module_path = ".".join(parts[:split])
|
|
50
|
+
try:
|
|
51
|
+
obj: Any = importlib.import_module(module_path)
|
|
52
|
+
except ImportError:
|
|
53
|
+
continue
|
|
54
|
+
|
|
55
|
+
try:
|
|
56
|
+
for attr in parts[split:]:
|
|
57
|
+
obj = getattr(obj, attr)
|
|
58
|
+
except AttributeError as exc:
|
|
59
|
+
raise MagicRegistryImportError(
|
|
60
|
+
f"'{path}' does not resolve -- '{module_path}' has no"
|
|
61
|
+
f" attribute chain {'.'.join(parts[split:])!r}"
|
|
62
|
+
) from exc
|
|
63
|
+
|
|
64
|
+
if not isinstance(obj, type):
|
|
65
|
+
raise MagicRegistryImportError(f"'{path}' does not resolve to a class")
|
|
66
|
+
|
|
67
|
+
return obj
|
|
68
|
+
|
|
69
|
+
raise MagicRegistryImportError(f"cannot import any module prefix of '{path}'")
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
class _MagicConfigBase(Configurable.Config):
|
|
73
|
+
model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def magic_wrap(path: str, field_keys: Iterable[str]) -> type[Configurable]:
|
|
77
|
+
"""Build a Configurable wrapper subclassing the class at `path`.
|
|
78
|
+
|
|
79
|
+
The wrapper's Config declares exactly `field_keys` as fields, typed from
|
|
80
|
+
the target's own __init__ signature where possible. Anything not in
|
|
81
|
+
field_keys is never a Config field -- pass it at build() time instead,
|
|
82
|
+
same as any other Configurable's owner kwargs (e.g. build(params=...)).
|
|
83
|
+
|
|
84
|
+
Stamped with `_magic_path_` (kept separate from Registry's
|
|
85
|
+
`_registry_path_` -- curated and magic entries are distinct namespaces,
|
|
86
|
+
see docs/registry.md) so resolve_config.to_dict() can round-trip it as
|
|
87
|
+
`_magic_registry_`.
|
|
88
|
+
|
|
89
|
+
No caching: each call builds a fresh class. importlib already memoizes
|
|
90
|
+
the underlying import; a small pydantic model + dynamic subclass is
|
|
91
|
+
cheap enough to build per config node.
|
|
92
|
+
|
|
93
|
+
Args:
|
|
94
|
+
path: Dotted import path to the target class, e.g. "torch.optim.Adam".
|
|
95
|
+
field_keys: Names of the config node's sibling keys -- the fields
|
|
96
|
+
this Config should declare.
|
|
97
|
+
|
|
98
|
+
Returns:
|
|
99
|
+
A new Configurable subclass of the target class.
|
|
100
|
+
|
|
101
|
+
Raises:
|
|
102
|
+
MagicRegistryImportError: If `path` can't be imported / isn't a class.
|
|
103
|
+
MagicRegistryUnknownFieldError: If a field key has no matching
|
|
104
|
+
__init__ parameter and the target has no **kwargs catch-all.
|
|
105
|
+
"""
|
|
106
|
+
target_cls = _import_dotted(path)
|
|
107
|
+
field_keys = list(field_keys)
|
|
108
|
+
field_set = set(field_keys)
|
|
109
|
+
fields = field_specs(target_cls, field_keys)
|
|
110
|
+
|
|
111
|
+
config_cls = create_model(
|
|
112
|
+
f"Magic{target_cls.__name__}Config",
|
|
113
|
+
__base__=_MagicConfigBase,
|
|
114
|
+
**fields,
|
|
115
|
+
) # ty: ignore[no-matching-overload]
|
|
116
|
+
|
|
117
|
+
def __init__(self: Any, cfg: Any, **kwargs: Any) -> None:
|
|
118
|
+
field_values = {name: getattr(cfg, name) for name in type(cfg).model_fields}
|
|
119
|
+
target_cls.__init__(self, **field_values, **kwargs)
|
|
120
|
+
|
|
121
|
+
# build() infers legal extra ("owner") kwargs from __init__'s *signature*,
|
|
122
|
+
# not from the **kwargs catch-all above -- expose target_cls's remaining
|
|
123
|
+
# real params (everything not already a Config field) so e.g.
|
|
124
|
+
# build(params=...) is recognized instead of rejected as unexpected.
|
|
125
|
+
try:
|
|
126
|
+
target_sig = inspect.signature(target_cls.__init__)
|
|
127
|
+
remaining = [
|
|
128
|
+
p
|
|
129
|
+
for name, p in target_sig.parameters.items()
|
|
130
|
+
if name not in field_set and name != "self"
|
|
131
|
+
]
|
|
132
|
+
__init__.__signature__ = inspect.Signature( # ty: ignore[unresolved-attribute]
|
|
133
|
+
[
|
|
134
|
+
inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
|
|
135
|
+
inspect.Parameter("cfg", inspect.Parameter.POSITIONAL_OR_KEYWORD),
|
|
136
|
+
*remaining,
|
|
137
|
+
]
|
|
138
|
+
)
|
|
139
|
+
except (TypeError, ValueError):
|
|
140
|
+
pass # unintrospectable __init__ -- fall back to bare **kwargs signature
|
|
141
|
+
|
|
142
|
+
wrapper_cls = type(
|
|
143
|
+
f"Magic{target_cls.__name__}",
|
|
144
|
+
(target_cls, Configurable),
|
|
145
|
+
{"Config": config_cls, "__init__": __init__, "_magic_path_": path},
|
|
146
|
+
)
|
|
147
|
+
return cast(type[Configurable], wrapper_cls)
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
def field_specs(
|
|
151
|
+
target_cls: type, field_keys: Iterable[str]
|
|
152
|
+
) -> dict[str, tuple[Any, Any]]:
|
|
153
|
+
"""Build create_model field specs for exactly the given keys.
|
|
154
|
+
|
|
155
|
+
Types come from target_cls.__init__'s own annotations when resolvable;
|
|
156
|
+
falls back to Any per-field, or wholesale if the signature itself can't
|
|
157
|
+
be introspected (some C-extension __init__s raise on inspect.signature).
|
|
158
|
+
|
|
159
|
+
Args:
|
|
160
|
+
target_cls: The class whose __init__ is introspected.
|
|
161
|
+
field_keys: Field names to build specs for -- exactly the keys
|
|
162
|
+
present in the config node, not the full signature.
|
|
163
|
+
|
|
164
|
+
Returns:
|
|
165
|
+
Dict mapping field name to a (type, default) tuple for
|
|
166
|
+
pydantic.create_model, all required (default is `...`).
|
|
167
|
+
|
|
168
|
+
Raises:
|
|
169
|
+
MagicRegistryUnknownFieldError: If a key has no matching parameter
|
|
170
|
+
on target_cls.__init__ and there's no **kwargs catch-all.
|
|
171
|
+
"""
|
|
172
|
+
try:
|
|
173
|
+
sig = inspect.signature(target_cls.__init__)
|
|
174
|
+
except (TypeError, ValueError):
|
|
175
|
+
return dict.fromkeys(field_keys, (Any, ...))
|
|
176
|
+
|
|
177
|
+
try:
|
|
178
|
+
hints = typing.get_type_hints(target_cls.__init__)
|
|
179
|
+
except Exception:
|
|
180
|
+
hints = {}
|
|
181
|
+
|
|
182
|
+
has_var_keyword = any(
|
|
183
|
+
p.kind is inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
|
|
184
|
+
)
|
|
185
|
+
|
|
186
|
+
specs: dict[str, tuple[Any, Any]] = {}
|
|
187
|
+
for key in field_keys:
|
|
188
|
+
param = sig.parameters.get(key)
|
|
189
|
+
|
|
190
|
+
if param is None:
|
|
191
|
+
if has_var_keyword:
|
|
192
|
+
specs[key] = (Any, ...)
|
|
193
|
+
continue
|
|
194
|
+
raise MagicRegistryUnknownFieldError(
|
|
195
|
+
f"{target_cls.__module__}.{target_cls.__qualname__} has no"
|
|
196
|
+
f" '{key}' parameter"
|
|
197
|
+
)
|
|
198
|
+
|
|
199
|
+
annotation = hints.get(key, param.annotation)
|
|
200
|
+
if annotation is inspect.Parameter.empty or isinstance(annotation, str):
|
|
201
|
+
annotation = Any
|
|
202
|
+
specs[key] = (annotation, ...)
|
|
203
|
+
|
|
204
|
+
return specs
|
|
@@ -39,8 +39,12 @@ class Registry[T]:
|
|
|
39
39
|
True
|
|
40
40
|
"""
|
|
41
41
|
|
|
42
|
-
def __init__(
|
|
42
|
+
def __init__(
|
|
43
|
+
self, *, prefix: str = "", root: "Registry[Any] | None" = None
|
|
44
|
+
) -> None:
|
|
43
45
|
self._entries: dict[str, type[Any] | Registry[Any]] = {}
|
|
46
|
+
self._prefix = prefix
|
|
47
|
+
self._root = root if root is not None else self
|
|
44
48
|
|
|
45
49
|
def subgroup[S](self, name: str) -> "Registry[S]":
|
|
46
50
|
"""Create a named child registry and store it under ``name``.
|
|
@@ -61,7 +65,8 @@ class Registry[T]:
|
|
|
61
65
|
if name in self._entries:
|
|
62
66
|
raise DuplicateRegistrationError(f"'{name}' is already registered")
|
|
63
67
|
|
|
64
|
-
|
|
68
|
+
full_prefix = f"{self._prefix}.{name}" if self._prefix else name
|
|
69
|
+
child: Registry[S] = Registry(prefix=full_prefix, root=self._root)
|
|
65
70
|
self._entries[name] = child
|
|
66
71
|
|
|
67
72
|
return child
|
|
@@ -73,20 +78,44 @@ class Registry[T]:
|
|
|
73
78
|
name: Single-segment name. Must not contain a period.
|
|
74
79
|
|
|
75
80
|
Returns:
|
|
76
|
-
A decorator that
|
|
81
|
+
A decorator that stamps the class's dotted registry path onto
|
|
82
|
+
``cls._registry_path_``, stores it, and returns it unchanged.
|
|
77
83
|
|
|
78
84
|
Raises:
|
|
79
85
|
InvalidRegistryNameError: If ``name`` contains a period.
|
|
80
86
|
DuplicateRegistrationError: If ``name`` is already registered in
|
|
81
87
|
this registry (raised at factory-call time, before the
|
|
82
|
-
decorated class is defined)
|
|
88
|
+
decorated class is defined), or if ``cls`` is already
|
|
89
|
+
reachable at another path within this registry's tree
|
|
90
|
+
(raised at decoration time). A stamp left over from an
|
|
91
|
+
unrelated, independently-constructed ``Registry()`` --
|
|
92
|
+
e.g. a fresh registry built per test -- does not count,
|
|
93
|
+
since it no longer resolves to anything.
|
|
83
94
|
"""
|
|
84
95
|
_validate_name(name)
|
|
85
96
|
|
|
86
97
|
if name in self._entries:
|
|
87
98
|
raise DuplicateRegistrationError(f"'{name}' is already registered")
|
|
88
99
|
|
|
100
|
+
full_path = f"{self._prefix}.{name}" if self._prefix else name
|
|
101
|
+
|
|
89
102
|
def decorator(cls: type[S]) -> type[S]:
|
|
103
|
+
existing_path = getattr(cls, "_registry_path_", None)
|
|
104
|
+
if existing_path is not None:
|
|
105
|
+
# existing_path may belong to an unrelated Registry() tree
|
|
106
|
+
# (e.g. a fresh registry built per test) that simply hasn't
|
|
107
|
+
# reached this path yet -- that's not a real duplicate.
|
|
108
|
+
try:
|
|
109
|
+
still_live = self._root.get(existing_path) is cls
|
|
110
|
+
except RegistryKeyError:
|
|
111
|
+
still_live = False
|
|
112
|
+
if still_live:
|
|
113
|
+
raise DuplicateRegistrationError(
|
|
114
|
+
f"{cls.__name__} is already registered at"
|
|
115
|
+
f" '{existing_path}' -- cannot also register it at"
|
|
116
|
+
f" '{full_path}'"
|
|
117
|
+
)
|
|
118
|
+
cast(Any, cls)._registry_path_ = full_path
|
|
90
119
|
self._entries[name] = cls
|
|
91
120
|
return cls
|
|
92
121
|
|
|
@@ -183,8 +212,9 @@ class Registry[T]:
|
|
|
183
212
|
def path_of(self, cls: type) -> str | None:
|
|
184
213
|
"""Return the dotted registry path for ``cls``, or ``None`` if not registered.
|
|
185
214
|
|
|
186
|
-
|
|
187
|
-
|
|
215
|
+
Reads the path stamped on ``cls`` by ``register()`` and confirms it
|
|
216
|
+
resolves back to ``cls`` within this registry -- O(path depth), not
|
|
217
|
+
O(registry size).
|
|
188
218
|
|
|
189
219
|
Args:
|
|
190
220
|
cls: The class to look up.
|
|
@@ -201,32 +231,14 @@ class Registry[T]:
|
|
|
201
231
|
>>> reg.path_of(MLP)
|
|
202
232
|
'models.mlp'
|
|
203
233
|
"""
|
|
204
|
-
|
|
205
|
-
|
|
206
|
-
|
|
207
|
-
"""Recursively search for a class in the registry tree.
|
|
208
|
-
|
|
209
|
-
Args:
|
|
210
|
-
cls: The class to look up.
|
|
211
|
-
prefix: Dot-separated path prefix accumulated by parent calls.
|
|
234
|
+
path = getattr(cls, "_registry_path_", None)
|
|
235
|
+
if path is None:
|
|
236
|
+
return None
|
|
212
237
|
|
|
213
|
-
|
|
214
|
-
|
|
215
|
-
|
|
216
|
-
|
|
217
|
-
full_name = f"{prefix}.{name}" if prefix else name
|
|
218
|
-
|
|
219
|
-
if entry is cls:
|
|
220
|
-
# found answer
|
|
221
|
-
return full_name
|
|
222
|
-
|
|
223
|
-
if isinstance(entry, Registry):
|
|
224
|
-
result = entry._find_path(cls, full_name)
|
|
225
|
-
if result is not None:
|
|
226
|
-
# pass solution up
|
|
227
|
-
return result
|
|
228
|
-
|
|
229
|
-
return None
|
|
238
|
+
try:
|
|
239
|
+
return path if self.get(path) is cls else None
|
|
240
|
+
except RegistryKeyError:
|
|
241
|
+
return None
|
|
230
242
|
|
|
231
243
|
|
|
232
244
|
def _validate_name(name: str) -> None:
|
|
@@ -4,7 +4,8 @@ from typing import Any
|
|
|
4
4
|
|
|
5
5
|
from pydantic import BaseModel
|
|
6
6
|
|
|
7
|
-
from mlbricks.configurable import Configurable
|
|
7
|
+
from mlbricks.configurable import Configurable, is_build_time_field
|
|
8
|
+
from mlbricks.magic_registry import magic_wrap
|
|
8
9
|
from mlbricks.registry import Registry
|
|
9
10
|
|
|
10
11
|
__all__ = [
|
|
@@ -61,8 +62,11 @@ def parse(
|
|
|
61
62
|
A top-level _registry_ dict is itself replaced by a ParsedNode.
|
|
62
63
|
|
|
63
64
|
Raises:
|
|
64
|
-
RegistryParseError: If a _registry_ value is not a string
|
|
65
|
+
RegistryParseError: If a _registry_ value is not a string, or a node
|
|
66
|
+
has both _registry_ and _magic_registry_ keys.
|
|
65
67
|
RegistryKeyError: If a _registry_ key is not found in the registry.
|
|
68
|
+
MagicRegistryError: If a _magic_registry_ path can't be imported or a
|
|
69
|
+
field doesn't match the target's __init__.
|
|
66
70
|
|
|
67
71
|
Example:
|
|
68
72
|
>>> result = parse({"_registry_": "models.mlp", "dim": 64}, registry=reg)
|
|
@@ -80,10 +84,24 @@ def _parse_dict(
|
|
|
80
84
|
node: dict[str, Any],
|
|
81
85
|
registry: Registry[Any],
|
|
82
86
|
) -> dict[str, Any] | ParsedNode:
|
|
83
|
-
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
87
|
+
match "_registry_" in node, "_magic_registry_" in node:
|
|
88
|
+
case True, True:
|
|
89
|
+
raise RegistryParseError(
|
|
90
|
+
"node cannot have both '_registry_' and '_magic_registry_' keys"
|
|
91
|
+
)
|
|
92
|
+
case False, False:
|
|
93
|
+
# plain dict -- recurse into values without creating a ParsedNode
|
|
94
|
+
return {k: parse(v, registry) for k, v in node.items()}
|
|
95
|
+
case False, True:
|
|
96
|
+
return _parse_magic_dict(node, registry)
|
|
97
|
+
case True, False:
|
|
98
|
+
return _parse_registry_dict(node, registry)
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
def _parse_registry_dict(
|
|
102
|
+
node: dict[str, Any],
|
|
103
|
+
registry: Registry[Any],
|
|
104
|
+
) -> ParsedNode:
|
|
87
105
|
key = node["_registry_"]
|
|
88
106
|
if not isinstance(key, str):
|
|
89
107
|
raise RegistryParseError(
|
|
@@ -99,6 +117,24 @@ def _parse_dict(
|
|
|
99
117
|
return ParsedNode(key=key, cls=cls, kwargs=kwargs)
|
|
100
118
|
|
|
101
119
|
|
|
120
|
+
def _parse_magic_dict(
|
|
121
|
+
node: dict[str, Any],
|
|
122
|
+
registry: Registry[Any],
|
|
123
|
+
) -> ParsedNode:
|
|
124
|
+
path = node["_magic_registry_"]
|
|
125
|
+
if not isinstance(path, str):
|
|
126
|
+
raise RegistryParseError(
|
|
127
|
+
f"_magic_registry_ value must be a string, got {type(path).__name__!r}"
|
|
128
|
+
)
|
|
129
|
+
|
|
130
|
+
kwargs = {k: parse(v, registry) for k, v in node.items() if k != "_magic_registry_"}
|
|
131
|
+
|
|
132
|
+
# MagicRegistryError propagates as-is
|
|
133
|
+
cls = magic_wrap(path, kwargs.keys())
|
|
134
|
+
|
|
135
|
+
return ParsedNode(key=path, cls=cls, kwargs=kwargs)
|
|
136
|
+
|
|
137
|
+
|
|
102
138
|
def instantiate(parsed: Any) -> Any:
|
|
103
139
|
"""Construct all ParsedNodes in a parsed value into Python objects.
|
|
104
140
|
|
|
@@ -232,6 +268,10 @@ def _config_to_dict(
|
|
|
232
268
|
"Config has no owner -- define it inside a Configurable subclass"
|
|
233
269
|
)
|
|
234
270
|
|
|
271
|
+
magic_path = getattr(owner, "_magic_path_", None)
|
|
272
|
+
if magic_path is not None:
|
|
273
|
+
return _dump_fields(cfg, registry, key="_magic_registry_", path=magic_path)
|
|
274
|
+
|
|
235
275
|
path = registry.path_of(owner)
|
|
236
276
|
if path is None:
|
|
237
277
|
raise ConfigSerializationError(
|
|
@@ -239,11 +279,17 @@ def _config_to_dict(
|
|
|
239
279
|
" -- cannot serialize"
|
|
240
280
|
)
|
|
241
281
|
|
|
242
|
-
|
|
282
|
+
return _dump_fields(cfg, registry, key="_registry_", path=path)
|
|
243
283
|
|
|
244
|
-
for name in type(cfg).model_fields:
|
|
245
|
-
result[name] = to_dict(getattr(cfg, name), registry)
|
|
246
284
|
|
|
285
|
+
def _dump_fields(
|
|
286
|
+
cfg: Configurable.Config, registry: Registry[Any], *, key: str, path: str
|
|
287
|
+
) -> dict[str, Any]:
|
|
288
|
+
result: dict[str, Any] = {key: path}
|
|
289
|
+
for name, field in type(cfg).model_fields.items():
|
|
290
|
+
if is_build_time_field(field):
|
|
291
|
+
continue
|
|
292
|
+
result[name] = to_dict(getattr(cfg, name), registry)
|
|
247
293
|
return result
|
|
248
294
|
|
|
249
295
|
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|