Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/registry.py: 99%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3# ----------------------------------------------------------------------------------------------------------- 

4# Copyright (c) 2026 Huawei Technologies Co., Ltd. 

5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 

6# CANN Open Software License Agreement Version 2.0 (the "License"). 

7# Please refer to the License for details. You may not use this file except in compliance with the License. 

8# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 

9# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. 

10# See LICENSE in the root of the software repository for the full text of the License. 

11# ----------------------------------------------------------------------------------------------------------- 

12 

13"""Python custom op implementation registry and decorators.""" 

14 

15import inspect 

16import threading 

17from dataclasses import dataclass, field 

18from typing import Any, Dict, List, Optional, Type 

19 

20INTERFACE_EAGER_EXECUTE = "eager_execute" 

21INTERFACE_COMPILABLE = "compilable" 

22INTERFACE_ANNOTATED_ARGS = "annotated_args" 

23_INTERFACE_SPECS = ( 

24 (INTERFACE_EAGER_EXECUTE, "execute"), 

25 (INTERFACE_COMPILABLE, "compile"), 

26 (INTERFACE_ANNOTATED_ARGS, "declare_launch_args"), 

27) 

28 

29 

30@dataclass(frozen=True) 

31class OpImplDescriptor: 

32 """Normalized Python custom op implementation descriptor.""" 

33 

34 descriptor_key: str 

35 op_type: str 

36 module_name: str 

37 class_name: str 

38 interfaces: List[str] = field(default_factory=list) 

39 cls: Type[Any] = field(compare=False, repr=False, default=object) 

40 

41 def to_bridge_dict(self) -> dict: 

42 return { 

43 "descriptor_key": self.descriptor_key, 

44 "op_type": self.op_type, 

45 "module_name": self.module_name, 

46 "class_name": self.class_name, 

47 "interfaces": list(self.interfaces), 

48 } 

49 

50 

51class _OpImplRegistry: 

52 def __init__(self) -> None: 

53 self._lock = threading.RLock() 

54 self._descriptor_key_to_desc: Dict[str, OpImplDescriptor] = {} 

55 self._op_type_to_desc: Dict[str, OpImplDescriptor] = {} 

56 

57 def clear(self) -> None: 

58 with self._lock: 

59 self._descriptor_key_to_desc.clear() 

60 self._op_type_to_desc.clear() 

61 

62 def register(self, descriptor: OpImplDescriptor) -> OpImplDescriptor: 

63 with self._lock: 

64 if descriptor.descriptor_key in self._descriptor_key_to_desc: 

65 raise ValueError( 

66 f"python op impl descriptor_key already exists: {descriptor.descriptor_key}" 

67 ) 

68 if descriptor.op_type in self._op_type_to_desc: 

69 raise ValueError( 

70 f"python op impl type already exists: {descriptor.op_type}" 

71 ) 

72 self._descriptor_key_to_desc[descriptor.descriptor_key] = descriptor 

73 self._op_type_to_desc[descriptor.op_type] = descriptor 

74 return descriptor 

75 

76 def get_by_descriptor_key(self, descriptor_key: str) -> Optional[OpImplDescriptor]: 

77 with self._lock: 

78 return self._descriptor_key_to_desc.get(descriptor_key) 

79 

80 def get_all(self) -> List[OpImplDescriptor]: 

81 with self._lock: 

82 return list(self._descriptor_key_to_desc.values()) 

83 

84 

85_OP_IMPL_REGISTRY = _OpImplRegistry() 

86 

87 

88def _build_descriptor_key(module_name: str, class_name: str, op_type: str) -> str: 

89 return f"{module_name}:{class_name}:{op_type}" 

90 

91 

92def _normalize_op_type(op_type: str) -> str: 

93 if not isinstance(op_type, str) or not op_type: 

94 raise TypeError("register_op_impl op_type must be a non-empty string") 

95 return op_type 

96 

97 

98def _collect_interfaces(cls: Type[Any]) -> List[str]: 

99 interfaces = [] 

100 for name, method_name in _INTERFACE_SPECS: 

101 method = getattr(cls, method_name, None) 

102 # Keep legacy interface discovery behavior for execute and 

103 # declare_launch_args. Compile is schema-bound and must reject an 

104 # explicitly declared non-callable callback at registration time. 

105 if ( 

106 name == INTERFACE_COMPILABLE 

107 and hasattr(cls, method_name) 

108 and not callable(method) 

109 ): 

110 raise TypeError(f"{method_name} must be callable") 

111 if callable(method): 

112 interfaces.append(name) 

113 return interfaces 

114 

115 

116def _get_interfaces(cls: Type[Any]) -> List[str]: 

117 interfaces = _collect_interfaces(cls) 

118 if not interfaces: 

119 supported_methods = ", ".join( 

120 method_name for _, method_name in _INTERFACE_SPECS 

121 ) 

122 class_name = f"{cls.__module__}.{cls.__qualname__}" 

123 raise TypeError( 

124 f"register_op_impl class '{class_name}' must implement at least one " 

125 f"supported method: {supported_methods}" 

126 ) 

127 return interfaces 

128 

129 

130def _register_op_impl_class(cls: Type[Any], *, op_type: str) -> Type[Any]: 

131 module_name = cls.__module__ 

132 class_name = cls.__name__ 

133 descriptor = OpImplDescriptor( 

134 descriptor_key=_build_descriptor_key(module_name, class_name, op_type), 

135 op_type=op_type, 

136 module_name=module_name, 

137 class_name=class_name, 

138 interfaces=_get_interfaces(cls), 

139 cls=cls, 

140 ) 

141 _OP_IMPL_REGISTRY.register(descriptor) 

142 setattr(cls, "__ge_op_impl_descriptor__", descriptor) 

143 return cls 

144 

145 

146def register_op_impl(*, op_type: str) -> callable: 

147 """Decorator for Python custom op implementation classes.""" 

148 

149 normalized_op_type = _normalize_op_type(op_type) 

150 

151 def decorator(cls: Type[Any]) -> Type[Any]: 

152 if not inspect.isclass(cls): 

153 raise TypeError("register_op_impl expects a class") 

154 if inspect.isabstract(cls): 

155 raise TypeError("register_op_impl expects a concrete class") 

156 return _register_op_impl_class(cls, op_type=normalized_op_type) 

157 

158 return decorator 

159 

160 

161def clear_registered_op_impls() -> None: 

162 _OP_IMPL_REGISTRY.clear() 

163 

164 

165def get_registered_op_impls() -> List[OpImplDescriptor]: 

166 return _OP_IMPL_REGISTRY.get_all() 

167 

168 

169def get_registered_op_impl_dicts() -> List[dict]: 

170 return [item.to_bridge_dict() for item in get_registered_op_impls()] 

171 

172 

173def get_registered_op_impl_by_descriptor_key( 

174 descriptor_key: str, 

175) -> Optional[OpImplDescriptor]: 

176 return _OP_IMPL_REGISTRY.get_by_descriptor_key(descriptor_key)