Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / source / resolver.py: 26%
235 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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"""Locate the real modeling source file for a model id.
17The generated modeling file is produced from the original ``modeling_*.py``
18that the checkpoint actually loads — for native HF models that is the file in
19the installed transformers package, for remote-code models it is the file
20shipped inside the checkpoint. Getting the right file matters: the generated
21bundle must stay importable in the exact environment the model runs in.
23The resolution chain mirrors what HF actually loads (``from_pretrained``):
25 pretrained_model_name_or_path -> AutoConfig -> model_type / architectures
26 -> modeling source file
28and prefers the *same* source HF would load. The deciding input is
29``trust_remote_code`` (forwarded from the YAML ``model:`` block): when it is
30set and the checkpoint declares ``auto_map``, HF loads the checkpoint's own
31modeling file — remote code wins even over an installed transformers module.
32Without it, HF uses the installed module (if any), and remote code is only
33reached as a fallback. ``None`` of the resolution helpers import
34``transformers`` at module scope: the codegen manager runs before the model is
35built, and source resolution must stay functional (degrading) in environments
36without torch/transformers.
37"""
38from __future__ import annotations
40import importlib.util
41import json
42import logging
43import os
44import re
45from dataclasses import dataclass
46from types import SimpleNamespace
47from typing import Any, Optional
49logger = logging.getLogger(__name__)
52@dataclass
53class ResolvedSource:
54 """A located modeling source file, ready to be copied into a bundle."""
56 module_name: str
57 file_path: str
58 architecture: Optional[str] = None
59 model_type: Optional[str] = None
61 def to_dict(self) -> dict[str, Any]:
62 """JSON-friendly form for ``CodegenMeta.source``."""
63 return {
64 "module_name": self.module_name,
65 "file_path": self.file_path,
66 "architecture": self.architecture,
67 "model_type": self.model_type,
68 }
71def resolve_model_source(
72 config: Any, model_id: str, *, hf_config: Any = None,
73) -> Optional[ResolvedSource]:
74 """Resolve the modeling source for a model id, degrading gracefully.
76 Resolves ``model_id -> AutoConfig -> architectures/model_type -> modeling
77 file``, with fallbacks when the environment cannot do the real lookup:
79 1. ``hf_config`` already resolved by the caller -> resolve against it.
80 2. ``AutoConfig.from_pretrained(model_id)`` (real HF config, lazy import;
81 None on any failure — no transformers / no network / bad id).
82 3. ``resolve_source_identity`` — installed module guessed from the
83 basename (covers locally-installed model types when AutoConfig is
84 unavailable).
85 4. Local checkpoint ``config.json`` -> remote code (``SimpleNamespace``
86 mirrors ``AutoConfig`` attribute access).
88 Returns ``None`` when no readable modeling source can be resolved.
89 """
90 if hf_config is not None:
91 return resolve_hf_source(config, hf_config)
92 cfg = _load_auto_config(model_id, config)
93 if cfg is not None:
94 try:
95 return resolve_hf_source(config, cfg)
96 except FileNotFoundError:
97 # AutoConfig resolved but no modeling file exists anywhere (not
98 # even remote code in the checkpoint) — keep degrading.
99 pass
100 source = resolve_source_identity(model_id)
101 if source is not None:
102 return source
103 return _resolve_source_from_yaml(config, model_id)
106def resolve_hf_source(config: Any, hf_config: Any) -> ResolvedSource:
107 """Resolve the modeling source for a resolved HF config.
109 Priority mirrors what the hyper trainer actually loads, which is plain
110 transformers ``from_pretrained`` semantics (see auto_model._init_model ->
111 _from_pretrained_parent_class, with ``trust_remote_code`` forwarded from the
112 YAML ``model:`` block). The deciding input is ``trust_remote_code``, NOT
113 the mere presence of ``auto_map``:
115 1. ``trust_remote_code`` AND the checkpoint ships remote code (``auto_map``)
116 -> remote code wins, even when an installed module also exists (this is
117 what HF loads). Must resolve to a file, else raise.
118 2. Otherwise an installed transformers module for ``model_type`` -> use it
119 (this is what HF loads when remote code is not trusted).
120 3. Otherwise remote code, if any (a checkpoint whose model_type has no
121 installed module; HF would require trust_remote_code to load it, so we
122 only reach here to still produce a source for generation).
123 """
124 arch = _first_architecture(hf_config)
125 model_type = getattr(hf_config, "model_type", None)
126 has_remote_code = bool(getattr(hf_config, "auto_map", None))
127 trust_remote_code = _trust_remote_code(config)
128 model_path = _model_path_of(config)
130 module_name = _module_for_model_type(model_type)
131 installed_path = import_module_file(module_name) if module_name else None
133 # 1. Trusted remote code is authoritative — HF loads the checkpoint's own
134 # modeling file over any installed module.
135 if has_remote_code and trust_remote_code:
136 remote_path = resolve_remote_code_file(model_path, hf_config)
137 if remote_path is not None:
138 logger.debug("source: resolved to trusted remote code %s", remote_path)
139 return _build_resolved_source(remote_path, "<remote>", arch, model_type)
140 raise FileNotFoundError(
141 f"codegen: trust_remote_code is set and the config for "
142 f"{model_path or model_type!r} declares auto_map "
143 f"({list(getattr(hf_config, 'auto_map', {}) or {})}), but no "
144 "modeling_*.py was found in the checkpoint dir or the loaded "
145 "config's module — remote code is the source HF would load"
146 )
148 # 2. Custom-registered model — what HyperParallel loads for a registered
149 # architecture. The registry entry is
150 # authoritative over an installed transformers module of the same
151 # model_type: model resolution prefers the registered class.
152 custom_module = _custom_registered_module(hf_config)
153 if custom_module is not None:
154 custom_path = import_module_file(custom_module)
155 if custom_path is None:
156 raise FileNotFoundError(
157 f"codegen: architecture {arch!r} is registered to module "
158 f"{custom_module!r} but its file could not be located"
159 )
160 logger.debug(
161 "source: resolved architecture=%r to registered module %s",
162 arch, custom_path,
163 )
164 return _build_resolved_source(custom_path, custom_module, arch, model_type)
166 # 3. Installed transformers module — what HF loads when remote code is not
167 # trusted (or not present).
168 if installed_path is not None:
169 logger.debug(
170 "source: resolved model_type=%r to installed module %s",
171 model_type, installed_path,
172 )
173 return _build_resolved_source(installed_path, module_name, arch, model_type)
175 # 4. Remote code without an installed module (HF would need
176 # trust_remote_code, but we can still locate the file for generation).
177 if has_remote_code:
178 remote_path = resolve_remote_code_file(model_path, hf_config)
179 if remote_path is not None:
180 logger.debug(
181 "source: resolved to remote code %s (no installed module)",
182 remote_path,
183 )
184 return _build_resolved_source(remote_path, "<remote>", arch, model_type)
186 raise FileNotFoundError(
187 f"codegen: cannot resolve modeling source for model_type={model_type!r}, "
188 f"architecture={arch!r} (no installed transformers module, no remote code)"
189 )
192def _custom_registered_module(hf_config: Any) -> Optional[str]:
193 """Consult hyper_parallel's custom-model registry for the modeling module.
195 ``MODEL_ARCH_MAPPING`` is filled lazily: a family's ``registration.py``
196 self-registers on first import (directory convention, see
197 :mod:`hyper_parallel.models.registry`). Call ``get_model_adapter`` first
198 so the lookup cannot miss an undiscovered family.
199 """
200 arch = _first_architecture(hf_config)
201 if not arch:
202 return None
203 try:
204 from hyper_parallel.models.registry import ( # pylint: disable=C0415
205 MODEL_ARCH_MAPPING,
206 get_model_adapter,
207 )
208 except ImportError:
209 return None
210 if MODEL_ARCH_MAPPING.get(arch) is None:
211 # Either spelling may be the key that discovers the family (mirrors
212 # get_is_hf_model's identity loop); stop once the architecture is
213 # registered. The architecture goes first: it is the registry key.
214 for identity in (arch, getattr(hf_config, "model_type", "")):
215 if not identity:
216 continue
217 get_model_adapter(identity)
218 if MODEL_ARCH_MAPPING.get(arch) is not None:
219 break
220 entry = MODEL_ARCH_MAPPING.get(arch)
221 return entry[0] if entry is not None else None
224def resolve_source_identity(model_id: str) -> Optional[ResolvedSource]:
225 """Best-effort source resolution keyed only on a model id.
227 Used as a degradation when no ``hf_config`` / ``AutoConfig`` is available.
228 Derives the module name from the model id's basename: model ids usually
229 carry the repo name the transformers module is keyed on (``DeepSeek-V3``
230 -> ``deepseek_v3``). ``None`` when nothing resolves so the caller falls
231 back further (or records the unresolved ``SourceSpec``).
232 """
233 if not model_id:
234 return None
235 name = os.path.basename(str(model_id).rstrip("/\\"))
236 safe = "".join(c if c.isalnum() or c == "_" else "_" for c in name).lower()
237 module_name = _module_for_model_type(safe)
238 if module_name is None:
239 return None
240 file_path = import_module_file(module_name)
241 if file_path is None:
242 return None
243 return ResolvedSource(
244 module_name=module_name,
245 file_path=file_path,
246 architecture=None,
247 model_type=safe,
248 )
251def resolve_remote_code_file(model_path: str, hf_config: Any) -> Optional[str]:
252 """Locate the remote-code modeling file for a checkpoint.
254 Remote-code checkpoints carry ``modeling_*.py`` + ``configuration_*.py``
255 next to the weights. Two layouts are handled:
257 * A local checkpoint directory — scan it.
258 * A Hub id — the files live in the HF cache, not under ``model_path``. A
259 config loaded with ``trust_remote_code`` was constructed from the cached
260 ``configuration_*.py``, so its class's ``__module__`` file locates the
261 cache dir to scan.
263 Within a directory, resolve in order of increasing specificity:
265 1. ``auto_map["AutoModel..."]`` value, e.g. ``"modeling_foo.FooForCausalLM"``
266 -> ``<dir>/modeling_foo.py``.
267 2. The architecture class name, e.g. ``FooForCausalLM`` -> the file whose
268 basename contains it.
269 3. The single ``modeling_*.py`` when there is exactly one.
271 Returns None when no remote code is reachable.
272 """
273 for directory in _remote_code_dirs(model_path, hf_config):
274 found = _scan_remote_code_dir(directory, hf_config)
275 if found is not None:
276 return found
277 return None
280def _remote_code_dirs(model_path: str, hf_config: Any) -> list[str]:
281 """Directories that may hold this checkpoint's remote code, best first."""
282 dirs: list[str] = []
283 if model_path and os.path.isdir(model_path):
284 dirs.append(model_path)
285 # A remote-code config instance is defined in the cached configuration
286 # module (``.../modules/transformers_modules/<repo>/configuration_x.py``),
287 # which sits next to the cached modeling file. This is how a Hub id
288 # resolves without downloading anything extra.
289 module_name = getattr(type(hf_config), "__module__", "") or ""
290 if module_name and module_name not in ("builtins", "types"):
291 origin = import_module_file(module_name)
292 if origin is None:
293 module = __import__("sys").modules.get(module_name)
294 origin = getattr(module, "__file__", None) if module else None
295 if origin:
296 cached_dir = os.path.dirname(origin)
297 if cached_dir and cached_dir not in dirs and os.path.isdir(cached_dir):
298 dirs.append(cached_dir)
299 return dirs
302def _scan_remote_code_dir(model_path: str, hf_config: Any) -> Optional[str]:
303 """Pick this checkpoint's modeling file out of one directory."""
304 try:
305 candidates = sorted(
306 name for name in os.listdir(model_path)
307 if name.startswith("modeling_") and name.endswith(".py")
308 )
309 except OSError:
310 return None
311 if not candidates:
312 return None
314 auto_map = getattr(hf_config, "auto_map", None) or {}
315 for value in auto_map.values():
316 if not isinstance(value, str) or "." not in value:
317 continue
318 # "modeling_foo.FooForCausalLM" -> "modeling_foo.py". Hub auto_map
319 # values can be "<repo>--modeling_foo.FooForCausalLM"; keep the tail.
320 module_ref = value.split("--")[-1]
321 module_file = module_ref.split(".", 1)[0] + ".py"
322 if module_file in candidates:
323 return os.path.join(model_path, module_file)
325 arch = _first_architecture(hf_config)
326 if arch:
327 for name in candidates:
328 if arch in name:
329 return os.path.join(model_path, name)
330 if len(candidates) == 1:
331 return os.path.join(model_path, candidates[0])
332 return None
335def import_module_file(module_name: str) -> Optional[str]:
336 """Resolve a Python module's ``__file__`` without importing it."""
337 try:
338 spec = importlib.util.find_spec(module_name)
339 except (ImportError, ModuleNotFoundError, ValueError, AttributeError):
340 return None
341 if spec is None or spec.origin is None:
342 return None
343 if not spec.origin.endswith(".py"):
344 return None
345 return spec.origin
348def _load_auto_config(model_id: str, config: Any) -> Optional[Any]:
349 """``AutoConfig.from_pretrained(model_id)`` with lazy import; None on failure.
351 ``trust_remote_code`` is forwarded so remote-code checkpoints can load
352 their own config classes — mirroring what ``HyperAutoModel.from_pretrained``
353 will do at build time. Any failure (no transformers installed, no network,
354 unknown id, remote config raising) returns ``None``; the caller degrades.
355 """
356 if not model_id:
357 return None
358 try:
359 # Transformers is optional for repository-only source resolution.
360 from transformers import AutoConfig # pylint: disable=import-outside-toplevel
361 except ImportError:
362 return None
363 kwargs: dict[str, Any] = {}
364 if _trust_remote_code(config):
365 kwargs["trust_remote_code"] = True
366 try:
367 return AutoConfig.from_pretrained(model_id, **kwargs)
368 except Exception:
369 return None
372def _resolve_source_from_yaml(config: Any, model_id: str) -> Optional[ResolvedSource]:
373 """Best-effort source resolution from a local checkpoint's ``config.json``.
375 Used when ``AutoConfig`` and the basename guess both came up empty: for a
376 local checkpoint directory, read its ``config.json`` and resolve against
377 it. ``SimpleNamespace`` attribute access mirrors what
378 ``AutoConfig.from_pretrained`` produces, without importing transformers.
379 Returns ``None`` on any failure so the caller degrades.
380 """
381 if not model_id or not os.path.isdir(model_id):
382 return None
383 config_path = os.path.join(model_id, "config.json")
384 if not os.path.isfile(config_path):
385 return None
386 try:
387 with open(config_path, "r", encoding="utf-8") as handle:
388 raw = json.load(handle)
389 hf_config = SimpleNamespace(**raw)
390 return resolve_hf_source(config, hf_config)
391 except (OSError, ValueError, FileNotFoundError, ImportError):
392 return None
395def _module_for_model_type(model_type: Optional[str]) -> Optional[str]:
396 """Map ``config.model_type`` to a transformers ``modeling_*`` module name."""
397 if not model_type:
398 return None
399 safe = "".join(c if c.isalnum() or c == "_" else "_" for c in model_type)
400 return f"transformers.models.{safe}.modeling_{safe}"
403def _first_architecture(hf_config: Any) -> Optional[str]:
404 architectures = getattr(hf_config, "architectures", None) or []
405 return architectures[0] if architectures else None
408def _model_path_of(config: Any) -> str:
409 """Extract the checkpoint dir / HF id from the model target.
411 Same key order as ``project_source_spec``: most YAMLs and
412 ``HyperAutoModel.from_pretrained`` configure the id under
413 ``pretrained_model_name_or_path``, with ``model_name_or_path`` and the
414 cropped/custom-builder ``config_path`` spelling as fallbacks.
415 """
416 model = getattr(config, "model", None)
417 if model is None:
418 return ""
419 path = getattr(model, "pretrained_model_name_or_path", None)
420 if path is None:
421 path = getattr(model, "model_name_or_path", None)
422 if path is None:
423 path = getattr(model, "config_path", None)
424 if path is None:
425 kwargs = getattr(model, "_kwargs", {})
426 path = kwargs.get("pretrained_model_name_or_path") or kwargs.get("config_path", "")
427 return path or ""
430def _trust_remote_code(config: Any) -> bool:
431 """Whether this model is loaded with trusted remote code.
433 The truth is the YAML's ``model.trust_remote_code`` — it is what the
434 trainer forwards into ``AutoConfig.from_pretrained`` / the HF
435 ``from_pretrained`` kwargs, and therefore what decides whether HF loads
436 the checkpoint's own modeling file or an installed transformers module.
437 The ``TrainerConfig`` object itself has no such field, so this must read
438 through the model target.
439 """
440 model = getattr(config, "model", None)
441 if model is None:
442 return False
443 value = getattr(model, "trust_remote_code", None)
444 if value is None:
445 value = getattr(model, "_kwargs", {}).get("trust_remote_code", False)
446 return bool(value)
449def _relative_imports_of(path: str) -> list[str]:
450 """Names a file relatively imports from its own package (one level).
452 Mirrors ``transformers.dynamic_module_utils.get_relative_imports`` without
453 importing transformers (source resolution must stay dependency-free). Only
454 single-dot relatives are considered: a sibling ``from .name import x`` /
455 ``import .name`` resolves within the same package directory, which is what
456 a ``<remote>`` modeling source ships. Multi-dot relatives (``..x``) do not
457 resolve against a flat remote code dir and are ignored.
458 """
459 try:
460 with open(path, "r", encoding="utf-8", errors="replace") as handle:
461 content = handle.read()
462 except OSError:
463 return []
464 names = re.findall(r"(?m)^\s*import\s+\.([A-Za-z_][A-Za-z0-9_]*)\s*$", content)
465 names += re.findall(r"(?m)^\s*from\s+\.([A-Za-z_][A-Za-z0-9_]*)\s+import", content)
466 return sorted(set(names))
469def resolve_remote_siblings(source: ResolvedSource) -> dict[str, str]:
470 """Transitive closure of a ``<remote>`` source's package siblings.
472 A trust-remote-code modeling file is compiled as the top-level module of a
473 synthetic package (loader.import_generated_module), so its single-dot
474 relative imports must resolve against *siblings in the same directory*.
475 ``source.module_name == "<remote>"`` or the Hugging Face dynamic-module
476 namespace ``transformers_modules.*`` marks this case; for an installed
477 transformers module the source uses absolute imports and this returns
478 empty.
480 Returns ``{basename: absolute_path}`` for every sibling reachable through
481 the source's relative-import closure, e.g. a ``configuration_deepseek.py``
482 that ``modeling_deepseek.py`` imports (and any module *it* relatively
483 imports). The caller copies these into the bundle so the artifact is
484 self-contained; the loader registers them under the synthetic package.
485 """
486 if isinstance(source, dict):
487 module_name = source.get("module_name")
488 file_path = source.get("file_path")
489 else:
490 module_name = getattr(source, "module_name", None)
491 file_path = getattr(source, "file_path", None)
492 is_remote = module_name == "<remote>" or (
493 isinstance(module_name, str)
494 and module_name.startswith("transformers_modules.")
495 )
496 if not is_remote or not file_path:
497 return {}
498 base_dir = os.path.dirname(os.path.abspath(file_path))
499 seen: dict[str, str] = {}
500 queue: list[str] = [os.path.abspath(file_path)]
501 while queue:
502 module_file = queue.pop(0)
503 for name in _relative_imports_of(module_file):
504 sibling = os.path.join(base_dir, f"{name}.py")
505 if not os.path.isfile(sibling):
506 continue
507 if sibling in seen:
508 continue
509 seen[sibling] = sibling
510 queue.append(sibling)
511 return seen
514def _build_resolved_source(
515 file_path: str, module_name: str, arch: Optional[str], model_type: Optional[str],
516) -> ResolvedSource:
517 return ResolvedSource(
518 module_name=module_name,
519 file_path=file_path,
520 architecture=arch,
521 model_type=model_type,
522 )
525__all__ = [
526 "ResolvedSource",
527 "import_module_file",
528 "resolve_hf_source",
529 "resolve_model_source",
530 "resolve_remote_code_file",
531 "resolve_remote_siblings",
532 "resolve_source_identity",
533]