ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Настроить mmap для embeddings

1. Цель задачи

Научиться использовать numpy.mmap для работы с большими наборами векторных эмбеддингов, которые не помещаются в оперативную память. Реализовать эффективный поиск ближайших соседей (kNN) по 100 ГБ векторов на сервере с 32 ГБ RAM, используя memory-mapped файлы для минимизации потребления памяти и утилизации дискового кэша ОС.

Ключевой результат Рабочий прототип поиска по 100 ГБ эмбеддингов на сервере с 32 ГБ RAM с временем ответа < 1 секунды на запрос (при последовательном доступе к страницам).

2. Исходные данные

Перед началом необходимо иметь:

Если нет реального набора эмбеддингов — симулируем:

1. Генерация синтетических данных

   • Размерность векторов: 768 (как у text-embedding-ada-002 или all-MiniLM-L6-v2)
   • Тип данных: float32 (4 байта на элемент)
   • Количество векторов: 100 ГБ / (768 * 4 байта) ≈ 34 000 000 векторов
   • Формат: бинарный файл (raw bytes) — embeddings.bin

2. Генерация тестовых запросов

   • 1000 случайных векторов той же размерности
   • Сохранить в отдельный файл queries.npy

3. Создание файла эмбеддингов

   import numpy as np
   import os
   
   # Параметры
   dim = 768
   dtype = np.float32
   n_vectors = 34_000_000  # ~100 ГБ
   file_size = n_vectors * dim * np.dtype(dtype).itemsize  # ~100 ГБ
   
   # Создаём файл и заполняем случайными данными (можно делать частями)
   mmap_file = 'embeddings.bin'
   with open(mmap_file, 'wb') as f:
       # Записываем нули, чтобы выделить место на диске
       f.write(b'\x00' * file_size)
   
   # Memory-map и заполнение случайными числами (по частям, чтобы не занять всю RAM)
   chunk_size = 100_000  # векторов за раз
   mmap = np.memmap(mmap_file, dtype=dtype, mode='r+', shape=(n_vectors, dim))
   for start in range(0, n_vectors, chunk_size):
       end = min(start + chunk_size, n_vectors)
       mmap[start:end] = np.random.randn(end - start, dim).astype(dtype)
       mmap.flush()
       print(f'Заполнено {end}/{n_vectors} векторов')
   del mmap  # закрываем mmap
   

3. Технологический стек

4. Этапы выполнения

Этап 1: Подготовка окружения и генерация данных (1-2 часа)

Действия

1. Создать виртуальное окружение и установить зависимости:

   python -m venv venv
   source venv/bin/activate  # или venv\Scripts\activate на Windows
   pip install numpy memory_profiler psutil matplotlib
   

2. Написать скрипт генерации данных

   • Создать файл generate_data.py
   • Реализовать генерацию 100 ГБ эмбеддингов (см. код в разделе 2)
   • Создать файл с тестовыми запросами (1000 векторов)

3. Проверить доступное дисковое пространство

   • Убедиться, что на диске есть > 120 ГБ свободного места
   • Если нет — уменьшить размер датасета (например, до 50 ГБ или 20 ГБ)

4. Запустить генерацию

   python generate_data.py
   

   Примечание: Генерация 100 ГБ может занять 30-60 минут. Можно запустить на ночь.

Ожидаемый результат этапа Файл embeddings.bin (~100 ГБ) и файл queries.npy (~3 МБ).

Этап 2: Реализация базового поиска с mmap (2-3 часа)

Действия

