Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/_bridge.py: 96%
174 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-28 11:24 +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 (
31 _compile_ctx_scope,
32 _declare_launch_args_ctx_scope,
33 _execute_ctx_scope,
34)
35from .registry import (
36 INTERFACE_ANNOTATED_ARGS,
37 INTERFACE_COMPILABLE,
38 INTERFACE_EAGER_EXECUTE,
39 get_registered_op_impl_by_descriptor_key,
40)
43@dataclass
44class _OpImplHolder:
45 descriptor_key: str
46 instance_id: str
47 instance: object
50_HOLDER_LOCK = threading.RLock()
51_OP_IMPL_HOLDERS: Dict[str, _OpImplHolder] = {}
54def load_and_get_op_impl_descriptors() -> list:
55 load_custom_op_plugins()
56 return get_registered_op_impls()
59def load_and_get_op_descriptors() -> dict:
60 load_custom_op_plugins()
61 return {
62 "protos": get_registered_op_protos(),
63 "impls": get_registered_op_impls(),
64 }
67def _get_holder(instance_id: str) -> _OpImplHolder:
68 with _HOLDER_LOCK:
69 holder = _OP_IMPL_HOLDERS.get(instance_id)
70 if holder is None:
71 raise KeyError(f"python op impl holder is not created: {instance_id}")
72 return holder
75def _get_eager_execute_holder(instance_id: str) -> _OpImplHolder:
76 holder = _get_holder(instance_id)
77 if not callable(getattr(holder.instance, "execute", None)):
78 raise TypeError(
79 f"python op impl does not implement callable execute: {instance_id}"
80 )
81 return holder
84def create_op_impl_holder(instance_id: str, descriptor_key: str) -> bool:
85 descriptor = get_registered_op_impl_by_descriptor_key(descriptor_key)
86 if descriptor is None:
87 raise KeyError(f"python op impl descriptor_key not found: {descriptor_key}")
88 with _HOLDER_LOCK:
89 if instance_id in _OP_IMPL_HOLDERS:
90 return True
91 _OP_IMPL_HOLDERS[instance_id] = _OpImplHolder(
92 descriptor_key=descriptor_key,
93 instance_id=instance_id,
94 instance=descriptor.cls(),
95 )
96 return True
99def destroy_op_impl_holder(instance_id: str) -> bool:
100 with _HOLDER_LOCK:
101 return _OP_IMPL_HOLDERS.pop(instance_id, None) is not None
104def _get_callback_for_signature(cls, method_name: str):
105 method = inspect.getattr_static(cls, method_name)
106 if isinstance(method, staticmethod):
107 return method.__func__
108 if isinstance(method, classmethod):
109 return method.__get__(None, cls)
110 if inspect.isfunction(method):
111 return method.__get__(object(), cls)
112 return getattr(cls, method_name)
115def validate_op_impl_descriptor(
116 descriptor_key: str,
117 ir_meta: Optional[dict],
118) -> bool:
119 descriptor = get_registered_op_impl_by_descriptor_key(descriptor_key)
120 if descriptor is None:
121 raise KeyError(f"python op impl descriptor_key not found: {descriptor_key}")
123 if INTERFACE_EAGER_EXECUTE in descriptor.interfaces:
124 if ir_meta is None:
125 raise RuntimeError(
126 f"canonical IR not found for schema-bound execute: {descriptor.op_type}"
127 )
128 method = _get_callback_for_signature(descriptor.cls, "execute")
129 _validate_args_signature(method, ir_meta, descriptor, method_name="execute")
131 if INTERFACE_COMPILABLE in descriptor.interfaces:
132 if ir_meta is None:
133 raise RuntimeError("canonical IR not found for schema-bound compile")
134 method = _get_callback_for_signature(descriptor.cls, "compile")
135 _validate_args_signature(method, ir_meta, descriptor, method_name="compile")
137 if INTERFACE_ANNOTATED_ARGS in descriptor.interfaces:
138 if ir_meta is None:
139 raise RuntimeError(
140 "canonical IR not found for schema-bound declare_launch_args"
141 )
142 method = _get_callback_for_signature(descriptor.cls, "declare_launch_args")
143 _validate_args_signature(
144 method, ir_meta, descriptor, method_name="declare_launch_args"
145 )
146 return True
149def _build_inputs(
150 ir_inputs: list,
151 get_required_input,
152 get_optional_input,
153 get_dynamic_input_num,
154 get_dynamic_input,
155) -> list:
156 args = []
157 for ir_index, item in enumerate(ir_inputs):
158 kind = item["kind"]
159 if kind == InputType.REQUIRED:
160 args.append(get_required_input(ir_index))
161 elif kind == InputType.OPTIONAL:
162 args.append(get_optional_input(ir_index))
163 elif kind == InputType.DYNAMIC:
164 instance_num = get_dynamic_input_num(ir_index)
165 args.append(
166 [
167 get_dynamic_input(ir_index, relative_index)
168 for relative_index in range(instance_num)
169 ]
170 )
171 else:
172 raise ValueError(
173 f"unsupported custom op IR input kind: {kind}, ir index: {ir_index}"
174 )
175 return args
178def _build_execute_inputs(ctx: EagerOpExecutionContext, ir_inputs: list) -> list:
179 return _build_inputs(
180 ir_inputs,
181 ctx.get_required_input_tensor,
182 ctx.get_optional_input_tensor,
183 ctx.get_dynamic_input_num,
184 ctx.get_dynamic_input_tensor,
185 )
188def _read_runtime_attr(attrs, index: int, ir_type: str):
189 getter_name, _ = _get_runtime_attr_spec(ir_type, index)
190 return getattr(attrs, getter_name)(index)
193def _build_execute_attrs(ctx: EagerOpExecutionContext, ir_attrs: list) -> dict:
194 if not ir_attrs:
195 return {}
196 attrs = ctx.get_attrs()
197 return {
198 item["name"]: _read_runtime_attr(attrs, index, item["type"])
199 for index, item in enumerate(ir_attrs)
200 }
203def _build_schema_inputs(ctx, ir_inputs: list) -> list:
204 return _build_inputs(
205 ir_inputs,
206 ctx._get_required_input_tensor,
207 ctx._get_optional_input_tensor,
208 ctx._get_dynamic_input_num,
209 ctx._get_dynamic_input_tensor,
210 )
213def _build_schema_outputs(ctx, ir_outputs: list) -> list:
214 args = []
215 for ir_index, item in enumerate(ir_outputs):
216 kind = item["kind"]
217 if kind == OutputType.REQUIRED:
218 args.append(ctx._get_required_output_tensor(ir_index))
219 elif kind == OutputType.DYNAMIC:
220 instance_num = ctx._get_dynamic_output_num(ir_index)
221 args.append(
222 [
223 ctx._get_dynamic_output_tensor(ir_index, relative_index)
224 for relative_index in range(instance_num)
225 ]
226 )
227 else:
228 raise ValueError(f"unsupported custom op IR output kind: {kind}")
229 return args
232def _build_schema_attrs(ctx, ir_attrs: list) -> dict:
233 if not ir_attrs:
234 return {}
235 attrs = ctx._get_attrs()
236 return {
237 item["name"]: _read_runtime_attr(attrs, index, item["type"])
238 for index, item in enumerate(ir_attrs)
239 }
242def call_execute(
243 instance_id: str,
244 ir_meta: Optional[dict],
245 ctx: EagerOpExecutionContext,
246) -> None:
247 try:
248 holder = _get_eager_execute_holder(instance_id)
249 custom_op = holder.instance
250 method = custom_op.execute
251 if ir_meta is None:
252 descriptor = custom_op.__ge_op_impl_descriptor__
253 raise RuntimeError(
254 f"canonical IR not found for schema-bound execute: {descriptor.op_type}"
255 )
256 args = _build_execute_inputs(ctx, ir_meta["inputs"])
257 kwargs = _build_execute_attrs(ctx, ir_meta["attrs"])
258 with _execute_ctx_scope(ctx):
259 result = method(*args, **kwargs)
260 if result is not None:
261 raise TypeError("execute must return None")
262 finally:
263 ctx._invalidate()
266def call_declare_launch_args(instance_id: str, ir_meta: Optional[dict], ctx) -> None:
267 try:
268 holder = _get_holder(instance_id)
269 method = getattr(holder.instance, "declare_launch_args", None)
270 if not callable(method):
271 raise TypeError(
272 f"python op impl does not implement declare_launch_args: {instance_id}"
273 )
274 if ir_meta is None:
275 raise RuntimeError(
276 "canonical IR not found for schema-bound declare_launch_args"
277 )
278 args = _build_schema_inputs(ctx, ir_meta["inputs"])
279 args.extend(_build_schema_outputs(ctx, ir_meta["outputs"]))
280 kwargs = _build_schema_attrs(ctx, ir_meta["attrs"])
281 with _declare_launch_args_ctx_scope(ctx):
282 result = method(*args, **kwargs)
283 if result is not None:
284 raise TypeError("declare_launch_args must return None")
285 finally:
286 ctx._invalidate()
289def call_compile(instance_id: str, ir_meta: Optional[dict], ctx) -> None:
290 """Invoke a schema-bound Python compile callback."""
292 try:
293 holder = _get_holder(instance_id)
294 method = getattr(holder.instance, "compile", None)
295 if not callable(method):
296 raise TypeError(f"python op impl does not implement compile: {instance_id}")
297 if ir_meta is None:
298 raise RuntimeError("canonical IR not found for schema-bound compile")
299 descriptor = holder.instance.__ge_op_impl_descriptor__
300 _validate_args_signature(method, ir_meta, descriptor, method_name="compile")
301 args = _build_schema_inputs(ctx, ir_meta["inputs"])
302 args.extend(_build_schema_outputs(ctx, ir_meta["outputs"]))
303 kwargs = _build_schema_attrs(ctx, ir_meta["attrs"])
304 with _compile_ctx_scope(ctx):
305 result = method(*args, **kwargs)
306 if result is not None:
307 raise TypeError("compile must return None")
308 finally:
309 ctx._invalidate()
312def clear_op_impl_holders() -> None:
313 with _HOLDER_LOCK:
314 _OP_IMPL_HOLDERS.clear()
317def clear_loaded_op_impl_modules() -> None:
318 """Clear all dynamically loaded op implementation modules from sys.modules to avoid test pollution."""
319 keys_to_remove = [key for key in sys.modules if key.startswith("_ge_py_custom_op_")]
320 for key in keys_to_remove:
321 del sys.modules[key]