Del print() a un logger thread-safe: decoradores para tracking de experimentos

Un modelo entrena seis horas y sale con la mejor métrica del mes. Nadie anotó el learning rate, el seed ni la versión del dataset que se usó esa tarde, porque el experimento se lanzó desde una notebook con cuatro celdas tocadas a mano. Tres semanas después, cuando alguien intenta reproducir ese resultado para el informe interno o para el despliegue, el experimento que funcionó ya no existe en ningún sitio verificable: queda un número suelto y la sospecha de que fue el batch_size más pequeño.

El problema no es falta de disciplina individual: anotar parámetros a mano no escala. Cada función de entrenamiento nueva exige repetir el mismo boilerplate de logging, y en algún punto alguien lo salta. Los decoradores en Python resuelven justo esta clase de problema porque envuelven la función de entrenamiento sin tocar su cuerpo: el registro de parámetros, duración y resultado deja de depender de que quien lanza el experimento se acuerde de hacerlo. Lo que sigue no es el primer decorador que se le ocurre a cualquiera —ese falla rápido—, sino cómo cambia en tres iteraciones sucesivas, cada una forzada por un fallo concreto de la anterior.

v1: un decorador que solo imprime por pantalla

La versión mínima envuelve la función de entrenamiento, guarda los parámetros de la llamada y mide cuánto tarda:

import time
from typing import Any, Callable, Dict


def experiment_tracker(func: Callable) -> Callable:
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        parameters: Dict[str, Any] = kwargs.copy()
        print(f"[Tracking Start] Parámetros: {parameters}")
        start_time = time.time()

        result = func(*args, **kwargs)

        duration = time.time() - start_time
        print(f"[Tracking End] Duración: {duration:.4f} segundos")
        return result
    return wrapper


@experiment_tracker
def train_model(epochs: int, lr: float) -> None:
    for _ in range(epochs):
        time.sleep(0.1)


train_model(epochs=5, lr=0.001)

Esto ya resuelve el olvido: los parámetros quedan visibles sin que el cuerpo de train_model sepa que existe un sistema de tracking. Pero el fallo aparece en cuanto hay más de un experimento a la vez. Si se lanzan dos entrenamientos en paralelo —dos pestañas de terminal, o un bucle que hace python train.py & dos veces seguidas— los print de ambos procesos se intercalan en el mismo stdout, o en el log agregado del scheduler que los recoja. No hay forma de saber a cuál de los dos pertenece cada línea de parámetros, y el resultado bueno de un experimento queda mezclado con las líneas del otro. Imprimir por pantalla ayuda mientras alguien mira la terminal, pero deja de servir en cuanto hace falta comparar corridas más tarde: es una nota que se borra sola.

v2: persistir con MLflow en vez de imprimir

MLflow persiste en un backend real en vez de un stream de texto que desaparece al cerrar la terminal. La instalación, el tracking server y el resto de fundamentos de MLflow están cubiertos en la guía de MLflow desde cero; aquí el interés es otro: cómo cambia la forma del decorador, cómo capturar también los argumentos posicionales, y qué falla en cuanto se paraleliza. El decorador queda así:

import inspect
import random
import time
from functools import wraps
from typing import Any, Callable

import mlflow


def experiment_mlflow(func: Callable) -> Callable:
    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        bound = inspect.signature(func).bind(*args, **kwargs)
        bound.apply_defaults()
        with mlflow.start_run():
            mlflow.log_params(dict(bound.arguments))
            result = func(*args, **kwargs)
            if isinstance(result, dict):
                for metric, value in result.items():
                    if isinstance(value, (int, float)) and not isinstance(value, bool):
                        mlflow.log_metric(metric, value)
            return result
    return wrapper


@experiment_mlflow
def train_model_mlflow(epochs: int, lr: float) -> dict:
    time.sleep(0.1)
    accuracy = 0.85 + random.random() * 0.1
    return {"accuracy": accuracy}


train_model_mlflow(epochs=5, lr=0.001)

