Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / config / replacement.py: 100%

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"""Replacement-only Target retaining family selectors for Python/CLI overrides.""" 

16 

17from __future__ import annotations 

18 

19from typing import Any, Callable 

20 

21from hyper_parallel.components.modules.family import ( 

22 SelectorSpec, 

23 get_family_selector_specs, 

24 get_module_family, 

25 normalize_family_selectors, 

26 resolve_module_family_target, 

27) 

28from hyper_parallel.components.quantization.config import LowPrecisionDtypeScheme 

29from hyper_parallel.components.modules.module_family import FamilyRequest, build_family 

30from hyper_parallel.trainer.config.target import Target, _build_nested_targets 

31 

32 

33class ReplacementTarget(Target): 

34 """Resolve the implementation without discarding the requested baseline.""" 

35 

36 def __init__(self, _target_: Callable, *, target_path: str, 

37 low_precision_dtype_scheme: str | None = None, **kwargs: Any) -> None: 

38 """Select a final callable and strictly validate constructor keywords. 

39 

40 Args: 

41 _target_: Requested baseline or ordinary replacement callable. 

42 target_path: Original import path retained in serialized configuration. 

43 low_precision_dtype_scheme: Catalog name resolved when the replacement plan is built. 

44 **kwargs: Explicit selectors and constructor arguments. 

45 """ 

46 self.requested_target = _target_ 

47 self.low_precision_dtype_scheme = low_precision_dtype_scheme 

48 kwargs = normalize_family_selectors(_target_, kwargs) 

49 selector_specs = get_family_selector_specs(_target_) 

50 self.selectors = {name: value for name, value in kwargs.items() if name in selector_specs} 

51 if low_precision_dtype_scheme is None: 

52 resolved, arguments = resolve_module_family_target(_target_, kwargs) 

53 else: 

54 if not isinstance(low_precision_dtype_scheme, str) or not low_precision_dtype_scheme.strip(): 

55 raise ValueError("low_precision_dtype_scheme must be a non-empty catalog name") 

56 # The catalog belongs to training configuration and is unavailable while parsing this Target. 

57 resolved = _target_ 

58 arguments = {name: value for name, value in kwargs.items() if name not in selector_specs} 

59 super().__init__(resolved, target_path=target_path, **arguments) 

60 

61 def bind_dtype_scheme(self, scheme: LowPrecisionDtypeScheme) -> Target: 

62 """Select from scheme contents while retaining its catalog name in the source config. 

63 

64 Args: 

65 scheme: Resolved catalog entry providing format selectors. 

66 """ 

67 formats = { 

68 "weight_format": scheme.weight_format, 

69 "act_format": scheme.act_format, 

70 "is_fake_quantize": scheme.is_fake_quantize, 

71 } 

72 for name, value in formats.items(): 

73 if name in self.selectors and ( 

74 type(self.selectors[name]) is not type(value) or self.selectors[name] != value): 

75 raise ValueError(f"{name} conflicts with dtype scheme {self.low_precision_dtype_scheme!r}") 

76 resolved, arguments = resolve_module_family_target(self.requested_target, { 

77 **self._kwargs, **self.selectors, **formats, 

78 }) 

79 return Target(resolved, target_path=self._target_path, **arguments) 

80 

81 def build(self, **runtime_kwargs: Any) -> Any: 

82 """Build only after a named scheme has been bound by the replacement planner.""" 

83 if self.low_precision_dtype_scheme is not None: 

84 raise ValueError("Bind low_precision_dtype_scheme through the replacement plan before building") 

85 family = get_module_family(self.requested_target) 

86 if family is None or family.default_variant is None: 

87 return super().build(**runtime_kwargs) 

88 arguments = {name: _build_nested_targets(value) for name, value in self._kwargs.items()} 

89 arguments.update(self.selectors) 

90 return build_family( 

91 FamilyRequest(self.requested_target, arguments), 

92 module=runtime_kwargs.get("module"), 

93 module_fqn=runtime_kwargs.get("module_fqn", ""), 

94 context=runtime_kwargs.get("context"), 

95 ) 

96 

97 def get_family_selector_spec(self, name: str) -> SelectorSpec | None: 

98 """Query the requested baseline's schema for Target-level overrides. 

99 

100 Args: 

101 name: Selector name supplied by the override path. 

102 

103 Returns: 

104 Its declaration, or None when it is not a selector for this target. 

105 """ 

106 if name == "low_precision_dtype_scheme": 

107 return SelectorSpec(str, "") 

108 return get_family_selector_specs(self.requested_target).get(name) 

109 

110 def replace(self, **changes: Any) -> "ReplacementTarget": 

111 """Re-resolve selectors after an override, without mutating the original. 

112 

113 Args: 

114 **changes: Explicit selectors or constructor arguments to replace. 

115 """ 

116 return type(self)( 

117 self.requested_target, target_path=self._target_path, 

118 **{**self._kwargs, **self.selectors, 

119 "low_precision_dtype_scheme": self.low_precision_dtype_scheme, 

120 **normalize_family_selectors(self.requested_target, changes)} 

121 ) 

122 

123 def to_dict(self) -> dict[str, Any]: 

124 """Serialize the baseline path and original selectors for round trips.""" 

125 result = {**super().to_dict(), **self.selectors} 

126 if self.low_precision_dtype_scheme is not None: 

127 result["low_precision_dtype_scheme"] = self.low_precision_dtype_scheme 

128 return result