examples/optimized_gpu_pool.py - master - nanocubit/ZeroLink

examples/optimized_gpu_pool.py
153 строки6 KB
slavam
feat: добавлена рабочая реализация GPU пула памяти, обновлены все компоненты ZeroLink
23 янв 2026, 20:56
23 янв 2026, 20:560106973
"""
Пример оптимизированного GPU Memory Pool для ZeroLink

Этот файл демонстрирует потенциальные оптимизации для
zerolink/core/gpu/vmm_pool.py
"""

import torch
import threading
from typing import Dict, Optional, List
import queue
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor


@dataclass
class AllocationRequest:
    """Запрос на аллокацию памяти"""
    size: int
    callback: callable
    request_id: str


class OptimizedGPUMemoryPool:
    """
    Оптимизированный GPU Memory Pool с использованием lock-free очередей
    и предварительного выделения памяти
    """
    
    def __init__(self, device_id: int = 0, total_size_gb: int = 4):
        self.device_id = device_id
        self.total_size_bytes = total_size_gb * 1024**3
        self.available_size_bytes = self.total_size_bytes
        
        # Lock-free очередь для запросов на аллокацию
        self.allocation_queue = queue.Queue()
        self.free_queue = queue.Queue()
        
        # Предварительно выделенные блоки памяти
        self.preallocated_blocks = []
        self.block_size = 256 * 1024 * 1024  # 256MB блоки
        
        # Пул потоков для обработки запросов
        self.executor = ThreadPoolExecutor(max_workers=4)
        
        # Инициализация предварительно выделенных блоков
        self._preallocate_memory()
        
        # Запуск обработчика очереди
        self._start_allocator_thread()
    
    def _preallocate_memory(self):
        """Предварительное выделение блоков памяти"""
        num_blocks = self.total_size_bytes // self.block_size
        for i in range(num_blocks):
            try:
                block = torch.cuda.FloatTensor(self.block_size // 4)  # 4 байта на float
                self.preallocated_blocks.append({
                    'tensor': block,
                    'offset': 0,
                    'size': self.block_size,
                    'available': True
                })
                self.free_queue.put(i)
            except RuntimeError:
                # Если не хватает памяти, прекращаем выделение
                break
    
    def _start_allocator_thread(self):
        """Запуск потока для обработки запросов на аллокацию"""
        def allocator_worker():
            while True:
                try:
                    request = self.allocation_queue.get(timeout=1)
                    if request is None:  # Сигнал остановки
                        break
                    
                    # Обработка запроса на аллокацию
                    result = self._process_allocation(request.size)
                    request.callback(result)
                    
                    self.allocation_queue.task_done()
                except queue.Empty:
                    continue
                except Exception as e:
                    print(f"Ошибка в обработчике аллокации: {e}")
        
        self.allocator_thread = threading.Thread(target=allocator_worker, daemon=True)
        self.allocator_thread.start()
    
    def _process_allocation(self, size: int):
        """Обработка запроса на аллокацию"""
        if size > self.block_size:
            # Для больших запросов используем обычную аллокацию
            try:
                tensor = torch.cuda.FloatTensor(size // 4)
                return {'tensor': tensor, 'type': 'direct'}
            except RuntimeError:
                return None
        
        # Для маленьких запросов используем предварительно выделенные блоки
        try:
            block_idx = self.free_queue.get_nowait()
            block = self.preallocated_blocks[block_idx]
            
            if block['available']:
                # Возвращаем подмножество предварительно выделенного тензора
                subset = block['tensor'][block['offset']:(block['offset'] + size//4)]
                block['offset'] += size//4
                
                if block['offset'] >= block['size'] // 4:
                    block['available'] = False
                
                return {
                    'tensor': subset,
                    'type': 'preallocated',
                    'block_idx': block_idx
                }
        except queue.Empty:
            return None
    
    def allocate_async(self, size: int, callback: callable):
        """Асинхронная аллокация с использованием очереди"""
        request = AllocationRequest(size, callback, f"req_{id(callback)}")
        self.allocation_queue.put(request)
    
    def cleanup(self):
        """Очистка ресурсов"""
        self.allocation_queue.put(None)  # Сигнал остановки
        self.executor.shutdown(wait=True)


# Пример использования оптимизированного пула
def example_usage():
    """Пример использования оптимизированного пула"""
    pool = OptimizedGPUMemoryPool(device_id=0, total_size_gb=2)
    
    def handle_allocation(result):
        if result:
            print(f"Аллокация выполнена: {result['type']}")
            # Использование тензора...
        else:
            print("Не удалось выполнить аллокацию")
    
    # Асинхронная аллокация
    pool.allocate_async(1024*1024, handle_allocation)  # 1MB
    
    # Очистка
    pool.cleanup()


if __name__ == "__main__":
    example_usage()