Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / tools / model_integration / structure_inventory.py: 77%
168 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Model/config/checkpoint inventory and classified structure differences."""
17from __future__ import annotations
19import inspect
20import json
21from dataclasses import asdict, dataclass, field
22from enum import Enum
23from pathlib import Path
24from typing import Any, Callable, Optional
27class DifferenceCategory(str, Enum):
28 """Action category for one authoritative-versus-candidate difference."""
30 GENERIC_HF = "generic_hf_supported"
31 MODEL_ADAPTER = "model_adapter"
32 COMPONENT_REPLACEMENT = "component_replacement"
33 CUSTOM_MODEL = "hf_component_custom_model"
34 FRAMEWORK_GAP = "framework_contract_gap"
35 OUT_OF_SCOPE = "decode_cache_or_quantization_out_of_scope"
38@dataclass(frozen=True)
39class InventoryDifference:
40 """One classified structural difference."""
42 path: str
43 reference: Any
44 candidate: Any
45 category: DifferenceCategory
46 reason: str
48 def to_dict(self) -> dict[str, Any]:
49 """Return a JSON-serializable representation."""
50 payload = asdict(self)
51 payload["category"] = self.category.value
52 return payload
55@dataclass
56class ModelInventory:
57 """Serializable inventory of a model's final module and state tree."""
59 source: dict[str, Any] = field(default_factory=dict)
60 config: dict[str, Any] = field(default_factory=dict)
61 modules: list[dict[str, Any]] = field(default_factory=list)
62 parameters: list[dict[str, Any]] = field(default_factory=list)
63 buffers: list[dict[str, Any]] = field(default_factory=list)
64 aliases: dict[str, list[str]] = field(default_factory=dict)
65 checkpoint: dict[str, Any] = field(default_factory=dict)
67 def to_dict(self) -> dict[str, Any]:
68 """Return a JSON-serializable representation."""
69 return asdict(self)
72def _qualified_type(value: Any) -> str:
73 value_type = type(value)
74 return f"{value_type.__module__}.{value_type.__qualname__}"
77def _shape(value: Any) -> list[int]:
78 return [int(size) for size in getattr(value, "shape", ())]
81def _layout(value: Any) -> dict[str, Any] | None:
82 layout = getattr(value, "layout", None)
83 if layout is None:
84 placements = getattr(value, "placements", None)
85 mesh = getattr(value, "device_mesh", None)
86 else:
87 placements = getattr(layout, "placements", None)
88 mesh = getattr(layout, "mesh", None)
89 if placements is None and mesh is None:
90 return None
91 return {
92 "mesh_dim_names": list(getattr(mesh, "mesh_dim_names", ()) or ()),
93 "mesh_shape": list(getattr(getattr(mesh, "mesh", None), "shape", ()) or ()),
94 "placements": [repr(placement) for placement in (placements or ())],
95 }
98def _forward_signature(module: Any) -> str | None:
99 forward = getattr(module, "forward", None)
100 if not callable(forward):
101 return None
102 try:
103 return str(inspect.signature(forward))
104 except (TypeError, ValueError):
105 return None
108def _config_dict(model: Any) -> dict[str, Any]:
109 config = getattr(model, "config", None)
110 if config is None:
111 return {}
112 to_dict = getattr(config, "to_dict", None)
113 if callable(to_dict):
114 return to_dict()
115 if hasattr(config, "__dict__"):
116 return {
117 name: value
118 for name, value in vars(config).items()
119 if isinstance(value, (str, int, float, bool, type(None), list, dict))
120 }
121 return {"repr": repr(config)}
124def build_model_inventory(
125 model: Any,
126 *,
127 source: Optional[dict[str, Any]] = None,
128 checkpoint: Optional[dict[str, Any]] = None,
129) -> ModelInventory:
130 """Inventory the instantiated final model tree, including aliases.
132 Args:
133 model: Torch-compatible module exposing named module/state iterators.
134 source: Recorded source/revision/import metadata.
135 checkpoint: Optional checkpoint index inventory.
137 Returns:
138 Serializable model inventory.
139 """
140 modules = []
141 for module_fqn, module in model.named_modules(remove_duplicate=False):
142 modules.append(
143 {
144 "fqn": module_fqn or "<root>",
145 "type": _qualified_type(module),
146 "forward_signature": _forward_signature(module),
147 "training": bool(getattr(module, "training", False)),
148 }
149 )
151 parameter_aliases: dict[int, list[str]] = {}
152 parameter_by_id = {}
153 for parameter_fqn, parameter in model.named_parameters(remove_duplicate=False):
154 parameter_aliases.setdefault(id(parameter), []).append(parameter_fqn)
155 parameter_by_id[id(parameter)] = parameter
156 parameters = []
157 for parameter_id, fqns in sorted(parameter_aliases.items(), key=lambda item: min(item[1])):
158 parameter = parameter_by_id[parameter_id]
159 parameters.append(
160 {
161 "fqn": min(fqns),
162 "aliases": sorted(fqns),
163 "shape": _shape(parameter),
164 "dtype": str(getattr(parameter, "dtype", None)),
165 "device": str(getattr(parameter, "device", None)),
166 "requires_grad": bool(getattr(parameter, "requires_grad", False)),
167 "is_meta": bool(getattr(parameter, "is_meta", False)),
168 "layout": _layout(parameter),
169 }
170 )
172 non_persistent_ids = set()
173 for module in model.modules():
174 for name in getattr(module, "_non_persistent_buffers_set", set()):
175 buffer_value = getattr(module, name, None)
176 if buffer_value is not None:
177 non_persistent_ids.add(id(buffer_value))
178 buffer_aliases: dict[int, list[str]] = {}
179 buffer_by_id = {}
180 for buffer_fqn, buffer in model.named_buffers(remove_duplicate=False):
181 buffer_aliases.setdefault(id(buffer), []).append(buffer_fqn)
182 buffer_by_id[id(buffer)] = buffer
183 buffers = []
184 for buffer_id, fqns in sorted(buffer_aliases.items(), key=lambda item: min(item[1])):
185 buffer = buffer_by_id[buffer_id]
186 buffers.append(
187 {
188 "fqn": min(fqns),
189 "aliases": sorted(fqns),
190 "shape": _shape(buffer),
191 "dtype": str(getattr(buffer, "dtype", None)),
192 "device": str(getattr(buffer, "device", None)),
193 "is_meta": bool(getattr(buffer, "is_meta", False)),
194 "persistent": buffer_id not in non_persistent_ids,
195 "layout": _layout(buffer),
196 }
197 )
198 aliases = {
199 "parameters": [fqns for fqns in parameter_aliases.values() if len(fqns) > 1],
200 "buffers": [fqns for fqns in buffer_aliases.values() if len(fqns) > 1],
201 }
202 return ModelInventory(
203 source=dict(source or {}),
204 config=_config_dict(model),
205 modules=modules,
206 parameters=parameters,
207 buffers=buffers,
208 aliases=aliases,
209 checkpoint=dict(checkpoint or {}),
210 )
213def inspect_local_model_assets(model_path: str | Path) -> dict[str, Any]:
214 """Inspect local config and checkpoint indices without loading tensors."""
215 root = Path(model_path).expanduser().resolve()
216 if not root.is_dir():
217 raise ValueError(f"model asset path is not a directory: {root}")
218 config_path = root / "config.json"
219 config = json.loads(config_path.read_text(encoding="utf-8")) if config_path.is_file() else {}
220 checkpoint = {
221 "keys": [],
222 "weight_map": {},
223 "metadata": {},
224 "lfs_pointers": {},
225 }
227 def lfs_pointer(path: Path) -> dict[str, Any] | None:
228 """Parse one small Git-LFS pointer without resolving its payload."""
229 if path.stat().st_size > 1024:
230 return None
231 lines = path.read_text(encoding="utf-8").splitlines()
232 if not lines or lines[0] != "version https://git-lfs.github.com/spec/v1":
233 return None
234 fields = {}
235 for line in lines[1:]:
236 name, _, value = line.partition(" ")
237 fields[name] = value
238 oid = fields.get("oid", "")
239 size = fields.get("size", "")
240 if not oid.startswith("sha256:") or not size.isdigit():
241 raise ValueError(f"invalid Git-LFS pointer: {path}")
242 return {
243 "oid": oid,
244 "payload_size": int(size),
245 "payload_available": False,
246 }
248 index_paths = sorted(root.glob("*.index.json"))
249 for index_path in index_paths:
250 pointer = lfs_pointer(index_path)
251 if pointer is not None:
252 checkpoint["lfs_pointers"][index_path.name] = pointer
253 checkpoint["metadata"][index_path.name] = {"git_lfs_pointer": True}
254 continue
255 try:
256 index = json.loads(index_path.read_text(encoding="utf-8"))
257 except json.JSONDecodeError as exc:
258 raise ValueError(f"checkpoint index is neither JSON nor a Git-LFS pointer: {index_path}") from exc
259 weight_map = index.get("weight_map", {})
260 checkpoint["weight_map"].update(weight_map)
261 checkpoint["metadata"][index_path.name] = index.get("metadata", {})
262 for weight_path in sorted(root.glob("*.safetensors")):
263 pointer = lfs_pointer(weight_path)
264 if pointer is not None:
265 checkpoint["lfs_pointers"][weight_path.name] = pointer
266 checkpoint["keys"] = sorted(checkpoint["weight_map"])
267 checkpoint["weights_available"] = not checkpoint["lfs_pointers"]
268 return {
269 "path": str(root),
270 "config": config,
271 "checkpoint": checkpoint,
272 "files": sorted(path.name for path in root.iterdir() if path.is_file()),
273 }
276def diff_inventories(
277 reference: ModelInventory,
278 candidate: ModelInventory,
279 classifier: Optional[Callable[[str, Any, Any], tuple[DifferenceCategory, str]]] = None,
280) -> list[InventoryDifference]:
281 """Compare module/state records and classify every difference."""
282 differences = []
284 def classify(path: str, expected: Any, actual: Any) -> tuple[DifferenceCategory, str]:
285 """Classify one inventory difference into its owning integration layer."""
286 if classifier is not None:
287 return classifier(path, expected, actual)
288 if path.startswith("parameters.") or path.startswith("buffers."):
289 return DifferenceCategory.MODEL_ADAPTER, "state layout/name requires explicit adapter mapping"
290 if path.startswith("modules."):
291 return DifferenceCategory.COMPONENT_REPLACEMENT, "module structure requires replacement review"
292 return DifferenceCategory.CUSTOM_MODEL, "configuration differs from the authoritative architecture"
294 for collection_name in ("modules", "parameters", "buffers"):
295 reference_entries = {
296 entry["fqn"]: entry for entry in getattr(reference, collection_name)
297 }
298 candidate_entries = {
299 entry["fqn"]: entry for entry in getattr(candidate, collection_name)
300 }
301 for fqn in sorted(set(reference_entries) | set(candidate_entries)):
302 expected = reference_entries.get(fqn)
303 actual = candidate_entries.get(fqn)
304 if expected == actual:
305 continue
306 path = f"{collection_name}.{fqn}"
307 category, reason = classify(path, expected, actual)
308 differences.append(
309 InventoryDifference(path, expected, actual, category, reason)
310 )
311 if reference.config != candidate.config:
312 category, reason = classify("config", reference.config, candidate.config)
313 differences.append(
314 InventoryDifference("config", reference.config, candidate.config, category, reason)
315 )
316 return differences
319__all__ = [
320 "DifferenceCategory",
321 "InventoryDifference",
322 "ModelInventory",
323 "build_model_inventory",
324 "diff_inventories",
325 "inspect_local_model_assets",
326]