Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / session.py: 86%
118 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"""Stateful collector for immutable codegen build snapshots."""
17from __future__ import annotations
19from collections.abc import Mapping, Sequence
20from typing import Any, Callable, Optional
22from hyper_parallel.codegen.plan.observation import observe_emission_plan
23from hyper_parallel.codegen.snapshot import (
24 CodegenSnapshot,
25 CodegenStatus,
26 FailureSnapshot,
27 ReplacementSnapshot,
28 copy_mapping,
29)
32class CodegenBuildSession:
33 """Collect native build observations without influencing the build."""
35 def __init__(self) -> None:
36 """Initialize an empty collection session."""
37 self._status = CodegenStatus.COLLECTING
38 self._source: dict[str, Any] = {}
39 self._replacements = ReplacementSnapshot()
40 self._plan: Optional[dict[str, Any]] = None
41 self._diagnostics: dict[str, Any] = {}
42 self._reason: Optional[str] = None
43 self._failure: Optional[FailureSnapshot] = None
44 self._source_recorded = False
45 self._replacements_recorded = False
46 self._diagnostics_recorded = False
47 self._finished = False
48 self._snapshot: Optional[CodegenSnapshot] = None
50 @property
51 def status(self) -> CodegenStatus:
52 """Return the current session status."""
53 return self._status
55 @property
56 def has_replacement_applications(self) -> bool:
57 """Whether native model construction installed any replacement."""
58 return bool(self._replacements.applications)
60 def record_source(self, source: Mapping[str, Any]) -> None:
61 """Record source and model identity once while collecting."""
62 self._require_open(CodegenStatus.COLLECTING)
63 if self._source_recorded:
64 raise RuntimeError("codegen source has already been recorded")
65 self._source = self._as_mapping(source, "source")
66 self._source_recorded = True
68 def record_replacements(
69 self,
70 *,
71 declarations: Sequence[Mapping[str, Any]] = (),
72 applications: Sequence[Mapping[str, Any]] = (),
73 unmatched: Sequence[Mapping[str, Any]] = (),
74 ) -> None:
75 """Record native replacement declarations and observed applications."""
76 self._require_open(CodegenStatus.COLLECTING)
77 if self._replacements_recorded:
78 raise RuntimeError("codegen replacements have already been recorded")
79 self._replacements = ReplacementSnapshot.from_data(
80 declarations=declarations,
81 applications=applications,
82 unmatched=unmatched,
83 )
84 self._replacements_recorded = True
86 def record_plan(self, plan: Mapping[str, Any]) -> None:
87 """Record the native planner result and enter the derived state."""
88 self._require_open(CodegenStatus.COLLECTING)
89 if not self._source_recorded:
90 raise RuntimeError("record_source() must run before record_plan()")
91 self._plan = self._as_mapping(plan, "plan")
92 self._status = CodegenStatus.DERIVED
94 def observe_plan(self, plan: Any, model: Any, mesh: Any) -> None:
95 """Observe a native planner result without exposing codegen to the builder."""
96 observed_plan = observe_emission_plan(plan, model).to_dict()
97 observed_plan["parallel_dims"] = {
98 "dp_size": getattr(mesh, "dp_size", 1),
99 "tp_size": mesh.tp_size,
100 "cp_size": mesh.cp_size,
101 "pp_size": getattr(mesh, "pp_size", 1),
102 "ep_size": mesh.ep_size,
103 "sequence_parallel": mesh.sequence_parallel,
104 "loss_parallel": mesh.loss_parallel,
105 }
106 self.record_plan(observed_plan)
108 def record_diagnostics(self, diagnostics: Mapping[str, Any]) -> None:
109 """Record optional counters and audit facts once."""
110 self._require_open()
111 if self._diagnostics_recorded:
112 raise RuntimeError("codegen diagnostics have already been recorded")
113 self._diagnostics = self._as_mapping(diagnostics, "diagnostics")
114 self._diagnostics_recorded = True
116 def mark_not_required(self, reason: str) -> None:
117 """Finish collection when the native build has no sharding plan."""
118 self._require_open(CodegenStatus.COLLECTING)
119 if not self._source_recorded:
120 raise RuntimeError("record_source() must run before mark_not_required()")
121 self._reason = self._require_reason(reason)
122 self._status = CodegenStatus.NOT_REQUIRED
124 def mark_generating(self) -> None:
125 """Mark the future side-effecting generation boundary."""
126 self._require_open(CodegenStatus.DERIVED)
127 self._status = CodegenStatus.GENERATING
129 def mark_generated(self) -> None:
130 """Mark successful generation."""
131 self._require_open(CodegenStatus.GENERATING)
132 self._status = CodegenStatus.GENERATED
134 def mark_skipped(self, reason: str) -> None:
135 """Mark generation as intentionally skipped with a stable reason."""
136 self._require_open(CodegenStatus.COLLECTING, CodegenStatus.DERIVED)
137 if not self._source_recorded:
138 raise RuntimeError("record_source() must run before mark_skipped()")
139 self._reason = self._require_reason(reason)
140 self._status = CodegenStatus.SKIPPED
142 def mark_failed(
143 self,
144 stage: str,
145 reason: str,
146 details: Optional[Mapping[str, Any]] = None,
147 ) -> None:
148 """Record a codegen-owned failure without retaining an exception."""
149 self._require_open(
150 CodegenStatus.COLLECTING,
151 CodegenStatus.DERIVED,
152 CodegenStatus.GENERATING,
153 CodegenStatus.NOT_REQUIRED,
154 )
155 self._failure = FailureSnapshot.from_data(stage, reason, details)
156 self._status = CodegenStatus.FAILED
158 def finish(
159 self,
160 generator: Optional[Callable[[CodegenSnapshot], Any]] = None,
161 ) -> CodegenSnapshot:
162 """Seal the session, optionally publishing a derived snapshot.
164 ``generator`` is invoked only for a fully derived native build. Its
165 failure is recorded as codegen-owned state and deliberately does not
166 escape into the caller's already successful model build.
168 Args:
169 generator: Optional side-effecting publisher that consumes the
170 immutable derived snapshot.
172 Returns:
173 The same snapshot instance on repeated calls.
175 Raises:
176 RuntimeError: If collection or generation is incomplete.
177 """
178 if self._finished:
179 if self._snapshot is None:
180 raise RuntimeError("finished codegen session has no snapshot")
181 return self._snapshot
182 if self._status == CodegenStatus.COLLECTING:
183 raise RuntimeError("codegen session is still collecting")
184 if self._status == CodegenStatus.GENERATING:
185 raise RuntimeError("codegen session is still generating")
186 if self._status == CodegenStatus.DERIVED and generator is not None:
187 publish_snapshot = self._build_snapshot(CodegenStatus.DERIVED)
188 self.mark_generating()
189 try:
190 generator(publish_snapshot)
191 except Exception as exc: # pylint: disable=broad-exception-caught
192 # Generation must not invalidate an already successful native build.
193 self.mark_failed(
194 "generation",
195 str(exc) or type(exc).__name__,
196 {"exception_type": type(exc).__name__},
197 )
198 else:
199 self.mark_generated()
200 elif self._status == CodegenStatus.NOT_REQUIRED and generator is not None:
201 publish_snapshot = self._build_snapshot(CodegenStatus.NOT_REQUIRED)
202 try:
203 generator(publish_snapshot)
204 except Exception as exc: # pylint: disable=broad-exception-caught
205 # Publication must not invalidate an already successful native build.
206 self.mark_failed(
207 "generation",
208 str(exc) or type(exc).__name__,
209 {"exception_type": type(exc).__name__},
210 )
211 self._snapshot = self._build_snapshot(self._status)
212 self._finished = True
213 return self._snapshot
215 def _build_snapshot(self, status: CodegenStatus) -> CodegenSnapshot:
216 """Build an immutable view without sealing the session."""
217 return CodegenSnapshot._from_owned_data( # pylint: disable=protected-access
218 status=status,
219 source=self._source,
220 replacements=self._replacements,
221 plan=self._plan,
222 diagnostics=self._diagnostics,
223 reason=self._reason,
224 failure=self._failure,
225 )
227 def _require_open(self, *statuses: CodegenStatus) -> None:
228 if self._finished:
229 raise RuntimeError("codegen session is already finished")
230 if statuses and self._status not in statuses:
231 expected = ", ".join(status.value for status in statuses)
232 raise RuntimeError(
233 f"codegen session is {self._status.value}; expected one of: {expected}"
234 )
236 @staticmethod
237 def _as_mapping(value: Mapping[str, Any], path: str) -> dict[str, Any]:
238 return copy_mapping(value, path)
240 @staticmethod
241 def _require_reason(reason: str) -> str:
242 if not isinstance(reason, str) or not reason.strip():
243 raise ValueError("codegen reason must be a non-empty string")
244 return reason
247__all__ = ["CodegenBuildSession"]