1. Написать класс MmapVectorSearch

   import numpy as np
   import time
   from typing import Tuple, List
   
   class MmapVectorSearch:
       def __init__(self, mmap_file: str, dim: int, dtype=np.float32, mode='r'):
           self.dim = dim
           self.dtype = dtype
           self.mmap_file = mmap_file
           
           # Определяем количество векторов по размеру файла
           import os
           file_size = os.path.getsize(mmap_file)
           self.n_vectors = file_size // (dim * np.dtype(dtype).itemsize)
           
           # Создаём memory-map (только чтение)
           self.mmap = np.memmap(mmap_file, dtype=dtype, mode=mode, shape=(self.n_vectors, dim))
           
           # Предварительно нормализуем векторы (если нужно для cosine similarity)
           # ВАЖНО: это потребует чтения всех данных! Для >RAM не подходит.
           # Вместо этого нормализуем на лету или храним нормы отдельно.
           
       def search(self, query: np.ndarray, k: int = 10) -> Tuple[np.ndarray, np.ndarray]:
           """
           Поиск k ближайших соседей для одного запроса.
           Возвращает индексы и расстояния.
           """
           # Нормализуем запрос
           query = query / np.linalg.norm(query)
           
           # Размер чанка для чтения (чтобы не загружать всё в RAM)
           chunk_size = 100_000  # настраивается
           
           all_distances = []
           all_indices = []
           
           for start in range(0, self.n_vectors, chunk_size):
               end = min(start + chunk_size, self.n_vectors)
               
               # Читаем чанк векторов (это загрузит страницы в page cache)
               chunk = self.mmap[start:end]
               
               # Нормализуем чанк (на лету)
               norms = np.linalg.norm(chunk, axis=1, keepdims=True)
               norms[norms == 0] = 1  # избегаем деления на ноль
               chunk_normalized = chunk / norms
               
               # Cosine similarity = dot product нормализованных векторов
               similarities = np.dot(chunk_normalized, query)
               
               # Находим top-k в этом чанке
               if len(similarities) > k:
                   top_k_indices_chunk = np.argpartition(similarities, -k)[-k:]
                   top_k_similarities = similarities[top_k_indices_chunk]
               else:
                   top_k_indices_chunk = np.arange(len(similarities))
                   top_k_similarities = similarities
               
               # Сохраняем с глобальными индексами
               all_distances.extend(top_k_similarities)
               all_indices.extend(start + top_k_indices_chunk)
           
           # Глобальный top-k
           all_distances = np.array(all_distances)
           all_indices = np.array(all_indices)
           
           global_top_k = np.argpartition(all_distances, -k)[-k:]
           global_top_k = global_top_k[np.argsort(all_distances[global_top_k])[::-1]]
           
           return all_indices[global_top_k], all_distances[global_top_k]
       
       def close(self):
           """Закрыть mmap."""
           del self.mmap
   

2. Оптимизация

   • Проблема Нормализация каждого чанка на лету требует чтения всех данных.
   • Решение 1 Хранить предварительно нормализованные векторы (требует перегенерации данных).
   • Решение 2 Использовать L2-расстояние вместо cosine similarity (не требует нормализации).
   • Решение 3 Хранить нормы отдельно в небольшом файле (8 байт на вектор → 34M * 8 = 272 МБ — помещается в RAM).

3. Реализовать хранение норм отдельно

   # При генерации данных
   norms = np.linalg.norm(embeddings, axis=1)
   np.save('norms.npy', norms)  # ~272 МБ для 34M векторов
   
   # В классе MmapVectorSearch
   self.norms = np.load('norms.npy', mmap_mode='r')  # тоже mmap!
   

Ожидаемый результат этапа Рабочий класс MmapVectorSearch, который может выполнять поиск по 100 ГБ данных, используя < 1 ГБ RAM.

Этап 3: Оптимизация производительности (2-3 часа)

Действия

1. Профилирование текущей реализации

   import time
   from memory_profiler import memory_usage
   
   search = MmapVectorSearch('embeddings.bin', dim=768)
   query = np.random.randn(768).astype(np.float32)
   
   # Замер времени
   start = time.time()
   indices, distances = search.search(query, k=10)
   elapsed = time.time() - start
   print(f'Время поиска: {elapsed:.3f} сек')
   
   # Замер памяти
   mem_usage = memory_usage((search.search, (query,), {'k': 10}))
   print(f'Пиковое потребление памяти: {max(mem_usage):.2f} МБ')
   

2. Оптимизация размера чанка

   • Экспериментировать с chunk_size от 10_000 до 1_000_000
   • Найти оптимальный баланс между временем и памятью
   • Построить график зависимости

3. Использование batch-запросов

   def search_batch(self, queries: np.ndarray, k: int = 10) -> List[Tuple[np.ndarray, np.ndarray]]:
       """Поиск для нескольких запросов одновременно (эффективнее)."""
       results = []
       for query in queries:
           results.append(self.search(query, k))
       return results
   

4. Предзагрузка page cache

   # Linux: предзагрузить файл в page cache (если достаточно места)
   vmtouch -t embeddings.bin
   

   Примечание: Это займёт ~100 ГБ page cache, но ОС может вытеснять страницы при нехватке памяти.

5. Сравнение с FAISS (опционально):

   import faiss
   
   # Создаём индекс FAISS (IVF для >RAM)
   index = faiss.index_factory(768, "IVF100,SQ8")
   # FAISS не поддерживает mmap напрямую для >RAM, но можно использовать
   # faiss.read_index_with_options с mmap=True для индекса на диске
   

Ожидаемый результат этапа Оптимизированная версия с временем поиска < 1 секунды на запрос.

Этап 4: Тестирование и валидация (1-2 часа)

Действия

1. Тест на корректность

   • Создать маленький датасет (1000 векторов, 10 МБ)
   • Сравнить результаты mmap-поиска с поиском в RAM (через np.dot)
   • Проверить, что индексы и расстояния совпадают