Dos correcciones respecto a la versión ingenua de este decorador: inspect.signature(func).bind(*args, **kwargs) con apply_defaults() captura también los parámetros que llegan posicionalmente —mlflow.log_params(kwargs) a secas pierde cualquier argumento que no se pase como keyword—; y el filtro isinstance(value, (int, float)) and not isinstance(value, bool) evita que log_metric() reviente en cuanto el diccionario de retorno trae un string, un None o un booleano en vez de un número. mlflow.start_run(), log_params() y log_metric() siguen siendo la API fluida vigente de MLflow, que en 2026 exige Python 3.10 o superior. Si el equipo ya trabaja con Weights & Biases o TensorBoard en lugar de MLflow, la forma del decorador no cambia, solo las tres llamadas de persistencia: con wandb.init() el patrón actual es with wandb.init(project=..., config=kwargs) as run: run.log(result); con TensorBoard, el equivalente es crear un SummaryWriter y llamar a writer.add_scalar(tag, scalar_value, global_step) por cada métrica, usando el nombre de la métrica como tag.

El fallo de v2 no es de persistencia, es de concurrencia dentro del mismo proceso. La propia documentación de la API fluida de MLflow lo advierte sin rodeos: "the fluent tracking API is not currently threadsafe. Any concurrent callers to the tracking API must implement mutual exclusion manually", y describe mlflow.active_run() como thread-local: cada hilo solo ve la corrida que él mismo abrió. Eso deja inútil la receta ingenua de anidar con nested=True a secas dentro de cada worker: una corrida padre abierta en el hilo principal no aparece como activa dentro de un ThreadPoolExecutor, así que no hay nada con qué anidar. La propia guía de patrones de tracking de MLflow tiene un ejemplo llamado justamente "child runs for thread-safe parallel execution" que hace exactamente eso —nested=True sin más, dentro de un ThreadPoolExecutor—, lo que contradice el comportamiento thread-local que documenta su propia referencia de API. Pero pasar parent_run_id de forma explícita tampoco resuelve el problema de fondo: seguiría siendo la API fluida (start_run, log_params, log_metric) la que se llama de forma concurrente desde varios hilos, y esa es exactamente la API que la documentación marca como no thread-safe, con exclusión mutua manual exigida sin importar el patrón de anidado que se use encima. Paralelizar de verdad las llamadas de tracking exige sacar la API fluida del código concurrente y usar MlflowClient, la API de bajo nivel, con un run_id explícito en cada llamada:

import random
import time
from concurrent.futures import ThreadPoolExecutor

import mlflow
from mlflow import MlflowClient

client = MlflowClient()
experiment_id = mlflow.set_experiment("hparam-sweep").experiment_id
configs = [{"lr": 0.01}, {"lr": 0.05}, {"lr": 0.1}, {"lr": 0.2}]


def train_worker(config: dict) -> float:
    run_id = client.create_run(experiment_id).info.run_id
    for key, value in config.items():
        client.log_param(run_id, key, value)
    time.sleep(0.1)
    accuracy = 0.85 + random.random() * 0.1
    client.log_metric(run_id, "accuracy", accuracy)
    client.set_terminated(run_id)
    return accuracy


with ThreadPoolExecutor(max_workers=4) as executor:
    scores = list(executor.map(train_worker, configs))

Cada llamada de client recibe su run_id explícito: no hay ninguna corrida "activa" que resolver por hilo, así que no hace falta ningún candado alrededor de estas llamadas para que el tracking sea seguro entre hilos —la exclusión mutua que exige la documentación es un problema específico de la API fluida sin run_id, no de MlflowClient. La alternativa sin cambiar de API es forzar la exclusión mutua a mano: envolver cada llamada fluida (start_run, log_params, log_metric) en un único threading.Lock global, serializando el acceso a la API fluida en vez de evitarla. Funciona, pero el paralelismo real queda limitado a lo que cada hilo hace fuera de ese candado —normalmente el entrenamiento en sí—, no al tracking.

v3: un logger propio y thread-safe para cuando sí hace falta paralelismo

Aquí es donde tiene sentido un logger de fichero propio, sin depender de un servidor de tracking, protegido con un threading.Lock para que los hilos que comparten esa instancia nunca escriban en el mismo archivo a la vez. La primera versión de este componente tenía seis bugs reales que conviene declarar en vez de solo corregir en silencio, y el más importante de los seis no es de código sino de qué datos entran al log:

