Approach: provide both a decorator and a context manager that measure elapsed time and emit a concise, structured log record (JSON string) including dataset_id, epoch, step, phase (train/eval), and optional metrics. Use standard logging; format messages as JSON for easy downstream parsing.python
import time
import json
import logging
from functools import wraps
from contextlib import ContextDecorator
# configure root logger once
logger = logging.getLogger("trainer")
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter("%(message)s")) # message will be JSON
logger.addHandler(handler)
logger.setLevel(logging.INFO)
def make_record(elapsed, dataset_id=None, epoch=None, step=None, phase=None, **metrics):
rec = {
"ts": int(time.time()*1000),
"elapsed_ms": int(elapsed*1000),
"dataset_id": dataset_id,
"epoch": epoch,
"step": step,
"phase": phase,
}
rec.update(metrics)
# remove None values for conciseness
return {k: v for k, v in rec.items() if v is not None}
class timed_context(ContextDecorator):
def __init__(self, dataset_id=None, epoch=None, step=None, phase=None, **metrics):
self.dataset_id = dataset_id
self.epoch = epoch
self.step = step
self.phase = phase
self.metrics = metrics
def __enter__(self):
self._start = time.perf_counter()
return self
def __exit__(self, exc_type, exc, tb):
elapsed = time.perf_counter() - self._start
logger.info(json.dumps(make_record(elapsed, **self.__dict__)))
return False # don't suppress exceptions
def timed(dataset_id=None, epoch=None, step=None, phase=None, **static_metrics):
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
start = time.perf_counter()
result = func(*args, **kwargs)
elapsed = time.perf_counter() - start
# allow function to return metrics (e.g., loss) that we include
metrics = {}
if isinstance(result, dict):
metrics = result
record = make_record(elapsed, dataset_id=dataset_id, epoch=epoch, step=step, phase=phase, **static_metrics, **metrics)
logger.info(json.dumps(record))
return result
return wrapper
return decorator
python
# context manager usage inside training loop
for epoch in range(1, 3):
for step, batch in enumerate(train_loader, 1):
with timed_context(dataset_id="mnist", epoch=epoch, step=step, phase="train"):
train_step(batch) # heavy compute
# decorator usage for evaluation function returning metrics
@timed(dataset_id="mnist", phase="eval")
def evaluate(data):
loss = 0.123 # compute
return {"loss": loss, "examples": len(data)}
