Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / manager.py: 92%

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"""Generate and atomically publish single-process codegen bundles. 

16 

17The manager is the only entry point the trainer / CLI talks to. It does not 

18perform AST patching or source emitting itself; those operations belong to 

19the ``emit`` layer. Every invocation emits and atomically replaces the bundle 

20in the current process. 

21""" 

22 

23from __future__ import annotations 

24 

25import logging 

26from typing import Optional 

27 

28from hyper_parallel.codegen.artifact import ( 

29 ArtifactLayout, 

30 artifact_generation_lock, 

31 clean_temp_artifacts, 

32 ensure_artifact_dir, 

33 resolve_default_artifact_layout, 

34 resolve_model_artifact_layout, 

35 write_bundle_atomic, 

36) 

37from hyper_parallel.codegen.emit import emit_bundle 

38from hyper_parallel.codegen.emit.replacement import replacement_records_from_snapshot 

39from hyper_parallel.codegen.emit.snapshot_patches import ( 

40 emission_patches_from_snapshot_payload, 

41) 

42from hyper_parallel.codegen.emit.context import EmissionContext 

43from hyper_parallel.codegen.snapshot import CodegenSnapshot, CodegenStatus 

44from hyper_parallel.models.registry import ( 

45 artifact_model_type, 

46 family_dir_for_model_type, 

47) 

48 

49logger = logging.getLogger(__name__) 

50 

51 

52def ensure_codegen_artifact_from_snapshot( 

53 snapshot: CodegenSnapshot, 

54 artifact_dir: Optional[str] = None, 

55 *, 

56 timeout_s: float = 600.0, 

57) -> ArtifactLayout: 

58 """Generate and replace a bundle in the current single process. 

59 

60 This entry must remain independent of config projection, meta-model 

61 construction, offline meshes, replacement matching, and planner calls. 

62 Distributed launchers must use a dedicated sidecar instead of invoking 

63 this function independently on every training rank. 

64 """ 

65 if not isinstance(snapshot, CodegenSnapshot): 

66 raise TypeError("snapshot must be a CodegenSnapshot") 

67 if snapshot.status not in (CodegenStatus.DERIVED, CodegenStatus.NOT_REQUIRED): 

68 raise ValueError( 

69 "snapshot-driven generation requires status 'derived' or 'not_required', " 

70 f"got {snapshot.status.value!r}" 

71 ) 

72 source = snapshot.source 

73 if not source.get("file_path"): 

74 raise ValueError("snapshot source must include a modeling file_path") 

75 model_type = source.get("model_type") 

76 public_model_type = artifact_model_type(str(model_type)) if model_type else None 

77 model_name = ( 

78 _model_type_to_modeling_token(str(public_model_type)) 

79 if model_type 

80 else source.get("architecture") 

81 ) 

82 if artifact_dir is None and source.get("model_type"): 

83 layout = resolve_model_artifact_layout( 

84 family_dir_for_model_type(str(source["model_type"])), 

85 model_name=str(model_name) if model_name else None, 

86 ) 

87 else: 

88 layout = resolve_default_artifact_layout( 

89 artifact_dir, 

90 model_name=str(model_name) if model_name else None, 

91 ) 

92 ensure_artifact_dir(layout) 

93 

94 with artifact_generation_lock(layout, timeout_s=timeout_s): 

95 clean_temp_artifacts(layout) 

96 logger.info("codegen: generating artifact %s", layout.artifact_dir) 

97 context = _build_emission_context(snapshot) 

98 files = emit_bundle(layout, context) 

99 write_bundle_atomic(layout, files) 

100 return layout 

101 

102 

103def _build_emission_context(snapshot: CodegenSnapshot) -> EmissionContext: 

104 """Map one immutable observation into an in-memory emission context.""" 

105 payload = snapshot._emission_view() # pylint: disable=protected-access 

106 source = dict(payload["source"]) 

107 status = snapshot.status 

108 frozen_plan = dict(payload["plan"] or {}) 

109 module_overrides = list(replacement_records_from_snapshot(payload)) 

110 is_derived = status is CodegenStatus.DERIVED 

111 emission_patches = ( 

112 emission_patches_from_snapshot_payload(payload) if is_derived else {} 

113 ) 

114 return EmissionContext( 

115 source=source, 

116 parallel_dims=dict(frozen_plan.get("parallel_dims") or {}), 

117 status=payload["status"], 

118 injections=list(frozen_plan.get("injections") or ()), 

119 param_plan=dict(frozen_plan.get("param_plan") or {}), 

120 boundary_classes=dict(frozen_plan.get("boundary_classes") or {}), 

121 nested_components=list(frozen_plan.get("nested_components") or ()), 

122 mesh_dim_names=list(frozen_plan.get("mesh_dim_names") or ()) or None, 

123 module_overrides=(module_overrides if is_derived else []), 

124 emission_patches=emission_patches, 

125 ) 

126 

127 

128def _model_type_to_modeling_token(model_type: str) -> str: 

129 """Convert HF ``model_type`` into the generated modeling file token.""" 

130 parts = [part for part in model_type.replace("-", "_").split("_") if part] 

131 if not parts: 

132 return "model" 

133 return "_".join(part[:1].upper() + part[1:] for part in parts) 

134 

135 

136__all__ = [ 

137 "ensure_codegen_artifact_from_snapshot", 

138]