Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / artifact.py: 27%
102 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"""Artifact-layout resolution and atomic writes for codegen bundles.
17The artifact bundle lives in a ``generated/`` directory and holds the
18generated modeling file, its diff against the original, and
19``codegen_meta.json``. CLI workflows can anchor it beside a YAML file; model
20construction workflows can use the current working directory or an explicit
21artifact directory. The bundle is written as one atomic unit: every file is first staged into a
22temporary *sibling* directory, then the whole directory is swapped into place,
23so a half-written bundle can never be imported.
24"""
26from __future__ import annotations
28import os
29import shutil
30import tempfile
31from dataclasses import dataclass
32from pathlib import Path
35@dataclass
36class ArtifactLayout:
37 """Describe the bundle path set for one YAML."""
39 yaml_path: str
40 artifact_dir: str
41 modeling_path: str
42 diff_path: str
43 meta_path: str
44 init_path: str
46 @classmethod
47 def build(cls, yaml_path: str, model_name: str | None = None) -> "ArtifactLayout":
48 """Create the layout from a YAML path and optional model name.
50 ``artifact_dir`` is ``<yaml_dir>/generated/``. Example directories
51 contain one Codegen YAML, so another per-YAML directory is redundant.
52 """
53 yaml_path = os.path.abspath(yaml_path)
54 artifact_dir = os.path.join(os.path.dirname(yaml_path), "generated")
55 modeling_name = _default_modeling_name(model_name)
56 return cls(
57 yaml_path=yaml_path,
58 artifact_dir=artifact_dir,
59 modeling_path=os.path.join(artifact_dir, modeling_name),
60 diff_path=os.path.join(artifact_dir, modeling_name + ".diff"),
61 meta_path=os.path.join(artifact_dir, "codegen_meta.json"),
62 init_path=os.path.join(artifact_dir, "__init__.py"),
63 )
65 @classmethod
66 def build_default(
67 cls,
68 artifact_dir: str | None = None,
69 model_name: str | None = None,
70 ) -> "ArtifactLayout":
71 """Create the layout for model-construction codegen.
73 ``artifact_dir`` defaults to ``Path.cwd()/generated``. ``yaml_path``
74 is kept as an empty string for backward-compatible metadata shape; the
75 signature is computed from the resolved config projection, not from a
76 YAML file path.
77 """
78 resolved_dir = os.path.abspath(artifact_dir or os.path.join(os.getcwd(), "generated"))
79 modeling_name = _default_modeling_name(model_name)
80 return cls(
81 yaml_path="",
82 artifact_dir=resolved_dir,
83 modeling_path=os.path.join(resolved_dir, modeling_name),
84 diff_path=os.path.join(resolved_dir, modeling_name + ".diff"),
85 meta_path=os.path.join(resolved_dir, "codegen_meta.json"),
86 init_path=os.path.join(resolved_dir, "__init__.py"),
87 )
90def _default_modeling_name(model_name: str | None) -> str:
91 base = model_name or "model"
92 # '' / '-' are not valid Python identifiers in module names.
93 safe = "".join(c if c.isalnum() or c == "_" else "_" for c in base)
94 return f"modeling_{safe}_gen_npu.py"
97def resolve_artifact_layout(
98 yaml_path: str, model_name: str | None = None
99) -> ArtifactLayout:
100 """Resolve the artifact layout for ``yaml_path``."""
101 return ArtifactLayout.build(yaml_path, model_name)
104def resolve_default_artifact_layout(
105 artifact_dir: str | None = None,
106 model_name: str | None = None,
107) -> ArtifactLayout:
108 """Resolve the default model-construction artifact layout."""
109 return ArtifactLayout.build_default(artifact_dir, model_name)
112def ensure_artifact_dir(layout: ArtifactLayout) -> None:
113 """Create the artifact directory (recursive, no-op if present)."""
114 Path(layout.artifact_dir).mkdir(parents=True, exist_ok=True)
117def write_bundle_atomic(
118 layout: ArtifactLayout, files: dict[str, str], meta: object = None
119) -> None:
120 """Write the whole bundle into place as one atomic directory switch.
122 ``files`` maps filename -> content; keys must be plain file names (not
123 absolute paths) so nothing bypasses the staging directory. ``meta``, when
124 given, is hashed against the *staged* files and written
125 ``codegen_meta.json`` LAST inside the staging directory, so meta being
126 present AND parseable always means the rest of the bundle is already there.
128 The bundle is first staged into a temporary *sibling* of ``artifact_dir``
129 (never inside it), then the whole directory is swapped in. On Windows
130 ``os.replace`` cannot replace a non-empty target directory, so the swap is
131 three steps: move the old bundle aside, rename the staged dir into place,
132 then delete the backup. A reader on ``artifact_dir`` therefore sees either
133 the old complete bundle or the new complete bundle — never a mix.
134 """
135 ensure_artifact_dir(layout)
136 parent = os.path.dirname(layout.artifact_dir) or "."
138 tmpdir = tempfile.mkdtemp(prefix=".codegen_tmp_", dir=parent)
139 old_dir = None
140 try:
141 _stage_bundle(tmpdir, layout, files, has_meta=meta is not None)
142 if meta is not None:
143 _stage_meta(tmpdir, layout, meta)
145 if os.path.exists(layout.artifact_dir):
146 old_dir = tempfile.mkdtemp(prefix=".codegen_old_", dir=parent)
147 os.rmdir(old_dir) # reserve a name; must be absent for os.replace
148 os.replace(layout.artifact_dir, old_dir)
149 os.replace(tmpdir, layout.artifact_dir)
150 if old_dir is not None:
151 shutil.rmtree(old_dir, ignore_errors=True)
152 old_dir = None
153 finally:
154 # If the staged dir never made it into place, put the old bundle back.
155 if old_dir is not None and os.path.isdir(old_dir):
156 if not os.path.exists(layout.artifact_dir):
157 os.replace(old_dir, layout.artifact_dir)
158 else:
159 shutil.rmtree(old_dir, ignore_errors=True)
160 if os.path.isdir(tmpdir):
161 shutil.rmtree(tmpdir, ignore_errors=True)
164def _stage_bundle(
165 tmpdir: str,
166 layout: ArtifactLayout,
167 files: dict[str, str],
168 *,
169 has_meta: bool,
170) -> None:
171 """Copy any existing bundle files plus ``files`` into ``tmpdir``.
173 Existing non-temp files are preserved so a partial update (e.g. adding
174 ``source_*.py``) never drops the rest of
175 the bundle; a key in ``files`` always wins over the preserved copy. When
176 ``has_meta`` is False the prior ``codegen_meta.json`` is also preserved —
177 otherwise the caller is going to write a fresh one into the staging dir.
178 """
180 def _is_tmp_name(name: str) -> bool:
181 return name.startswith(".codegen_tmp_") or name.startswith(".codegen_old_")
183 normalized = dict(files)
184 for name in normalized:
185 _validate_bundle_filename(name)
187 init_present = "__init__.py" in normalized
188 if not has_meta and os.path.isdir(layout.artifact_dir):
189 for name in os.listdir(layout.artifact_dir):
190 if _is_tmp_name(name):
191 continue
192 if name == "codegen_meta.json":
193 continue # preserved below, after the loop
194 if name in normalized:
195 continue
196 src = os.path.join(layout.artifact_dir, name)
197 if not os.path.isfile(src):
198 continue
199 if name == "__init__.py":
200 init_present = True
201 shutil.copy2(src, os.path.join(tmpdir, name))
203 if not init_present:
204 normalized["__init__.py"] = ""
206 for name, content in normalized.items():
207 # newline="" disables the platform's default CRLF translation so the
208 # staged bytes equal the emitted text bytes exactly. On Windows the
209 # default (newline=None) would turn every LF into CRLF, inflating each
210 # file by one byte per line and breaking the drift/hash checks that
211 # compare on-disk bytes against the re-emitted LF text.
212 with open(
213 os.path.join(tmpdir, name), "w", encoding="utf-8", newline=""
214 ) as handle:
215 handle.write(content)
217 if not has_meta and os.path.isfile(layout.meta_path):
218 shutil.copy2(
219 layout.meta_path, os.path.join(tmpdir, os.path.basename(layout.meta_path))
220 )
223def _validate_bundle_filename(name: str) -> None:
224 """Reject absolute, nested, and parent-relative bundle paths."""
225 if (
226 not name
227 or os.path.isabs(name)
228 or os.path.basename(name) != name
229 or "/" in name
230 or "\\" in name
231 or name in {".", ".."}
232 ):
233 raise ValueError(f"bundle file keys must be plain names, got {name!r}")
236def _stage_meta(tmpdir: str, layout: ArtifactLayout, meta: object) -> None:
237 """Hash the staged content files then write ``codegen_meta.json`` last."""
238 from hyper_parallel.codegen.meta import (
239 record_output_hashes,
240 write_codegen_meta,
241 )
243 record_output_hashes(meta, layout, base_dir=tmpdir)
244 meta_path = os.path.join(tmpdir, os.path.basename(layout.meta_path))
245 write_codegen_meta(meta, meta_path)
248def clean_temp_artifacts(layout: ArtifactLayout) -> None:
249 """Remove stale ``.codegen_tmp_*`` / ``.codegen_old_*`` dirs.
251 Staging now happens in the artifact dir's *parent*, so these leftovers are
252 scanned there rather than inside the bundle.
253 """
254 parent = os.path.dirname(layout.artifact_dir) or "."
255 if not os.path.isdir(parent):
256 return
257 for name in os.listdir(parent):
258 if name.startswith(".codegen_tmp_") or name.startswith(".codegen_old_"):
259 shutil.rmtree(os.path.join(parent, name), ignore_errors=True)
262__all__ = [
263 "ArtifactLayout",
264 "clean_temp_artifacts",
265 "ensure_artifact_dir",
266 "resolve_default_artifact_layout",
267 "resolve_artifact_layout",
268 "write_bundle_atomic",
269]