Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / emit / replacement.py: 0%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 

4# ============================================================================ 

5"""Compile ``replace_module`` overrides against the generation-time meta model. 

6 

7The trainer's hf backend applies ``plan_overrides.replace_module`` by 

8desugaring each entry into a :class:`ModuleReplacementSpec` and running 

9``apply_module_replacements`` on the live model. Generation must reject the 

10same bad rules up front: the manager compiles the spec against the 

11generation-time meta model so conflicts, type mismatches and unmatched 

12patterns fail here rather than at runtime inside the built model, and records 

13the resulting FQN set in ``meta.module_overrides``. 

14 

15The inline pipeline lowers the surviving replacements into the generated 

16file's module bodies, and ``meta.covered["module_overrides"]`` tells the 

17trainer to skip its own replacement step. 

18""" 

19from __future__ import annotations 

20 

21from typing import Any, Sequence 

22 

23 

24# --------------------------------------------------------------------------- 

25# generation-time compile (fail-fast) 

26# --------------------------------------------------------------------------- 

27 

28def compile_overrides_for_meta( 

29 meta_model: Any, 

30 specs: Sequence[Any], 

31 *, 

32 factory_paths: Sequence[str | None] = (), 

33 configs_by_spec: dict[int, Any] | None = None, 

34) -> tuple[dict[str, Any], ...]: 

35 """Match replacement specs against the generation-time meta model. 

36 

37 Runs :func:`compile_module_replacements` on the meta (empty-weights) model 

38 so a conflicting or type-mismatched replacement rule fails here, at 

39 generation time, rather than at runtime inside the built model. Returns 

40 one JSON-safe record per matched target. 

41 

42 Fail-fast behavior (a pattern matching no module, a type mismatch, a 

43 conflict) is owned by the native :func:`compile_module_replacements` 

44 itself — the same checks the trainer's hf backend runs — so this wrapper 

45 keeps no mirrored pre-check; it only re-raises the native 

46 :class:`ValueError` with the codegen context prefix. Every rule reaching 

47 this module on the codegen path comes from the user's YAML 

48 ``plan_overrides.replace_module`` — there is no plan- or internally-derived 

49 rule source here — so a no-match is a config error (a typo, or a rule 

50 written for a different architecture) and must fail generation rather than 

51 silently drop an intended replacement from the artifact. 

52 

53 A record carries: 

54 

55 .. code-block:: python 

56 

57 { 

58 "match": ["model.layers.0.mlp"], # spec.match, the patterns 

59 "fqn": "model.layers.0.mlp", # first aliased FQN 

60 "fqns": ["model.layers.0.mlp"], # every registered alias 

61 "module_type": "<dotted path>", # source module type 

62 "factory": "<dotted path>", # the @module_replacement fn 

63 "exact_type": False, 

64 "target_config": {...}, # optional YAML Target static args 

65 } 

66 

67 The ``factory`` (and ``module_type``) are stored as import paths, not the 

68 live objects — the replacement factory is a per-entry closure that cannot 

69 be written as a Python literal, and the runtime re-imports the *raw* 

70 decoration at those paths to rebuild equal specs. When the YAML entry 

71 bound extra static args to its replace-module Target, those must ride 

72 along in ``configs_by_spec`` (keyed by ``id(spec)``) so the runtime can 

73 re-pre-bind them onto the rebuilt factory instead of dropping to the 

74 factory's defaults (or failing a required arg). 

75 """ 

76 if not specs: 

77 return () 

78 if configs_by_spec is None: 

79 configs_by_spec = {} 

80 from hyper_parallel.models.replacement import compile_module_replacements 

81 

82 paths_by_spec = { 

83 id(spec): path 

84 for spec, path in zip(specs, factory_paths) 

85 if path is not None 

86 } 

87 

88 try: 

89 plan = compile_module_replacements(meta_model, specs) 

90 except ValueError as exc: 

91 raise ValueError( 

92 f"codegen: plan_overrides replace_module: {exc}" 

93 ) from exc 

94 return tuple( 

95 _record_for_target(target, paths_by_spec, configs_by_spec) 

96 for target in plan.targets 

97 ) 

98 

99 

100def _record_for_target( 

101 target: Any, 

102 paths_by_spec: dict[int, str], 

103 configs_by_spec: dict[int, Any] | None = None, 

104) -> dict[str, Any]: 

105 """Serialize one replacement target into its meta record.""" 

106 spec = target.spec 

107 factory_path = paths_by_spec.get(id(spec)) 

108 if factory_path is None: 

109 raise TypeError( 

110 f"codegen: replacement factory for {spec.match!r} has no serializable " 

111 "import path — the gen backend needs a YAML Target-backed " 

112 "@module_replacement factory" 

113 ) 

114 record = { 

115 "match": list(spec.match), 

116 "fqn": target.module_fqns[0], 

117 "fqns": list(target.module_fqns), 

118 "module_type": _type_path(spec.module_type), 

119 "factory": factory_path, 

120 "exact_type": bool(spec.exact_type), 

121 } 

122 config = (configs_by_spec or {}).get(id(spec)) 

123 if config: 

124 record["target_config"] = config 

125 return record 

126 

127 

128def _type_path(module_type: type) -> str: 

129 return f"{module_type.__module__}.{module_type.__qualname__}" 

130 

131 

132__all__ = ["compile_overrides_for_meta"]