|
| 1 | +# Copyright (c) Microsoft Corporation. |
| 2 | +# Licensed under the MIT License. |
| 3 | + |
| 4 | +"""Canonical metadata for the components of a model package.""" |
| 5 | + |
| 6 | +from __future__ import annotations |
| 7 | + |
| 8 | +__all__ = [ |
| 9 | + "ComponentDescriptor", |
| 10 | + "ComponentManifest", |
| 11 | + "get_hf_component_sources", |
| 12 | + "resolve_component_manifest", |
| 13 | +] |
| 14 | + |
| 15 | +import dataclasses |
| 16 | +from collections.abc import Iterator, Mapping |
| 17 | +from types import MappingProxyType |
| 18 | +from typing import TYPE_CHECKING |
| 19 | + |
| 20 | +if TYPE_CHECKING: |
| 21 | + pass |
| 22 | + |
| 23 | + |
| 24 | +@dataclasses.dataclass(frozen=True) |
| 25 | +class ComponentDescriptor: |
| 26 | + """One package component and all metadata needed to address it. |
| 27 | +
|
| 28 | + Attributes: |
| 29 | + name: Key used by :class:`~mobius.ModelPackage`. |
| 30 | + module_attribute_path: Dotted Python attribute path from the root |
| 31 | + :class:`onnxscript.nn.Module` passed to ``task.build()`` to the |
| 32 | + sub-module that constructs this component. The empty string means |
| 33 | + the root module itself. This is not a package key or checkpoint |
| 34 | + prefix. |
| 35 | + role: Task-defined optimization category. Current roles include |
| 36 | + ``decoder``, ``encoder``, ``vision``, ``embedding``, and ``glue``. |
| 37 | + source_paths: Runtime HuggingFace ``named_modules()`` paths whose |
| 38 | + weights belong to this component. |
| 39 | + source_path_aliases: Pairs of ``(local_prefix, source_prefix)`` for |
| 40 | + component paths that cannot be aligned by a shared anchor segment. |
| 41 | + """ |
| 42 | + |
| 43 | + name: str |
| 44 | + module_attribute_path: str |
| 45 | + role: str |
| 46 | + source_paths: tuple[str, ...] = () |
| 47 | + source_path_aliases: tuple[tuple[str, str], ...] = () |
| 48 | + |
| 49 | + def __post_init__(self) -> None: |
| 50 | + if not self.name: |
| 51 | + raise ValueError("component name must not be empty") |
| 52 | + if not self.role: |
| 53 | + raise ValueError(f"component {self.name!r} must declare a role") |
| 54 | + if any(not path for path in self.source_paths): |
| 55 | + raise ValueError( |
| 56 | + f"component {self.name!r} source_paths must not contain empty paths" |
| 57 | + ) |
| 58 | + if any(not local or not source for local, source in self.source_path_aliases): |
| 59 | + raise ValueError( |
| 60 | + f"component {self.name!r} source_path_aliases must contain " |
| 61 | + "non-empty local/source prefixes" |
| 62 | + ) |
| 63 | + |
| 64 | + def source_module_names(self, local_module_path: str) -> tuple[str, ...]: |
| 65 | + """Candidate HuggingFace names for a component-local module path. |
| 66 | +
|
| 67 | + Source roots and Mobius paths commonly share an anchor segment even |
| 68 | + when their prefixes differ. For example, source root |
| 69 | + ``model.language_model.layers`` and local path |
| 70 | + ``model.layers.0.self_attn.q_proj`` share ``layers`` and resolve to |
| 71 | + ``model.language_model.layers.0.self_attn.q_proj``. |
| 72 | + """ |
| 73 | + if not local_module_path: |
| 74 | + return self.source_paths |
| 75 | + |
| 76 | + local_parts = local_module_path.split(".") |
| 77 | + candidates = [local_module_path] |
| 78 | + for local_prefix, source_prefix in self.source_path_aliases: |
| 79 | + if local_module_path == local_prefix: |
| 80 | + candidates.append(source_prefix) |
| 81 | + elif local_module_path.startswith(f"{local_prefix}."): |
| 82 | + suffix = local_module_path[len(local_prefix) + 1 :] |
| 83 | + candidates.append(f"{source_prefix}.{suffix}") |
| 84 | + for source_path in self.source_paths: |
| 85 | + source_parts = source_path.split(".") |
| 86 | + anchor = source_parts[-1] |
| 87 | + anchor_indices = [ |
| 88 | + index for index, part in enumerate(local_parts) if part == anchor |
| 89 | + ] |
| 90 | + if anchor_indices: |
| 91 | + for index in anchor_indices: |
| 92 | + suffix = local_parts[index + 1 :] |
| 93 | + candidates.append(".".join((*source_parts, *suffix))) |
| 94 | + return tuple(dict.fromkeys(candidates)) |
| 95 | + |
| 96 | + |
| 97 | +@dataclasses.dataclass(frozen=True) |
| 98 | +class ComponentManifest(Mapping[str, ComponentDescriptor]): |
| 99 | + """Ordered, immutable component metadata keyed by package component name.""" |
| 100 | + |
| 101 | + components: tuple[ComponentDescriptor, ...] |
| 102 | + _by_name: Mapping[str, ComponentDescriptor] = dataclasses.field( |
| 103 | + init=False, |
| 104 | + repr=False, |
| 105 | + compare=False, |
| 106 | + ) |
| 107 | + |
| 108 | + def __post_init__(self) -> None: |
| 109 | + by_name: dict[str, ComponentDescriptor] = {} |
| 110 | + for component in self.components: |
| 111 | + if component.name in by_name: |
| 112 | + raise ValueError( |
| 113 | + f"component manifest declares {component.name!r} more than once" |
| 114 | + ) |
| 115 | + by_name[component.name] = component |
| 116 | + object.__setattr__(self, "_by_name", MappingProxyType(by_name)) |
| 117 | + |
| 118 | + def __getitem__(self, name: str) -> ComponentDescriptor: |
| 119 | + return self._by_name[name] |
| 120 | + |
| 121 | + def __iter__(self) -> Iterator[str]: |
| 122 | + return iter(self._by_name) |
| 123 | + |
| 124 | + def __len__(self) -> int: |
| 125 | + return len(self._by_name) |
| 126 | + |
| 127 | + @property |
| 128 | + def names(self) -> tuple[str, ...]: |
| 129 | + """Component names in task declaration order.""" |
| 130 | + return tuple(self._by_name) |
| 131 | + |
| 132 | + |
| 133 | +def get_hf_component_sources( |
| 134 | + module_class: type, |
| 135 | + model_type: str, |
| 136 | + hf_config: object, |
| 137 | +) -> dict[str, tuple[str, ...]]: |
| 138 | + """Read runtime HuggingFace component paths from a registered model class.""" |
| 139 | + resolver = getattr(module_class, "get_hf_component_sources", None) |
| 140 | + if resolver is not None: |
| 141 | + resolved = resolver(model_type=model_type, hf_config=hf_config) |
| 142 | + else: |
| 143 | + resolved = getattr(module_class, "HF_COMPONENT_SOURCES", {}) |
| 144 | + return {name: tuple(paths) for name, paths in resolved.items()} |
| 145 | + |
| 146 | + |
| 147 | +def resolve_component_manifest( |
| 148 | + task: object, |
| 149 | + *, |
| 150 | + module_class: type | None = None, |
| 151 | + model_type: str | None = None, |
| 152 | + hf_config: object | None = None, |
| 153 | +) -> ComponentManifest: |
| 154 | + """Combine task roles/paths and model source ownership into one manifest.""" |
| 155 | + roles = dict(getattr(task, "model_roles", {}) or {}) |
| 156 | + component_spec = getattr(task, "components", None) |
| 157 | + module_paths = dict(component_spec.items()) if component_spec is not None else {} |
| 158 | + |
| 159 | + component_sources: dict[str, tuple[str, ...]] = {} |
| 160 | + component_aliases: dict[str, tuple[tuple[str, str], ...]] = {} |
| 161 | + if module_class is not None and model_type is not None and hf_config is not None: |
| 162 | + component_sources = get_hf_component_sources( |
| 163 | + module_class, |
| 164 | + model_type, |
| 165 | + hf_config, |
| 166 | + ) |
| 167 | + raw_aliases = getattr(module_class, "HF_COMPONENT_MODULE_ALIASES", {}) |
| 168 | + component_aliases = { |
| 169 | + name: tuple(aliases.items()) for name, aliases in raw_aliases.items() |
| 170 | + } |
| 171 | + |
| 172 | + ordered_names = tuple(dict.fromkeys((*roles, *module_paths))) |
| 173 | + descriptors = tuple( |
| 174 | + ComponentDescriptor( |
| 175 | + name=name, |
| 176 | + module_attribute_path=module_paths.get( |
| 177 | + name, |
| 178 | + "" if name == "model" else name, |
| 179 | + ), |
| 180 | + role=roles.get(name, "decoder"), |
| 181 | + source_paths=component_sources.get(name, ()), |
| 182 | + source_path_aliases=component_aliases.get(name, ()), |
| 183 | + ) |
| 184 | + for name in ordered_names |
| 185 | + ) |
| 186 | + return ComponentManifest(descriptors) |
0 commit comments