Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/_capi/_allocator_callback_adapter.py: 27%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3# ------------------------------------------------------------------- 

4# ----------------------------------------------------------------------------------------------------------- 

5# Copyright (c) 2026 Huawei Technologies Co., Ltd. 

6# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 

7# CANN Open Software License Agreement Version 2.0 (the "License"). 

8# Please refer to the License for details. You may not use this file except in compliance with the License. 

9# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 

10# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. 

11# See LICENSE in the root of the software repository for the full text of the License. 

12# ----------------------------------------------------------------------------------------------------------- 

13 

14"""Private adapter for wiring Python Allocator objects to C callbacks.""" 

15 

16from typing import Dict 

17 

18from ge.allocator import Allocator, MemBlock 

19 

20from .pysession_wrapper import ( 

21 c_func_t_free, 

22 c_func_t_get_addr, 

23 c_func_t_malloc, 

24 c_func_t_on_destroy, 

25) 

26 

27# 防止 Python 回调对象被 GC 回收，C++ 析构时回调清理 

28_prevent_gc: Dict[int, "_AllocatorCCallbacks"] = {} 

29 

30 

31class _AllocatorCCallbacks: 

32 """Adapters a Python Allocator to C function-pointer callbacks via ctypes.""" 

33 

34 def __init__(self, allocator: Allocator): 

35 self._allocator = allocator 

36 # 持有引用防止 GC 回收；C++ 回调 free 时释放。 

37 self._blocks: Dict[int, MemBlock] = {} 

38 

39 def _malloc(_py_obj, size): 

40 blk = self._allocator.malloc(size) 

41 key = id(blk) 

42 self._blocks[key] = blk 

43 return key 

44 

45 def _free(_py_obj, block_key): 

46 blk = self._blocks.pop(block_key, None) 

47 if blk is not None: 

48 self._allocator.free(blk) 

49 

50 def _get_addr(block_key): 

51 blk = self._blocks.get(block_key) 

52 return blk.addr if blk else 0 

53 

54 self.c_malloc = c_func_t_malloc(_malloc) 

55 self.c_free = c_func_t_free(_free) 

56 self.c_get_addr = c_func_t_get_addr(_get_addr) 

57 

58 

59def _on_allocator_destroy(prevent_gc_key): 

60 """Called from C++ ~PyCallbackAllocator to release Python callback refs.""" 

61 cb = _prevent_gc.pop(prevent_gc_key, None) 

62 if cb is not None: 

63 cb.c_malloc = None 

64 cb.c_free = None 

65 cb.c_get_addr = None 

66 

67 

68_c_on_allocator_destroy = c_func_t_on_destroy(_on_allocator_destroy) 

69 

70 

71def create_allocator_c_callbacks(allocator: Allocator): 

72 """Create and retain ctypes callbacks for a Python allocator.""" 

73 cb = _AllocatorCCallbacks(allocator) 

74 prevent_gc_key = id(cb) 

75 _prevent_gc[prevent_gc_key] = cb 

76 return cb, prevent_gc_key, _c_on_allocator_destroy 

77 

78 

79def rollback_allocator_c_callbacks(prevent_gc_key: int) -> None: 

80 """Rollback retained ctypes callbacks when C++ registration fails.""" 

81 _prevent_gc.pop(prevent_gc_key, None)