2. Тест на потребление памяти

   import psutil
   
   process = psutil.Process()
   mem_before = process.memory_info().rss / 1024 / 1024  # МБ
   
   search = MmapVectorSearch('embeddings.bin', dim=768)
   mem_after_open = process.memory_info().rss / 1024 / 1024
   print(f'Память после открытия mmap: {mem_after_open - mem_before:.2f} МБ')
   
   # Выполнить поиск
   indices, distances = search.search(query, k=10)
   mem_after_search = process.memory_info().rss / 1024 / 1024
   print(f'Память после поиска: {mem_after_search - mem_before:.2f} МБ')
   

3. Тест на производительность

   • Выполнить 100 запросов, замерить среднее время
   • Построить график времени от количества векторов (10%, 25%, 50%, 75%, 100% данных)
   • Проверить, что время растёт линейно с размером данных

4. Стресс-тест

   • Запустить 10 параллельных потоков поиска
   • Проверить, что потребление памяти не превышает 2-3 ГБ
   • Проверить, что нет взаимных блокировок (mmap thread-safe для чтения)

Ожидаемый результат этапа Подтверждение, что решение работает корректно и укладывается в ограничения по памяти.

Этап 5: Документирование и оформление (1 час)

Действия

1. Написать README

   • Описание задачи и подхода
   • Инструкция по запуску
   • Результаты бенчмарков

2. Оформить код

   • Добавить docstrings
   • Добавить type hints
   • Разделить на модули: data_generator.py, mmap_search.py, benchmark.py

3. Создать скрипт для воспроизведения

   # reproduce.sh
   python generate_data.py --size 10  # 10 ГБ для быстрого теста
   python benchmark.py --data embeddings_10gb.bin --queries queries.npy
   

Ожидаемый результат этапа Чистый, документированный код в Git-репозитории.

5. Критерии приемки (Definition of Done)

• Сгенерирован файл эмбеддингов размером ~100 ГБ (или меньше, если ограничено дисковое пространство)
• Реализован класс MmapVectorSearch с методами search() и search_batch()
• Потребление RAM при открытии mmap не превышает 100 МБ (только метаданные)
• Потребление RAM при выполнении одного запроса не превышает 1 ГБ
• Время поиска одного запроса (k=10) по 100 ГБ данных < 1 секунды (при прогретом page cache)
• Результаты поиска через mmap совпадают с результатами поиска в RAM (на маленьком датасете)
• Код покрыт type hints и docstrings
• Создан скрипт для воспроизведения бенчмарка
• Написан README с описанием подхода и результатами

6. Ожидаемый результат

Основной артефакт Git-репозиторий со структурой:

mmap-embeddings/
├── README.md
├── requirements.txt
├── generate_data.py          # Генерация синтетических данных
├── mmap_search.py            # Класс MmapVectorSearch
├── benchmark.py              # Бенчмарки
├── tests/
│   ├── test_correctness.py   # Тест на корректность
│   └── test_memory.py        # Тест на потребление памяти
└── results/
    ├── benchmark_results.csv  # Результаты замеров
    └── memory_profile.png     # График потребления памяти

Содержание README

• Описание проблемы (100 ГБ данных на 32 ГБ RAM)
• Подход (memory-mapped файлы)
• Инструкция по запуску
• Результаты бенчмарков (таблица с временем и памятью)
• Выводы и рекомендации

Опциональные дополнительные результаты

• Сравнение с FAISS (IVF, HNSW) на том же датасете
• Анализ влияния размера чанка на производительность
• Исследование поведения page cache при разных паттернах доступа

7. Возможные сложности и их решение

8. Бюджет времени (оценка)

Примечание Для первого раза рекомендуется выделить 2 полных дня (с учётом времени на генерацию данных и отладку). Генерацию 100 ГБ данных можно запустить на ночь.

9. Связанные вопросы из базы знаний

10. Чек-лист самопроверки

• Я сгенерировал тестовые данные и убедился, что они корректны (проверил размер файла, тип данных)
• Я реализовал поиск через mmap и проверил, что результаты совпадают с эталоном (на маленьком датасете)
• Я замерил потребление памяти и убедился, что оно не превышает 1 ГБ при поиске
• Я провёл бенчмарк и получил время поиска < 1 секунды на запрос (при прогретом кэше)
• Я задокументировал код и написал README с инструкцией по воспроизведению
• Я проверил, что решение работает на разных ОС (Linux, macOS, Windows)
• Я рассмотрел альтернативные подходы (FAISS, Dask) и могу обосновать выбор mmap
• Я добавил обработку граничных случаев (пустой файл, запрос с нулевой нормой, k > n_vectors)