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

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"""Bridge-facing Python runtime helpers for GE custom op implementations.""" 

14 

15import inspect 

16import sys 

17import threading 

18from dataclasses import dataclass 

19from typing import Dict, Optional 

20 

21from .base import BaseCustomOp, EagerExecuteOp, EagerOpExecutionContext 

22from .bootstrap import get_registered_op_impls, load_custom_op_plugins 

23from .context import _execute_ctx_scope 

24from .registry import get_registered_op_impl_by_descriptor_key 

25 

26 

27@dataclass 

28class _OpImplHolder: 

29 descriptor_key: str 

30 instance_id: str 

31 instance: BaseCustomOp 

32 

33 

34_HOLDER_LOCK = threading.RLock() 

35_OP_IMPL_HOLDERS: Dict[str, _OpImplHolder] = {} 

36 

37_IR_INPUT_REQUIRED = 0 

38_IR_INPUT_OPTIONAL = 1 

39_IR_INPUT_DYNAMIC = 2 

40 

41_RUNTIME_ATTR_GETTERS = { 

42 "VT_INT": "get_int", 

43 "VT_FLOAT": "get_float", 

44 "VT_BOOL": "get_bool", 

45 "VT_STRING": "get_str", 

46 "VT_DATA_TYPE": "get_data_type", 

47 "VT_TENSOR": "get_tensor", 

48 "VT_LIST_INT": "get_list_int", 

49 "VT_LIST_FLOAT": "get_list_float", 

50 "VT_LIST_BOOL": "get_list_bool", 

51 "VT_LIST_STRING": "get_list_str", 

52 "VT_LIST_DATA_TYPE": "get_list_data_type", 

53 "VT_LIST_LIST_INT": "get_list_list_int", 

54} 

55 

56 

57def load_and_get_op_impl_descriptors() -> list: 

58 load_custom_op_plugins() 

59 return get_registered_op_impls() 

60 

61 

62def _get_holder(instance_id: str) -> _OpImplHolder: 

63 with _HOLDER_LOCK: 

64 holder = _OP_IMPL_HOLDERS.get(instance_id) 

65 if holder is None: 

66 raise KeyError(f"python op impl holder is not created: {instance_id}") 

67 return holder 

68 

69 

70def _get_eager_execute_op(instance_id: str) -> EagerExecuteOp: 

71 instance = _get_holder(instance_id).instance 

72 if not isinstance(instance, EagerExecuteOp): 

73 raise TypeError( 

74 f"python op impl does not implement EagerExecuteOp: {instance_id}" 

75 ) 

76 return instance 

77 

78 

79def create_op_impl_holder(instance_id: str, descriptor_key: str) -> bool: 

80 descriptor = get_registered_op_impl_by_descriptor_key(descriptor_key) 

81 if descriptor is None: 

82 raise KeyError(f"python op impl descriptor_key not found: {descriptor_key}") 

83 with _HOLDER_LOCK: 

84 if instance_id in _OP_IMPL_HOLDERS: 

85 return True 

86 _OP_IMPL_HOLDERS[instance_id] = _OpImplHolder( 

87 descriptor_key=descriptor_key, 

88 instance_id=instance_id, 

89 instance=descriptor.cls(), 

90 ) 

91 return True 

92 

93 

94def destroy_op_impl_holder(instance_id: str) -> bool: 

95 with _HOLDER_LOCK: 

96 return _OP_IMPL_HOLDERS.pop(instance_id, None) is not None 

97 

98 

99def _is_legacy_execute(method) -> bool: 

100 params = list(inspect.signature(method).parameters.values()) 

101 return ( 

102 len(params) == 1 

103 and params[0].name == "ctx" 

104 and params[0].kind 

105 in ( 

106 inspect.Parameter.POSITIONAL_ONLY, 

107 inspect.Parameter.POSITIONAL_OR_KEYWORD, 

108 ) 

109 ) 

110 

111 

112def _build_execute_inputs(ctx: EagerOpExecutionContext, ir_inputs: list) -> list: 

113 args = [] 

114 for ir_index, item in enumerate(ir_inputs): 

115 kind = item["kind"] 

116 if kind == _IR_INPUT_REQUIRED: 

117 args.append(ctx.get_required_input_tensor(ir_index)) 

118 elif kind == _IR_INPUT_OPTIONAL: 

119 args.append(ctx.get_optional_input_tensor(ir_index)) 

120 elif kind == _IR_INPUT_DYNAMIC: 

121 instance_num = ctx.get_dynamic_input_num(ir_index) 

122 args.append( 

123 [ 

124 ctx.get_dynamic_input_tensor(ir_index, relative_index) 

125 for relative_index in range(instance_num) 

126 ] 

127 ) 

128 else: 

129 raise ValueError( 

130 f"unsupported custom op IR input kind: {kind}, ir index: {ir_index}" 

131 ) 

132 return args 

133 

134 

135def _read_runtime_attr(attrs, index: int, ir_type: str): 

136 getter_name = _RUNTIME_ATTR_GETTERS.get(ir_type) 

137 if getter_name is None: 

138 raise ValueError( 

139 f"unsupported custom op runtime attr type: {ir_type}, attr index: {index}" 

140 ) 

141 return getattr(attrs, getter_name)(index) 

142 

143 

144def _build_execute_attrs(ctx: EagerOpExecutionContext, ir_attrs: list) -> dict: 

145 if not ir_attrs: 

146 return {} 

147 attrs = ctx.get_attrs() 

148 return { 

149 item["name"]: _read_runtime_attr(attrs, index, item["type"]) 

150 for index, item in enumerate(ir_attrs) 

151 } 

152 

153 

154def call_execute( 

155 instance_id: str, 

156 ir_meta: Optional[dict], 

157 ctx: EagerOpExecutionContext, 

158) -> None: 

159 try: 

160 custom_op = _get_eager_execute_op(instance_id) 

161 method = custom_op.execute 

162 if _is_legacy_execute(method): 

163 method(ctx) 

164 return 

165 if ir_meta is None: 

166 descriptor = custom_op.__ge_op_impl_descriptor__ 

167 raise RuntimeError( 

168 f"canonical IR not found for schema-bound execute: {descriptor.op_type}" 

169 ) 

170 args = _build_execute_inputs(ctx, ir_meta["inputs"]) 

171 kwargs = _build_execute_attrs(ctx, ir_meta["attrs"]) 

172 with _execute_ctx_scope(ctx): 

173 method(*args, **kwargs) 

174 finally: 

175 ctx._invalidate() 

176 

177 

178def clear_op_impl_holders() -> None: 

179 with _HOLDER_LOCK: 

180 _OP_IMPL_HOLDERS.clear() 

181 

182 

183def clear_loaded_op_impl_modules() -> None: 

184 """Clear all dynamically loaded op implementation modules from sys.modules to avoid test pollution.""" 

185 keys_to_remove = [key for key in sys.modules if key.startswith("_ge_py_custom_op_")] 

186 for key in keys_to_remove: 

187 del sys.modules[key]