kalfa 3.0.0.dev0__py3-none-any.whl

kalfa/core/record.py

@@ -0,0 +1,347 @@
+import importlib.util
+import json
+import os
+import platform
+import shutil
+import socket
+import subprocess
+import sys
+import tempfile
+import threading
+from contextlib import contextmanager
+from datetime import datetime
+from io import StringIO
+from pathlib import Path
+
+import torch
+from ruamel.yaml import YAML
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
+
+from .errors import RunError
+from .refs import tokenize
+
+
+FILES = ("manifest.json", "resolved.yaml", "plan.json", "history.jsonl", "updates.jsonl", "events.jsonl")
+STOP_FILE = "stop.json"
+LOG_FILE = "log.txt"
+FINAL = ("done", "failed", "stopped")
+BEAT = 10.0
+
+
+def file_mode():
+    mask = os.umask(0)
+    os.umask(mask)
+    return 0o666 & ~mask
+
+
+MODE = file_mode()
+
+
+@contextmanager
+def atomic(path):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    handle, temporary = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp")
+    os.close(handle)
+    try:
+        yield Path(temporary)
+        os.chmod(temporary, MODE)
+        os.replace(temporary, path)
+    finally:
+        if os.path.exists(temporary):
+            os.unlink(temporary)
+
+
+def now():
+    return datetime.now().isoformat(timespec="seconds")
+
+
+def version():
+    from kalfa import __version__
+
+    return __version__
+
+
+class RecordWriter:
+    def __init__(self, directory):
+        self.directory = Path(directory)
+        self.lock = threading.RLock()
+        self.within = None
+
+    def path(self, name):
+        return self.directory / name
+
+    def write_json(self, name, value):
+        with atomic(self.path(name)) as temporary:
+            temporary.write_text(json.dumps(value, indent=1, default=str), encoding="utf-8")
+
+    def write_text(self, name, text):
+        with atomic(self.path(name)) as temporary:
+            temporary.write_text(text, encoding="utf-8")
+
+    def read_json(self, name):
+        path = self.path(name)
+        return json.loads(path.read_text(encoding="utf-8")) if path.exists() else None
+
+    def append(self, name, line):
+        path = self.path(name)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("a", encoding="utf-8") as handle:
+            handle.write(json.dumps(line, default=str) + "\n")
+
+    def event(self, kind, turn=None, **fields):
+        self.append("events.jsonl", {"at": now(), "turn": turn, "kind": kind, **fields})
+
+    def manifest(self, **changes):
+        with self.lock:
+            note = self.read_json("manifest.json") or {}
+            note.update(changes)
+            note["updated"] = now()
+            if note.get("status") in FINAL:
+                note.pop("within", None)
+            elif self.within is not None:
+                note["within"] = dict(self.within)
+            self.write_json("manifest.json", note)
+            return note
+
+    def note(self, turn, phase, done=None, total=None, unit=None):
+        before = self.within
+        started = before["started"] if before is not None and before.get("turn") == turn else now()
+        self.within = {"turn": turn, "phase": phase, "done": done, "total": total, "unit": unit, "started": started}
+
+    def beat(self):
+        with self.lock:
+            try:
+                note = self.read_json("manifest.json")
+            except (OSError, ValueError):
+                return False
+            if not isinstance(note, dict) or note.get("status") != "running":
+                return False
+            note["updated"] = now()
+            if self.within is not None:
+                note["within"] = dict(self.within)
+            try:
+                self.write_json("manifest.json", note)
+            except OSError:
+                return False
+            return True
+
+    def log(self, text):
+        path = self.path(LOG_FILE)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        stamp = now()
+        with self.lock, path.open("a", encoding="utf-8") as handle:
+            for line in str(text).splitlines() or [""]:
+                handle.write(f"{stamp} {line}\n")
+
+    def stop_requested(self):
+        return self.path(STOP_FILE).exists()
+
+    def stop_note(self):
+        return self.read_json(STOP_FILE)
+
+
+class Pulse:
+    def __init__(self, writer, every=None):
+        self.writer = writer
+        self.every = float(BEAT if every is None else every)
+        self.ended = threading.Event()
+        self.thread = None
+
+    def __enter__(self):
+        self.ended.clear()
+        self.thread = threading.Thread(target=self.run, name="kalfa-pulse", daemon=True)
+        self.thread.start()
+        return self
+
+    def __exit__(self, kind, error, trace):
+        self.ended.set()
+        thread, self.thread = self.thread, None
+        if thread is not None:
+            thread.join()
+        return False
+
+    def run(self):
+        while not self.ended.wait(self.every):
+            self.writer.beat()
+
+
+def commented(value):
+    if isinstance(value, dict):
+        out = CommentedMap()
+        for key, item in value.items():
+            out[key] = commented(item)
+        return out
+    if isinstance(value, (list, tuple)):
+        out = CommentedSeq()
+        for item in value:
+            out.append(commented(item))
+        return out
+    return value
+
+
+def locate(document, tokens):
+    holder = document
+    for token in tokens[:-1]:
+        try:
+            holder = holder[token]
+        except (KeyError, IndexError, TypeError):
+            return None, None
+    if not tokens or not isinstance(holder, (CommentedMap, CommentedSeq)):
+        return None, None
+    try:
+        holder[tokens[-1]]
+    except (KeyError, IndexError, TypeError):
+        return None, None
+    return holder, tokens[-1]
+
+
+def annotate(document, overrides):
+    for path, label in (overrides or {}).items():
+        try:
+            tokens = tokenize(path)
+        except RunError:
+            continue
+        holder, key = locate(document, tokens)
+        if holder is None:
+            continue
+        try:
+            holder.yaml_add_eol_comment(str(label), key=key)
+        except (KeyError, IndexError, TypeError, AttributeError):
+            try:
+                holder.yaml_add_eol_comment(str(label), key=key, column=0)
+            except (KeyError, IndexError, TypeError, AttributeError):
+                continue
+    return document
+
+
+def resolved_mapping(config, record=None):
+    data = {key: value for key, value in (config.data or {}).items() if key != "sweep"}
+    if record is not None:
+        data["record"] = str(record)
+    plugins = [str(name) for name in (getattr(config, "plugins", None) or [])]
+    if plugins and "plugins" not in data:
+        return {"plugins": plugins, **data}
+    return data
+
+
+def resolved_text(config, record=None):
+    yaml = YAML()
+    yaml.width = 120
+    document = annotate(commented(resolved_mapping(config, record)), getattr(config, "overrides", None))
+    stream = StringIO()
+    yaml.dump(document, stream)
+    return stream.getvalue()
+
+
+def host_note():
+    return {"hostname": socket.gethostname(), "platform": platform.platform(), "python": platform.python_version(),
+            "torch": torch.__version__, "cuda": torch.version.cuda or None, "cwd": str(Path.cwd())}
+
+
+def git_note(root):
+    try:
+        commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=str(root), capture_output=True, text=True,
+                                check=True, timeout=5)
+        status = subprocess.run(["git", "status", "--porcelain"], cwd=str(root), capture_output=True, text=True,
+                                check=True, timeout=5)
+    except (OSError, ValueError, subprocess.SubprocessError):
+        return None
+    return {"commit": commit.stdout.strip(), "dirty": bool(status.stdout.strip())}
+
+
+def start_record(ctx, directory, paths=(), plan=None, parent=None, sweep=None):
+    directory = Path(directory)
+    if directory.exists() and any(directory.iterdir()) and parent is None:
+        raise RunError(f"record directory {directory} exists and is not empty; name another record or resume")
+    directory.mkdir(parents=True, exist_ok=True)
+    writer = RecordWriter(directory)
+    config = ctx.config
+    train = getattr(ctx, "train", None)
+    device = getattr(ctx, "device", None)
+    writer.manifest(format=1, kalfa=version(), name=directory.resolve().name, started=now(), finished=None,
+                    status="running", config_files=[str(Path(path).resolve()) for path in paths],
+                    params=dict(getattr(config, "params", None) or {}),
+                    unit=train.unit if train is not None else "epochs",
+                    turns_planned=train.total_turns() if train is not None else 0, host=host_note(),
+                    device=str(device) if device is not None else None, git=git_note(Path.cwd()), sweep=sweep,
+                    resumed_from=str(parent) if parent is not None else None, failure=None)
+    writer.write_text("resolved.yaml", resolved_text(config, record=directory))
+    if plan is not None:
+        writer.write_json("plan.json", plan)
+    copy_plugins(config, directory, imports=getattr(getattr(ctx, "tree", None), "imports", None) or ())
+    ctx.record = writer
+    return writer
+
+
+def module_file(name, search=False):
+    module = sys.modules.get(name)
+    file = getattr(module, "__file__", None) if module is not None else None
+    if file is None and search and module is None:
+        try:
+            spec = importlib.util.find_spec(name)
+        except (ImportError, ValueError, AttributeError):
+            return None
+        file = spec.origin if spec is not None and spec.has_location else None
+    if not file or not str(file).endswith(".py"):
+        return None
+    return Path(file).resolve()
+
+
+def plugin_source(name, roots, search=False):
+    parts = str(name).split(".")
+    for cut in range(1, len(parts) + 1):
+        file = module_file(".".join(parts[:cut]), search=search)
+        if file is None:
+            continue
+        home = next((root for root in roots if file.is_relative_to(root)), None)
+        if home is not None:
+            return file, home
+    return (module_file(name, search=search) if search else None), None
+
+
+def copy_module(file, home, destination):
+    relative = file.relative_to(home) if home is not None else Path(file.name)
+    if file.name == "__init__.py":
+        target = destination / (relative.parent if home is not None else Path(file.parent.name))
+        shutil.copytree(file.parent, target, ignore=shutil.ignore_patterns("__pycache__"), dirs_exist_ok=True)
+        return target
+    target = destination / relative
+    target.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(file, target)
+    return target
+
+
+def manifest_files(directory):
+    try:
+        files = json.loads((Path(directory) / "manifest.json").read_text(encoding="utf-8")).get("config_files") or []
+    except (OSError, ValueError, AttributeError):
+        return []
+    return [str(file) for file in files if isinstance(file, str)]
+
+
+def plugin_roots(config):
+    roots = []
+    for item in getattr(config, "directories", None) or []:
+        root = Path(item).resolve()
+        if (root / "resolved.yaml").is_file():
+            if (root / "plugins").is_dir():
+                roots.append(root / "plugins")
+            roots.extend(Path(file).resolve().parent for file in manifest_files(root))
+        roots.append(root)
+    return list(dict.fromkeys(roots))
+
+
+def copy_plugins(config, directory, imports=()):
+    roots = plugin_roots(config)
+    named = [str(name) for name in (getattr(config, "plugins", None) or [])]
+    destination = Path(directory) / "plugins"
+    copied = []
+    seen = set()
+    for name, search in [(name, True) for name in named] + [(str(text), False) for text in imports or ()]:
+        file, home = plugin_source(name, roots, search=search)
+        if file is None or file in seen:
+            continue
+        seen.add(file)
+        copied.append(copy_module(file, home, destination))
+    return copied

