ZeroLink/test_basic_functionality.py at master · nanocubit/ZeroLink · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
"""
test_basic_functionality.py

Простой тест основной функциональности ZeroLink v2.0.
"""

import torch
from zerolink.runtime.unified import UnifiedRuntime


def test_unified_runtime_creation():
    """Тест создания UnifiedRuntime."""
    print("Тест 1: Создание UnifiedRuntime...")

    # Создаем runtime с минимальными параметрами
    runtime = UnifiedRuntime(
        gpu_device_id=0,
        gpu_pool_size_gb=1,
        num_cpu_workers=1
    )

    print("✓ UnifiedRuntime успешно создан")

    # Проверяем, что атрибуты установлены
    assert hasattr(runtime, 'gpu_pool')
    assert hasattr(runtime, 'cpu_executor')

    print("✓ Атрибуты runtime проверены")

    # Проверяем статус GPU пула (может быть None если CUDA недоступна)
    print(f"✓ GPU Pool: {'Available' if runtime.gpu_pool is not None else 'Not available (CUDA)'}")

    # Останавливаем runtime
    runtime.stop()
    print("✓ UnifiedRuntime успешно остановлен")


def test_cpu_execution():
    """Тест выполнения CPU задач."""
    print("\nТест 2: Выполнение CPU задач...")

    runtime = UnifiedRuntime(
        gpu_device_id=0,
        gpu_pool_size_gb=1,
        num_cpu_workers=2
    )

    # Простая CPU задача
    def simple_cpu_task(x):
        return x * 2

    # Выполняем задачу
    result = runtime.execute_cpu(simple_cpu_task, 21)
    assert result == 42
    print(f"✓ CPU задача выполнена: 21 * 2 = {result}")

    # Тест map
    inputs = [1, 2, 3, 4]
    results = runtime.map_cpu(simple_cpu_task, inputs)
    expected = [2, 4, 6, 8]
    assert results == expected
    print(f"✓ CPU map задачи выполнены: {inputs} -> {results}")

    runtime.stop()
    print("✓ CPU тесты пройдены")


def test_gpu_execution_if_available():
    """Тест выполнения GPU задач (если доступна)."""
    print("\nТест 3: Выполнение GPU задач...")

    runtime = UnifiedRuntime(
        gpu_device_id=0,
        gpu_pool_size_gb=1,
        num_cpu_workers=1
    )

    if torch.cuda.is_available():
        print("  CUDA доступна, тестируем GPU функциональность...")

        # Простая GPU задача
        def simple_gpu_task(size):
            t = torch.randn(size, size, device='cuda')
            return t.sum().item()

        try:
            result = runtime.execute_gpu(simple_gpu_task, 32)  # Маленькая матрица для теста
            print(f"  ✓ GPU задача выполнена, результат: {result:.4f}")
        except Exception as e:
            print(f"  ⚠ GPU задача не выполнена (ожидаемо): {e}")
    else:
        print("  CUDA недоступна, пропускаем GPU тесты")

    runtime.stop()
    print("✓ GPU тесты завершены")


def test_get_pool_stats():
    """Тест получения статистики пула."""
    print("\nТест 4: Получение статистики пула...")

    runtime = UnifiedRuntime(
        gpu_device_id=0,
        gpu_pool_size_gb=1,
        num_cpu_workers=1
    )

    stats = runtime.get_pool_stats()
    assert isinstance(stats, dict)
    assert 'cpu' in stats
    print("✓ Статистика пула успешно получена")
    print(f"  CPU статистика: {stats['cpu']}")

    runtime.stop()
    print("✓ Тест статистики завершен")


if __name__ == "__main__":
    print("=== Тестирование основной функциональности ZeroLink v2.0 ===\n")

    try:
        test_unified_runtime_creation()
        test_cpu_execution()
        test_gpu_execution_if_available()
        test_get_pool_stats()

        print("\n=== Все тесты пройдены успешно! ===")
        print("ZeroLink v2.0 работает корректно.")
        print("Основные компоненты функционируют как ожидалось.")

    except Exception as e:
        print(f"\n❌ Ошибка в тестах: {e}")
        import traceback
        traceback.print_exc()