V2: tiny ml-experiment framework

experiment_v2.py

from __future__ import annotations

import base64
import dataclasses
import hashlib
import importlib
import json
import os
import re
import tempfile
import traceback
from dataclasses import asdict, field
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path, PosixPath, PurePosixPath
from typing import (
    Any,
    Callable,
    Iterator,
    Protocol,
    Self,
    Type,
    TypeVar,
    assert_type,
    cast,
    dataclass_transform,
    final,
    get_type_hints,
    overload,
)

import jsonlines
import numpy as np
import serde as serde_mod
import torch
from plum import dispatch
from serde import from_dict, serde, to_dict
from serde.json import from_json, to_json
from torch import Tensor

from symd.logging import error, info, warn


class PyTSerializer:
    @dispatch
    def serialize(self, value: torch.Tensor) -> Any:
        return value.tolist()


class PyTDeserializer:
    @dispatch
    def deserialize(self, cls: type[torch.Tensor], value: Any) -> torch.Tensor:
        return torch.tensor(value)


class PathSerializer:
    @dispatch
    def serialize(self, value: PosixPath) -> str:
        return str(value)


class PathDeserializer:
    @dispatch
    def deserialize(self, cls: type[PosixPath], value: Any) -> PosixPath:
        return PosixPath(value)


serde_mod.add_serializer(PyTSerializer())
serde_mod.add_deserializer(PyTDeserializer())

serde_mod.add_serializer(PathSerializer())
serde_mod.add_deserializer(PathDeserializer())


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


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


class Serdeable(Protocol):
    def serialize(self, as_dict: bool = False) -> str | dict[str, Any]: ...

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


class IDableAndSerdeable(IDable, Serdeable, Protocol):
    pass


TYPE_TAG = '_type'

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


def deserialize(cls: type[T], d: str | dict[str, Any]) -> T:
    if isinstance(d, dict):
        # _ = d.pop(TYPE_TAG, None)
        return from_dict(cls, d)

    # Now the only indirection is that serialize method
    # where I add the _type tag, requireing a obj -> dict -> str
    # pipeline and using stdlib js vs. whaqtever pystre ues
    # hope that is ok
    return from_json(cls, d)

    ## Pyserde will stip the _type tag for us automatcally so all is well.
    # d1 = json.loads(d)
    # _c = d1.pop(TYPE_TAG, None)
    # if _c is not None and str_to_class(_c) != cls:
    #     warn(
    #         (
    #             'Type tag does not match provided class type.'
    #             'This will probably error',
    #             'Type-Tag: ',
    #             class_to_str(_c),
    #             _c,
    #             'Class Type:',
    #             class_to_str(cls),
    #             cls,
    #         )
    #     )
    #
    # # For instance lam = 0.0  will serioalize to json float
    # # bu t lam 0 will go to json int
    # return from_dict(
    #     cls, d1
    # )  # TODO: consider setting coerce for floating point


# Todo maybe overload this
def serialize(self: T, as_dict: bool = False) -> str | dict[str, Any]:
    # Theres def a better way to do this, lik hook into the serializer
    # without having to convert it to a dict
    # safeer too
    d = to_dict(self)
    # tag = getattr(type(self), '__type_tag__', None) or class_to_str(type(self))
    tag = class_to_str(type(self))
    d[TYPE_TAG] = tag
    # d = to_json(self)
    # d = json.loads(d)
    # d['_type'] = class_to_str(type(self))
    if as_dict:
        return d
    d = json.dumps(d, separators=(',', ':'), sort_keys=True)
    return d


def class_to_str(cls: type) -> str:
    return f'{cls.__module__}.{cls.__qualname__}'


def str_to_class(path: str) -> type[T]:
    module_name, _, class_name = path.rpartition('.')
    module = importlib.import_module(module_name)
    return cast(type[T], getattr(module, class_name))


@overload
def make_serdeable(cls: type[T], **pyserde_args: Any) -> type[T]: ...


@overload
def make_serdeable(
    cls: Any = None, **pyserde_args: Any
) -> Callable[[type[T]], type[T]]: ...


