Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/_bridge.py: 94%

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# ----------------------------------------------------------------------------------------------------------- 

12 

13"""Bridge-facing Python runtime helpers for GE custom op implementations.""" 

14 

15import sys 

16import threading 

17from dataclasses import dataclass 

18from typing import Dict 

19 

20from .base import BaseCustomOp, EagerExecuteOp, EagerOpExecutionContext 

21from .bootstrap import get_registered_op_impls, load_custom_op_plugins 

22from .registry import get_registered_op_impl_by_descriptor_key 

23 

24 

25@dataclass 

26class _OpImplHolder: 

27 descriptor_key: str 

28 instance_id: str 

29 instance: BaseCustomOp 

30 

31 

32_HOLDER_LOCK = threading.RLock() 

33_OP_IMPL_HOLDERS: Dict[str, _OpImplHolder] = {} 

34 

35 

36def load_and_get_op_impl_descriptors() -> list: 

37 load_custom_op_plugins() 

38 return get_registered_op_impls() 

39 

40 

41def _get_holder(instance_id: str) -> _OpImplHolder: 

42 with _HOLDER_LOCK: 

43 holder = _OP_IMPL_HOLDERS.get(instance_id) 

44 if holder is None: 

45 raise KeyError(f"python op impl holder is not created: {instance_id}") 

46 return holder 

47 

48 

49def _get_eager_execute_op(instance_id: str) -> EagerExecuteOp: 

50 instance = _get_holder(instance_id).instance 

51 if not isinstance(instance, EagerExecuteOp): 

52 raise TypeError( 

53 f"python op impl does not implement EagerExecuteOp: {instance_id}" 

54 ) 

55 return instance 

56 

57 

58def create_op_impl_holder(instance_id: str, descriptor_key: str) -> bool: 

59 descriptor = get_registered_op_impl_by_descriptor_key(descriptor_key) 

60 if descriptor is None: 

61 raise KeyError(f"python op impl descriptor_key not found: {descriptor_key}") 

62 with _HOLDER_LOCK: 

63 if instance_id in _OP_IMPL_HOLDERS: 

64 return True 

65 _OP_IMPL_HOLDERS[instance_id] = _OpImplHolder( 

66 descriptor_key=descriptor_key, 

67 instance_id=instance_id, 

68 instance=descriptor.cls(), 

69 ) 

70 return True 

71 

72 

73def destroy_op_impl_holder(instance_id: str) -> bool: 

74 with _HOLDER_LOCK: 

75 return _OP_IMPL_HOLDERS.pop(instance_id, None) is not None 

76 

77 

78def call_execute(instance_id: str, ctx: EagerOpExecutionContext) -> None: 

79 try: 

80 custom_op = _get_eager_execute_op(instance_id) 

81 custom_op.execute(ctx) 

82 finally: 

83 ctx._invalidate() 

84 

85 

86def clear_op_impl_holders() -> None: 

87 with _HOLDER_LOCK: 

88 _OP_IMPL_HOLDERS.clear() 

89 

90 

91def clear_loaded_op_impl_modules() -> None: 

92 """Clear all dynamically loaded op implementation modules from sys.modules to avoid test pollution.""" 

93 keys_to_remove = [key for key in sys.modules if key.startswith("_ge_py_custom_op_")] 

94 for key in keys_to_remove: 

95 del sys.modules[key]