• Candado a nivel de clase: declarar _lock = threading.Lock() como atributo de clase hace que todas las instancias de ExperimentLogger compartan el mismo candado, aunque escriban en ficheros distintos. Eso serializa escrituras que no tienen ninguna relación entre sí, justo lo contrario de lo que se busca al añadir paralelismo. El candado debe crearse por instancia, dentro de __init__.
• Vuelca datos crudos, y truncar por longitud no basta: guardar args/kwargs completos y str(result) completo puede reventar con TypeError en cuanto aparece un array de NumPy o un DataFrame —parcheable con default=str—, y puede volcar PII, tokens de API o la representación íntegra de un dataset al fichero de log. Limitar por tamaño no resuelve el problema de fondo: un token o una contraseña cortos pasan igual de crudos si el único criterio es la longitud. La solución no es un límite de caracteres, es no loguear ningún string por defecto salvo que se declare explícitamente qué campos son loggeables, y redactar por nombre de campo —token, password, secret, key, authorization, email y variantes— incluso si el campo se declaró loggeable. Este es el fallo que de verdad justifica diseñar la captura con cuidado, no solo envolver la función.
• Sin manejo de errores en la función decorada: si la función decorada lanza una excepción, el logger.log() nunca se ejecuta y el experimento fallido no deja ningún rastro, ni siquiera de que se intentó. Envolver la llamada en try/except/finally permite registrar también los fallos, con su duración hasta el momento del error.
• El propio log() sin proteger dentro del finally: si escribir a disco falla —disco lleno, permisos, ruta inválida— justo en ese finally, la nueva excepción sustituye a la que se estaba propagando y la oculta. El logging necesita su propio try/except que nunca deje que un fallo de I/O tape el error real del entrenamiento.
• datetime.utcnow() deprecado: Python lo marca obsoleto desde la versión 3.12 en favor de datetime.now(timezone.utc); sigue funcionando hoy, pero cualquier código nuevo debería evitarlo.
• Sin @wraps(func): el wrapper interno no preservaba el nombre ni los metadatos de la función decorada, así que train_model_advanced.__name__ devolvía "wrapper" en vez del nombre real —rompe cualquier herramienta que dependa de introspección: depuradores, otros decoradores, frameworks de testing—.

Con los seis corregidos, el componente queda así:

import datetime
import inspect
import json
import sys
import threading
from contextlib import contextmanager
from functools import wraps
from typing import Any, Callable, Dict, FrozenSet, TypeVar

F = TypeVar("F", bound=Callable[..., Any])

MAX_STR_LEN = 200
SENSITIVE_NAME_MARKERS = ("token", "password", "secret", "key", "authorization", "email", "credential")


def _is_sensitive_name(name: str) -> bool:
    lowered = name.lower()
    return any(marker in lowered for marker in SENSITIVE_NAME_MARKERS)


def _safe_value(name: str, value: Any, log_fields: FrozenSet[str]) -> Any:
    """Reduce un valor a algo corto y seguro. Redacta por defecto: ningún string
    sale salvo que su campo esté en la allowlist, y un nombre sensible se
    redacta aunque esté en la allowlist."""
    if _is_sensitive_name(name):
        return ""
    if value is None or isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        return value
    if isinstance(value, str):
        if name not in log_fields:
            return ""
        return value if len(value) <= MAX_STR_LEN else value[:MAX_STR_LEN] + "...(truncado)"
    return f"<{type(value).__name__}>"  # nunca el repr completo: puede ser un dataset, un tensor o PII


def _safe_params(
    func: Callable, args: Any, kwargs: Dict[str, Any], log_fields: FrozenSet[str]
) -> Dict[str, Any]:
    bound = inspect.signature(func).bind(*args, **kwargs)
    bound.apply_defaults()
    return {name: _safe_value(name, value, log_fields) for name, value in bound.arguments.items()}


def _safe_result(result: Any) -> Any:
    if isinstance(result, dict):
        return {
            key: value for key, value in result.items()
            if isinstance(value, (int, float)) and not isinstance(value, bool)
        }
    return f"<{type(result).__name__}>"


