Управление памятью в NumPy

Исторический обзор

С версии 1.7.0 NumPy предоставляет набор функций PyDataMem_* (PyDataMem_NEW, PyDataMem_FREE, PyDataMem_RENEW), которые поддерживаются alloc, free, realloc соответственно.

С тех пор Python также улучшил свои возможности управления памятью и начал предоставлять различные политики управления управления памятью начиная с версии 3.4. Эти процедуры разделены на набор областей, каждая область имеет структуру PyMemAllocatorEx для управления памятью. Python также добавил модуль tracemalloc для отслеживания вызовов различных процедур. Эти фиксации отслеживания были добавлены в NumPy PyDataMem_* процедуры.

NumPy добавил небольшой кэш выделенной памяти в свои внутренние функции npy_alloc_cache, npy_alloc_cache_zero, и npy_free_cache. Они оборачивают alloc, alloc-and-memset(0) и free соответственно, но при вызове npy_free_cache, он добавляет указатель в короткий список доступных блоков, помеченных размером. Эти блоки могут быть повторно использованы последующими вызовами npy_alloc*, избегая избыточного использования памяти.

Настраиваемые процедуры управления памятью в NumPy (NEP 49)

Пользователи могут захотеть переопределить внутренние процедуры управления памятью данных своими собственными. Так как NumPy не использует стратегию домена Python для управления памятью данных, он предоставляет альтернативный набор C-API для изменения процедур управления памятью. Нет стратегий управления памятью на уровне домена Python для больших фрагментов данных объектов, поэтому они менее подходят для нужд NumPy. Пользователи, которые хотят изменить процедуры управления памятью данных NumPy, могут использовать PyDataMem_SetHandler, который использует структуру PyDataMem_Handler для хранения указателей на функции, используемые для управления памятью данных. Вызовы все еще оборачиваются внутренними процедурами для вызова PyTraceMalloc_Track, PyTraceMalloc_Untrack. Поскольку функции могут изменяться в течение времени существования процесса, каждый ndarray содержит функции, используемые во время его создания, и они будут использоваться для перераспределения или освобождения памяти данных экземпляра.

typePyDataMem_Handler

Структура для хранения указателей на функции, используемые для управления памятью

typedef struct {
    char name[127];  /* multiple of 64 to keep the struct aligned */
    uint8_t version; /* currently 1 */
    PyDataMemAllocator allocator;
} PyDataMem_Handler;

где структура аллокатора находится

/* The declaration of free differs from PyMemAllocatorEx */
typedef struct {
    void *ctx;
    void* (*malloc) (void *ctx, size_t size);
    void* (*calloc) (void *ctx, size_t nelem, size_t elsize);
    void* (*realloc) (void *ctx, void *ptr, size_t new_size);
    void (*free) (void *ctx, void *ptr, size_t size);
} PyDataMemAllocator;

PyObject*PyDataMem_SetHandler(PyObject*handler)

Установить новую политику выделения. Если входное значение NULL, политика будет сброшена до значения по умолчанию. Возвращает предыдущую политику или NULL, если произошла ошибка. Мы оборачиваем функции, предоставленные пользователем, чтобы они по-прежнему вызывали обратные вызовы управления памятью Python и NumPy.

PyObject*PyDataMem_GetHandler()

Возвращает текущую политику, которая будет использоваться для выделения данных для следующего PyArrayObject. При ошибке возвращает NULL.

Пример настройки и использования PyDataMem_Handler см. в тесте numpy/_core/tests/test_mem_policy.py

Что происходит при освобождении памяти, если политика не задана

Редкий, но полезный метод — выделить буфер вне NumPy, использовать PyArray_NewFromDescr для упаковки буфера в ndarray, а затем установить флаг OWNDATA в значение true. Когда ndarray будет освобождён, соответствующая функция из ndarray’s PyDataMem_Handler должна быть вызвана для освобождения буфера. Но поскольку поле PyDataMem_Handler никогда не устанавливалось, оно будет NULL. Для обратной совместимости NumPy вызовет free() для освобождения буфера. Если NUMPY_WARN_IF_NO_MEM_POLICY будет установлено в 1, будет выведено предупреждение. Текущий вариант по умолчанию — не выводить предупреждение, но это может измениться в будущей версии NumPy.

Более подходящим методом будет использование PyCapsule в качестве базового объекта:

/* define a PyCapsule_Destructor, using the correct deallocator for buff */
void free_wrap(void *capsule){
    void * obj = PyCapsule_GetPointer(capsule, PyCapsule_GetName(capsule));
    free(obj);
};

/* then inside the function that creates arr from buff */
...
arr = PyArray_NewFromDescr(... buf, ...);
if (arr == NULL) {
    return NULL;
}
capsule = PyCapsule_New(buf, "my_wrapped_buffer",
                        (PyCapsule_Destructor)&free_wrap);
if (PyArray_SetBaseObject(arr, capsule) == -1) {
    Py_DECREF(arr);
    return NULL;
}
...

Пример отслеживания памяти с помощью np.lib.tracemalloc_domain

Обратите внимание, что с Python 3.6 (или новее) встроенный модуль tracemalloc можно использовать для отслеживания выделений внутри NumPy. NumPy размещает свои выделения памяти процессора в области np.lib.tracemalloc_domain. Дополнительную информацию можно найти по адресу: https://docs.python.org/3/library/tracemalloc.html.

Вот пример использования np.lib.tracemalloc_domain:

"""
   The goal of this example is to show how to trace memory
   from an application that has NumPy and non-NumPy sections.
   We only select the sections using NumPy related calls.
"""

import tracemalloc
import numpy as np

# Flag to determine if we select NumPy domain
use_np_domain = True

nx = 300
ny = 500

# Start to trace memory
tracemalloc.start()

# Section 1
# ---------

# NumPy related call
a = np.zeros((nx,ny))

# non-NumPy related call
b = [i**2 for i in range(nx*ny)]

snapshot1 = tracemalloc.take_snapshot()
# We filter the snapshot to only select NumPy related calls
np_domain = np.lib.tracemalloc_domain
dom_filter = tracemalloc.DomainFilter(inclusive=use_np_domain,
                                      domain=np_domain)
snapshot1 = snapshot1.filter_traces([dom_filter])
top_stats1 = snapshot1.statistics('traceback')

print("================ SNAPSHOT 1 =================")
for stat in top_stats1:
    print(f"{stat.count} memory blocks: {stat.size / 1024:.1f} KiB")
    print(stat.traceback.format()[-1])

# Clear traces of memory blocks allocated by Python
# before moving to the next section.
tracemalloc.clear_traces()

# Section 2
#----------

# We are only using NumPy
c = np.sum(a*a)

snapshot2 = tracemalloc.take_snapshot()
top_stats2 = snapshot2.statistics('traceback')

print()
print("================ SNAPSHOT 2 =================")
for stat in top_stats2:
    print(f"{stat.count} memory blocks: {stat.size / 1024:.1f} KiB")
    print(stat.traceback.format()[-1])

tracemalloc.stop()

print()
print("============================================")
print("\nTracing Status : ", tracemalloc.is_tracing())

try:
    print("\nTrying to Take Snapshot After Tracing is Stopped.")
    snap = tracemalloc.take_snapshot()
except Exception as e:
    print("Exception : ", e)