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

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 ONNX Plugin descriptor registry.""" 

14 

15import threading 

16from dataclasses import dataclass, field 

17from typing import Callable, Dict, List, Optional, Tuple 

18 

19PARSE_NODE = "parse_node" 

20PARSE_OPERATOR = "parse_operator" 

21 

22 

23@dataclass(frozen=True) 

24class OnnxPluginDescriptor: 

25 """Normalized descriptor for an ONNX parser plugin.""" 

26 

27 descriptor_key: str 

28 source: str 

29 domain: str 

30 opsets: Tuple[int, ...] 

31 target: str 

32 origin_types: Tuple[str, ...] 

33 module_name: str 

34 parser_node: Optional[Callable[..., None]] = field( 

35 default=None, compare=False, repr=False 

36 ) 

37 parser_operator: Optional[Callable[..., None]] = field( 

38 default=None, compare=False, repr=False 

39 ) 

40 

41 @property 

42 def callback_kinds(self) -> Tuple[str, ...]: 

43 kinds = [] 

44 if self.parser_node is not None: 

45 kinds.append(PARSE_NODE) 

46 if self.parser_operator is not None: 

47 kinds.append(PARSE_OPERATOR) 

48 return tuple(kinds) 

49 

50 @property 

51 def callback_kind(self) -> str: 

52 """Compatibility view for descriptors with one callback.""" 

53 return ( 

54 self.callback_kinds[0] 

55 if len(self.callback_kinds) == 1 

56 else ",".join(self.callback_kinds) 

57 ) 

58 

59 def to_bridge_dict(self) -> dict: 

60 descriptor = { 

61 "descriptor_key": self.descriptor_key, 

62 "source": self.source, 

63 "domain": self.domain, 

64 "opsets": list(self.opsets), 

65 "target": self.target, 

66 "origin_types": list(self.origin_types), 

67 "module_name": self.module_name, 

68 } 

69 if len(self.callback_kinds) == 1: 

70 descriptor["callback_kind"] = self.callback_kinds[0] 

71 else: 

72 descriptor["callback_kinds"] = list(self.callback_kinds) 

73 return descriptor 

74 

75 

76class _OnnxPluginRegistry: 

77 def __init__(self) -> None: 

78 self._lock = threading.RLock() 

79 self._descriptor_key_to_desc: Dict[str, OnnxPluginDescriptor] = {} 

80 self._origin_type_to_desc: Dict[str, OnnxPluginDescriptor] = {} 

81 

82 def clear(self) -> None: 

83 with self._lock: 

84 self._descriptor_key_to_desc.clear() 

85 self._origin_type_to_desc.clear() 

86 

87 def register(self, descriptor: OnnxPluginDescriptor) -> OnnxPluginDescriptor: 

88 with self._lock: 

89 if descriptor.descriptor_key in self._descriptor_key_to_desc: 

90 raise ValueError( 

91 "python ONNX Plugin descriptor_key already exists: " 

92 f"{descriptor.descriptor_key}" 

93 ) 

94 for origin_type in descriptor.origin_types: 

95 if origin_type in self._origin_type_to_desc: 

96 raise ValueError( 

97 f"python ONNX Plugin origin type already exists: {origin_type}" 

98 ) 

99 self._descriptor_key_to_desc[descriptor.descriptor_key] = descriptor 

100 for origin_type in descriptor.origin_types: 

101 self._origin_type_to_desc[origin_type] = descriptor 

102 return descriptor 

103 

104 def get_all(self) -> List[OnnxPluginDescriptor]: 

105 with self._lock: 

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

107 

108 def get_by_origin_type(self, origin_type: str) -> Optional[OnnxPluginDescriptor]: 

109 # Registration finishes before parsing; parse-time lookup is read-only. 

110 return self._origin_type_to_desc.get(origin_type) 

111 

112 def replace(self, descriptor: OnnxPluginDescriptor) -> OnnxPluginDescriptor: 

113 with self._lock: 

114 if descriptor.descriptor_key not in self._descriptor_key_to_desc: 

115 raise ValueError( 

116 "python ONNX Plugin descriptor_key does not exist: " 

117 f"{descriptor.descriptor_key}" 

118 ) 

119 self._descriptor_key_to_desc[descriptor.descriptor_key] = descriptor 

120 for origin_type in descriptor.origin_types: 

121 self._origin_type_to_desc[origin_type] = descriptor 

122 return descriptor 

123 

124 

125_ONNX_PLUGIN_REGISTRY = _OnnxPluginRegistry() 

126 

127 

128def register_onnx_plugin( 

129 descriptor: OnnxPluginDescriptor, 

130) -> OnnxPluginDescriptor: 

131 return _ONNX_PLUGIN_REGISTRY.register(descriptor) 

132 

133 

134def replace_registered_onnx_plugin( 

135 descriptor: OnnxPluginDescriptor, 

136) -> OnnxPluginDescriptor: 

137 return _ONNX_PLUGIN_REGISTRY.replace(descriptor) 

138 

139 

140def clear_registered_onnx_plugins() -> None: 

141 _ONNX_PLUGIN_REGISTRY.clear() 

142 

143 

144def get_registered_onnx_plugins() -> List[OnnxPluginDescriptor]: 

145 return _ONNX_PLUGIN_REGISTRY.get_all() 

146 

147 

148def get_registered_onnx_plugin_dicts() -> List[dict]: 

149 return [item.to_bridge_dict() for item in get_registered_onnx_plugins()] 

150 

151 

152def get_registered_onnx_plugin_by_origin_type( 

153 origin_type: str, 

154) -> Optional[OnnxPluginDescriptor]: 

155 return _ONNX_PLUGIN_REGISTRY.get_by_origin_type(origin_type)