Оптимизация Django: или как мы избежали ошибок памяти

(Кросс-пост из моего блога на Medium)
Работа в технологическом стартапе сродни тушению серии пожаров, которые возникают время от времени (быстрее, если вы работаете с головокружительной скоростью). Каждый раз, когда вы тушите огонь, достаточно, чтобы сохранять спокойствие в течение следующих нескольких месяцев, но в глубине души вы знаете, что это еще не конец. Таким образом, становится важным выбери свои битвы.

Проблема под рукой

Во-первых, некоторый контекст — у нас есть эта задача Celery, которая запускается каждую ночь и выполняет некоторые трудоемкие, но важные вычисления, которые сохраняют «свежесть» определенной таблицы в базе данных. Это одна из важнейших частей нашей кодовой базы, и важно, чтобы эта задача успешно выполнялась каждый день. Раньше он доставлял нам проблемы из-за относительно больших объемов данных, которые он обрабатывает, и к нему применялось более одной оптимизации. Недавно задача перестала полностью выполняться, и я обнаружил сообщения о segfault при исследовании производственных журналов на предмет другой проблемы.
Рассмотрим ультра упрощенную версию кода задачи следующим образом:

books = Book.objects.filter(some_field=some_condition)
authors = Author.objects.filter(some_other_field=some_condition)
results = {}
for x in authors:
    for y in books:
        # Insert computations here
        results[(x, y)] = True # Not that simple, but still
return results

books а также authors оба являются наборами запросов Django, хотя и не представляют наши реальные модели. У последнего было 2012 объектов, а у первого около 17 тысяч объектов на момент написания. С момента создания этой задачи этот двойной цикл for хорошо служил нам. Затем мы столкнулись с первыми ошибками памяти в прошлом году и недавними ошибками сегментации. Я решил сократить код до самого необходимого, аналогично тому, что я показал выше, и запустить несколько тестов с помощью команды htop. Вот то же самое, когда я исследовал проблему на нашей промежуточной машине, которая имеет 4 гигабайта оперативной памяти, а не 10 гигабайт на производственной (согласно упрощенному коду):

Использование памяти растет довольно быстро, и, прежде чем мы это осознаем, мы достигаем предела в 4 гигабайта на машине. Если бы мне нужно было измерять с точки зрения пространственной сложности, это было бы O (MN), где M и N — обе (приблизительно?) линейные функции, которые отслеживают рост обоих наборов запросов с течением времени. M явно медленнее, чем N.

Пакетная обработка набора запросов

Затем я повторно применил первую оптимизацию, примененную год назад, — использование пакетных наборов запросов. Есть два основных места¹, где потребляется память в нашем скрипте, первое — это коннектор базы данных python (в данном случае коннектор Python MySQL DB), который выполняет задачу извлечения результатов, а второе — кеш набора запросов. Пакетная обработка Queryset в первую очередь относится к оптимизации.
Вот код с пакетной обработкой, примененной ко второму набору запросов:

def batch_qs(qs, batch_size=1000):
    """
    Returns a (start, end, total, queryset) tuple for each batch in the given
    queryset. Useful when memory is an issue. Picked from djangosnippets.
    """
    if isinstance(qs, QuerySet):
        total = qs.count()
    else:
        total = len(qs)
    for start in range(0, total, batch_size):
        end = min(start + batch_size, total)
        yield (start, end, total, qs[start:end])

books = Book.objects.filter(some_field=some_condition)
authors = Author.objects.filter(some_other_field=some_condition)
results = {}
for x in authors:
    for _, _, _, qs in batch_qs(books, batch_size=1000):
        for y in qs:
            # Insert computations here
            results[(x, y)] = True # Not that simple, but still
return results

Выполняем задачу еще раз и смотрим на htop.

Хм. Скорость, с которой растет использование памяти, медленнее, чем в прошлый раз, но она неизбежно достигает предела и segfaults. Это имеет смысл, поскольку пакетная обработка наборов запросов не уменьшает количество результатов, извлекаемых из базы данных. Он только стремится сохранить столько, сколько указано в памяти за раз, настроив размер пакета. Это приводит нас к нашему нынешнему затруднительному положению и второму месту для оптимизации.

