modulearn 0.1.0__py3-none-any.whl

modulearn/registry.py

@@ -0,0 +1,283 @@
+#!/usr/bin/env python3
+"""The registry: how a host application declares its nodes to ModuLearn.
+
+WHAT THIS REPLACES
+------------------
+ModuLearn is a visual, "Unreal-Blueprints"-style editor for building training runs:
+you drop Dataset / Model / Hyperparameter / Loss / Transform nodes on a canvas,
+wire them into a single Train sink, and launch. The canvas is just a JSON graph.
+
+The *only* thing tying that generic editor to your project is this registry. Your
+app creates a :class:`Registry`, calls ``add_dataset`` / ``add_model`` / … to
+describe what it can train, and hands the registry to :func:`modulearn.server.create_app`.
+Everything downstream — the node palette (:mod:`modulearn.catalog`), the type-checked
+graph compiler (:mod:`modulearn.compiler`), and the web editor — is built from it.
+There is no hardcoded node list anywhere, so the palette can never drift from what
+your backend actually accepts.
+
+THE VOCABULARY
+--------------
+* **Port** — a typed socket on a node. Two ports connect only if their types match
+  exactly. Families: ``dataset``, ``model``, ``loss``, ``scalar`` and ``run``.
+  Scalars carry a field-specific subtype (``scalar/lr``) so a learning-rate value
+  only fits the ``lr`` input, never ``epochs``.
+* **Param** — an in-node widget the user edits directly (float / int / enum / bool /
+  int_list), or a read-only descriptor drawn as a subtitle.
+* **NodeSpec** — one node type: its ports, its params, and where it lives in the
+  palette. This module builds these for you via the ``add_*`` helpers.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+# The link type-families. A port type is either a family ("model") or a subtype
+# ("scalar/lr"); the base before "/" must be one of these. Kept in lockstep with
+# the compiler and the frontend's wire colors.
+PORT_TYPES: dict[str, str] = {
+    "dataset": "A resolved dataset: its feature columns, target columns and kind.",
+    "model": "A model specification (architecture + params), ready to train.",
+    "loss": "A training-loss specification.",
+    "scalar": "A single numeric value from a hyperparameter node, on the "
+              "field-specific subtype 'scalar/<field>' so it only fits its match.",
+    "run": "The terminal handle of a launched training run (its run_id).",
+}
+
+# Palette groupings, in display order.
+CATEGORIES: list[str] = [
+    "dataset", "model", "hyperparameter", "loss", "transform", "train",
+]
+
+
+@dataclass
+class Param:
+    """One in-node value. ``kind`` is float|int|enum|bool|int_list.
+
+    ``readonly`` marks a descriptor (e.g. a dataset's feature count) — the editor
+    renders it as a compact subtitle instead of an editable widget. ``choices`` is
+    required when ``kind == "enum"``.
+    """
+    name: str
+    label: str
+    kind: str
+    default: Any = None
+    help: str = ""
+    min: float | None = None
+    max: float | None = None
+    choices: list[str] | None = None
+    readonly: bool = False
+
+    def to_json(self) -> dict:
+        if self.kind == "enum" and not self.choices:
+            raise ValueError(f"enum param {self.name!r} needs choices")
+        d: dict = {"name": self.name, "label": self.label, "kind": self.kind,
+                   "default": self.default, "help": self.help,
+                   "readonly": self.readonly}
+        if self.min is not None:
+            d["min"] = self.min
+        if self.max is not None:
+            d["max"] = self.max
+        if self.choices is not None:
+            d["choices"] = list(self.choices)
+        return d
+
+
+@dataclass
+class Port:
+    name: str
+    type: str
+
+    def to_json(self) -> dict:
+        base = self.type.split("/", 1)[0]
+        if base not in PORT_TYPES:
+            raise ValueError(f"unknown port type {self.type!r}; base must be one "
+                             f"of {sorted(PORT_TYPES)}")
+        return {"name": self.name, "type": self.type}
+
+
+@dataclass
+class NodeSpec:
+    """One node type in the catalog. ``id`` is ``"<category>.<name>"``."""
+    id: str
+    category: str
+    title: str
+    help: str = ""
+    title_short: str | None = None
+    inputs: list[Port] = field(default_factory=list)
+    outputs: list[Port] = field(default_factory=list)
+    params: list[Param] = field(default_factory=list)
+    experimental: bool = False
+    meta: dict = field(default_factory=dict)
+
+    def to_json(self) -> dict:
+        d = {
+            "id": self.id, "category": self.category, "title": self.title,
+            "help": self.help,
+            "inputs": [p.to_json() for p in self.inputs],
+            "outputs": [p.to_json() for p in self.outputs],
+            "params": [p.to_json() for p in self.params],
+            "experimental": self.experimental,
+        }
+        if self.title_short:
+            d["title_short"] = self.title_short
+        if self.meta:
+            d["meta"] = self.meta
+        return d
+
+
+# A semantic check runs after structural/type validation succeeds. It receives the
+# compiled graph and returns a list of human-readable error strings ([] == ok).
+SemanticCheck = Callable[["Any"], list[str]]
+
+
+class Registry:
+    """The declarative surface a host app fills in, then hands to ``create_app``.
+
+    Call the ``add_*`` methods in any order. Every hyperparameter you add becomes a
+    matching typed input on the Train sink automatically, so you never wire the sink
+    by hand. Call :meth:`catalog` (the ``/api/nodes`` payload) once you're done —
+    it lazily appends the Train node from the hyperparameters registered so far.
+    """
+
+    def __init__(self) -> None:
+        self._nodes: list[NodeSpec] = []
+        self._by_id: dict[str, NodeSpec] = {}
+        self._hparam_ids: list[str] = []          # order-preserving, for the sink
+        self._checks: list[SemanticCheck] = []
+        self._live_transforms: set[str] = set()
+
+    # -- low-level -----------------------------------------------------------
+    def add_node(self, spec: NodeSpec) -> NodeSpec:
+        if spec.id in self._by_id:
+            raise ValueError(f"duplicate node id {spec.id!r}")
+        if spec.category not in CATEGORIES:
+            raise ValueError(f"unknown category {spec.category!r}")
+        self._nodes.append(spec)
+        self._by_id[spec.id] = spec
+        return spec
+
+    def add_semantic_check(self, fn: SemanticCheck) -> None:
+        """Register a project-specific validity rule (e.g. model/dataset kind must
+        match). Runs only once the graph is structurally sound."""
+        self._checks.append(fn)
+
+    # -- ergonomic builders --------------------------------------------------
+    def add_dataset(self, name: str, *, title: str, help: str = "",
+                    features: list[str] | None = None,
+                    targets: list[str] | None = None,
+                    kind: str = "", trainable: bool = True,
+                    meta: dict | None = None) -> NodeSpec:
+        """A training source. Emits a single ``dataset`` port, no inputs.
+
+        ``kind`` is a free-form tag (e.g. "tabular", "sequence") you can match
+        against a model's required kind in a semantic check. ``trainable=False``
+        lists a feature-only set that the compiler will refuse to train."""
+        features = features or []
+        targets = targets or []
+        m = {"trainable": trainable, "kind": kind,
+             "features": list(features), "targets": list(targets)}
+        m.update(meta or {})
+        return self.add_node(NodeSpec(
+            id=f"dataset.{name}", category="dataset", title=title,
+            title_short=name, help=help,
+            outputs=[Port("dataset", "dataset")],
+            params=[
+                Param("n_features", "features", "int", len(features),
+                      help=", ".join(features) or "(none)", readonly=True),
+                Param("n_targets", "targets", "int", len(targets),
+                      help=", ".join(targets) or "(none)", readonly=True),
+            ] + ([Param("kind", "kind", "enum", kind, choices=[kind],
+                        help="Dataset kind.", readonly=True)] if kind else []),
+            meta=m,
+        ))
+
+    def add_model(self, name: str, *, title: str, title_short: str | None = None,
+                  help: str = "", params: list[Param] | None = None,
+                  requires_kind: str = "") -> NodeSpec:
+        """A trainable architecture. Takes a ``dataset`` in, emits a ``model`` out.
+
+        ``requires_kind`` (optional) is matched against the dataset's ``kind`` by
+        the built-in kind check, so an incompatible pairing fails to compile."""
+        return self.add_node(NodeSpec(
+            id=f"model.{name}", category="model", title=title,
+            title_short=title_short or name, help=help,
+            inputs=[Port("dataset", "dataset")],
+            outputs=[Port("model", "model")],
+            params=list(params or []),
+            meta={"requires_kind": requires_kind} if requires_kind else {},
+        ))
+
+    def add_hyperparameter(self, name: str, *, label: str, kind: str = "float",
+                           default: Any = None, help: str = "",
+                           min: float | None = None,
+                           max: float | None = None) -> NodeSpec:
+        """A single scalar knob. Emits ``scalar/<name>`` so it only plugs into the
+        Train sink's matching ``<name>`` input (added automatically)."""
+        self._hparam_ids.append(name)
+        return self.add_node(NodeSpec(
+            id=f"hyperparameter.{name}", category="hyperparameter", title=label,
+            help=help, outputs=[Port("value", f"scalar/{name}")],
+            params=[Param(name, label, kind, default, help=help, min=min, max=max)],
+        ))
+
+    def add_loss(self, params: list[Param], *, title: str = "Loss",
+                 help: str = "Training objective.") -> NodeSpec:
+        """The Loss node. Emits a ``loss`` port into the Train sink."""
+        return self.add_node(NodeSpec(
+            id="loss.loss", category="loss", title=title, help=help,
+            outputs=[Port("loss", "loss")], params=list(params),
+        ))
+
+    def add_transform(self, name: str, *, title: str, help: str = "",
+                      live: bool = True) -> NodeSpec:
+        """A chainable feature transform: ``dataset`` in, ``dataset`` out.
+
+        ``live=False`` marks it experimental — shown in the palette but rejected by
+        the compiler until your backend can honor it."""
+        if live:
+            self._live_transforms.add(name)
+        return self.add_node(NodeSpec(
+            id=f"transform.{name}", category="transform", title=title,
+            help=help if live else f"(not yet wireable) {help}",
+            inputs=[Port("dataset", "dataset")],
+            outputs=[Port("dataset", "dataset")],
+            experimental=not live,
+        ))
+
+    # -- assembly ------------------------------------------------------------
+    def _train_node(self) -> NodeSpec:
+        """The terminal Train sink: model + loss + one scalar input per registered
+        hyperparameter, each on its field-specific subtype."""
+        scalar_inputs = [Port(h, f"scalar/{h}") for h in self._hparam_ids]
+        return NodeSpec(
+            id="train.train", category="train", title="Train",
+            help="Compiles the wired graph and launches a background training job. "
+                 "Exactly one per graph.",
+            inputs=[Port("model", "model"), Port("loss", "loss"), *scalar_inputs],
+            outputs=[Port("run", "run")],
+        )
+
+    def nodes(self) -> list[NodeSpec]:
+        """All registered nodes plus the auto-built Train sink."""
+        return [*self._nodes, self._train_node()]
+
+    def catalog(self) -> dict:
+        """The ``GET /api/nodes`` payload: port glossary, ordered categories, nodes."""
+        return {
+            "port_types": PORT_TYPES,
+            "categories": CATEGORIES,
+            "nodes": [n.to_json() for n in self.nodes()],
+        }
+
+    # accessors used by the compiler
+    @property
+    def hyperparameter_ids(self) -> list[str]:
+        return list(self._hparam_ids)
+
+    @property
+    def semantic_checks(self) -> list[SemanticCheck]:
+        return list(self._checks)
+
+    @property
+    def live_transforms(self) -> set[str]:
+        return set(self._live_transforms)

modulearn/server.py

@@ -0,0 +1,195 @@
+#!/usr/bin/env python3
+"""The ModuLearn web server: a FastAPI app factory around your :class:`Registry`.
+
+:func:`create_app` gives you the whole visual editor — palette, canvas, live
+type-checking, one-click launch and on-canvas run tracking — wired to *your*
+training code through two things you provide:
+
+  * a :class:`~modulearn.registry.Registry` describing your nodes, and
+  * an ``on_train(compiled, reporter)`` callback that actually trains.
+
+ModuLearn owns the run directory, the background thread, and the state/metrics files
+the editor polls; your callback just trains and calls ``reporter`` to report
+progress. Training runs in a daemon thread owned by this process and persists to
+disk, so the page is a stateless live window: closing the browser never stops a run.
+
+THE STATE CONTRACT (what the editor's live panel + chart read)
+--------------------------------------------------------------
+``reporter.state(...)`` merges keys into ``state.json``; ``reporter.metric(...)``
+appends a row to ``metrics.jsonl``. The frontend understands these keys:
+  state: ``phase`` (running|done|error|paused), ``epoch``, ``epochs``,
+         ``best_val``, ``test_score``, ``error``.
+  metric row: ``epoch``, ``train``, ``val`` — drives the live learning curve.
+Anything else you write is preserved and returned by the API, just not charted.
+"""
+from __future__ import annotations
+
+import json
+import threading
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable
+
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import FileResponse, Response
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+
+from modulearn import compiler
+from modulearn.registry import Registry
+
+STATIC = Path(__file__).resolve().parent / "static"
+
+
+@dataclass
+class RunReporter:
+    """The handle your ``on_train`` uses to publish progress. Writes are atomic
+    enough for a poller: state is a merged JSON doc, metrics an append-only log."""
+    run_dir: Path
+    _state: dict = field(default_factory=dict)
+
+    def state(self, **kw) -> None:
+        self._state.update(kw)
+        (self.run_dir / "state.json").write_text(json.dumps(self._state, indent=2))
+
+    def metric(self, **row) -> None:
+        with (self.run_dir / "metrics.jsonl").open("a") as f:
+            f.write(json.dumps(row) + "\n")
+
+
+# on_train(compiled, reporter) -> None. Raise to mark the run errored.
+TrainFn = Callable[[compiler.CompiledGraph, RunReporter], None]
+
+
+class _GraphReq(BaseModel):
+    nodes: list[dict] = []
+    links: list[dict] = []
+
+
+class _Job:
+    def __init__(self, app_state: "_AppState", compiled: compiler.CompiledGraph):
+        self.compiled = compiled
+        self._app = app_state
+        self.thread = threading.Thread(target=self._run, daemon=True)
+
+    def _run(self):
+        run_dir = self._app.run_dir(self.compiled.run_id)
+        reporter = RunReporter(run_dir)
+        reporter.state(phase="running", epoch=0)
+        try:
+            self._app.on_train(self.compiled, reporter)
+            if reporter._state.get("phase") == "running":
+                reporter.state(phase="done")
+        except Exception as e:  # never let a job crash silently
+            reporter.state(phase="error", error=str(e))
+        finally:
+            self._app.jobs.pop(self.compiled.run_id, None)
+
+    def start(self):
+        self.thread.start()
+
+
+@dataclass
+class _AppState:
+    registry: Registry
+    on_train: TrainFn
+    runs: Path
+    jobs: dict = field(default_factory=dict)
+
+    def run_dir(self, run_id: str) -> Path:
+        d = self.runs / run_id
+        d.mkdir(parents=True, exist_ok=True)
+        return d
+
+
+def create_app(registry: Registry, on_train: TrainFn, *,
+               runs_dir: str | Path = "runs",
+               title: str = "ModuLearn", subtitle: str = "") -> FastAPI:
+    """Build the FastAPI app. ``title``/``subtitle`` brand the editor's header."""
+    state = _AppState(registry=registry, on_train=on_train, runs=Path(runs_dir))
+    app = FastAPI(title=title)
+
+    def _read(run_id: str, fname: str) -> dict:
+        p = state.runs / run_id / fname
+        return json.loads(p.read_text()) if p.exists() else {}
+
+    @app.get("/api/nodes")
+    def api_nodes():
+        return registry.catalog()
+
+    @app.get("/api/meta")
+    def api_meta():
+        return {"title": title, "subtitle": subtitle}
+
+    @app.get("/api/runs")
+    def api_runs():
+        runs = []
+        if state.runs.exists():
+            for d in sorted(state.runs.iterdir(),
+                            key=lambda p: p.stat().st_mtime, reverse=True):
+                if not d.is_dir():
+                    continue
+                st = _read(d.name, "state.json")
+                runs.append({"run_id": d.name, "phase": st.get("phase", "?"),
+                             "epoch": st.get("epoch"), "epochs": st.get("epochs"),
+                             "best_val": st.get("best_val"),
+                             "test_score": st.get("test_score"),
+                             "live": d.name in state.jobs})
+        return {"runs": runs}
+
+    @app.get("/api/run/{run_id}")
+    def api_run(run_id: str):
+        d = state.runs / run_id
+        if not d.exists():
+            raise HTTPException(404, "no such run")
+        metrics = []
+        mp = d / "metrics.jsonl"
+        if mp.exists():
+            metrics = [json.loads(l) for l in mp.read_text().splitlines() if l.strip()]
+        return {"run_id": run_id, "state": _read(run_id, "state.json"),
+                "config": _read(run_id, "config.json"), "metrics": metrics,
+                "live": run_id in state.jobs}
+
+    @app.post("/api/graph/compile")
+    def api_compile(req: _GraphReq):
+        try:
+            c = compiler.compile_graph({"nodes": req.nodes, "links": req.links}, registry)
+        except compiler.GraphError as e:
+            raise HTTPException(400, {"errors": e.errors})
+        return {"ok": True, "config": c.as_dict()}
+
+    @app.post("/api/graph/start")
+    def api_start(req: _GraphReq):
+        try:
+            c = compiler.compile_graph({"nodes": req.nodes, "links": req.links}, registry)
+        except compiler.GraphError as e:
+            raise HTTPException(400, {"errors": e.errors})
+        d = state.run_dir(c.run_id)
+        (d / "graph.json").write_text(json.dumps(
+            {"nodes": req.nodes, "links": req.links}, indent=2))
+        (d / "config.json").write_text(json.dumps(c.as_dict(), indent=2))
+        job = _Job(state, c)
+        state.jobs[c.run_id] = job
+        job.start()
+        return {"run_id": c.run_id, "config": c.as_dict()}
+
+    @app.get("/api/graph/{run_id}")
+    def api_graph_get(run_id: str):
+        p = state.runs / run_id / "graph.json"
+        if not p.exists():
+            raise HTTPException(404, "no saved graph for this run")
+        return json.loads(p.read_text())
+
+    @app.get("/")
+    def index():
+        return FileResponse(STATIC / "graph.html")
+
+    @app.get("/favicon.ico")
+    def favicon():
+        svg = ('<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 32 32">'
+               '<circle cx="16" cy="16" r="13" fill="#a8542f"/></svg>')
+        return Response(content=svg, media_type="image/svg+xml")
+
+    app.mount("/static", StaticFiles(directory=STATIC), name="static")
+    return app