Elegant tiny ml-experiment framework

experiment.py

import hashlib
import importlib
import json
import logging
import os
import random
import time
from contextlib import contextmanager
from dataclasses import asdict, dataclass, field, make_dataclass
from datetime import datetime, timezone
from enum import Enum
from os import makedirs
from os.path import exists, join
from pathlib import Path
from typing import Any, Iterator, Protocol, Self, TypeVar

import msgspec.json
import numpy as np
import torch
import torch.nn as nn
from polars import DataFrame
from send2trash import send2trash
from torch import Tensor

from symd.logging import (
    LogLevel,
    debug,
    info,
    init_logger,
    init_wandb_handler,
    scoped_jsonl_handler,
    warn,
)


def class_to_str(cls: type) -> str:
    """Serialize a class into a fully qualified string."""
    return f'{cls.__module__}.{cls.__qualname__}'


def str_to_class(path: str) -> type:
    """Deserialize a class string back into the class object."""
    module_name, _, class_name = path.rpartition('.')
    module = importlib.import_module(module_name)
    return getattr(module, class_name)


class URLable(Protocol):
    def url(self) -> Path: ...


class IDable(Protocol):
    def id(self) -> str: ...


class Serdeable(Protocol):
    def serialize(self) -> dict[str, Any]: ...

    @classmethod
    def deserialize(cls: type[Self], d: dict[str, Any]) -> Self: ...


I = TypeVar('I', bound=Serdeable)
O = TypeVar('O', bound=Serdeable)


class Experiment(Protocol[I, O]):
    def __call__(self, cfg: I, root: Path) -> O: ...
    def url(self) -> Path: ...


def run_hash() -> str:
    """Timestamp + short hash to avoid collisions."""
    ts = datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S')
    entropy = f'{time.time_ns()}:{os.getpid()}:{random.getrandbits(64)}'
    digest = hashlib.blake2b(
        entropy.encode('utf-8'), digest_size=5
    ).hexdigest()
    return f'{ts}-{digest}'


def _json_default(obj: Any) -> Any:
    if isinstance(obj, Enum):
        return obj.name
    if isinstance(obj, Path):
        return str(obj)
    if isinstance(obj, torch.device):
        return str(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, (np.number,)):
        return obj.item()
    if isinstance(obj, Tensor):
        return obj.detach().cpu().tolist()
    if isinstance(obj, type):
        return class_to_str(obj)
    serialize = getattr(obj, 'serialize', None)
    if callable(serialize):
        return serialize()
    return str(obj)


def run(exp: Experiment[I, O], cfg: I, root: Path | None = None) -> O:
    """Run exp(cfg), persist inputs/outputs, return output."""
    root = root if root is not None else Path(os.getcwd())
    exp_dir = root / exp.url()

    log_filepath = exp_dir / 'log.jsonl'
    exp_dir.mkdir(parents=True, exist_ok=True)

    with scoped_jsonl_handler(log_filepath):
        info({'event': 'experiment_start', 'exp_dir': exp_dir})

        input_cfg_filepath = exp_dir / 'input.json'
        debug('Serializing input config ...')
        input_cfg_dict = cfg.serialize()
        with input_cfg_filepath.open('w', encoding='utf-8') as f:
            json.dump(
                input_cfg_dict,
                f,
                indent=2,
                sort_keys=True,
                ensure_ascii=False,
                default=_json_default,
            )
        info({'event': 'wrote_input', 'path': input_cfg_filepath})

        output = exp(cfg, root)

        debug('Serializing output config ...')
        output_cfg_filepath = exp_dir / 'output.json'
        output_cfg_dict = output.serialize()
        with output_cfg_filepath.open('w', encoding='utf-8') as f:
            json.dump(
                output_cfg_dict,
                f,
                indent=2,
                sort_keys=True,
                ensure_ascii=False,
                default=_json_default,
            )
        info({'event': 'wrote_output', 'path': output_cfg_filepath})

        return output


def load_experiment(
    exp: Experiment[I, O],
    input_cls: type[I],
    output_cls: type[O],
    root: Path | None = None,
) -> tuple[I, O]:
    root = root if root is not None else Path(os.getcwd())
    exp_dir = root / exp.url()
    input_path = exp_dir / 'input.json'
    output_path = exp_dir / 'output.json'
    with input_path.open('r') as f:
        input_cfg_dict: dict[str, Any] = json.load(f)
    input_cfg = input_cls.deserialize(input_cfg_dict)
    with output_path.open('r') as f:
        output_cfg_dict: dict[str, Any] = json.load(f)
    output_cfg = output_cls.deserialize(output_cfg_dict)
    # for run_dir in sorted([p for p in exp_dir.iterdir() if p.is_dir()]):
    #     if not input_path.exists() or not output_path.exists():
    #         continue
    #
    #     with input_path.open('r') as f:
    #         input_cfg_dict: dict[str, Any] = json.load(f)
    #     input_cfg = input_cls.deserialize(input_cfg_dict)

    return input_cfg, output_cfg