Cómo acelerar tus scripts Python CPU-bound en 3 pasos

Esta guía rápida y práctica te muestra cómo identificar cuellos de botella y aplicar tres técnicas efectivas para acelerar código CPU-bound en Python: vectorización (NumPy), JIT (Numba) y paralelización por procesos. Incluye estructura de proyecto, código completo y mediciones simples.

Cuando usar cada técnica (regla rápida)

• Si puedes expresar el trabajo como operaciones sobre arrays, usa NumPy (menos código, muy eficiente).
• Si tienes loops complejos con ramas y operaciones numéricas, prueba Numba (JIT) para mantener la API Python.
• Si la tarea es inherentemente paralelizables pero cada unidad hace trabajo pesado y no comparte mucho estado, usa multiprocessing (evita GIL).

Estructura del proyecto

py-speedup-example/
├─ requirements.txt
├─ profiler.py         # detecta hotspots con cProfile
├─ worker.py           # implementación pura Python (baseline)
├─ numba_impl.py       # versión acelerada con Numba
├─ numpy_impl.py       # versión vectorizada con NumPy
└─ bench.py            # script que mide tiempos y compara

Problema de ejemplo

Vamos a calcular la función f(x) = sum_{k=1..N} sin(x*k) / k para muchos valores de x (coste por valor). Es representativo: operaciones trigonométricas y bucles — workload CPU-bound.

1) Perfilado (siempre primero)

# profiler.py
import cProfile
import pstats
import worker

if __name__ == '__main__':
    pr = cProfile.Profile()
    pr.enable()
    worker.main()   # ejecuta la versión baseline durante un tiempo representativo
    pr.disable()
    ps = pstats.Stats(pr).sort_stats('cumtime')
    ps.print_stats(20)

Ejecuta: python profiler.py. Mira las funciones que consumen más tiempo (cumtime).

2) Implementaciones

Baseline: worker.py (pure Python)

# worker.py
import math

def f_scalar(x, N=100000):
    s = 0.0
    for k in range(1, N+1):
        s += math.sin(x*k)/k
    return s

def main():
    xs = [i*0.0001 for i in range(1000)]
    out = [f_scalar(x) for x in xs]
    print('done', len(out))

if __name__ == '__main__':
    main()

NumPy vectorizado: numpy_impl.py

# numpy_impl.py
import numpy as np

def f_vec(xs, N=100000):
    k = np.arange(1, N+1)
    # broadcasting: xs[:,None] * k[None,:]
    # precaución con memoria: procesar en batches
    out = np.empty(len(xs), dtype=np.float64)
    batch = 100  # ajustar según RAM
    for i in range(0, len(xs), batch):
        xb = xs[i:i+batch]
        # shape (batch, N)
        s = np.sin(np.multiply.outer(xb, k)) / k
        out[i:i+batch] = s.sum(axis=1)
    return out

if __name__ == '__main__':
    xs = np.linspace(0, 0.1, 1000)
    print(len(f_vec(xs)))

Numba JIT: numba_impl.py

# numba_impl.py
import math
from numba import njit, prange

@njit(fastmath=True)
def f_scalar_numba(x, N=100000):
    s = 0.0
    for k in range(1, N+1):
        s += math.sin(x*k)/k
    return s

@njit(parallel=True, fastmath=True)
def compute_numba(xs, N=100000):
    out = [0.0]*len(xs)
    for i in prange(len(xs)):
        out[i] = f_scalar_numba(xs[i], N)
    return out

if __name__ == '__main__':
    xs = [i*0.0001 for i in range(1000)]
    print(len(compute_numba(xs)))

Multiprocessing: bench parallel map (bench.py)

# bench.py
import time
from concurrent.futures import ProcessPoolExecutor
import worker

def chunked_map(fn, xs, workers=None, chunk_size=10):
    with ProcessPoolExecutor(max_workers=workers) as ex:
        # submit in chunks to reduce submit overhead
        futures = []
        for i in range(0, len(xs), chunk_size):
            sub = xs[i:i+chunk_size]
            futures.append(ex.submit(lambda s: [fn(x) for x in s], sub))
        res = []
        for f in futures:
            res.extend(f.result())
    return res

if __name__ == '__main__':
    xs = [i*0.0001 for i in range(1000)]

    t0 = time.time()
    worker.main()  # baseline
    print('baseline', time.time()-t0)

    from numba_impl import compute_numba
    t0 = time.time()
    compute_numba(xs)
    print('numba', time.time()-t0)

    import numpy as np
    from numpy_impl import f_vec
    t0 = time.time()
    f_vec(np.array(xs))
    print('numpy', time.time()-t0)

    t0 = time.time()
    def wrap(x):
        return worker.f_scalar(x)
    chunked_map(wrap, xs, workers=4, chunk_size=25)
    print('multiprocessing', time.time()-t0)

Notas clave y por qué funcionan

• GIL: los bucles Python puros están limitados por el GIL. Multiprocessing crea procesos independientes para evitarlo.
• NumPy mueve el trabajo a código C vectorizado y reduce overhead por elemento.
• Numba compila tus loops a código nativo (LLVM), manteniendo la lógica en Python y evitando el overhead de Python por iteración.
• La paralelización tiene overhead de serialización: agrupa tareas en chunks para amortiguarlo.

Buenas prácticas y consejos

• Mide antes y después: usa time.perf_counter() o timeit, y perfila con cProfile o pyinstrument.
• Evita crear arrays gigantescos con NumPy sin pensar en memoria: procesa por batches.
• Con Numba, habilita fastmath=True solo si toleras pequeñas diferencias numéricas.
• En multiprocessing, evita pasar grandes estructuras frecuentemente; usa inicializadores o memoria compartida si hace falta (multiprocessing.shared_memory o sharedctypes).

Dependencias

# requirements.txt
numpy
numba

Cómo ejecutar

1. Instala dependencias: pip install -r requirements.txt
2. Corre el profiler: python profiler.py
3. Compara implementaciones: python bench.py

Errores comunes

• Esperar mejoras sin perfilar: optimizar código que no es hotspot no ayuda.
• Usar multiprocessing para tareas muy cortas (alto overhead de IPC).
• Ignorar el coste de memoria al vectorizar: puede degradar el rendimiento por swapping.

Si quieres profundizar: prueba pybind11 para extender con C++ si Numba/NumPy no cubren tu caso, o cython si prefieres control fino. Ten en cuenta la reproducibilidad numérica y detalles de inicialización al paralelizar: semillas aleatorias y orden de reducción pueden cambiar resultados.

Consejo avanzado: mide con contadores de hardware (perf, linux perf_events) para identificar misses de caché y vectorizabilidad; muchas veces reorganizar datos (AoS -> SoA) produce ganancias significativas.