Разделяемая память в Python: multiprocessing под нагрузкой

Задача простая: посчитать статистику по массиву NumPy размером 2 ГБ, распараллелив работу на восемь процессов. Наивный код через multiprocessing.Pool.map отрабатывал заметно дольше, чем однопоточный вариант, — при этом пик потребления памяти уходил за 16 ГБ на машине с 32 ГБ. Причина не в вычислениях, а в том, что каждый воркер получал свою копию массива через pickle. Ниже — где именно теряется время и память и что даёт переход на multiprocessing.shared_memory (доступен с Python 3.8).

Почему Pool.map копирует данные

В CPython у каждого процесса свой GIL и своё адресное пространство. Когда вы передаёте объект в Pool.map или как аргумент в Process, объект сериализуется через модуль pickle, отправляется в дочерний процесс через pipe и там десериализуется. Для массива NumPy это означает полную побайтовую копию буфера — плюс временные объекты на стороне отправителя и получателя.

На Linux по умолчанию используется метод старта fork¹, и глобальные объекты, объявленные до создания пула, наследуются дочерними процессами через copy-on-write. Но как только воркер начинает читать массив, ядро всё равно копирует затронутые страницы, а счётчик ссылок Python, который лежит рядом с объектом, портит COW почти сразу. На Windows и macOS метод старта по умолчанию spawn, и там наследования нет вовсе — всё едет через pickle.

Разделяемая память не ускоряет вычисления. Она убирает то, что вычислениями не является: сериализацию, передачу через pipe и дублирование буфера на каждый процесс.

Как считали

Условия замеров: массив float64 формы (268_435_456,), это ровно 2 ГБ. Восемь воркеров, каждый суммирует свой срез. Машина — 8 физических ядер, 32 ГБ RAM, Python 3.11, NumPy 1.26, Linux с методом старта fork. Время — среднее по пяти запускам через time.perf_counter(), память — пик RSS процесса-родителя и суммарно по группе через /usr/bin/time -v.

Наивный вариант через Pool.map

import numpy as np
from multiprocessing import Pool
import time

# 2 ГБ данных
data = np.arange(268_435_456, dtype=np.float64)

def partial_sum(chunk):
    # chunk сюда приезжает как копия через pickle
    return float(np.sum(chunk))

if __name__ == "__main__":
    parts = np.array_split(data, 8)
    start = time.perf_counter()
    with Pool(8) as pool:
        results = pool.map(partial_sum, parts)
    total = sum(results)
    print(f"sum={total}, elapsed={time.perf_counter() - start:.2f}s")

Каждый срез сериализуется и уезжает в воркер отдельной копией. Родитель держит исходный массив плюс временные буферы под pickle, дочерние процессы — свои копии срезов. Основное время съедает не суммирование, а упаковка и передача 2 ГБ через pipe.

Вариант через shared_memory

Идея: создать блок разделяемой памяти один раз, положить туда массив, а воркерам передать только имя блока, форму и dtype — несколько десятков байт. Каждый воркер сам подключается к блоку и строит поверх него view NumPy без копирования.

import numpy as np
from multiprocessing import Pool
from multiprocessing import shared_memory
import time

def worker(args):
    name, shape, dtype_str, start_idx, end_idx = args
    existing = shared_memory.SharedMemory(name=name)
    # view поверх общего буфера, копии нет
    arr = np.ndarray(shape, dtype=np.dtype(dtype_str), buffer=existing.buf)
    result = float(np.sum(arr[start_idx:end_idx]))
    existing.close()  # detach, но не unlink
    return result