Природа наборов запросов

Говорят, что наборы запросов Django лениво загружаются и кэшируются¹ ². Ленивая загрузка означает, что пока вы не выполните определенные действия с набором запросов, например повторите его, соответствующий запрос БД не будет выполнен. Кэширование означает, что при повторном использовании одного и того же набора запросов не будет выполняться несколько запросов к БД.

qs = Book.objects.filter(id=3)
# The DB query has not been executed at this point
x = qs
# Just assigning variables doesn't do anything
for x in qs:
    print x
# The query is executed at this point, on iteration
for x in qs:
    print "%d" % x.id
# The query is not executed this time, due to caching

Теперь, как кеширование входит в смесь здесь? Получается, что из-за кеша мы не можем «выбросить» (сборщик мусора) уже использованные партии. В нашем случае мы используем каждый пакет только один раз, и их кэширование — это расточительное использование памяти. Кэши не будут очищены до конца функции. Чтобы бороться с этим, мы используем функцию iterator()³ в наборе запросов.

...
for x in authors:
    for _, _, _, qs in books:
        for y in qs.iterator():
            # Insert computations here
            results[(x, y)] = True # Not that simple, but still
...

Давайте попробуем. Рост использования памяти должен замедлиться до минимума.

Но это не так!

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

results[(x, y)] = True

Простой тест, который я провел, заключался в том, чтобы закомментировать эту строку и посмотреть, сколько памяти занимает задача. Использование памяти больше не росло бешено, а вместо этого было заморожено на определенном значении.
Для каждого автора из внешнего цикла мы перебирали разные «копии» книг во внутреннем цикле. Для двух авторов A1 и A2 кортежи (A1, B) и (A2, B) указывали на разные копии B в памяти, из-за использования итератора. Таким образом, нам придется снова ввести кэширование, но на наших условиях.

books = Book.objects.filter(some_field=some_condition)
authors = Author.objects.filter(some_other_field=some_condition)
results = {}
book_cache = {}
for x in authors:
    for _, _, _, qs in batch_qs(books, batch_size=1000):
        for y in qs:
            # Insert computations here
            if y.id not in book_cache:
                book_cache[y.id] = y
            cached_y = book_cache[y.id]
            results[(x, cached_y)] = True # Not that simple, but still
return results

И вот мы идем. Использование памяти растет медленно, колеблется вокруг определенных значений и через некоторое время значительно падает (не показано). Цель book_cache заключается в сохранении каждого объекта книги, который просматривается при одном полном проходе всех объектов книги, для повторного использования при последующих полных проходах. Объекты, на которые указывает y в этих последующих проходах будет собираться мусор, и будут использоваться кэшированные версии. Теперь для (A1, B) и (A2, B) оба B указывают на один и тот же объект в памяти. В конце концов, iterator() позволяет нам контролировать, как мы тратим нашу свободную память. Объемная сложность этого окончательного кода будет O(M + N), так как у нас есть только одна копия каждого объекта книги в памяти.

Обратите внимание, что это решает нашу проблему на данный момент. По мере роста базы данных будет больше проблем. Это означает, что код не идеален, но цель никогда не «совершенство», а «оптимизация на данный момент».
PS: чтобы ответить на некоторые комментарии Medium — я действительно могу хранить идентификаторы объектов, пары из них, в качестве ключей. Это было бы идеальным решением. Но в данной ситуации мне это не подходит, потому что приведенный выше код представляет собой небольшую часть более крупного рабочего процесса, где требуются экземпляры объектов.

[1] Запросы Django с эффективным использованием памяти (www.poeschko.com/2012/02/memory-efficient-django-queries/)
[2] Наборы запросов Django (https://docs.djangoproject.com/en/2.0/topics/db/optimization/#understand-queryset-evaluation)
[3] итератор() (