Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / plan / placements.py: 81%

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"""Convert planner placements to and from source-emission data.""" 

16 

17from __future__ import annotations 

18 

19import re 

20from typing import Any 

21 

22from hyper_parallel.core.dtensor.placement_types import ( 

23 Partial, 

24 Replicate, 

25 Shard, 

26 StridedShard, 

27) 

28 

29 

30def placement_to_string(placement: Any) -> str: 

31 """Serialize one placement to its canonical source representation.""" 

32 if isinstance(placement, StridedShard): 

33 if placement.uneven_shard: 

34 raise TypeError("cannot observe StridedShard(uneven_shard=True)") 

35 return f"SS({placement.dim},{placement.split_factor})" 

36 if isinstance(placement, Shard): 

37 if placement.uneven_shard: 

38 raise TypeError("cannot observe Shard(uneven_shard=True)") 

39 return f"S({placement.dim})" 

40 if isinstance(placement, Replicate): 

41 return "R" 

42 if isinstance(placement, Partial): 

43 return f"P({placement.reduce_op})" 

44 raise TypeError( 

45 f"cannot serialize placement {placement!r} of unknown type " 

46 f"{type(placement).__name__}" 

47 ) 

48 

49 

50def named_placement_to_dict(named: Any) -> dict[str, Any]: 

51 """Serialize a nested named-placement contract to plain dictionaries.""" 

52 if not named: 

53 return {} 

54 return { 

55 str(getattr(name, "value", name)): _value_to_data(value) 

56 for name, value in named.items() 

57 } 

58 

59 

60def parse_placement(text: str) -> Any: 

61 """Parse a canonical placement string.""" 

62 if not isinstance(text, str) or not text: 

63 raise ValueError(f"cannot parse placement from {text!r}") 

64 if text == "R": 

65 return Replicate() 

66 match = re.fullmatch(r"([SP]|SS)\((.*)\)", text) 

67 if match is None: 

68 raise ValueError( 

69 f"unknown placement form {text!r} (expected S(dim) / R / " 

70 "P(reduce_op) / SS(dim,factor))" 

71 ) 

72 kind, arguments = match.groups() 

73 if kind == "P": 

74 if not arguments: 

75 raise ValueError(f"malformed partial placement {text!r}: missing reduce_op") 

76 return Partial(arguments) 

77 parts = [part.strip() for part in arguments.split(",")] 

78 try: 

79 dim = int(parts[0]) 

80 except (IndexError, ValueError) as exc: 

81 raise ValueError(f"malformed shard placement {text!r}: invalid dim") from exc 

82 if kind == "S": 

83 if len(parts) != 1: 

84 raise ValueError(f"malformed shard placement {text!r}: too many args") 

85 return Shard(dim) 

86 if len(parts) != 2: 

87 raise ValueError(f"malformed strided shard placement {text!r}: expected SS(dim,factor)") 

88 try: 

89 split_factor = int(parts[1]) 

90 except ValueError as exc: 

91 raise ValueError( 

92 f"malformed strided shard placement {text!r}: invalid split_factor" 

93 ) from exc 

94 return StridedShard(dim, split_factor) 

95 

96 

97def parse_named_placement(data: dict[str, Any]) -> Any: 

98 """Parse a nested named-placement contract from plain dictionaries.""" 

99 if not data: 

100 return {} 

101 return _value_from_data(data) 

102 

103 

104def value_to_data(value: Any) -> Any: 

105 """Project a planner value to deterministic source-emission data.""" 

106 return _value_to_data(value) 

107 

108 

109def _value_to_data(value: Any) -> Any: 

110 """Recursively project placements and planner values into plain data.""" 

111 if isinstance(value, (Partial, Replicate, Shard, StridedShard)): 

112 return placement_to_string(value) 

113 if hasattr(value, "to_dict"): 

114 return value.to_dict() 

115 if isinstance(value, dict): 

116 return { 

117 str(getattr(key, "value", key)): _value_to_data(item) 

118 for key, item in value.items() 

119 } 

120 if isinstance(value, (list, tuple)): 

121 return [_value_to_data(item) for item in value] 

122 if callable(value): 

123 return getattr(value, "__name__", str(value)) 

124 return value 

125 

126 

127def _value_from_data(value: Any) -> Any: 

128 """Recursively reconstruct placements from nested plain-data contracts.""" 

129 if isinstance(value, str): 

130 return parse_placement(value) 

131 if isinstance(value, dict): 

132 return { 

133 str(getattr(key, "value", key)): _value_from_data(item) 

134 for key, item in value.items() 

135 } 

136 if isinstance(value, list): 

137 return [_value_from_data(item) for item in value] 

138 return value 

139 

140 

141__all__ = [ 

142 "named_placement_to_dict", 

143 "parse_named_placement", 

144 "parse_placement", 

145 "placement_to_string", 

146 "value_to_data", 

147]