Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/allocator/allocator.py: 67%

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"""Allocator module for external memory allocation in GraphEngine.""" 

15 

16from abc import ABC, abstractmethod 

17 

18 

19class MemBlock: 

20 """Represents a block of memory allocated by an Allocator. 

21 

22 Attributes: 

23 addr: Device memory address (integer). 

24 size: Size in bytes. 

25 """ 

26 

27 def __init__(self, addr: int, size: int): 

28 self._addr = addr 

29 self._size = size 

30 

31 @property 

32 def addr(self) -> int: 

33 return self._addr 

34 

35 @property 

36 def size(self) -> int: 

37 return self._size 

38 

39 

40class Allocator(ABC): 

41 """Base class for external allocators. 

42 

43 Subclass this to implement custom device memory allocation strategies. 

44 The allocator is registered to a stream via Session.register_external_allocator(). 

45 """ 

46 

47 @abstractmethod 

48 def malloc(self, size: int) -> MemBlock: 

49 """Allocate device memory of the given size. 

50 

51 Args: 

52 size: Number of bytes to allocate. 

53 

54 Returns: 

55 A MemBlock whose addr is a valid device memory address. 

56 

57 Raises: 

58 MemoryError: If allocation fails. 

59 """ 

60 raise NotImplementedError 

61 

62 @abstractmethod 

63 def free(self, block: MemBlock) -> None: 

64 """Free memory previously allocated by malloc().""" 

65 raise NotImplementedError