Como um único .delete() do Django estourou a memória de uma instância de 4 GB | Recortando Dados de Produção

ORMs podem trazer muita simplicidade e são mais fáceis de manter do que comandos SQL escritos na mão, mas, como tudo na vida, isso tem um custo. No meu caso, o custo foi memória, mais precisamente, um erro de OOM (Out Of Memory). Então, para registrar isso para o futuro e também ajudar outras pessoas a evitar os mesmos erros, resolvi escrever mais sobre o assunto: suas causas e possíveis soluções.

Contexto

O contexto em que o OOM aconteceu foi a etapa de transformação de um pipeline que tem como objetivo gerar um recorte dos dados de produção. Eu sei que hoje existem algumas ferramentas para implementar isso, mas tínhamos alguns requisitos específicos (como deixar certas linhas de certas tabelas intocadas) e, para reaproveitar os mesmos models já definidos no Django e simplificar a manutenção, decidimos usar um Django command.

Para criar o recorte, decidimos escolher uma porcentagem e deletar o resto. O problema é que o primeiro delete envolve uma cascata de regras CASCADE (sim) por conta do uso do método .delete(). Esse método (e outros erros de implementação) foi a causa do erro de OOM.

Causa

Bom, dizer que a causa do erro foi exclusivamente o método não é toda a verdade, já que isso só apareceu quando fiz o primeiro teste de QA com dados de produção, e essa é a causa da maioria dos problemas de backend: escala e volume. Localmente, tentei criar uma boa massa de dados para testar minha etapa de transformação, mas nada daquilo foi suficiente para representar os dados de produção (mais de 150 GB). Além disso, em produção os dados são muito mais concentrados nas folhas do que na raiz: para dar uma noção de escala (números ilustrativos), o formato é mais ou menos assim: avô → 20 mil pais → 100 mil netos. Isso criou o cenário perfeito para expor o meu erro de implementação: o uso errado do .delete().

Como o delete funciona

Quando você chama .delete() em um queryset, o trabalho de verdade é delegado a uma instância de Collector. A função dele é descobrir tudo que precisa ser deletado antes de deletar qualquer coisa: ele percorre recursivamente (collect()) os relacionamentos em cascata do model base, puxando os objetos afetados para um dict em memória, e só então dispara o SQL de verdade, um DELETE ... IN (...) por tabela. (Ele consegue pular o carregamento de uma relação quando é capaz de fazer um "fast delete" com uma única subquery, mas qualquer coisa que tenha suas próprias cascatas, signals ou filhos ainda é trazida para a memória.) Esse design de "primeiro coleta a árvore inteira, depois deleta" é exatamente para onde a memória vai. Três partes do Collector contam essa história: onde os objetos ficam guardados, como a coleta recursiona e quando o SQL finalmente roda.

Primeiro, o armazenamento (todos os trechos vêm de django/db/models/deletion.py, Django 5.2, encurtados para facilitar a leitura). Os objetos a serem deletados são acumulados em um único dict em memória:

class Collector:
    def __init__(self, using, origin=None):
        ...
        # Inicialmente, {model: {instances}}; depois os valores viram listas.
        self.data = defaultdict(set)   # os objetos a serem deletados são coletados aqui, em memória

A recursão vem do handler de CASCADE: para cada objeto relacionado que ele encontra, chama collect() de novo, que por sua vez coleta os objetos relacionados a esses novos objetos, e assim por diante árvore abaixo:

def CASCADE(collector, field, sub_objs, using):
    collector.collect(          # recursão: coleta os filhos dos filhos ...
        sub_objs,
        source=field.remote_field.model,
        source_attr=field.name,
        nullable=field.null,
        fail_on_restricted=False,
    )

    def collect(self, objs, ...):
        new_objs = self.add(objs, ...)      # guarda essas instâncias em self.data
        ...
        for related in get_candidate_relations_to_delete(model._meta):
            field = related.field
            on_delete = field.remote_field.on_delete
            ...
            for batch in batches:
                sub_objs = self.related_objects(related_model, [field], batch)  # busca os filhos
                if getattr(on_delete, "lazy_sub_objs", False) or sub_objs:
                    on_delete(self, field, sub_objs, self.using)   # CASCADE -> collect() de novo

Só depois que tudo foi coletado é que o SQL roda, um DELETE ... IN (...) por model, construído a partir das chaves primárias de cada instância que está em self.data:

    def delete(self):
        ...
        # deleta as instâncias
        for model, instances in self.data.items():
            query = sql.DeleteQuery(model)
            pk_list = [obj.pk for obj in instances]          # todas as pks coletadas, ainda em memória
            count = query.delete_batch(pk_list, self.using)  # DELETE FROM ... WHERE pk IN (...)

Ou seja, a combinação do método collect sendo chamado para cada objeto relacionado em cascata e sendo colocado em memória pela instância do collector, mais o volume de dados, mais a estrutura da distribuição desses dados, consome rapidamente a capacidade máxima de memória da instância da máquina (4 GB, para ser exato).

A solução inicial

A solução inicial para isso, em vez de simplesmente chamar o método delete no queryset, foi deletar os descendentes relacionados em lotes, de modo a manter o uso de memória constante, a clássica ideia de resolver primeiro os problemas pequenos. Aqui está a implementação:

Os models envolvidos formam uma cascata Blog → Post → Comment (models fictícios), ou seja, avô → pai → neto, onde o volume vive nas folhas. Primeiro calculamos as chaves primárias que queremos remover com SQL puro (assim os ids nunca caem todos de uma vez no Python) e, em vez de um único queryset.delete() gigante, deletamos um lote limitado por vez:

from django.db import transaction

# A versão ingênua: uma chamada para o recorte inteiro. queryset.delete() executa
# o Collector do Django, que puxa os descendentes em cascata (Post -> Comment -> ...)
# para a memória antes de disparar os DELETEs. Com dados em escala de produção,
# é isso que esgota uma instância de 4 GB.
Post.objects.filter(id__in=drop_ids).delete()

# A versão em lotes: deleta as mesmas linhas um lote limitado por vez, de forma que
# o Collector nunca segura em memória mais do que os descendentes de um único lote.
def batched_delete(model, ids, *, size=1000):
    # Deleta através do _base_manager (o manager sem filtros), e não do .objects:
    # um manager padrão customizado pode esconder linhas (por exemplo, as com
    # soft delete) que a lista de ids ainda mira, o que faria a deleção ficar
    # silenciosamente incompleta.
    manager = model._base_manager
    for start in range(0, len(ids), size):
        with transaction.atomic():
            manager.filter(id__in=ids[start:start + size]).delete()
    return len(ids)

batched_delete(Post, drop_ids)

E depois que a nova implementação foi mergeada, fizemos uma nova rodada de QA, e o erro de OOM estava resolvido. Então está tudo certo, né? Acontece que não. Naquele momento, mesmo tendo alguma noção da estrutura dos dados, eu não esperava que essas deleções usando esses lotes definidos levassem mais de 2 horas para (não) serem processadas, o que fez o pipeline inteiro estourar o tempo limite e, com isso, gerou a necessidade de repensar tudo de novo para resolver esse novo problema.

Conclusão

Mesmo que essa primeira tentativa não tenha resolvido todos os problemas do pipeline nem conseguido entregá-lo de vez, ela me ajudou a entender melhor como o Django funciona por baixo dos panos com o seu ORM. Então, no próximo post, vou tentar explicar quais foram as minhas tentativas de resolver isso definitivamente (espero), assim que possível.