Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / meta.py: 42%
98 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"""``codegen_meta.json`` read/write and schema validation.
17The meta file is the single source of truth for an artifact bundle's
18identity: it records the signature, YAML path + hash, source file the
19modeling output was generated from, parallel dims, and the resolved
20replacements/injections/param plan so a later preflight pass can detect
21drift without re-running the whole generation.
22"""
24from __future__ import annotations
26import json
27import os
28from dataclasses import asdict, dataclass, field
29from typing import Any, Optional
31from hyper_parallel.codegen.artifact import ArtifactLayout
32from hyper_parallel.codegen.hash import file_line_count, sha256_file
35@dataclass
36class CodegenMeta:
37 """Serializable artifact metadata for one generated modeling bundle."""
39 codegen_version: str
40 signature: str
41 yaml_path: str
42 yaml_sha256: str
43 source: dict[str, Any]
44 parallel_dims: dict[str, Any]
45 model_name: Optional[str] = None
46 outputs: dict[str, Any] = field(default_factory=dict)
47 injections: list[dict[str, Any]] = field(default_factory=list)
48 skipped_overrides: list[dict[str, Any]] = field(default_factory=list)
49 param_plan: dict[str, Any] = field(default_factory=dict)
50 frozen_sharded_params: list[str] = field(default_factory=list)
51 tied_pairs: list[list[str]] = field(default_factory=list)
52 special_handlers: dict[str, str] = field(default_factory=dict)
53 #: FQN -> owning class name for every boundary. The emitter
54 #: groups boundary FQNs by this to rewrite one ``class.forward`` per class.
55 boundary_classes: dict[str, str] = field(default_factory=dict)
56 entrypoints: dict[str, str] = field(default_factory=dict)
57 #: Sunk ``replace_module`` overrides: one JSON-safe record per
58 #: matched target, ``{match, fqn, fqns, module_type, factory, exact_type}``.
59 #: The inline pipeline lowers them into the generated file, and the record
60 #: is what lets the trainer skip its own HF replacement step
61 #: (``covered["module_overrides"]``).
62 module_overrides: list[dict[str, Any]] = field(default_factory=list)
63 #: Entry model class name (the architecture name, e.g. ``AnthropicV3ForCausalLM``).
64 #: The inline pipeline resolves the model's adapter from it.
65 model_class: Optional[str] = None
66 #: Generated class names whose ``forward`` an inline strategy body replaces
67 #: (the EP MoE shell), so it orchestrates through instance attributes bound
68 #: at install time (``hp_bind_inline_state``) instead of a compiled
69 #: boundary. Derived from structure at generation time and copied here so
70 #: the runtime reads the artifact's own record rather than a second copy
71 #: that has to be kept in sync by hand. A *rendered* component class keeps
72 #: its own forward plus the emitted boundary form, so it is not listed.
73 external_state_classes: list[str] = field(default_factory=list)
74 #: ``covered`` must always carry a bool ``sharding_plan``,
75 #: so the default mirrors ``COVERED_DEFAULTS`` rather than ``{}`` — a bare
76 #: ``CodegenMeta`` round-trips through ``load_codegen_meta`` instead of
77 #: failing schema validation. The manager overwrites this explicitly.
78 covered: dict[str, Any] = field(default_factory=lambda: dict(COVERED_DEFAULTS))
79 not_covered: list[str] = field(default_factory=list)
80 #: Plan-coordinate active axes (``FrozenPlan.mesh_dim_names``), JSON form.
81 #: The runtime slices the live mesh down to these axes so placements never
82 #: land on dp axes the plan never shards on. ``None`` means the runtime
83 #: uses the full mesh for compatibility with older artifacts.
84 mesh_dim_names: Optional[list[str]] = None
85 #: ``<remote>`` source sibling modules copied into the bundle:
86 #: ``["configuration_deepseek.py", ...]``
87 #: listing every package sibling the generated file's relative imports
88 #: resolve against. Empty for installed transformers sources (they use
89 #: absolute imports) and for artifacts generated before this field existed.
90 remote_siblings: list[str] = field(default_factory=list)
93#: Stages the generated artifact declares it covers. ``sharding_plan`` starts
94#: as the default here (False, kept for metas that predate the field), but
95#: ``_fill_plan_fields`` (manager.py) sets it True once the plan is frozen:
96#: the artifact carries the parallel logic and the runtime shards from
97#: ``codegen_meta.json``. When True, the trainer short-circuits the
98#: planner+applier and runs the artifact's own parallelize step.
99#: ``module_overrides`` turns True once generation sunk ``replace_module``, so
100#: the trainer's HF replacement step is skipped.
101COVERED_DEFAULTS = {
102 "sharding_plan": False,
103 "module_overrides": False,
104 "axes": [],
105}
107#: Stages that remain trainer-owned.
108NOT_COVERED_DEFAULT = [
109 "mesh",
110 "fsdp2",
111 "pp",
112 "activation_checkpoint",
113 "compile",
114]
117def meta_to_dict(meta: CodegenMeta) -> dict[str, Any]:
118 """Serialize a CodegenMeta to a JSON-friendly dict."""
119 return asdict(meta)
122def meta_from_dict(data: dict[str, Any]) -> CodegenMeta:
123 """Rebuild a CodegenMeta from a JSON dict."""
124 validate_meta_schema(data)
125 return CodegenMeta(
126 **{k: v for k, v in data.items() if k in CodegenMeta.__dataclass_fields__}
127 )
130def load_codegen_meta(meta_path: str) -> Optional[CodegenMeta]:
131 """Read the meta file; ``None`` if absent or malformed."""
132 if not os.path.isfile(meta_path):
133 return None
134 try:
135 with open(meta_path, "r", encoding="utf-8") as handle:
136 return meta_from_dict(json.load(handle))
137 except (json.JSONDecodeError, TypeError, ValueError, KeyError):
138 return None
141def write_codegen_meta(meta: CodegenMeta, path: str) -> None:
142 """Write the meta as pretty-printed JSON."""
143 with open(path, "w", encoding="utf-8") as handle:
144 json.dump(meta_to_dict(meta), handle, indent=2, sort_keys=True)
147def validate_meta_schema(data: dict[str, Any]) -> None:
148 """Fail fast if required meta fields are missing or mistyped."""
149 required = ("codegen_version", "signature", "yaml_path", "yaml_sha256")
150 for key in required:
151 if key not in data:
152 raise ValueError(f"codegen_meta missing required field {key!r}")
153 if not isinstance(data[key], str):
154 raise TypeError(f"codegen_meta field {key!r} must be a string")
155 for key in ("source", "parallel_dims", "outputs"):
156 if key in data and not isinstance(data[key], dict):
157 raise TypeError(f"codegen_meta field {key!r} must be a dict")
158 if "entrypoints" in data:
159 entrypoints = data["entrypoints"]
160 if not isinstance(entrypoints, dict) or not all(
161 isinstance(k, str) and isinstance(v, str) for k, v in entrypoints.items()
162 ):
163 raise TypeError("codegen_meta field 'entrypoints' must be a dict[str, str]")
164 if "covered" in data:
165 covered = data["covered"]
166 if not isinstance(covered, dict):
167 raise TypeError("codegen_meta field 'covered' must be a dict")
168 if not isinstance(covered.get("sharding_plan"), bool):
169 raise TypeError("codegen_meta field 'covered.sharding_plan' must be a bool")
170 if "not_covered" in data:
171 not_covered = data["not_covered"]
172 if not isinstance(not_covered, list) or not all(
173 isinstance(item, str) for item in not_covered
174 ):
175 raise TypeError("codegen_meta field 'not_covered' must be a list[str]")
176 if "module_overrides" in data:
177 overrides = data["module_overrides"]
178 if not isinstance(overrides, list) or not all(
179 isinstance(item, dict) for item in overrides
180 ):
181 raise TypeError(
182 "codegen_meta field 'module_overrides' must be a list[dict]"
183 )
184 if "remote_siblings" in data:
185 siblings = data["remote_siblings"]
186 if not isinstance(siblings, list) or not all(
187 isinstance(item, str) for item in siblings
188 ):
189 raise TypeError("codegen_meta field 'remote_siblings' must be a list[str]")
190 if "external_state_classes" in data:
191 class_names = data["external_state_classes"]
192 if not isinstance(class_names, list) or not all(
193 isinstance(item, str) for item in class_names
194 ):
195 raise TypeError(
196 "codegen_meta field 'external_state_classes' must be a list[str]"
197 )
200def record_output_hashes(
201 meta: CodegenMeta,
202 layout: ArtifactLayout,
203 *,
204 base_dir: str = "",
205) -> CodegenMeta:
206 """Populate ``meta.outputs`` with per-file hashes and line counts.
208 When ``base_dir`` is given every regular file is hashed from *there* (the
209 staging directory); otherwise the artifact directory is used. Hashing must
210 happen where the bytes actually live at write time: the atomic bundle
211 writes ``codegen_meta.json`` into the staging dir, so ``base_dir`` is the
212 staging dir and the layout paths still point at the old bundle.
214 The meta file itself is deliberately NOT recorded: its sha256 cannot be
215 stored in its own bytes (the value changes the file, so the hash could
216 never be a fixed point), and the previous record-then-write order hashed
217 the *old* meta and made preflight report drift on every second
218 generation. The meta file's integrity is covered by the signature check,
219 the schema validation, and the hashes of every other bundle file — it does
220 not need a self-hash.
221 """
222 output_dir = base_dir or layout.artifact_dir
223 role_by_name = {
224 os.path.basename(layout.modeling_path): "modeling",
225 os.path.basename(layout.diff_path): "diff",
226 os.path.basename(layout.init_path): "init",
227 }
228 meta.outputs = {}
229 for name in sorted(os.listdir(output_dir)):
230 target = os.path.join(output_dir, name)
231 if name == os.path.basename(layout.meta_path) or not os.path.isfile(target):
232 continue
233 key = role_by_name.get(name, name)
234 meta.outputs[key] = {
235 "sha256": sha256_file(target),
236 "lines": file_line_count(target),
237 }
238 return meta
241__all__ = [
242 "CodegenMeta",
243 "load_codegen_meta",
244 "meta_from_dict",
245 "meta_to_dict",
246 "record_output_hashes",
247 "validate_meta_schema",
248 "write_codegen_meta",
249]