`lru_cache`: aumentando dramaticamente a performance de chamadas a funções no Python (em alguns casos)

Há uns tempos atrás, num encontro do Python Floripa, dei uma lightning talk sobre uma coisa que raramente vejo ser utilizada: o lru_cache. Apesar de pouco utilizado, é uma coisinha útil e pode ser bastante útil; é um decorador que controla o cache ao nível de funções. Ele pode drasticamente aumentar a performance do seu programa, com praticamente nenhuma complicação ou trabalho a mais. Assim, pensei que podia ser um bom assunto pro primeiro post.

Mas o que é cache?

Na computação, o cache é um espaço destinado ao armazenamento de informações recentemente usadas, a fim de tornar o acesso a essas informações mais rápido. Isso acontece em diversos níveis de abstração: por exemplo, o processador tem um cache de memória recentemente usada, porque o acesso à chamada memória L1 interna a ele (e seus níveis um pouco menos rápidos, L2, L3...) é muito mais rápido do que buscar a memória em RAM. Seu navegador da internet, de modo parecido, também usa cache, mas em um nível mais alto: ao invés de baixar arquivos que não se modificam com frequência a cada vez que uma página é acessada, eles baixam esses arquivos (JavaScript, CSS, até o HTML pra páginas não dinâmicas) uma vez só e mantém uma cópia disso localmente, porque é mais rápido pegar do disco ou memória do que baixar de novo.

Como entra o cache no contexto de uma função escrita em Python, então? Simples: pros casos em que rodar uma função com os mesmos argumentos retorna sempre os mesmos valores, podemos guardar o resultado da função e usá-lo novamente depois. Isso é útil principalmente pros casos em que é computacionalmente caro re-rodar uma função várias vezes. O "lru" do nome da função vem de "Least Recently Used"; quer dizer, ele descarta da memória os dados que foram usados menos recentemente (mantendo sempre os mais recentes).

Na prática

Pra entender como funciona o lru_cache, vamos implementar uma versão simples dele na mão primeiro. Vamos criar uma função que retorna o cubo do número passado a ela:

def cubo(x):
    return x ** 3

Quanto tempo leva pra calcular os cubos de 0 a 100, 1000 vezes?

%%timeit
for i in range(1000):
    for j in range(100):
        cubo(j)
# 46 ms ± 846 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Por volta de 46 milissegundos na minha máquina.

Mas calcular o valor ao cubo de um número é relativamente devagar, mais devagar do que retornar um valor da memória. E se, ao invés de recalcularmos toda vez, calcularmos só uma vez cada número, e depois se o mesmo x for pedido novamente, só retornarmos ele da memória? Vamos ver:

valores_cubo = {}
def cubo_com_cache(x):
    # Se já tivermos visto o valor antes, retornamos ele imediatamente
    if x in valores_cubo:
        return valores_cubo[x]
    # Se não...
    else:
        cubo = x ** 3  # Calculamos o cubo à moda antiga
        valores_cubo[x] = cubo  # Guardamos o resultado no nosso dicionário
        return cubo  # Retornamos o cubo calculado normalmente

Testando a velocidade...

%%timeit
for i in range(1000):
    for j in range(100):
        cubo_com_cache(j)
17.7 ms ± 167 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Levamos menos da metade do tempo pra fazer a mesma coisa!

lru_cache

Com isso, escrevemos a nossa versão simplificada do lru_cache. Agora que entendemos o funcionamento e benefícios do cache ao nível de funções, vamos comparar o que fizemos acima com o que o Python nos traz pronto. lru_cache é um decorador que deve ser importado do módulo built-in functools. Estando importado, basta decorar uma função qualquer com ele:

from functools import lru_cache

@lru_cache()
def cubo_lru(x):
    return x ** 3

E o Python cuida de toda a criação do dicionário e retorno de valores já vistos automagicamente! Vamos à comparação de velocidade:

%%timeit
for i in range(1000):
    for j in range(100):
        cubo_lru(j)
