|
| 1 | +# Copyright OpenSearch Contributors |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | +# |
| 4 | +# The OpenSearch Contributors require contributions made to |
| 5 | +# this file be licensed under the Apache-2.0 license or a |
| 6 | +# compatible open source license. |
| 7 | + |
| 8 | +import threading |
| 9 | + |
| 10 | + |
| 11 | +class ResourceManager: |
| 12 | + """ |
| 13 | + A thread-safe resource manager that tracks and manages GPU and CPU memory allocations. |
| 14 | +
|
| 15 | + This class provides mechanisms to safely allocate and release memory resources |
| 16 | + in a multi-threaded environment, ensuring that memory usage doesn't exceed |
| 17 | + the specified limits. |
| 18 | +
|
| 19 | + Attributes: |
| 20 | + _total_gpu_memory (float): Total available GPU memory in bytes |
| 21 | + _total_cpu_memory (float): Total available CPU memory in bytes |
| 22 | + _available_gpu_memory (float): Currently available GPU memory in bytes |
| 23 | + _available_cpu_memory (float): Currently available CPU memory in bytes |
| 24 | + _lock (threading.Lock): Thread lock for synchronization |
| 25 | + """ |
| 26 | + |
| 27 | + def __init__(self, total_gpu_memory: float, total_cpu_memory: float): |
| 28 | + """ |
| 29 | + Initialize the ResourceManager with specified GPU and CPU memory limits. |
| 30 | +
|
| 31 | + Args: |
| 32 | + total_gpu_memory (float): Total GPU memory available for allocation, in bytes |
| 33 | + total_cpu_memory (float): Total CPU memory available for allocation, in bytes |
| 34 | + """ |
| 35 | + self._total_gpu_memory = total_gpu_memory |
| 36 | + self._total_cpu_memory = total_cpu_memory |
| 37 | + self._available_gpu_memory = total_gpu_memory |
| 38 | + self._available_cpu_memory = total_cpu_memory |
| 39 | + self._lock = threading.Lock() |
| 40 | + |
| 41 | + def allocate(self, gpu_memory: float, cpu_memory: float) -> bool: |
| 42 | + """ |
| 43 | + Attempt to allocate the specified amount of GPU and CPU memory. |
| 44 | +
|
| 45 | + Args: |
| 46 | + gpu_memory (float): Amount of GPU memory to allocate, in bytes |
| 47 | + cpu_memory (float): Amount of CPU memory to allocate, in bytes |
| 48 | +
|
| 49 | + Returns: |
| 50 | + bool: True if allocation was successful, False if insufficient resources |
| 51 | + """ |
| 52 | + with self._lock: |
| 53 | + if not ( |
| 54 | + self._available_gpu_memory >= gpu_memory |
| 55 | + and self._available_cpu_memory >= cpu_memory |
| 56 | + ): |
| 57 | + return False |
| 58 | + self._available_gpu_memory -= gpu_memory |
| 59 | + self._available_cpu_memory -= cpu_memory |
| 60 | + return True |
| 61 | + |
| 62 | + def release(self, gpu_memory: float, cpu_memory: float) -> None: |
| 63 | + """ |
| 64 | + Release previously allocated GPU and CPU memory back to the pool. |
| 65 | +
|
| 66 | + Args: |
| 67 | + gpu_memory (float): Amount of GPU memory to release, in bytes |
| 68 | + cpu_memory (float): Amount of CPU memory to release, in bytes |
| 69 | + """ |
| 70 | + with self._lock: |
| 71 | + self._available_gpu_memory += gpu_memory |
| 72 | + self._available_cpu_memory += cpu_memory |
| 73 | + |
| 74 | + def get_available_gpu_memory(self) -> float: |
| 75 | + """ |
| 76 | + Get the current amount of available GPU memory. |
| 77 | +
|
| 78 | + Returns: |
| 79 | + float: Amount of available GPU memory in bytes |
| 80 | + """ |
| 81 | + with self._lock: |
| 82 | + return self._available_gpu_memory |
| 83 | + |
| 84 | + def get_available_cpu_memory(self) -> float: |
| 85 | + """ |
| 86 | + Get the current amount of available GPU memory. |
| 87 | +
|
| 88 | + Returns: |
| 89 | + float: Amount of available GPU memory in bytes |
| 90 | + """ |
| 91 | + with self._lock: |
| 92 | + return self._available_cpu_memory |
0 commit comments