mlbricks-studio 1.0.0b1__py3-none-any.whl

mlb_studio/__init__.py

@@ -0,0 +1,5 @@
+from .builder import Builder, BuilderWidget
+from .graph import new_project
+from .version import __version__
+
+__all__ = ["Builder", "BuilderWidget", "new_project", "__version__"]

mlb_studio/api_graph_runtime.py

@@ -0,0 +1,347 @@
+from __future__ import annotations
+
+"""Declarative runtime contracts for MLBricks components.
+
+Studio is an orchestrator: contracts describe how a visual node maps to the
+original MLBricks constructor and forward API.  The graph executor itself does
+not need per-component branches.
+"""
+
+from copy import deepcopy
+from dataclasses import dataclass
+import json
+from pathlib import Path
+from typing import Any, Mapping
+
+from .import_pool import IMPORT_POOL
+
+_SCHEMA_PATH = Path(__file__).with_name("mlbricks_api_schema.json")
+_SCHEMA_CACHE: dict[str, Any] | None = None
+
+
+def _schema() -> dict[str, Any]:
+    global _SCHEMA_CACHE
+    if _SCHEMA_CACHE is None:
+        try:
+            _SCHEMA_CACHE = json.loads(_SCHEMA_PATH.read_text(encoding="utf-8"))
+        except Exception:
+            _SCHEMA_CACHE = {}
+    return _SCHEMA_CACHE or {}
+
+
+def _bool(value: Any) -> bool:
+    if isinstance(value, bool):
+        return value
+    return str(value).strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _coerce(value: Any, annotation: str | None) -> Any:
+    kind = str(annotation or "").lower()
+    if value is None:
+        return None
+    if "bool" in kind:
+        return _bool(value)
+    if "int" in kind:
+        return int(float(value))
+    if "float" in kind:
+        return float(value)
+    return value
+
+
+@dataclass(frozen=True)
+class APIComponentContract:
+    """Mapping between one Studio node and an original MLBricks API.
+
+    ``input_ports`` maps Studio input-port names to forward argument names.
+    ``output_ports`` maps Studio output-port names to result positions/keys.
+    For the common one-input/one-output case no custom code is required.
+    """
+
+    component_type: str
+    import_key: str
+    input_ports: Mapping[str, str]
+    output_ports: Mapping[str, Any]
+    parameter_aliases: Mapping[str, tuple[str, ...]] = None
+    runtime_sources: Mapping[str, str] = None
+    input_initializers: Mapping[str, Mapping[str, Any]] = None
+
+    def _parameter_value(self, key: str, spec: dict[str, Any], params: dict[str, Any], runtime: dict[str, Any]):
+        aliases = (self.parameter_aliases or {}).get(key, ())
+        candidates = (key, *aliases)
+        for candidate in candidates:
+            if candidate in params and params[candidate] not in (None, ""):
+                return params[candidate], True
+        runtime_key = (self.runtime_sources or {}).get(key)
+        if runtime_key and runtime.get(runtime_key) not in (None, ""):
+            return runtime[runtime_key], True
+        if spec.get("value") is not None:
+            return spec.get("value"), True
+        if spec.get("default") is not None:
+            return spec.get("default"), True
+        return None, False
+
+    def constructor_kwargs(self, node: dict[str, Any], runtime: dict[str, Any]) -> dict[str, Any]:
+        params = deepcopy(node.get("params") or {})
+        component_schema = (_schema().get("components") or _schema()).get(self.component_type) or {}
+        kwargs: dict[str, Any] = {}
+        for spec in component_schema.get("parameters") or []:
+            key = str(spec.get("key") or spec.get("name") or "").strip()
+            if not key:
+                continue
+            value, present = self._parameter_value(key, spec, params, runtime)
+            if not present:
+                if spec.get("required"):
+                    raise ValueError(f"{node.get('name', self.component_type)} requires {key}.")
+                continue
+            # UI/runtime values override constructor defaults, but 'auto' backend
+            # is intentionally preserved because it is part of the original API.
+            kwargs[key] = _coerce(value, spec.get("annotation") or spec.get("type"))
+        return kwargs
+
+    def instantiate(self, node: dict[str, Any], runtime: dict[str, Any]):
+        cls = IMPORT_POOL.resolve_component(self.import_key)
+        return cls(**self.constructor_kwargs(node, runtime))
+
+
+    def initialize_input(self, port: str, inputs: Mapping[str, Any], node: Mapping[str, Any], module):
+        """Create a declarative default for an unconnected logical input.
+
+        This keeps first-depth initialization in the component contract instead
+        of adding component-specific branches to the graph executor. Supported
+        initializers are intentionally small and tensor-shape driven.
+        """
+        spec = dict((self.input_initializers or {}).get(port) or {})
+        if not spec:
+            return None, False
+
+        source_name = str(spec.get("source") or "")
+        source = inputs.get(source_name)
+        if source is None:
+            return None, False
+
+        kind = str(spec.get("kind") or "").lower()
+        if kind == "zeros_like":
+            return source.new_zeros(source.shape), True
+
+        if kind == "module_method":
+            method_name = str(spec.get("method") or "")
+            method = getattr(module, method_name, None)
+            if not callable(method):
+                return None, False
+            return method(source), True
+
+        if kind == "zeros_feature":
+            width = None
+            width_attr = str(spec.get("width_attr") or "")
+            if width_attr and hasattr(module, width_attr):
+                width = getattr(module, width_attr)
+            if width in (None, ""):
+                width_param = str(spec.get("width_param") or "")
+                if width_param:
+                    width = (node.get("params") or {}).get(width_param)
+            if width in (None, ""):
+                return None, False
+            shape = (*source.shape[:-1], int(width))
+            return source.new_zeros(shape), True
+
+        raise ValueError(f"Unsupported input initializer kind {kind!r} for {self.component_type}.{port}.")
+
+    def execute(self, module, inputs: Mapping[str, Any]):
+        kwargs = {
+            api_arg: inputs[port]
+            for port, api_arg in self.input_ports.items()
+            if port in inputs
+        }
+        result = module(**kwargs)
+        if len(self.output_ports) == 1 and "main" in self.output_ports:
+            selector = self.output_ports["main"]
+            if selector is None:
+                return {"main": result}
+        outputs: dict[str, Any] = {}
+        for port, selector in self.output_ports.items():
+            if selector is None:
+                outputs[port] = result
+            elif isinstance(selector, int):
+                outputs[port] = result[selector]
+            else:
+                outputs[port] = result[selector] if isinstance(result, dict) else getattr(result, str(selector))
+        return outputs
+
+
+class VisionClassificationContract(APIComponentContract):
+    # TensorGraph contract for image-classification vision engines.
+    _CLASSIFIER_ENGINES = {"serpentine", "vit", "visiontransformer", "cnn"}
+
+    def constructor_kwargs(self, node, runtime):
+        kwargs = super().constructor_kwargs(node, runtime)
+        engine = str(kwargs.get("engine") or "Serpentine").strip()
+        key = engine.lower().replace("-", "").replace("_", "").replace(" ", "")
+        if key not in self._CLASSIFIER_ENGINES:
+            raise ValueError(
+                f"{node.get('name', self.component_type)} engine={engine!r} is "
+                "not an image-classification TensorGraph mode. Studio currently "
+                "supports Serpentine, ViT/VisionTransformer and CNN here; "
+                "Diffusion requires timesteps and AR requires visual token IDs."
+            )
+        return kwargs
+
+
+class APIComponentRegistry:
+    def __init__(self):
+        self._contracts: dict[str, APIComponentContract] = {}
+
+    def register(self, contract: APIComponentContract) -> APIComponentContract:
+        self._contracts[contract.component_type] = contract
+        return contract
+
+    def get(self, component_type: str | None) -> APIComponentContract | None:
+        return self._contracts.get(str(component_type or ""))
+
+    def __contains__(self, component_type: str) -> bool:
+        return component_type in self._contracts
+
+
+API_COMPONENTS = APIComponentRegistry()
+
+# RoPE is a standard one-input/one-output MLBricks module. The original API
+# expects a Q/K-like tensor [B,H,T,D]; start_pos defaults to zero in MLBricks,
+# so Studio can use the universal execution path without a TensorGraph branch.
+API_COMPONENTS.register(APIComponentContract(
+    component_type="rope",
+    import_key="rope",
+    input_ports={"main": "x"},
+    output_ports={"main": None},
+))
+
+# First proof of the universal path.  No BOLT-specific branch is required in
+# TensorGraph: this declaration maps the visual node to Bolt's original API.
+API_COMPONENTS.register(APIComponentContract(
+    component_type="bolt",
+    import_key="bolt",
+    input_ports={"main": "x"},
+    output_ports={"main": None},
+    parameter_aliases={
+        "d_model": ("dim", "hidden_size"),
+        "num_heads": ("heads", "head"),
+        "backend": ("kernel",),
+    },
+    runtime_sources={
+        "d_model": "model_dim",
+        "num_heads": "heads",
+        "backend": "backend",
+    },
+))
+
+# Multi-input proof: the existing Studio Main/Skip lanes map directly to the
+# original ResController API.  Main carries the update tensor; Skip carries
+# the residual stream.  TensorGraph does not need a ResController branch.
+API_COMPONENTS.register(APIComponentContract(
+    component_type="rescontroller",
+    import_key="rescontroller",
+    input_ports={"main": "update", "skip": "residual"},
+    output_ports={"main": None},
+    runtime_sources={
+        "backend": "backend",
+    },
+))
+# Stateful multi-input/multi-output proof: SAFFN exposes its original forward
+# signature as named graph ports.  The executor maps API return tuple[0] to
+# ``main`` and tuple[1] to ``state``; no SAFFN branch is required in TensorGraph.
+API_COMPONENTS.register(APIComponentContract(
+    component_type="saffn",
+    import_key="saffn",
+    input_ports={
+        "x": "x",
+        "esa_update": "esa_update",
+        "previous_esa": "previous_esa",
+        "previous_state": "previous_state",
+    },
+    output_ports={"main": 0, "state": 1},
+    parameter_aliases={
+        "d_model": ("dim", "hidden_size"),
+        "depth_embedding_dim": ("depth_dim",),
+    },
+    runtime_sources={
+        "d_model": "model_dim",
+        "backend": "backend",
+    },
+    input_initializers={
+        # At physical depth 0 there is no earlier signal/state. Studio supplies
+        # correctly shaped zeros only when these sockets are left unconnected.
+        "previous_esa": {"kind": "zeros_like", "source": "esa_update"},
+        "previous_state": {
+            # Use the original MLBricks helper so Studio does not duplicate
+            # StateAwareFFN's state-shape policy.
+            "kind": "module_method",
+            "source": "x",
+            "method": "initial_state",
+        },
+    },
+))
+
+# VirtualStateAwareFFN preserves the exact StateAwareFFN forward contract and
+# adds internal virtual refinements. It therefore reuses the same named graph
+# wiring and first-depth initialization policy without a compiler branch.
+API_COMPONENTS.register(APIComponentContract(
+    component_type="virtual_saffn",
+    import_key="virtual_saffn",
+    input_ports={
+        "x": "x",
+        "esa_update": "esa_update",
+        "previous_esa": "previous_esa",
+        "previous_state": "previous_state",
+    },
+    output_ports={"main": 0, "state": 1},
+    parameter_aliases={
+        "d_model": ("dim", "hidden_size"),
+        "depth_embedding_dim": ("depth_dim",),
+    },
+    runtime_sources={
+        "d_model": "model_dim",
+        "backend": "backend",
+    },
+    input_initializers={
+        "previous_esa": {"kind": "zeros_like", "source": "esa_update"},
+        "previous_state": {
+            "kind": "module_method",
+            "source": "x",
+            "method": "initial_state",
+        },
+    },
+))
+
+# MicroVirtualFFN is a standard one-input/one-output original MLBricks API.
+# The default forward pass validates one configured refinement pass; callers
+# can still use the full ``refine`` API through a custom API Function node.
+API_COMPONENTS.register(APIComponentContract(
+    component_type="micro_ffn",
+    import_key="micro_ffn",
+    input_ports={"main": "x"},
+    output_ports={"main": None},
+    parameter_aliases={
+        "d_model": ("dim", "hidden_size"),
+    },
+    runtime_sources={
+        "d_model": "model_dim",
+        "backend": "backend",
+    },
+))
+
+# Vision classification path. Both current MLBricks Kit 1.0.0b2 families take
+# raw image tensors [B,C,H,W] for Serpentine / ViT / CNN and return class logits.
+API_COMPONENTS.register(VisionClassificationContract(
+    component_type="vesa",
+    import_key="vesa",
+    input_ports={"main": "images"},
+    output_ports={"main": None},
+    runtime_sources={"backend": "backend"},
+))
+
+API_COMPONENTS.register(VisionClassificationContract(
+    component_type="visualbolt",
+    import_key="visualbolt",
+    input_ports={"main": "images"},
+    output_ports={"main": None},
+    runtime_sources={"backend": "backend"},
+))
+

