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