kalfa/core/refs.py

@@ -0,0 +1,88 @@
+import re
+
+from .errors import RunError
+
+
+TABLES = ("models", "losses", "metrics", "optimizers")
+SINGULAR = {"models": "model", "losses": "loss", "metrics": "metric", "optimizers": "optimizer"}
+SEGMENT = re.compile(r"(?P<name>[^.\[\]\s]+)(?:\[\d+\])*")
+INDEX = re.compile(r"\[(\d+)\]")
+
+
+def tokenize(path):
+    if not isinstance(path, str) or not path.strip():
+        raise RunError(f"a path is dotted like optimizers.main.lr, got {path!r}")
+    tokens = []
+    for piece in path.strip().split("."):
+        match = SEGMENT.fullmatch(piece)
+        if match is None:
+            raise RunError(f"malformed path {path!r}: {piece!r} is not a name with an optional [index]")
+        tokens.append(match.group("name"))
+        tokens.extend(int(index) for index in INDEX.findall(piece))
+    return tokens
+
+
+def join_tokens(tokens):
+    text = ""
+    for token in tokens:
+        if isinstance(token, int):
+            text += f"[{token}]"
+        else:
+            text += f".{token}" if text else token
+    return text
+
+
+def ref_parts(asked):
+    if isinstance(asked, str) and asked.startswith("@"):
+        return asked[1:]
+    if isinstance(asked, dict) and "ref" in asked:
+        attr = asked.get("attr")
+        return f"{asked['ref']}.{attr}" if attr else str(asked["ref"])
+    return None
+
+
+def ref_json(asked):
+    tokens = tokenize(ref_parts(asked))
+    split = 2 if tokens[0] in TABLES else 1
+    out = {"ref": join_tokens(tokens[:split])}
+    if tokens[split:]:
+        out["attr"] = join_tokens(tokens[split:])
+    return out
+
+
+def ref_json_of(ctx, value):
+    for section in TABLES:
+        for name, entry in getattr(ctx, section).items():
+            if entry is value or getattr(entry, "module", None) is value:
+                return {"ref": f"{section}.{name}"}
+            if section == "models" and getattr(entry, "ema", None) is value:
+                return {"ref": f"{section}.{name}", "attr": "ema"}
+    for section in ("data", "generate"):
+        if getattr(ctx, section, None) is value:
+            return {"ref": section}
+    for node in ctx.tree.nodes.values():
+        if node.obj is value:
+            return {"ref": node.path}
+    return None
+
+
+def jsonable(value, ctx=None):
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    from .build import Node, Ref
+
+    if isinstance(value, Ref):
+        return {"ref": value.path} if value.attr is None else {"ref": value.path, "attr": value.attr}
+    if isinstance(value, Node):
+        return {"node": value.path}
+    if isinstance(value, dict):
+        return {str(key): jsonable(item, ctx) for key, item in value.items()}
+    if isinstance(value, (list, tuple, set, frozenset)):
+        return [jsonable(item, ctx) for item in value]
+    if ctx is not None:
+        found = ref_json_of(ctx, value)
+        if found is not None:
+            return found
+    if hasattr(value, "tolist"):
+        return value.tolist()
+    return str(value)

