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_ANNOTATED_ARGS = "annotated_args" 

22_INTERFACE_SPECS = ( 

23 (INTERFACE_EAGER_EXECUTE, "execute"), 

24 (INTERFACE_ANNOTATED_ARGS, "declare_launch_args"), 

25) 

26 

27 

28@dataclass(frozen=True) 

29class OpImplDescriptor: 

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

31 

32 descriptor_key: str 

33 op_type: str 

34 module_name: str 

35 class_name: str 

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

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

38 

39 def to_bridge_dict(self) -> dict: 

40 return { 

41 "descriptor_key": self.descriptor_key, 

42 "op_type": self.op_type, 

43 "module_name": self.module_name, 

44 "class_name": self.class_name, 

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

46 } 

47 

48 

49class _OpImplRegistry: 

50 def __init__(self) -> None: 

51 self._lock = threading.RLock() 

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

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

54 

55 def clear(self) -> None: 

56 with self._lock: 

57 self._descriptor_key_to_desc.clear() 

58 self._op_type_to_desc.clear() 

59 

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

61 with self._lock: 

62 if descriptor.descriptor_key in self._descriptor_key_to_desc: 

63 raise ValueError( 

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

65 ) 

66 if descriptor.op_type in self._op_type_to_desc: 

67 raise ValueError( 

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

69 ) 

70 self._descriptor_key_to_desc[descriptor.descriptor_key] = descriptor 

71 self._op_type_to_desc[descriptor.op_type] = descriptor 

72 return descriptor 

73 

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

75 with self._lock: 

76 return self._descriptor_key_to_desc.get(descriptor_key) 

77 

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

79 with self._lock: 

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

81 

82 

83_OP_IMPL_REGISTRY = _OpImplRegistry() 

84 

85 

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

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

88 

89 

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

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

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

93 return op_type 

94 

95 

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

97 return [ 

98 name 

99 for name, method_name in _INTERFACE_SPECS 

100 if callable(getattr(cls, method_name, None)) 

101 ] 

102 

103 

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

105 interfaces = _collect_interfaces(cls) 

106 if not interfaces: 

107 supported_methods = ", ".join( 

108 method_name for _, method_name in _INTERFACE_SPECS 

109 ) 

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

111 raise TypeError( 

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

113 f"supported method: {supported_methods}" 

114 ) 

115 return interfaces 

116 

117 

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

119 module_name = cls.__module__ 

120 class_name = cls.__name__ 

121 descriptor = OpImplDescriptor( 

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

123 op_type=op_type, 

124 module_name=module_name, 

125 class_name=class_name, 

126 interfaces=_get_interfaces(cls), 

127 cls=cls, 

128 ) 

129 _OP_IMPL_REGISTRY.register(descriptor) 

130 setattr(cls, "__ge_op_impl_descriptor__", descriptor) 

131 return cls 

132 

133 

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

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

136 

137 normalized_op_type = _normalize_op_type(op_type) 

138 

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

140 if not inspect.isclass(cls): 

141 raise TypeError("register_op_impl expects a class") 

142 if inspect.isabstract(cls): 

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

144 return _register_op_impl_class(cls, op_type=normalized_op_type) 

145 

146 return decorator 

147 

148 

149def clear_registered_op_impls() -> None: 

150 _OP_IMPL_REGISTRY.clear() 

151 

152 

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

154 return _OP_IMPL_REGISTRY.get_all() 

155 

156 

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

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

159 

160 

161def get_registered_op_impl_by_descriptor_key( 

162 descriptor_key: str, 

163) -> Optional[OpImplDescriptor]: 

164 return _OP_IMPL_REGISTRY.get_by_descriptor_key(descriptor_key)