@dataclass_transform(field_specifiers=(field,))
def make_serdeable(cls: Any = None, **pyserde_args: Any) -> Any:
    def wrap(C: type[T]) -> type[T]:
        C = cast(
            type[T], serde(C, **pyserde_args)
        )  # ensure pyserde codegen on the class

        @classmethod
        def _de(cls: type[T], d: str | dict[str, Any]) -> T:
            return deserialize(cls, d)

        def _se(self: T, as_dict: bool = False) -> str | dict[str, Any]:
            return serialize(self, as_dict)

        C.serialize = _se  # ty:ignore[unresolved-attribute]
        C.deserialize = _de  # ty:ignore[unresolved-attribute]
        return C

    return wrap if cls is None else wrap(cls)


def fingerprint(obj: Any, *, n_chars=16) -> str:
    d = serialize(obj)
    h = hashlib.sha256(d.encode('utf-8')).digest()
    tok = base64.urlsafe_b64encode(h).decode('ascii').rstrip('=')
    return tok[:16]


# make the narrow case first so the checkers that are depdne on order
# can noarrow first
@overload
def make_idable(cls: type[T]) -> type[T]: ...


@overload
def make_idable(cls: Any = None) -> Callable[[type[T]], type[T]]: ...


def make_idable(cls: Any = None) -> Any:
    def wrap(C: type[T]) -> type[T]:
        def _id(self: T) -> str | PurePosixPath:
            return fingerprint(self)

        C.id = _id  # ty:ignore[unresolved-attribute]
        return C

    return wrap if cls is None else wrap(cls)


@overload
def experiment_input(cls: type[T]) -> type[T]: ...


@overload
def experiment_input(cls: Any = None) -> Callable[[type[T]], type[T]]: ...


# TODO .....
@dataclass_transform(field_specifiers=(field,))
def experiment_input(cls: Any = None) -> Any:
    def wrap(C: type[T]) -> type[T]:
        C = make_serdeable(C)  # ty:ignore[invalid-assignment]
        C = make_idable(C)
        return C

    return wrap if cls is None else wrap(cls)


def _default_run_grp_id(fn: Callable[..., Any]) -> str:
    module = getattr(fn, '__module__', fn.__class__.__module__)
    qualname = getattr(fn, '__qualname__', None)
    name = getattr(fn, '__name__', None)
    ident = qualname or name or fn.__class__.__qualname__
    return f'{module}:{ident}'


class BaseContext:
    __slots__ = ('_root_dir', '_run_dir', '_run_id')

    def __init__(self) -> None:
        self._root_dir: Path | None = None
        self._run_dir: Path | None = None
        self._run_id: PurePosixPath | None = None

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        forbidden = {'set_root', 'set_run_dir', 'root', 'run_dir'}
        for name in forbidden:
            if name in cls.__dict__:
                raise TypeError(
                    f"{cls.__name__} must not override final member '{name}'"
                )

    @final
    def set_run_id(self, run_id: PurePosixPath):
        self._run_id = run_id

    @property
    @final
    def run_id(self) -> PurePosixPath:
        assert self._run_id is not None, 'Context not initialized by runner'
        return self._run_id

    @property
    @final
    def datasets_dir(self) -> Path:
        return self.root_dir / 'datasets'

    @property
    @final
    def root_dir(self) -> Path:
        assert self._root_dir is not None, 'Context not initialized by runner'
        return self._root_dir

    @property
    @final
    def run_dir(self) -> Path:
        assert self._run_dir is not None, 'Context not initialized by runner'
        return self._run_dir

    def embody(self, root: Path) -> None:
        if not root.is_absolute():
            warn(
                (
                    'Root is not resolved by the surrounding caller'
                    'That is a fine pattern if you want us to control the hydration'
                    'Of the environment, but it might make more sense to do it higher'
                    'for your own sake'
                )
            )
        # Hydrate those absolute paths
        self._root_dir = root.resolve()
        self._run_dir = (self.root_dir / self.run_id).resolve()

        self.root_dir.mkdir(parents=True, exist_ok=True)
        self.run_dir.mkdir(parents=True, exist_ok=True)
        self.datasets_dir.mkdir(parents=True, exist_ok=True)


name_part_re = re.compile(r'^[a-zA-Z0-9._-]+$')
run_part_re = re.compile(r'^[a-zA-Z0-9._=-]+$')


