Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/_bridge.py: 96%
151 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 10:22 +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 ._native 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 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
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
99def _get_callback_for_signature(cls, method_name: str):
100 method = inspect.getattr_static(cls, method_name)
101 if isinstance(method, staticmethod):
102 return method.__func__
103 if isinstance(method, classmethod):
104 return method.__get__(None, cls)
105 if inspect.isfunction(method):
106 return method.__get__(object(), cls)
107 return getattr(cls, method_name)
110def validate_op_impl_descriptor(descriptor_key: str, ir_meta: Optional[dict]) -> bool:
111 descriptor = get_registered_op_impl_by_descriptor_key(descriptor_key)
112 if descriptor is None:
113 raise KeyError(f"python op impl descriptor_key not found: {descriptor_key}")
115 if INTERFACE_EAGER_EXECUTE in descriptor.interfaces:
116 if ir_meta is None:
117 raise RuntimeError(
118 f"canonical IR not found for schema-bound execute: {descriptor.op_type}"
119 )
120 method = _get_callback_for_signature(descriptor.cls, "execute")
121 _validate_args_signature(method, ir_meta, descriptor, method_name="execute")
123 if INTERFACE_ANNOTATED_ARGS in descriptor.interfaces:
124 if ir_meta is None:
125 raise RuntimeError(
126 "canonical IR not found for schema-bound declare_launch_args"
127 )
128 method = _get_callback_for_signature(descriptor.cls, "declare_launch_args")
129 _validate_args_signature(
130 method, ir_meta, descriptor, method_name="declare_launch_args"
131 )
132 return True
135def _build_inputs(
136 ir_inputs: list,
137 get_required_input,
138 get_optional_input,
139 get_dynamic_input_num,
140 get_dynamic_input,
141) -> list:
142 args = []
143 for ir_index, item in enumerate(ir_inputs):
144 kind = item["kind"]
145 if kind == InputType.REQUIRED:
146 args.append(get_required_input(ir_index))
147 elif kind == InputType.OPTIONAL:
148 args.append(get_optional_input(ir_index))
149 elif kind == InputType.DYNAMIC:
150 instance_num = get_dynamic_input_num(ir_index)
151 args.append(
152 [
153 get_dynamic_input(ir_index, relative_index)
154 for relative_index in range(instance_num)
155 ]
156 )
157 else:
158 raise ValueError(
159 f"unsupported custom op IR input kind: {kind}, ir index: {ir_index}"
160 )
161 return args
164def _build_execute_inputs(ctx: EagerOpExecutionContext, ir_inputs: list) -> list:
165 return _build_inputs(
166 ir_inputs,
167 ctx.get_required_input_tensor,
168 ctx.get_optional_input_tensor,
169 ctx.get_dynamic_input_num,
170 ctx.get_dynamic_input_tensor,
171 )
174def _read_runtime_attr(attrs, index: int, ir_type: str):
175 getter_name, _ = _get_runtime_attr_spec(ir_type, index)
176 return getattr(attrs, getter_name)(index)
179def _build_execute_attrs(ctx: EagerOpExecutionContext, ir_attrs: list) -> dict:
180 if not ir_attrs:
181 return {}
182 attrs = ctx.get_attrs()
183 return {
184 item["name"]: _read_runtime_attr(attrs, index, item["type"])
185 for index, item in enumerate(ir_attrs)
186 }
189def _build_declare_inputs(ctx, ir_inputs: list) -> list:
190 return _build_inputs(
191 ir_inputs,
192 ctx._get_required_input_tensor,
193 ctx._get_optional_input_tensor,
194 ctx._get_dynamic_input_num,
195 ctx._get_dynamic_input_tensor,
196 )
199def _build_declare_outputs(ctx, ir_outputs: list) -> list:
200 args = []
201 for ir_index, item in enumerate(ir_outputs):
202 kind = item["kind"]
203 if kind == OutputType.REQUIRED:
204 args.append(ctx._get_required_output_tensor(ir_index))
205 elif kind == OutputType.DYNAMIC:
206 instance_num = ctx._get_dynamic_output_num(ir_index)
207 args.append(
208 [
209 ctx._get_dynamic_output_tensor(ir_index, relative_index)
210 for relative_index in range(instance_num)
211 ]
212 )
213 else:
214 raise ValueError(f"unsupported custom op IR output kind: {kind}")
215 return args
218def _build_declare_attrs(ctx, ir_attrs: list) -> dict:
219 if not ir_attrs:
220 return {}
221 attrs = ctx._get_attrs()
222 return {
223 item["name"]: _read_runtime_attr(attrs, index, item["type"])
224 for index, item in enumerate(ir_attrs)
225 }
228def call_execute(
229 instance_id: str,
230 ir_meta: Optional[dict],
231 ctx: EagerOpExecutionContext,
232) -> None:
233 try:
234 holder = _get_eager_execute_holder(instance_id)
235 custom_op = holder.instance
236 method = custom_op.execute
237 if ir_meta is None:
238 descriptor = custom_op.__ge_op_impl_descriptor__
239 raise RuntimeError(
240 f"canonical IR not found for schema-bound execute: {descriptor.op_type}"
241 )
242 args = _build_execute_inputs(ctx, ir_meta["inputs"])
243 kwargs = _build_execute_attrs(ctx, ir_meta["attrs"])
244 with _execute_ctx_scope(ctx):
245 result = method(*args, **kwargs)
246 if result is not None:
247 raise TypeError("execute must return None")
248 finally:
249 ctx._invalidate()
252def call_declare_launch_args(instance_id: str, ir_meta: Optional[dict], ctx) -> None:
253 try:
254 holder = _get_holder(instance_id)
255 method = getattr(holder.instance, "declare_launch_args", None)
256 if not callable(method):
257 raise TypeError(
258 f"python op impl does not implement declare_launch_args: {instance_id}"
259 )
260 if ir_meta is None:
261 raise RuntimeError(
262 "canonical IR not found for schema-bound declare_launch_args"
263 )
264 args = _build_declare_inputs(ctx, ir_meta["inputs"])
265 args.extend(_build_declare_outputs(ctx, ir_meta["outputs"]))
266 kwargs = _build_declare_attrs(ctx, ir_meta["attrs"])
267 with _declare_launch_args_ctx_scope(ctx):
268 result = method(*args, **kwargs)
269 if result is not None:
270 raise TypeError("declare_launch_args must return None")
271 finally:
272 ctx._invalidate()
275def clear_op_impl_holders() -> None:
276 with _HOLDER_LOCK:
277 _OP_IMPL_HOLDERS.clear()
280def clear_loaded_op_impl_modules() -> None:
281 """Clear all dynamically loaded op implementation modules from sys.modules to avoid test pollution."""
282 keys_to_remove = [key for key in sys.modules if key.startswith("_ge_py_custom_op_")]
283 for key in keys_to_remove:
284 del sys.modules[key]