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 Dict, List, Optional, Type 

19 

20from .base import BaseCustomOp, EagerExecuteOp 

21 

22INTERFACE_EAGER_EXECUTE = "eager_execute" 

23_INTERFACE_SPECS = ((INTERFACE_EAGER_EXECUTE, EagerExecuteOp),) 

24 

25 

26@dataclass(frozen=True) 

27class OpImplDescriptor: 

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

29 

30 descriptor_key: str 

31 op_type: str 

32 module_name: str 

33 class_name: str 

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

35 cls: Type[BaseCustomOp] = field(compare=False, repr=False, default=BaseCustomOp) 

36 

37 def to_bridge_dict(self) -> dict: 

38 return { 

39 "descriptor_key": self.descriptor_key, 

40 "op_type": self.op_type, 

41 "module_name": self.module_name, 

42 "class_name": self.class_name, 

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

44 } 

45 

46 

47class _OpImplRegistry: 

48 def __init__(self) -> None: 

49 self._lock = threading.RLock() 

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

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

52 

53 def clear(self) -> None: 

54 with self._lock: 

55 self._descriptor_key_to_desc.clear() 

56 self._op_type_to_desc.clear() 

57 

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

59 with self._lock: 

60 if descriptor.descriptor_key in self._descriptor_key_to_desc: 

61 raise ValueError( 

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

63 ) 

64 if descriptor.op_type in self._op_type_to_desc: 

65 raise ValueError( 

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

67 ) 

68 self._descriptor_key_to_desc[descriptor.descriptor_key] = descriptor 

69 self._op_type_to_desc[descriptor.op_type] = descriptor 

70 return descriptor 

71 

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

73 with self._lock: 

74 return self._descriptor_key_to_desc.get(descriptor_key) 

75 

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

77 with self._lock: 

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

79 

80 

81_OP_IMPL_REGISTRY = _OpImplRegistry() 

82 

83 

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

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

86 

87 

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

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

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

91 return op_type 

92 

93 

94def _collect_interfaces(cls: Type[BaseCustomOp]) -> List[str]: 

95 return [name for name, base_cls in _INTERFACE_SPECS if issubclass(cls, base_cls)] 

96 

97 

98def _get_interfaces(cls: Type[BaseCustomOp]) -> List[str]: 

99 interfaces = _collect_interfaces(cls) 

100 if not interfaces: 

101 raise TypeError("register_op_impl expects a supported BaseCustomOp subclass") 

102 return interfaces 

103 

104 

105def _register_op_impl_class( 

106 cls: Type[BaseCustomOp], *, op_type: str 

107) -> Type[BaseCustomOp]: 

108 module_name = cls.__module__ 

109 class_name = cls.__name__ 

110 descriptor = OpImplDescriptor( 

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

112 op_type=op_type, 

113 module_name=module_name, 

114 class_name=class_name, 

115 interfaces=_get_interfaces(cls), 

116 cls=cls, 

117 ) 

118 _OP_IMPL_REGISTRY.register(descriptor) 

119 setattr(cls, "__ge_op_impl_descriptor__", descriptor) 

120 return cls 

121 

122 

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

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

125 

126 normalized_op_type = _normalize_op_type(op_type) 

127 

128 def decorator(cls: Type[BaseCustomOp]) -> Type[BaseCustomOp]: 

129 if not inspect.isclass(cls) or not issubclass(cls, BaseCustomOp): 

130 raise TypeError("register_op_impl expects a BaseCustomOp subclass") 

131 return _register_op_impl_class(cls, op_type=normalized_op_type) 

132 

133 return decorator 

134 

135 

136def clear_registered_op_impls() -> None: 

137 _OP_IMPL_REGISTRY.clear() 

138 

139 

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

141 return _OP_IMPL_REGISTRY.get_all() 

142 

143 

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

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

146 

147 

148def get_registered_op_impl_by_descriptor_key( 

149 descriptor_key: str, 

150) -> Optional[OpImplDescriptor]: 

151 return _OP_IMPL_REGISTRY.get_by_descriptor_key(descriptor_key)