class DefaultContext(BaseContext):
    def append_to_json_lines_file(self, name: str | Path, obj: Any) -> Path:
        rel = self.validate_name(name, default_suffix='.jsonl')
        p = self.run_dir / rel
        with jsonlines.open(p, 'a') as f:
            f.write(serialize(obj, as_dict=True))
        return p

    def save_json(self, name: str | Path, obj: Any) -> Path:
        rel = self.validate_name(name, default_suffix='.json')
        p = self.run_dir / rel
        self.atomic_write_text(
            p,
            json.dumps(
                serialize(obj, as_dict=True),
                indent=2,
                sort_keys=True,
                ensure_ascii=False,
            ),
        )
        return p

    def save_text(self, name: str | Path, text: str) -> Path:
        rel = self.validate_name(name, default_suffix='.txt')
        p = self.run_dir / rel
        self.atomic_write_text(p, text)
        return p

    def save_dataset(self, name: str | Path, obj: Any) -> Path:
        rel = self.validate_name(name, default_suffix='.pt')
        p = self.datasets_dir / rel
        if p.exists():
            raise ValueError(f'Path at {p.resolve()} already exists.')
        p.parent.mkdir(parents=True, exist_ok=True)
        torch.save(obj, p)
        return p

    def load_dataset(
        self, name: str | Path, map_location: str | torch.device = 'cpu'
    ) -> Path:
        rel = self.validate_name(name, default_suffix='.pt')
        p = self.datasets_dir / rel
        return torch.load(p, map_location=map_location)

    def save_torch(self, name: str | Path, obj: Any) -> Path:
        rel = self.validate_name(name, default_suffix='.pt')
        p = self.run_dir / rel
        p.parent.mkdir(parents=True, exist_ok=True)
        torch.save(obj, p)
        return p

    def load_torch_data(
        self,
        rel_data_path: str | Path,
        *,
        map_location: str | torch.device = 'cpu',
    ) -> Any:
        rel = self.validate_name(rel_data_path)
        p = (self.datasets_dir / rel).resolve()
        return torch.load(p, map_location=map_location)

    @staticmethod
    def atomic_write_text(path: Path, text: str) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        fd, tmp = tempfile.mkstemp(
            dir=str(path.parent),
            prefix=path.name + '.',
            suffix='.tmp',
        )
        try:
            with os.fdopen(fd, 'w', encoding='utf-8') as f:
                f.write(text)
            os.replace(tmp, path)
        finally:
            try:
                os.remove(tmp)
            except FileNotFoundError:
                pass

    @staticmethod
    def validate_name(
        name: str | Path, *, default_suffix: str | None = None
    ) -> Path:
        p = Path(name)
        if p.is_absolute():
            raise ValueError(f'name must be relative, got absolute path: {p}')
        if any(part in ('..', '', '.') for part in p.parts):
            raise ValueError(f'illegal name components in: {p}')
        for part in p.parts:
            if not name_part_re.match(part):
                raise ValueError(
                    f"illegal chars in component '{part}' (name={p})"
                )
        if default_suffix is not None and p.suffix == '':
            p = p.with_suffix(default_suffix)
        return p


CtxT = TypeVar('CtxT', bound=BaseContext)
ExpFn = Callable[[I, CtxT], O]


# -----------------------------------------------------------------------------
# 6) Type-tagged payloads for input/output
# -----------------------------------------------------------------------------


def set_seed(seed: int):
    torch.manual_seed(seed)
    # random.??
    # numpy.??