if __name__ == "__main__":
    data = np.arange(268_435_456, dtype=np.float64)

    shm = shared_memory.SharedMemory(create=True, size=data.nbytes)
    shared = np.ndarray(data.shape, dtype=data.dtype, buffer=shm.buf)
    shared[:] = data[:]  # единственное копирование, в общий блок

    n = data.shape[0]
    step = n // 8
    tasks = []
    for i in range(8):
        lo = i * step
        hi = n if i == 7 else (i + 1) * step
        tasks.append((shm.name, data.shape, str(data.dtype), lo, hi))

    start = time.perf_counter()
    with Pool(8) as pool:
        results = pool.map(worker, tasks)
    total = sum(results)
    print(f"sum={total}, elapsed={time.perf_counter() - start:.2f}s")

    shm.close()
    shm.unlink()  # освободить блок, обязательно один раз

Воркеры получают строку-имя, а не 250 МБ данных. Копирование происходит ровно один раз — при заполнении общего блока в родителе. Дальше все восемь процессов читают один и тот же физический буфер.

Цифры до и после

На описанных выше условиях (8 ядер, 32 ГБ, Python 3.11, NumPy 1.26, массив 2 ГБ, метод старта fork). Ваши числа будут отличаться — снимайте свои на своём железе и объёме, порядок величин важнее конкретных значений.

Смысл в том, что при чистой сериализации время растёт линейно с объёмом данных и почти не зависит от того, сколько вы реально считаете. С разделяемой памятью объём передачи фиксирован, и время определяется полезной работой. Чем больше массив и чем легче вычисление на элемент, тем сильнее выигрыш от shared_memory.

Где подход ломается

• Утечка блоков. Если процесс упал до unlink(), блок остаётся в /dev/shm на Linux и продолжает занимать память. Начиная с Python 3.8 за это отвечает resource_tracker, но при жёстком падении по SIGKILL он не всегда успевает подчистить. Проверяйте /dev/shm и удаляйте осиротевшие блоки.
• Двойной unlink. unlink() должен вызвать ровно один процесс — обычно тот, что создал блок. Воркеры делают только close(). Повторный unlink уронит с ошибкой, что блока уже нет.
• Запись из нескольких процессов. Если воркеры не только читают, но и пишут в общий буфер, вам нужна синхронизация: разделяемая память сама по себе не даёт атомарности. Разбивайте на непересекающиеся диапазоны или ставьте блокировки.
• Ограничение размера /dev/shm. На Linux tmpfs под /dev/shm по умолчанию ограничен, часто половиной RAM. Блок больше этого лимита создать не получится — проверьте df -h /dev/shm.
• spawn на Windows и macOS. При spawn дочерний процесс не наследует глобальные объекты, поэтому имя блока обязательно передавать явно как аргумент — что мы и делаем в примере. Код выше работает и на fork, и на spawn.

Когда shared_memory не нужна

Если данные маленькие (единицы мегабайт) или процессов один-два, накладные расходы на управление блоком и синхронизацию перевешивают экономию. То же, если задача упирается в CPU, а данные и так копируются редко. Для сложных структур, а не сырых буферов, посмотрите на multiprocessing.Manager — он медленнее, но избавляет от ручного управления памятью.

Что мерить и что делать дальше

Прежде чем переписывать код на shared_memory, замерьте, сколько времени уходит именно на передачу данных, а не на вычисление. Проще всего — засечь время пустого воркера, который принимает массив и сразу возвращает его длину: если оно сопоставимо с полным прогоном, вы платите за сериализацию, и разделяемая память окупится.

• Снимите пик RSS группы процессов через /usr/bin/time -v или psutil — до и после.
• Проверьте /dev/shm после аварийных завершений на утечку блоков.
• Если структура сложнее плоского массива, посмотрите в сторону Arrow или разделяемых файлов с mmap перед тем, как городить ручную сериализацию.

¹ Метод старта процесса — способ, которым multiprocessing создаёт дочерний процесс. fork копирует адресное пространство родителя (Linux), spawn запускает чистый интерпретатор и заново импортирует модуль (Windows, macOS, и он же по умолчанию на macOS с Python 3.8+). От метода зависит, наследуются ли глобальные объекты.