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