# Think of it like this
# We start of with a word like
# --
# apple
# bananna
# --
# So the universe right now is U = {apple, bannana}
# If that is all we have, we cannot proceed. But what if we added more
# information into this world, where apple and bannana had an attribute.
# So now the world is U' ~= U x A, where A = {fruit, vegetable}
# What we see here on the left and right handside are actually two ways to represent
# the universe. On the left hand side think of it as the flattened view
# so that we have
# --
# fruit-apple
# fruit-bannana
# --
# On the right hand side, think of it as a quotienting or factorization of the universe.
# By adding a new attribute set, we now may speak of elements in terms of equivalency
# relations we can define.
# And so in the above example fruit-apple != fruit-bannana becuase the symbols are literally
# different, but if we define the relation u1 ~ u2 <=> (food-cat1, symbol1) (food-cat2, symbol2)
# where food-cat1 == footcat2, we indicue an equivalency relation.
#
# When we add the folder "fruit", we induce an equivalency relation.
# In order to introduce such a relation, we need to
# --
# fruit/apple
# fruit/banana
# --
# So now these two items are equal in the quotient.
# By giving identifiers a common prefix, we induce an equivalency relation.
# In other words, if I was walking through this tree, and i was at the first level,
# all I would see is a single element in my universe, the folder "fruit".
#
# So the upshot is, say an experiment was defined by two attributes. the set of allowed
# values of attr 1 is {x, y} and attr 2 is {a, b}. Baring you inject more attributes on which
# you assign your own meeting (precisely the rungrp id, You essentially have two choices of
# representation / identity  of two instances of x a and  x b.
# 1.
# x__a
# x__b
# or 2.
# x/a
# x/b
#
# Whats the advantage of each? in the first representation your world is flat.
# When is this advantageous? When you wish to think of just one notion of equiality.
# But say x were the experiment type. there it is natural to talk about the points
# x__a, x__b (which represent a (computation, input) tuple), as *equivalent
# as experiments*.
# Whats great about this is that / quite literally induces a equivalence relation
# and we represent that in the filesystem. In practice this means that
# the fs at level 0 (x) is actually a view of the quoteitn space in terms of the coset
# reps.
#
#
type EquivalencyClassRelation[I] = Callable[[I], PurePosixPath]


def run_id(
    *,
    cfg: I,
    run_grp_id: PurePosixPath | None,
) -> PurePosixPath:
    ts = datetime.now().strftime('%Y_%m_%d-%I_%M_%S_%p')
    key = cfg.id()
    return (
        (run_grp_id / f'{ts}__{key}__run')
        if run_grp_id is not None
        else PurePosixPath(f'{ts}__{key}__run')
    )


def execute(
    *,
    fn: ExpFn[I, CtxT, O],
    cfg: I,
    root: Path | None = None,
    run_grp_id: PurePosixPath | str | None = None,
    input_key: str | PurePosixPath | None = None,
    ctx_cls: type[CtxT] = DefaultContext,  # type: ignore[assignment]
) -> tuple[Path, O]:
    """Run an experiment and persist input/output.

    Writes:
      - input.json
      - output.json
    """
    rid = run_id(
        cfg=cfg,
        run_grp_id=PurePosixPath(run_grp_id)
        if isinstance(run_grp_id, str)
        else run_grp_id,
    )
    root_dir = (root if root is not None else Path(os.getcwd())).resolve()
    info(f'Resolved root dir to {root_dir}')

    ctx = ctx_cls()
    ctx.set_run_id(rid)
    ctx.embody(root_dir)

    run_dir = ctx.run_dir
    with open(run_dir / 'input-pre.json', 'w') as f:
        (
            json.dump(
                cfg.serialize(as_dict=True),
                f,
                indent=2,
            ),
        )

    try:
        out = fn(cfg, ctx)
    except Exception:
        with open(run_dir / 'traceback.txt', 'w') as f:
            f.write(traceback.format_exc())
            error('Found exception, see traceback')
            error(f'Run directory: {run_dir}')
        raise

    with open(run_dir / 'input.json', 'w') as f:
        (
            json.dump(
                cfg.serialize(as_dict=True),
                f,
                indent=2,
            ),
        )

    with open(run_dir / 'output.json', 'w') as f:
        (
            json.dump(
                out.serialize(as_dict=True),
                f,
                indent=2,
            ),
        )
    info(out)
    info(f'Run_directory: {run_dir}')
    return run_dir, out


def run(
    fn: ExpFn[I, CtxT, O],
    cfg: I,
    *,
    root: Path | None = None,
    run_grp_id: PurePosixPath | str | None = None,
    ctx_cls: type[CtxT] = DefaultContext,  # type: ignore[assignment]
) -> tuple[Path, O]:
    return execute(
        fn=fn,
        cfg=cfg,
        root=root,
        run_grp_id=run_grp_id,
        ctx_cls=ctx_cls,
    )


def deserialize_tagged_dict(d: dict[str, Any]) -> Any:
    c = d.pop('_type', None)
    return from_dict(c, d)


def deserialize_tagged_json(s: str | Path) -> Any:
    s = s if isinstance(s, str) else s.read_text(encoding='utf-8')
    d = json.loads(s)  # s to get ty to be quite
    c = d.pop('_type', None)
    c = str_to_class(c)
    return from_dict(c, d)