mlb_studio/api_registry.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import inspect
+import json
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Iterable
+
+from .import_pool import (
+    COMPONENT_IMPORTS,
+    CONFIG_KEYS,
+    IMPORT_POOL,
+)
+
+_SCHEMA_PATH = Path(__file__).with_name("mlbricks_api_schema.json")
+
+
+# These components intentionally keep the richer source-derived Builder schema
+# instead of replacing it with a plain constructor signature at runtime.
+# SOUP accepts scalar-or-per-layer sequences/config mappings, while the primary
+# ElasticBit 4-32 UI represents ElasticBit.RuntimeMatrix.from_auto.
+SOURCE_DEFINED_FIELDS = {"soup", "elasticbit_runtime", "lm_head"}
+
+CHOICES = {
+    "backend": ["auto", "native", "pytorch"],
+    "precision": ["fp32", "fp16", "bf16"],
+    "activation": ["gelu", "gelu_tanh", "relu", "silu", "swish", "tanh"],
+    "device": ["auto", "cpu", "cuda", "None"],
+    "engine": ["Serpentine", "ViT", "CNN", "Diffusion", "AR"],
+    "ffn": ["standard", "ffnbrick", "virtual_ffnbrick", "micro_ffnbrick"],
+    "residual": ["standard", "rescontroller"],
+    "norm": ["rmsnorm", "layernorm"],
+}
+
+COMPONENT_CHOICES = {
+    "bolt": {
+        "position": ["none", "rope"],
+    },
+    "vesa": {
+        "position": ["auto", "none", "2d_sincos", "learned"],
+        "scan": ["cross", "horizontal", "vertical", "raster"],
+    },
+    "visualbolt": {
+        "position": ["auto", "none", "2d_sincos", "learned"],
+        "scan": ["cross", "horizontal", "vertical", "raster"],
+    },
+}
+
+
+def _fallback_schema():
+    payload = json.loads(_SCHEMA_PATH.read_text(encoding="utf-8"))
+    return deepcopy(payload["components"])
+
+
+def _safe_default(value: Any):
+    if value is inspect._empty:
+        return None
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    return str(value)
+
+
+def _field_type(name, annotation, default, choices=None):
+    choices = CHOICES if choices is None else choices
+    if name in choices:
+        return "select"
+    if isinstance(default, bool):
+        return "bool"
+    if isinstance(default, (int, float)):
+        return "number"
+    text = "" if annotation is inspect._empty else str(annotation).lower()
+    if "bool" in text:
+        return "bool"
+    if "int" in text or "float" in text:
+        return "number"
+    return "text"
+
+
+def _fields(obj, *, choices=None):
+    choices = CHOICES if choices is None else choices
+    sig = inspect.signature(obj)
+    out = []
+    for name, p in sig.parameters.items():
+        if name in {"self", "args", "kwargs"}:
+            continue
+        default = _safe_default(p.default)
+        item = {
+            "key": name,
+            "label": name.replace("_", " ").title(),
+            "type": _field_type(name, p.annotation, default, choices),
+            "required": p.default is inspect._empty,
+            "value": default,
+            "annotation": "" if p.annotation is inspect._empty else str(p.annotation),
+        }
+        if name in choices:
+            item["options"] = choices[name]
+        out.append(item)
+    return sig, out
+
+
+def _canonical_metadata(component_type: str, fallback: dict[str, Any]) -> dict[str, Any]:
+    info = IMPORT_POOL.import_info(component_type)
+    canonical_path = info.get("canonical_path")
+    public_name = info.get("canonical_symbol") or fallback.get("public_name")
+    payload = {
+        **fallback,
+        "available": True,
+        "runtime_available": None,
+        "source": "MLBricks source schema + lazy import pool",
+        "public_name": public_name,
+        "import_path": canonical_path or fallback.get("import_path"),
+        "import_module": info.get("canonical_module"),
+        "import_symbol": info.get("canonical_symbol"),
+        "import_candidates": [candidate.path for candidate in COMPONENT_IMPORTS.get(component_type, ())],
+        "loaded": bool(info.get("loaded")),
+        "resolved_from": info.get("resolved_from"),
+        "runtime_error": info.get("error"),
+    }
+    if info.get("config_path"):
+        config = deepcopy(payload.get("config_api") or {})
+        config.update({
+            "public_name": info.get("config_symbol"),
+            "import_path": info.get("config_path"),
+            "import_module": info.get("config_module"),
+            "import_symbol": info.get("config_symbol"),
+        })
+        payload["config_api"] = config
+    return payload
+
+
+def _inspect_one(component_type: str, fallback: dict[str, Any]) -> dict[str, Any]:
+    result = _canonical_metadata(component_type, fallback)
+    try:
+        obj = IMPORT_POOL.resolve_component(component_type)
+        component_choices = {**CHOICES, **COMPONENT_CHOICES.get(component_type, {})}
+        sig, inspected_fields = _fields(obj, choices=component_choices)
+        source_defined = component_type in SOURCE_DEFINED_FIELDS
+        fields = deepcopy(fallback.get("parameters", [])) if source_defined else inspected_fields
+
+        config_info = result.get("config_api")
+        if component_type in CONFIG_KEYS:
+            cfg_obj = IMPORT_POOL.resolve_config(component_type)
+            cfg_choices = {**CHOICES, **COMPONENT_CHOICES.get(component_type, {})}
+            cfg_sig, cfg_fields = _fields(cfg_obj, choices=cfg_choices)
+            if len(cfg_fields) > 1:
+                fields = cfg_fields
+                config_info = {
+                    **(config_info or {}),
+                    "public_name": cfg_obj.__name__,
+                    "signature": f"{cfg_obj.__name__}{cfg_sig}",
+                    "parameters": cfg_fields,
+                }
+
+        doc = inspect.getdoc(obj) or ""
+        runtime_available = True
+        runtime_error = None
+        if component_type == "elasticbit_runtime":
+            checker = getattr(obj, "native_runtime_available", None)
+            if callable(checker):
+                try:
+                    runtime_available = bool(checker())
+                except Exception as exc:
+                    runtime_available = False
+                    runtime_error = f"{type(exc).__name__}: {exc}"
+            if not runtime_available and runtime_error is None:
+                runtime_error = "ElasticBit runtime is not built in this environment."
+
+        info = IMPORT_POOL.import_info(component_type)
+        return {
+            **result,
+            "available": True,
+            "runtime_available": runtime_available,
+            "source": "runtime inspection + MLBricks source schema" if source_defined else "runtime inspection",
+            "public_name": fallback.get("public_name", obj.__name__) if source_defined else obj.__name__,
+            "signature": fallback.get("signature", f"{obj.__name__}{sig}") if source_defined else f"{obj.__name__}{sig}",
+            "description": fallback.get("description", "") if source_defined else (doc.splitlines()[0] if doc else fallback.get("description", "")),
+            "parameters": fields if fields else fallback.get("parameters", []),
+            "config_api": config_info,
+            "runtime_error": runtime_error,
+            "loaded": True,
+            "resolved_from": info.get("resolved_from"),
+        }
+    except Exception as exc:
+        info = IMPORT_POOL.import_info(component_type)
+        return {
+            **result,
+            "available": True,
+            "runtime_available": False,
+            "loaded": False,
+            "source": "MLBricks source schema + lazy import pool",
+            "runtime_error": f"{type(exc).__name__}: {exc}",
+            "resolved_from": info.get("resolved_from"),
+        }
+
+
+def discover_mlbricks_api(
+    component_types: Iterable[str] | None = None,
+    *,
+    eager: bool = False,
+):
+    """Return Builder API metadata without requiring eager MLBricks imports.
+
+    By default, the supplied source-derived schema is used and canonical import
+    routes are attached to each component.  If ``eager=True`` (or a specific
+    ``component_types`` iterable is supplied), only those requested components
+    are imported and inspected.  This is the import-pool behavior used by the
+    Builder UI: adding a component can warm just that component's API instead of
+    importing the whole MLBricks package.
+    """
+    result = _fallback_schema()
+    requested = set(str(x) for x in (component_types or ()))
+    inspect_all = bool(eager and component_types is None)
+
+    for component_type in COMPONENT_IMPORTS:
+        fallback = result.get(component_type, {})
+        if inspect_all or component_type in requested:
+            result[component_type] = _inspect_one(component_type, fallback)
+        else:
+            result[component_type] = _canonical_metadata(component_type, fallback)
+
+    return result
+
+
+def refresh_component_api(component_type: str) -> dict[str, Any] | None:
+    component_type = str(component_type)
+    if component_type not in COMPONENT_IMPORTS:
+        return None
+    fallback = _fallback_schema().get(component_type, {})
+    return _inspect_one(component_type, fallback)