# 12.2 ms ± 92.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Além de mais limpo e conciso, usar o lru_cache é também mais rápido do que nossa versão caseira. Isso acontece porque nossa versão é um tanto direta e ignora algumas possíveis otimizações que são feitas no lru_cache. Então por ho-

%%timeit
for i in range(1000):
    for j in range(129):
        cubo_lru(j)
# 75.8 ms ± 441 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Mas o que é isso? Ao aumentarmos o número de cálculos em ~30% o tempo de execução da mesma função com cache aumentou mais de 6 vezes! Por que isso aconteceu?

Os tradeoffs de cache

Nem tudo nessa vida é almoço de graça. Os leitores mais observadores terão notado que o benefício que o uso de cache traz em velocidade de execução vem com um custo: o de memória consumida pra armazenamento dos resultados.

Isso em geral não é problema, principalmente pra funções com tipos simples de retorno (como nossos ints), mas é algo a se manter em mente. Nosso algoritmo caseiro de caching ignora essa ressalva e cresce infinitamente. Se chamássemos a função "cacheada" com argumentos diferentes um número suficiente de vezes, eventualmente ficaríamos sem memória o suficiente pra continuar.

Pensando nisso, a implementação do Python de lru_cache traz o mecanismo de segurança (e controle do quanto de memória você está disposto a sacrificar) de limitar o tamanho máximo do cache que você usa. Por padrão, esse tamanho é de 128 resultados.

Vê agora o que aconteceu ali em cima? O lru_cache mantém em memória os resultados mais recentemente usados. Como o tamanho máximo era de 128 resultados, ao chegar no cálculo do 128, o 0 era descartado da memória, e precisava ser recalculado de novo da próxima vez. O resultado da chamada de 0 era então guardado em memória, mas pra isso o 1 era descartado, e assim por diante, obrigando a função a recalcular o resultado em todos os casos.

O lru_cache só é útil quando as chamadas recentes são um bom indicativo das chamadas que estão por vir.

Você pode tranquilamente mudar o tamanho máximo armazenado, porém. É o parâmetro maxsize:

@lru_cache(maxsize=256)
def cubo_lru(x):
    return x ** 3

Nosso problema acima agora está resolvido (ao menos para até 256 resultados). Os programadores intrépidos também podem usar None como maxsize pra desabilitar o limite máximo. A referência recomenda o uso de maxsize potência de 2 (32, 64, 128) pra melhor performance.

Falando de parâmetros, o lru_cache toma mais um: typed, que pode ser True ou False. Quando True, o cache trata argumentos com tipos diferentes (por exemplo, 3 (int) e 3.0 (float)) como sendo dois argumentos distintos, e quando False ele os considera iguais.

Mais ressalvas

Você pode já ter percebido quando implementamos nossa versão simples, mas vale a pena reiterar: não adianta usar cache pra funções que dependam de estado externo, ou que tenham algum efeito colateral, ou que precisem retornar objetos únicos. Lembremos que quando o cache funciona, a função não é executada; meramente lembramos do resultado de quando a chamamos pela última vez com esses argumentos.

Vamos ver o que cada uma dessas ressalvas significa com exemplos curtos que podem ser inesperados se não lembrarmos do que estamos fazendo:

Funções que dependem de estado externo

O cache, quando vê argumentos que já estão na memória, só retorna o resultado prévio, sem se importar se alguma coisa mudou ou não. Isso é perigoso quando o resultado da função depende de algo além dos argumentos passados:

# Temos uma variável x externa à função
x = 2

# Definimos uma função com cache que depende de x
@lru_cache()
def mult_a_x(a):
    return a * x

# Usamos a função
print(mult_a_x(2))  # printa 4

# Modificamos x
x = 3

# O cache não sabe que o resultado não é mais válido!
print(mult_a_x(2))  # printa 4

Uma possível solução pra esse problema (desde que se tenha cuidado e se saiba o que está fazendo) é usar o cache_clear que o decorador providencia. Isso descarta todo o cache e começa do zero; pode ser útil pra quando realmente se quer usar cache numa função que dependa de estado externo que mude raramente e previsivelmente. Se for fazer isso, tome cuidado pra lembrar de limpar o cache quando o estado mudar!

