Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / artifact.py: 72%
194 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"""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 copied remote-code
19dependencies. Two layout anchors exist: the model-anchored default
20(``hyper_parallel/models/<model_type>/generated``, one bundle per model type)
21and an explicit directory for isolated model-construction workflows. The bundle is
22written as one atomic unit: every file is first staged into a temporary
23*sibling* directory, then the whole directory is swapped into place, so a
24half-written bundle can never be imported.
25"""
27from __future__ import annotations
29from contextlib import contextmanager
30import ctypes
31from dataclasses import dataclass
32import errno
33import logging
34import os
35from pathlib import Path
36import shutil
37import tempfile
38import time
39from typing import Iterator
41# Bundle directory name shared by every layout constructor.
42_GENERATED_DIR = "generated"
44_LOGGER = logging.getLogger(__name__)
45_LOCK_INITIALIZATION_GRACE_S = 1.0
48@dataclass
49class ArtifactLayout:
50 """Describe the paths belonging to one generated bundle."""
52 artifact_dir: str
53 modeling_path: str
54 diff_path: str
55 init_path: str
57 @classmethod
58 def build_default(
59 cls,
60 artifact_dir: str | None = None,
61 model_name: str | None = None,
62 ) -> "ArtifactLayout":
63 """Create the layout for model-construction codegen.
65 ``artifact_dir`` defaults to ``Path.cwd()/generated``.
66 """
67 resolved_dir = os.path.abspath(artifact_dir or os.path.join(os.getcwd(), _GENERATED_DIR))
68 modeling_name = _default_modeling_name(model_name)
69 return cls(
70 artifact_dir=resolved_dir,
71 modeling_path=os.path.join(resolved_dir, modeling_name),
72 diff_path=os.path.join(resolved_dir, modeling_name + ".diff"),
73 init_path=os.path.join(resolved_dir, "__init__.py"),
74 )
76 @classmethod
77 def build_model_anchored(
78 cls,
79 model_type: str,
80 model_name: str | None = None,
81 ) -> "ArtifactLayout":
82 """Create the model-anchored layout for one model type.
84 ``artifact_dir`` is ``hyper_parallel/models/<model_type>/generated``
85 — exactly one bundle per model type, colocated with the model's
86 adapter code. When ``models/<model_type>/`` does not exist yet, a
87 warning is emitted and the directory is created on demand, so a
88 model without a family directory still gets its own bundle home
89 instead of silently falling back to another location.
90 """
91 # ``artifact.py`` lives in ``hyper_parallel/codegen/``; the models
92 # tree is its sibling — package-relative, never cwd-dependent.
93 models_root = Path(__file__).resolve().parents[1] / "models"
94 family_root = models_root / model_type
95 if not family_root.is_dir():
96 _LOGGER.warning(
97 "codegen: model type %r has no directory under "
98 "hyper_parallel/models yet; the bundle will be created at "
99 "%s (add an adapter package there to colocate it with the "
100 "model's conversion code)",
101 model_type,
102 os.path.join(str(family_root), _GENERATED_DIR),
103 )
104 artifact_dir = os.path.join(str(family_root), _GENERATED_DIR)
105 modeling_name = _default_modeling_name(model_name)
106 return cls(
107 artifact_dir=artifact_dir,
108 modeling_path=os.path.join(artifact_dir, modeling_name),
109 diff_path=os.path.join(artifact_dir, modeling_name + ".diff"),
110 init_path=os.path.join(artifact_dir, "__init__.py"),
111 )
114def _default_modeling_name(model_name: str | None) -> str:
115 base = model_name or "model"
116 # '' / '-' are not valid Python identifiers in module names.
117 safe = "".join(c if c.isalnum() or c == "_" else "_" for c in base)
118 return f"modeling_{safe}_gen_npu.py"
121def resolve_default_artifact_layout(
122 artifact_dir: str | None = None,
123 model_name: str | None = None,
124) -> ArtifactLayout:
125 """Resolve the default model-construction artifact layout."""
126 return ArtifactLayout.build_default(artifact_dir, model_name)
129def resolve_model_artifact_layout(
130 model_type: str,
131 model_name: str | None = None,
132) -> "ArtifactLayout":
133 """Resolve the model-anchored layout for ``model_type``.
135 The bundle directory is always ``hyper_parallel/models/<model_type>/
136 generated`` — a strict one-to-one mapping keyed by the model type, with
137 no aliasing between model types. When ``models/<model_type>/`` does not
138 exist yet, ``ArtifactLayout.build_model_anchored`` warns and the directory is
139 created on demand by :func:`ensure_artifact_dir`.
140 """
141 return ArtifactLayout.build_model_anchored(model_type, model_name=model_name)
144def ensure_artifact_dir(layout: ArtifactLayout) -> None:
145 """Create the artifact directory (recursive, no-op if present).
147 Raises a readable error when the directory cannot be created: the
148 model-anchored default lives inside the installed ``hyper_parallel``
149 package, which is read-only under a non-editable install.
150 """
151 try:
152 Path(layout.artifact_dir).mkdir(parents=True, exist_ok=True)
153 except OSError as exc:
154 raise RuntimeError(
155 f"codegen: cannot create artifact directory {layout.artifact_dir}: {exc}. "
156 "The default bundle location lives inside the hyper_parallel package; "
157 "for a non-editable install either reinstall with 'pip install -e .' "
158 "or point 'model.codegen_artifact_dir' at a writable directory in the YAML."
159 ) from exc
162def write_bundle_atomic(layout: ArtifactLayout, files: dict[str, str]) -> None:
163 """Write the whole bundle into place as one atomic directory switch.
165 ``files`` maps filename -> content; keys must be plain file names (not
166 absolute paths) so nothing bypasses the staging directory.
168 The bundle is first staged into a temporary *sibling* of ``artifact_dir``
169 (never inside it), then the whole directory is swapped in. On Windows
170 ``os.replace`` cannot replace a non-empty target directory, so the swap is
171 three steps: move the old bundle aside, rename the staged dir into place,
172 then delete the backup. A reader on ``artifact_dir`` therefore sees either
173 the old complete bundle or the new complete bundle — never a mix.
174 """
175 ensure_artifact_dir(layout)
176 parent = os.path.dirname(layout.artifact_dir) or "."
178 tmpdir = tempfile.mkdtemp(prefix=".codegen_tmp_", dir=parent)
179 old_dir = None
180 try:
181 _stage_bundle(tmpdir, files)
183 if os.path.exists(layout.artifact_dir):
184 old_dir = tempfile.mkdtemp(prefix=".codegen_old_", dir=parent)
185 os.rmdir(old_dir) # reserve a name; must be absent for os.replace
186 os.replace(layout.artifact_dir, old_dir)
187 os.replace(tmpdir, layout.artifact_dir)
188 if old_dir is not None:
189 shutil.rmtree(old_dir, ignore_errors=True)
190 old_dir = None
191 finally:
192 # If the staged dir never made it into place, put the old bundle back.
193 if old_dir is not None and os.path.isdir(old_dir):
194 if not os.path.exists(layout.artifact_dir):
195 os.replace(old_dir, layout.artifact_dir)
196 else:
197 shutil.rmtree(old_dir, ignore_errors=True)
198 if os.path.isdir(tmpdir):
199 shutil.rmtree(tmpdir, ignore_errors=True)
202def _stage_bundle(tmpdir: str, files: dict[str, str]) -> None:
203 """Write the complete newly emitted file set into ``tmpdir``."""
204 normalized = dict(files)
205 for name in normalized:
206 _validate_bundle_filename(name)
208 init_present = "__init__.py" in normalized
209 if not init_present:
210 normalized["__init__.py"] = ""
212 for name, content in normalized.items():
213 # newline="" disables the platform's default CRLF translation so the
214 # staged bytes equal the emitted text bytes exactly. On Windows the
215 # default (newline=None) would turn every LF into CRLF and make bundle
216 # bytes platform-dependent.
217 with open(
218 os.path.join(tmpdir, name), "w", encoding="utf-8", newline=""
219 ) as handle:
220 handle.write(content)
222def _validate_bundle_filename(name: str) -> None:
223 """Reject absolute, nested, and parent-relative bundle paths."""
224 if (
225 not name
226 or os.path.isabs(name)
227 or os.path.basename(name) != name
228 or "/" in name
229 or "\\" in name
230 or name in {".", ".."}
231 ):
232 raise ValueError(f"bundle file keys must be plain names, got {name!r}")
235def clean_temp_artifacts(layout: ArtifactLayout) -> None:
236 """Remove stale ``.codegen_tmp_*`` / ``.codegen_old_*`` dirs.
238 Staging now happens in the artifact dir's *parent*, so these leftovers are
239 scanned there rather than inside the bundle.
240 """
241 parent = os.path.dirname(layout.artifact_dir) or "."
242 if not os.path.isdir(parent):
243 return
244 for name in os.listdir(parent):
245 if name.startswith(".codegen_tmp_") or name.startswith(".codegen_old_"):
246 shutil.rmtree(os.path.join(parent, name), ignore_errors=True)
249@contextmanager
250def artifact_generation_lock(
251 layout: ArtifactLayout,
252 *,
253 timeout_s: float = 600.0,
254) -> Iterator[None]:
255 """Serialize bundle publication and recover locks left by dead owners."""
256 parent = os.path.dirname(layout.artifact_dir) or "."
257 bundle_name = os.path.basename(layout.artifact_dir.rstrip(os.sep)) or "generated"
258 lock_path = os.path.join(parent, f".{bundle_name}.codegen.lock")
259 deadline = time.monotonic() + timeout_s
260 descriptor = None
261 owner_identity = None
262 while descriptor is None:
263 try:
264 descriptor = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
265 try:
266 os.write(descriptor, str(os.getpid()).encode("ascii"))
267 except BaseException:
268 os.close(descriptor)
269 descriptor = None
270 try:
271 os.unlink(lock_path)
272 except FileNotFoundError:
273 pass
274 raise
275 owner_identity = _file_identity(lock_path)
276 except FileExistsError:
277 if _remove_stale_lock(lock_path):
278 continue
279 if time.monotonic() >= deadline:
280 raise TimeoutError(
281 f"codegen: timed out waiting for artifact lock {lock_path}"
282 ) from None
283 time.sleep(0.1)
284 try:
285 yield
286 finally:
287 os.close(descriptor)
288 if owner_identity is not None and _file_identity(lock_path) == owner_identity:
289 try:
290 os.unlink(lock_path)
291 except FileNotFoundError:
292 pass
295def _remove_stale_lock(lock_path: str) -> bool:
296 """Remove ``lock_path`` when its recorded owner process is dead."""
297 try:
298 identity = _file_identity(lock_path)
299 stat_result = os.stat(lock_path)
300 with open(lock_path, "r", encoding="ascii") as handle:
301 value = handle.read().strip()
302 except (FileNotFoundError, OSError, UnicodeError):
303 return False
305 try:
306 pid = int(value)
307 except ValueError:
308 # There is a tiny interval between O_EXCL creation and writing the PID.
309 # Only reclaim malformed locks after that initialization window.
310 age_s = max(0.0, time.time() - stat_result.st_mtime)
311 stale = age_s >= _LOCK_INITIALIZATION_GRACE_S
312 else:
313 stale = not _pid_is_alive(pid)
314 if not stale or _file_identity(lock_path) != identity:
315 return False
316 try:
317 os.unlink(lock_path)
318 except FileNotFoundError:
319 return True
320 except OSError:
321 return False
322 _LOGGER.warning("codegen: removed stale artifact lock %s", lock_path)
323 return True
326def _pid_is_alive(pid: int) -> bool:
327 """Return whether ``pid`` still identifies a live process."""
328 if pid <= 0:
329 return False
330 if os.name == "nt":
331 return _windows_pid_is_alive(pid)
332 try:
333 os.kill(pid, 0)
334 except ProcessLookupError:
335 return False
336 except PermissionError:
337 return True
338 except OSError as exc:
339 if exc.errno == errno.ESRCH:
340 return False
341 if exc.errno == errno.EPERM:
342 return True
343 return True
344 return True
347def _windows_pid_is_alive(pid: int) -> bool:
348 """Query a Windows process without using ``os.kill(pid, 0)``.
350 Windows does not implement the POSIX signal-zero probe; using it may
351 terminate the target process. ``OpenProcess`` plus ``GetExitCodeProcess``
352 is read-only and reports ``STILL_ACTIVE`` for a live owner.
353 """
354 kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
355 open_process = kernel32.OpenProcess
356 open_process.argtypes = (ctypes.c_ulong, ctypes.c_int, ctypes.c_ulong)
357 open_process.restype = ctypes.c_void_p
358 get_exit_code = kernel32.GetExitCodeProcess
359 get_exit_code.argtypes = (ctypes.c_void_p, ctypes.POINTER(ctypes.c_ulong))
360 get_exit_code.restype = ctypes.c_int
361 close_handle = kernel32.CloseHandle
362 close_handle.argtypes = (ctypes.c_void_p,)
363 close_handle.restype = ctypes.c_int
364 process = open_process(0x1000, False, pid)
365 if not process:
366 # Access denied proves that the process exists. Other errors are
367 # treated conservatively except ERROR_INVALID_PARAMETER (no such PID).
368 return ctypes.get_last_error() != 87
369 try:
370 exit_code = ctypes.c_ulong()
371 if not get_exit_code(process, ctypes.byref(exit_code)):
372 return True
373 return exit_code.value == 259 # STILL_ACTIVE
374 finally:
375 close_handle(process)
378def _file_identity(path: str) -> tuple[int, int, int] | None:
379 """Return enough stat identity to avoid unlinking a replaced lock file."""
380 try:
381 result = os.stat(path)
382 except FileNotFoundError:
383 return None
384 return result.st_ino, result.st_mtime_ns, result.st_size
387__all__ = [
388 "ArtifactLayout",
389 "artifact_generation_lock",
390 "clean_temp_artifacts",
391 "ensure_artifact_dir",
392 "resolve_default_artifact_layout",
393 "resolve_model_artifact_layout",
394 "write_bundle_atomic",
395]