Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/_bridge.py: 96%
157 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-20 19:14 +0800
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# -----------------------------------------------------------------------------------------------------------
13"""Bridge-facing Python runtime helpers for GE custom op implementations."""
15import inspect
16import sys
17import threading
18from dataclasses import dataclass
19from typing import Dict, Optional
21from ._ir_types import InputType, OutputType
22from ._signature import _get_runtime_attr_spec, _validate_args_signature
23from .base import EagerOpExecutionContext
24from .bootstrap import (
25 get_registered_op_impls,
26 get_registered_op_protos,
27 load_custom_op_plugins,
28)
29from .context import _declare_launch_args_ctx_scope, _execute_ctx_scope
30from .registry import (
31 INTERFACE_ANNOTATED_ARGS,
32 INTERFACE_EAGER_EXECUTE,
33 get_registered_op_impl_by_descriptor_key,
34)
37@dataclass
38class _OpImplHolder:
39 descriptor_key: str
40 instance_id: str
41 instance: object
44_HOLDER_LOCK = threading.RLock()
45_OP_IMPL_HOLDERS: Dict[str, _OpImplHolder] = {}
48def load_and_get_op_impl_descriptors() -> list:
49 load_custom_op_plugins()
50 return get_registered_op_impls()
53def load_and_get_op_descriptors() -> dict:
54 load_custom_op_plugins()
55 return {
56 "protos": get_registered_op_protos(),
57 "impls": get_registered_op_impls(),
58 }
61def _get_holder(instance_id: str) -> _OpImplHolder:
62 with _HOLDER_LOCK:
63 holder = _OP_IMPL_HOLDERS.get(instance_id)
64 if holder is None:
65 raise KeyError(f"python op impl holder is not created: {instance_id}")
66 return holder
69def _get_eager_execute_holder(instance_id: str) -> _OpImplHolder:
70 holder = _get_holder(instance_id)
71 if not callable(getattr(holder.instance, "execute", None)):
72 raise TypeError(
73 f"python op impl does not implement callable execute: {instance_id}"
74 )
75 return holder
78def _get_eager_execute_op(instance_id: str) -> object:
79 return _get_eager_execute_holder(instance_id).instance
82def create_op_impl_holder(instance_id: str, descriptor_key: str) -> bool:
83 descriptor = get_registered_op_impl_by_descriptor_key(descriptor_key)
84 if descriptor is None:
85 raise KeyError(f"python op impl descriptor_key not found: {descriptor_key}")
86 with _HOLDER_LOCK:
87 if instance_id in _OP_IMPL_HOLDERS:
88 return True
89 _OP_IMPL_HOLDERS[instance_id] = _OpImplHolder(
90 descriptor_key=descriptor_key,
91 instance_id=instance_id,
92 instance=descriptor.cls(),
93 )
94 return True
97def destroy_op_impl_holder(instance_id: str) -> bool:
98 with _HOLDER_LOCK:
99 return _OP_IMPL_HOLDERS.pop(instance_id, None) is not None
102def _is_legacy_execute(method) -> bool:
103 params = list(inspect.signature(method).parameters.values())
104 return (
105 len(params) == 1
106 and params[0].name == "ctx"
107 and params[0].kind
108 in (
109 inspect.Parameter.POSITIONAL_ONLY,
110 inspect.Parameter.POSITIONAL_OR_KEYWORD,
111 )
112 )
115def _get_callback_for_signature(cls, method_name: str):
116 method = inspect.getattr_static(cls, method_name)
117 if isinstance(method, staticmethod):
118 return method.__func__
119 if isinstance(method, classmethod):
120 return method.__get__(None, cls)
121 if inspect.isfunction(method):
122 return method.__get__(object(), cls)
123 return getattr(cls, method_name)
126def validate_op_impl_descriptor(descriptor_key: str, ir_meta: Optional[dict]) -> bool:
127 descriptor = get_registered_op_impl_by_descriptor_key(descriptor_key)
128 if descriptor is None:
129 raise KeyError(f"python op impl descriptor_key not found: {descriptor_key}")
131 if INTERFACE_EAGER_EXECUTE in descriptor.interfaces:
132 method = _get_callback_for_signature(descriptor.cls, "execute")
133 if not _is_legacy_execute(method):
134 if ir_meta is None:
135 raise RuntimeError(
136 "canonical IR not found for schema-bound execute: "
137 f"{descriptor.op_type}"
138 )
139 _validate_args_signature(method, ir_meta, descriptor, method_name="execute")
141 if INTERFACE_ANNOTATED_ARGS in descriptor.interfaces:
142 if ir_meta is None:
143 raise RuntimeError(
144 "canonical IR not found for schema-bound declare_launch_args"
145 )
146 method = _get_callback_for_signature(descriptor.cls, "declare_launch_args")
147 _validate_args_signature(
148 method, ir_meta, descriptor, method_name="declare_launch_args"
149 )
150 return True
153def _build_inputs(
154 ir_inputs: list,
155 get_required_input,
156 get_optional_input,
157 get_dynamic_input_num,
158 get_dynamic_input,
159) -> list:
160 args = []
161 for ir_index, item in enumerate(ir_inputs):
162 kind = item["kind"]
163 if kind == InputType.REQUIRED:
164 args.append(get_required_input(ir_index))
165 elif kind == InputType.OPTIONAL:
166 args.append(get_optional_input(ir_index))
167 elif kind == InputType.DYNAMIC:
168 instance_num = get_dynamic_input_num(ir_index)
169 args.append(
170 [
171 get_dynamic_input(ir_index, relative_index)
172 for relative_index in range(instance_num)
173 ]
174 )
175 else:
176 raise ValueError(
177 f"unsupported custom op IR input kind: {kind}, ir index: {ir_index}"
178 )
179 return args
182def _build_execute_inputs(ctx: EagerOpExecutionContext, ir_inputs: list) -> list:
183 return _build_inputs(
184 ir_inputs,
185 ctx.get_required_input_tensor,
186 ctx.get_optional_input_tensor,
187 ctx.get_dynamic_input_num,
188 ctx.get_dynamic_input_tensor,
189 )
192def _read_runtime_attr(attrs, index: int, ir_type: str):
193 getter_name, _ = _get_runtime_attr_spec(ir_type, index)
194 return getattr(attrs, getter_name)(index)
197def _build_execute_attrs(ctx: EagerOpExecutionContext, ir_attrs: list) -> dict:
198 if not ir_attrs:
199 return {}
200 attrs = ctx.get_attrs()
201 return {
202 item["name"]: _read_runtime_attr(attrs, index, item["type"])
203 for index, item in enumerate(ir_attrs)
204 }
207def _build_declare_inputs(ctx, ir_inputs: list) -> list:
208 return _build_inputs(
209 ir_inputs,
210 ctx._get_required_input_tensor,
211 ctx._get_optional_input_tensor,
212 ctx._get_dynamic_input_num,
213 ctx._get_dynamic_input_tensor,
214 )
217def _build_declare_outputs(ctx, ir_outputs: list) -> list:
218 args = []
219 for ir_index, item in enumerate(ir_outputs):
220 kind = item["kind"]
221 if kind == OutputType.REQUIRED:
222 args.append(ctx._get_required_output_tensor(ir_index))
223 elif kind == OutputType.DYNAMIC:
224 instance_num = ctx._get_dynamic_output_num(ir_index)
225 args.append(
226 [
227 ctx._get_dynamic_output_tensor(ir_index, relative_index)
228 for relative_index in range(instance_num)
229 ]
230 )
231 else:
232 raise ValueError(f"unsupported custom op IR output kind: {kind}")
233 return args
236def _build_declare_attrs(ctx, ir_attrs: list) -> dict:
237 if not ir_attrs:
238 return {}
239 attrs = ctx._get_attrs()
240 return {
241 item["name"]: _read_runtime_attr(attrs, index, item["type"])
242 for index, item in enumerate(ir_attrs)
243 }
246def call_execute(
247 instance_id: str,
248 ir_meta: Optional[dict],
249 ctx: EagerOpExecutionContext,
250) -> None:
251 try:
252 holder = _get_eager_execute_holder(instance_id)
253 custom_op = holder.instance
254 method = custom_op.execute
255 if _is_legacy_execute(method):
256 method(ctx)
257 return
258 if ir_meta is None:
259 descriptor = custom_op.__ge_op_impl_descriptor__
260 raise RuntimeError(
261 f"canonical IR not found for schema-bound execute: {descriptor.op_type}"
262 )
263 args = _build_execute_inputs(ctx, ir_meta["inputs"])
264 kwargs = _build_execute_attrs(ctx, ir_meta["attrs"])
265 with _execute_ctx_scope(ctx):
266 method(*args, **kwargs)
267 finally:
268 ctx._invalidate()
271def call_declare_launch_args(instance_id: str, ir_meta: Optional[dict], ctx) -> None:
272 try:
273 holder = _get_holder(instance_id)
274 method = getattr(holder.instance, "declare_launch_args", None)
275 if not callable(method):
276 raise TypeError(
277 f"python op impl does not implement declare_launch_args: {instance_id}"
278 )
279 if ir_meta is None:
280 raise RuntimeError(
281 "canonical IR not found for schema-bound declare_launch_args"
282 )
283 args = _build_declare_inputs(ctx, ir_meta["inputs"])
284 args.extend(_build_declare_outputs(ctx, ir_meta["outputs"]))
285 kwargs = _build_declare_attrs(ctx, ir_meta["attrs"])
286 with _declare_launch_args_ctx_scope(ctx):
287 result = method(*args, **kwargs)
288 if result is not None:
289 raise TypeError("declare_launch_args must return None")
290 finally:
291 ctx._invalidate()
294def clear_op_impl_holders() -> None:
295 with _HOLDER_LOCK:
296 _OP_IMPL_HOLDERS.clear()
299def clear_loaded_op_impl_modules() -> None:
300 """Clear all dynamically loaded op implementation modules from sys.modules to avoid test pollution."""
301 keys_to_remove = [key for key in sys.modules if key.startswith("_ge_py_custom_op_")]
302 for key in keys_to_remove:
303 del sys.modules[key]