kalfa/core/rng.py

@@ -0,0 +1,100 @@
+import hashlib
+import random
+
+import numpy
+import torch
+
+
+class Rng:
+    def __init__(self, seed=0):
+        self.seed = int(seed)
+        self.streams = {}
+        self.pending = {}
+
+    def seed_everything(self):
+        random.seed(self.seed)
+        numpy.random.seed(self.seed % (2**32))
+        torch.manual_seed(self.seed)
+
+    def seed_for(self, name):
+        digest = hashlib.sha256(f"{self.seed}:{name}".encode()).digest()
+        return int.from_bytes(digest[:8], "little") % (2**63 - 1)
+
+    @staticmethod
+    def place(device=None, like=None):
+        if device is None and like is not None:
+            device = like.device
+        if device is None:
+            return torch.device("cpu")
+        if isinstance(device, torch.device):
+            return device
+        found = getattr(device, "torch", None)
+        if isinstance(found, torch.device):
+            return found
+        return torch.device(str(device))
+
+    @staticmethod
+    def size(shape, like=None):
+        if shape is None:
+            return tuple(like.shape)
+        if isinstance(shape, int):
+            return (shape,)
+        return tuple(shape)
+
+    @staticmethod
+    def float_dtype(dtype, like):
+        if dtype is not None:
+            return dtype
+        if like is not None and like.dtype.is_floating_point:
+            return like.dtype
+        return None
+
+    def stream(self, name, device=None):
+        where = self.place(device)
+        key = f"{name}@{where.type}"
+        generator = self.streams.get(key)
+        if generator is None:
+            generator = torch.Generator(device=where)
+            generator.manual_seed(self.seed_for(name))
+            state = self.pending.pop(key, None)
+            if state is not None:
+                generator.set_state(state)
+            self.streams[key] = generator
+        return generator
+
+    def normal(self, name, shape, like=None, device=None, dtype=None):
+        where = self.place(device, like)
+        return torch.randn(self.size(shape, like), generator=self.stream(name, where), device=where,
+                           dtype=self.float_dtype(dtype, like))
+
+    def uniform(self, name, shape, like=None, device=None, dtype=None):
+        where = self.place(device, like)
+        return torch.rand(self.size(shape, like), generator=self.stream(name, where), device=where,
+                          dtype=self.float_dtype(dtype, like))
+
+    def integers(self, name, low, high, shape, device=None):
+        where = self.place(device)
+        return torch.randint(int(low), int(high), self.size(shape), generator=self.stream(name, where), device=where)
+
+    def permutation(self, name, count, device=None):
+        where = self.place(device)
+        return torch.randperm(int(count), generator=self.stream(name, where), device=where)
+
+    def state_dict(self):
+        streams = {key: generator.get_state().detach().cpu().clone() for key, generator in self.streams.items()}
+        for key, state in self.pending.items():
+            streams.setdefault(key, state)
+        return {"seed": self.seed, "streams": streams}
+
+    def load_state_dict(self, state):
+        if not state:
+            return
+        if state.get("seed") is not None:
+            self.seed = int(state["seed"])
+        for key, saved in (state.get("streams") or {}).items():
+            tensor = (saved if isinstance(saved, torch.Tensor) else torch.as_tensor(saved)).cpu()
+            generator = self.streams.get(key)
+            if generator is not None:
+                generator.set_state(tensor)
+            else:
+                self.pending[key] = tensor