# Temos uma variável x externa à função
x = 2

# Definimos uma função com cache que depende de x
@lru_cache()
def mult_a_x(a):
    return a * x

# Usamos a função
print(mult_a_x(2))  # printa 4

# Modificamos x
x = 3

# LEMBRAMOS DE LIMPAR O CACHE
mult_a_x.cache_clear()
print(mult_a_x(2))  # printa 6

Funções que tenham efeitos colaterais

O cache não executa o bloco interior da função quando vê argumentos que já estão em memória. Isso significa que ele pulará qualquer modificação do estado externo, não recalculará números aleatórios, não executará prints e nem dormirá em sleep:

import time

@lru_cache()
def dormir(x):
    time.sleep(x)

print('A')
dormir(1)  # Para por 1 segundo
print('B')
dormir(1)  # Não para!
print('C')

Funções que precisem retornar objetos únicos

Essa é um pouco mais difícil de enxergar de primeira. Tem a ver com o fato de que objetos como listas, dicionários, e instâncias de classes criadas, ao contrário de tipos primitivos como ints e floats, são guardados e passados por referência, e não por valor. Isso merece um post inteiro, mas pros leitores que já estejam familiarizados com o conceito, vai um exemplo pra demonstrar o perigo:

import time

@lru_cache()
def retorna_dict(x):
    return {'chave': x * 2}

d = retorna_dict(3)
print(d)  # printa {'chave': 6}

# imagine aqui que o programa segue e 
# modificamos um pouco o dicionário d
d['outra_chave'] = 5
d['chave'] = 2

# O que acontecerá aqui?
outro_d = retorna_dict(3)
print(outro_d)  # printa {'chave': 2, 'outra_chave': 5}

O que aconteceu? No dicionário interno do lru_cache, o que foi guardado como resultado da chamada da função pro argumento x=3 é (como acontece sempre no Python) uma referência pra um dicionário. Essa referência é passada pra variável d, e modificar os valores de d (a não ser atribuir a outra coisa) modifica o dicionário pro qual a referência aponta. O retorno do cache pra x=3 sempre vai ser a referência desse mesmo dicionário, e isso pode pegar de surpresa.

As possíveis soluções pra esse caso são:

  1. Não usar lru_cache pra funções que retornem tipos de referência
  2. Fazer um .copy() do resultado e modificar essa cópia (ex.: d = retorna_dict(3).copy() resolve)
  3. Não modificar objetos passados por referência resultantes de lru_cache

Informação e limpeza de cache

Tendo isso tudo em mente, pra te ajudar a decidir se vale a pena usar cache ou se o seu maxsize está muito grande ou muito pequeno, o nosso querido decorador também oferece uma função que mostra estatísticas acumuladas durante o uso da função decorada, o cache_info:

import random

for _ in range(100000):
    cubo_lru(random.randint(0, 400))
    
print(cubo_lru.cache_info())
# CacheInfo(hits=127623, misses=72377, maxsize=256, currsize=256)

Se você tem muito mais cache misses do que hits, talvez seja uma boa ideia aumentar o tamanho máximo ou esquecer a ideia de usar cache. Se o seu currsize (current size, tamanho atual) não chega perto do maxsize, é porque talvez seu maxsize esteja grande demais. Se você tem consideravelmente mais hits que misses, parabéns, sua função está sendo acelerada!

Bonus round

Só como um extra, vou mostrar um caso em que o lru_cache realmente brilha: funções recursivas. Case in point: fatorial.

def fact(n):
    if n == 1:
        return 1
    return fact(n - 1) * n

@lru_cache()
def fact_lru(n):
    if n == 1:
        return 1
    return fact_lru(n - 1) * n
%%timeit

fact(100)
# 21.4 µs ± 188 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%%timeit

fact_lru(100)
# 115 ns ± 2.94 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

Com uma linha e nenhuma biblioteca externa aumentamos a performance da nossa função em 186 vezes!

Show Comments

Get the latest posts delivered right to your inbox.