def try_deserialize_tagged_json(s: str | Path) -> Any | None:
    try:
        return deserialize_tagged_json(s)
    except (
        FileNotFoundError,
        IsADirectoryError,
        NotADirectoryError,
        PermissionError,
        json.JSONDecodeError,
        KeyError,
        ModuleNotFoundError,
        AttributeError,
        TypeError,
        ValueError,
    ):
        return None


def get_run_dirs(
    root_dir: Path, run_grp_id: str | None | PurePosixPath = None
) -> Iterator[Path]:
    """Iterate all run directories under a root.

    A "run directory" is any directory whose name begins with `run-`.
    """
    base = root_dir.resolve()
    if run_grp_id is not None:
        base = base / run_grp_id

    if not base.exists():
        return

    run_dirs = [p for p in base.rglob('run-*') if p.is_dir()]
    run_dirs.sort(key=lambda p: str(p))
    yield from run_dirs


def get_completed_run_dirs(
    root_dir: Path, run_grp_id: PurePosixPath | str | None = None
) -> Iterator[Path]:
    """Iterate only completed run directories.

    A run is "completed" if it has both `input.json` and `output.json`.
    """
    for run_dir in get_run_dirs(root_dir, run_grp_id=run_grp_id):
        if (run_dir / 'input.json').is_file() and (
            run_dir / 'output.json'
        ).is_file():
            yield run_dir


def try_load_single_run(run_dir: Path) -> tuple[Any, Any, Path]:
    run_dir = run_dir.resolve()
    # inp_str = (run_dir / 'input.json').read_text(encoding='utf-8')
    # out_str = (run_dir / 'output.json').read_text(encoding='utf-8')
    return (
        try_deserialize_tagged_json((run_dir / 'input.json')),
        try_deserialize_tagged_json((run_dir / 'output.json')),
        run_dir,
    )


def load_single_run(run_dir: Path) -> tuple[Any, Any, Path]:
    run_dir = run_dir.resolve()
    # inp_str = (run_dir / 'input.json').read_text(encoding='utf-8')
    # out_str = (run_dir / 'output.json').read_text(encoding='utf-8')
    return (
        deserialize_tagged_json((run_dir / 'input.json')),
        deserialize_tagged_json((run_dir / 'output.json')),
        run_dir,
    )


def load_runs(
    root_dir: Path, run_grp_id: str | None = None
) -> list[tuple[I | None, O | None, Path]]:
    """Load all runs found under `root_dir`.

    Returns a list of triples `(input, output, run_dir)`. The `input` and/or
    `output` objects may be `None` if the corresponding file is missing or
    cannot be deserialized.
    """
    runs: list[tuple[I | None, O | None, Path]] = []
    for run_dir in get_run_dirs(root_dir, run_grp_id=run_grp_id):
        inp = deserialize_tagged_json(run_dir / 'input.json')
        out = deserialize_tagged_json(run_dir / 'output.json')
        runs.append((inp, out, run_dir))
    return runs


def load_completed_runs(
    root_dir: Path, run_grp_id: str | None = None
) -> list[tuple[Serdeable, Serdeable, Path]]:
    """Load only completed runs found under `root_dir`."""
    runs: list[tuple[Serdeable, Serdeable, Path]] = []
    for run_dir in get_completed_run_dirs(root_dir, run_grp_id=run_grp_id):
        inp, out, _ = load_single_run(run_dir)
        runs.append((inp, out, run_dir))
    return runs


#######################################################################
# Deprecated Things
#######################################################################
def _to_jsonable(obj: Any) -> Any:
    if obj is None or isinstance(obj, (str, int, bool)):
        return obj
    serialize = getattr(obj, 'serialize', None)
    if callable(serialize):
        return _to_jsonable(serialize())
    if isinstance(obj, float):
        if not (obj == obj and obj not in (float('inf'), float('-inf'))):
            raise ValueError(f'Non-finite float in payload: {obj}')
        return obj
    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.number,)):
        return obj.item()
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, Tensor):
        return obj.detach().cpu().tolist()
    if isinstance(obj, type):
        return class_to_str(obj)
    if dataclasses.is_dataclass(obj):
        return {k: _to_jsonable(v) for k, v in asdict(obj).items()}
    if isinstance(obj, dict):
        return {str(k): _to_jsonable(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_to_jsonable(v) for v in obj]
    return str(obj)