class ExperimentLogger:
    """Persiste entradas de tracking en un fichero JSON Lines.

    Thread-safe si todos los hilos comparten esta misma instancia: el candado
    es por instancia, no por fichero. Dos instancias distintas apuntando al
    mismo filename no se coordinan entre sí.
    """

    def __init__(self, filename: str):
        self.filename = filename
        self._lock = threading.Lock()  # por instancia: cada fichero tiene su propio candado

    @contextmanager
    def open_log(self):
        with self._lock, open(self.filename, "a") as f:
            yield f

    def log(self, data: Dict[str, Any]) -> None:
        with self.open_log() as f:
            json.dump(data, f, default=str)  # red de seguridad final, no la defensa principal
            f.write("\n")


def experiment_tracker_advanced(
    logger: ExperimentLogger, log_fields: FrozenSet[str] = frozenset()
) -> Callable[[F], F]:
    """log_fields es la allowlist explícita de nombres de parámetro que sí se
    loguean como string; cualquier string fuera de esa lista se redacta, y un
    nombre sensible se redacta aunque esté en la lista."""

    def decorator(func: F) -> F:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            start_time = datetime.datetime.now(datetime.timezone.utc)
            status = "success"
            result = None
            try:
                result = func(*args, **kwargs)
                return result
            except Exception:
                status = "error"
                raise
            finally:
                end_time = datetime.datetime.now(datetime.timezone.utc)
                entry = {
                    "function": func.__name__,
                    "params": _safe_params(func, args, kwargs, log_fields),
                    "status": status,
                    "result": _safe_result(result) if result is not None else None,
                    "start": start_time.isoformat(),
                    "end": end_time.isoformat(),
                    "duration_sec": (end_time - start_time).total_seconds(),
                }
                try:
                    logger.log(entry)
                except Exception as log_error:
                    print(f"[ExperimentLogger] fallo al escribir el log: {log_error!r}", file=sys.stderr)
        return wrapper
    return decorator


logger = ExperimentLogger("experiments.jsonl")


@experiment_tracker_advanced(logger, log_fields=frozenset({"epochs", "lr"}))
def train_model_advanced(epochs: int, lr: float) -> dict:
    import random
    import time
    time.sleep(0.1)
    accuracy = 0.85 + random.random() * 0.1
    return {"accuracy": accuracy}


train_model_advanced(epochs=5, lr=0.001)

La garantía de thread-safety depende de compartir una única instancia de ExperimentLogger entre los hilos que escriben en el mismo fichero, como hace el ejemplo con la variable logger global: si en vez de eso se crean dos instancias distintas apuntando al mismo filename, cada una trae su propio candado y no se coordinan entre sí, con lo que las escrituras pueden interleavearse igual que en v1. El diseño de _safe_value es redactar por defecto: ningún string se loguea salvo que su nombre de parámetro esté en log_fields —la allowlist que declara quien usa el decorador—, y aun así se redacta si el nombre coincide con un patrón sensible (token, password, secret, key, authorization, email...). Es intencional: la sensibilidad de un dato depende de qué campo es, no de cuántos caracteres tiene, así que un token corto como "sk-abc123" se redacta igual que uno largo. El default=str en json.dump queda como red de seguridad para el caso residual —un tipo verdaderamente inesperado que ni siquiera pasó por _safe_value—, no como la defensa principal. Y conviene ser honesto sobre el coste que queda incluso compartiendo instancia: la escritura a disco ocurre dentro del candado, así que un hilo bloquea a los demás mientras dura el open() y el write(). Para volúmenes de logging altos esto se convierte en cuello de botella; la salida es sacar la escritura del camino caliente con una cola y un hilo dedicado, no seguir apilando lógica dentro del candado. Separar el ExperimentLogger del decorador también deja abierta la puerta a cambiar el backend —de fichero a base de datos, o a un bucket— sin tocar la firma del decorador ni el código de entrenamiento.

Qué versión usar según el caso

Las tres versiones no son necesariamente excluyentes: es razonable seguir escribiendo a MLflow con la API de cliente y run_id explícito por hilo, y en paralelo mantener un ExperimentLogger propio como registro de respaldo sin dependencias externas.

El hilo común de las tres versiones es el mismo: ningún resultado de un experimento vale nada si no queda escrito en algún sitio antes de que alguien intente reproducirlo. El decorador solo decide dónde se escribe y qué tan bien sobrevive a la concurrencia; la disciplina de no lanzar un entrenamiento sin él es lo que realmente evita perder el experimento irrepetible.