benchmark_optimizations.py - master - nanocubit/ZeroLink

benchmark_optimizations.py
132 строки5 KB
slavam
feat: добавлена рабочая реализация GPU пула памяти, обновлены все компоненты ZeroLink
23 янв 2026, 20:56
23 янв 2026, 20:560106973
#!/usr/bin/env python3
"""
Скрипт для бенчмаркинга оптимизаций ZeroLink
"""

import time
import torch
import numpy as np
from examples.optimized_gpu_pool import OptimizedGPUMemoryPool
from examples.optimized_protocol import OptimizedProtocol, BatchMessageProcessor


def benchmark_gpu_pool():
    """Бенчмарк оптимизированного GPU пула"""
    print("=== Бенчмарк GPU Memory Pool ===")
    
    # Создаем оптимизированный пул
    start_time = time.time()
    pool = OptimizedGPUMemoryPool(device_id=0, total_size_gb=1)
    pool_init_time = time.time() - start_time
    print(f"Инициализация пула: {pool_init_time:.4f}s")
    
    # Тест аллокации
    allocations = []
    start_time = time.time()
    
    for i in range(100):
        def callback(result):
            pass  # Заглушка для асинхронного вызова
        
        pool.allocate_async(1024*1024, callback)  # 1MB аллокации
        allocations.append(time.time())
    
    alloc_time = time.time() - start_time
    print(f"100 аллокаций по 1MB: {alloc_time:.4f}s")
    print(f"Среднее время аллокации: {alloc_time/100*1000:.2f}ms")
    
    pool.cleanup()
    print()


def benchmark_protocol():
    """Бенчмарк оптимизированного протокола"""
    print("=== Бенчмарк Protocol ===")
    
    protocol = OptimizedProtocol()
    
    # Тест упаковки/распаковки
    test_payload = b"Hello, ZeroLink!" * 1000  # 13KB
    
    # Измеряем упаковку
    start_time = time.time()
    for i in range(1000):
        packed = protocol.pack_message_optimized(1, test_payload)
    pack_time = time.time() - start_time
    print(f"1000 упаковок: {pack_time:.4f}s")
    print(f"Среднее время упаковки: {pack_time/1000*1000:.2f}ms")
    
    # Измеряем распаковку
    start_time = time.time()
    for i in range(1000):
        version, unpacked = protocol.unpack_message_optimized(packed)
    unpack_time = time.time() - start_time
    print(f"1000 распаковок: {unpack_time:.4f}s")
    print(f"Среднее время распаковки: {unpack_time/1000*1000:.2f}ms")
    
    # Тест сериализации тензора
    tensor = torch.randn(1000, 1000)  # 1000x1000 тензор
    start_time = time.time()
    for i in range(100):
        serialized = protocol.serialize_tensor_optimized(tensor)
    serialize_time = time.time() - start_time
    print(f"100 сериализаций тензора 1000x1000: {serialize_time:.4f}s")
    print(f"Среднее время сериализации: {serialize_time/100*1000:.2f}ms")
    
    print()


def benchmark_batch_processing():
    """Бенчмарк пакетной обработки"""
    print("=== Бенчмарк Batch Processing ===")
    
    batch_processor = BatchMessageProcessor(max_batch_size=10)
    
    # Тест пакетной обработки
    start_time = time.time()
    for i in range(50):
        result = batch_processor.add_message(i % 10, f"Message {i}" * 100)
        if result:
            # Пакет обработан
            pass
    
    # Обрабатываем оставшиеся сообщения
    remaining = batch_processor.process_batch()
    
    batch_time = time.time() - start_time
    print(f"50 сообщений с пакетной обработкой: {batch_time:.4f}s")
    print()


def compare_with_baseline():
    """Сравнение с базовой реализацией (гипотетической)"""
    print("=== Сравнение с базовой реализацией ===")
    print("Примечание: Реальные улучшения зависят от конкретной реализации")
    print("- Lock-free очереди уменьшают contention")
    print("- Memoryview уменьшает копирование данных")
    print("- Предварительное выделение уменьшает вызовы CUDA API")
    print("- Пакетная обработка уменьшает накладные расходы")
    print()


if __name__ == "__main__":
    print("Бенчмарк оптимизаций ZeroLink v2.0\n")
    
    try:
        benchmark_gpu_pool()
    except Exception as e:
        print(f"Ошибка в бенчмарке GPU пула: {e}")
    
    try:
        benchmark_protocol()
    except Exception as e:
        print(f"Ошибка в бенчмарке протокола: {e}")
    
    try:
        benchmark_batch_processing()
    except Exception as e:
        print(f"Ошибка в бенчмарке пакетной обработки: {e}")
    
    compare_with_baseline()
    
    print("Бенчмарк завершен. Результаты можно использовать для дальнейшей оптимизации.")