Project-Chronos 0.1.1__py3-none-any.whl

chronos/__init__.py

@@ -0,0 +1,16 @@
+# Project Chronos: On-device low-latency lookahead dual-layer MoE inference
+__version__ = "0.1.0"
+
+# Bootstrap minimind dependency before any submodule imports
+from chronos.deps import ensure_minimind as _ensure_minimind
+_ensure_minimind()
+
+# Register with HuggingFace AutoXXX factories so
+# AutoModelForCausalLM.from_pretrained("path/to/chronos") works out of the box.
+try:
+    from chronos.model import hf_io as _hf_io  # noqa: F401
+except Exception as _e:
+    # Non-fatal: HF registration is a convenience, not a hard requirement.
+    import warnings
+    warnings.warn(f"Chronos HF registration skipped: {_e}")
+

chronos/_backend_legacy.py

@@ -0,0 +1,64 @@
+"""
+chronos/backend.py — Unified backend selector.
+
+Usage:
+    from chronos.backend import build_model, get_backend
+
+    backend = get_backend()   # "mlx" | "cuda" | "cpu"
+    model, engine = build_model(config, backend=backend)
+"""
+from __future__ import annotations
+import importlib
+from typing import Literal
+
+BackendType = Literal["mlx", "cuda", "mps", "cpu"]
+
+
+def get_backend() -> BackendType:
+    """
+    Auto-detect best available backend.
+    Priority: MLX (Apple Silicon) > CUDA > MPS > CPU
+    """
+    # Try MLX first (Apple Silicon unified memory)
+    if importlib.util.find_spec("mlx") is not None:
+        try:
+            import mlx.core as mx
+            if mx.metal.is_available():
+                return "mlx"
+        except Exception:
+            pass
+
+    import torch
+    if torch.cuda.is_available():
+        return "cuda"
+    try:
+        if torch.backends.mps.is_available():
+            return "mps"
+    except Exception:
+        pass
+    return "cpu"
+
+
+def build_model(config, backend: BackendType = None, ssd_dir: str = "./expert_cache"):
+    """
+    Build ChronosModel + InferenceEngine for the given backend.
+
+    Returns (model, inference_engine).
+    """
+    if backend is None:
+        backend = get_backend()
+
+    if backend == "mlx":
+        from chronos.mlx.model import ChronosMLXModel
+        from chronos.mlx.inference import ChronosMLXInferenceEngine
+        model = ChronosMLXModel(config)
+        engine = ChronosMLXInferenceEngine(model, config, ssd_dir=ssd_dir)
+        return model, engine
+    else:
+        import torch
+        from chronos.model.model_chronos import ChronosForCausalLM
+        from chronos.runtime.inference_engine import ChronosInferenceEngine
+        device_map = {"cuda": "cuda", "mps": "mps", "cpu": "cpu"}[backend]
+        model = ChronosForCausalLM(config).to(device_map)
+        engine = ChronosInferenceEngine(model, config, ssd_dir=ssd_dir)
+        return model, engine

chronos/app.py

@@ -0,0 +1,122 @@
+"""
+chronos/app.py — Project Chronos Web UI entry point (package-installable).
+
+Exposed as `chronos-ui` CLI entry point via pyproject.toml.
+
+Usage:
+    chronos-ui
+    chronos-ui --port 7861 --share
+    python -m chronos.app
+"""
+import sys
+import os
+import argparse
+
+# When run from the installed package, Project_Chronos root may not be on
+# sys.path. We add it so that ui.* imports resolve correctly.
+_pkg_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if _pkg_root not in sys.path:
+    sys.path.insert(0, _pkg_root)
+
+from ui.gradio_compat import gr
+
+import chronos.deps  # auto-bootstrap minimind
+from ui.i18n import (
+    LANGUAGE_CHOICES, DEFAULT_LANGUAGE,
+    t, set_current_lang, build_language_update, get_registered_components,
+)
+from ui.tabs.config_tab    import build_config_tab
+from ui.tabs.train_tab     import build_train_tab
+from ui.tabs.inference_tab import build_inference_tab
+from ui.tabs.benchmark_tab import build_benchmark_tab
+from ui.tabs.autotune_tab  import build_autotune_tab
+from ui.tabs.iomon_tab     import build_iomon_tab
+from ui.tabs.pipeline_tab  import build_pipeline_tab
+from ui.tabs.data_tab      import build_data_tab
+
+
+APP_CSS = """
+.app-header { padding: 24px 0 16px; border-bottom: 1px solid var(--border-color-primary); margin-bottom: 16px; }
+.app-header-title { font-size: 2rem; font-weight: 700; margin: 4px 0; }
+.app-header-subtitle { color: var(--body-text-color-subdued); margin: 0; }
+.app-meta-line { display: flex; gap: 12px; margin-top: 10px; flex-wrap: wrap; }
+.app-meta-chip { background: var(--background-fill-secondary); border-radius: 6px; padding: 3px 10px; font-size: 12px; }
+.toolbar-strip { margin-bottom: 8px; }
+"""
+
+
+def _hero_html() -> str:
+    return (
+        "<div class='app-header'>"
+        "<div style='font-size:12px;color:var(--body-text-color-subdued);letter-spacing:.08em;text-transform:uppercase;'>On-Device MoE Inference Workbench</div>"
+        "<h1 class='app-header-title'>Project Chronos</h1>"
+        "<p class='app-header-subtitle'>"
+        "Lookahead Routing · MLA + Sliding Window Attention · Async SSD Prefetch · "
+        "Pretrain / SFT / RL / ORPO · Optuna Auto-Tune"
+        "</p>"
+        "<div class='app-meta-line'>"
+        "<span class='app-meta-chip'>Apache-2.0</span>"
+        "<span class='app-meta-chip'>Python 3.10+</span>"
+        "<span class='app-meta-chip'>PyTorch 2.1+</span>"
+        "<span class='app-meta-chip'><a href='https://github.com/jingyaogong/minimind' target='_blank'>MoE kernel: minimind</a></span>"
+        "</div>"
+        "</div>"
+    )
+
+
+def build_app() -> gr.Blocks:
+    with gr.Blocks(
+        title="Project Chronos",
+        css=APP_CSS,
+        fill_width=True,
+    ) as app:
+        gr.HTML(_hero_html())
+
+        with gr.Row(elem_classes="toolbar-strip"):
+            language = gr.Dropdown(
+                choices=LANGUAGE_CHOICES,
+                value=DEFAULT_LANGUAGE,
+                label="Language / 语言 / 言語",
+                scale=1,
+            )
+
+        config_state, config_inputs, cfg_save_dir = build_config_tab()
+        build_data_tab()
+        build_train_tab(config_state, cfg_save_dir)
+        build_pipeline_tab()
+        build_inference_tab(config_state)
+        build_benchmark_tab()
+        build_autotune_tab(config_state, config_inputs)
+        build_iomon_tab()
+
+        def on_language_change(lang):
+            set_current_lang(lang)
+            return build_language_update(lang)
+
+        language.change(
+            fn=on_language_change,
+            inputs=[language],
+            outputs=get_registered_components(),
+        )
+
+    return app
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Project Chronos Web UI")
+    parser.add_argument("--port",  type=int,  default=7860)
+    parser.add_argument("--host",  type=str,  default="127.0.0.1")
+    parser.add_argument("--share", action="store_true")
+    args = parser.parse_args()
+
+    app = build_app()
+    app.launch(
+        server_name=args.host,
+        server_port=args.port,
+        share=args.share,
+        show_error=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

chronos/backend/__init__.py

@@ -0,0 +1,40 @@
+"""
+chronos/backend/__init__.py
+
+Public re-exports for the backend dispatcher. Also preserves the legacy
+``get_backend()`` / ``build_model()`` API from the pre-M5 single-file
+``chronos/backend.py`` so existing UI and CLI code keeps working.
+"""
+from chronos.backend.dispatcher import (
+    BackendDispatcher,
+    BackendInfo,
+    available,
+    select,
+    training_available,
+    select_training,
+    device_str,
+    training_device_str,
+    resolve_training_device,
+    describe,
+    describe_training,
+    AUTO_PRIORITY,
+    TRAINING_AUTO_PRIORITY,
+)
+
+# Back-compat shim for the pre-M5 API that code elsewhere still calls.
+from chronos._backend_legacy import build_model, BackendType  # type: ignore
+
+
+def get_backend():
+    """Back-compat alias for ``select()`` (auto-detected default backend)."""
+    return select()
+
+
+__all__ = [
+    "BackendDispatcher", "BackendInfo",
+    "available", "select", "training_available", "select_training",
+    "device_str", "training_device_str", "resolve_training_device",
+    "describe", "describe_training",
+    "AUTO_PRIORITY", "TRAINING_AUTO_PRIORITY",
+    "get_backend", "build_model", "BackendType",
+]

chronos/backend/dispatcher.py

@@ -0,0 +1,347 @@
+"""
+chronos/backend/dispatcher.py
+
+Unified compute-backend dispatcher.
+
+Chronos dispatches across these backend names:
+
+  cpu     — PyTorch CPU (always available)
+  cuda    — PyTorch CUDA (NVIDIA GPU)
+  mps     — PyTorch Metal Performance Shaders (Apple Silicon via torch)
+  mlx     — Apple MLX (Apple Silicon, non-torch inference path in this repo)
+  xpu     — PyTorch Intel XPU
+  vulkan  — PyTorch Vulkan (only if torch was built with USE_VULKAN=ON)
+  opencl  — third-party extension hook (no upstream backend; plug-in only)
+
+Training support in the current repo: cpu, cuda, mps, xpu
+Inference support (stock paths):      cpu, cuda, mps, mlx, xpu, (vulkan if built-in)
+Inference via ext plugin:             opencl
+
+Important: although MLX is available as an inference backend, this repository
+does not currently implement a full MLX-native training stack comparable to
+``chronos.trainer.*``. Training resolvers therefore exclude ``mlx`` until a
+real MLX trainer exists.
+"""
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import os
+from dataclasses import dataclass
+from typing import List, Optional
+
+
+BACKENDS = ("cuda", "mlx", "mps", "xpu", "vulkan", "opencl", "cpu")
+
+# General runtime auto-detect priority. This covers inference and legacy
+# backend selection, where MLX is meaningful on Apple Silicon.
+AUTO_PRIORITY = ("mlx", "cuda", "xpu", "mps", "vulkan", "opencl", "cpu")
+
+# Training must only pick backends with an actual training implementation in
+# this repository. Keep this separate from AUTO_PRIORITY so inference can
+# still prefer MLX while training stays honest.
+TRAINING_AUTO_PRIORITY = ("cuda", "xpu", "mps", "cpu")
+TRAINING_BACKENDS = ("cuda", "xpu", "mps", "cpu")
+
+
+@dataclass
+class BackendInfo:
+    name: str
+    available: bool
+    supports_training: bool
+    supports_amp: bool
+    torch_device: Optional[str]
+    notes: str = ""
+
+
+def _probe_cuda() -> BackendInfo:
+    try:
+        import torch
+        avail = bool(torch.cuda.is_available())
+    except Exception:
+        avail = False
+    return BackendInfo(
+        name="cuda", available=avail, supports_training=avail,
+        supports_amp=avail, torch_device="cuda" if avail else None,
+    )
+
+
+def _probe_mps() -> BackendInfo:
+    try:
+        import torch
+        avail = bool(torch.backends.mps.is_available())
+    except Exception:
+        avail = False
+    return BackendInfo(
+        name="mps", available=avail, supports_training=avail,
+        # fp16 autocast on MPS is still patchy across torch versions; bf16
+        # generally OK. We report True but document the caveat.
+        supports_amp=avail,
+        torch_device="mps" if avail else None,
+        notes="bf16 autocast preferred over fp16 on MPS" if avail else "",
+    )
+
+
+def _probe_xpu() -> BackendInfo:
+    """Intel XPU (torch.xpu). Only present on torch with oneAPI."""
+    try:
+        import torch
+        avail = hasattr(torch, "xpu") and bool(torch.xpu.is_available())
+    except Exception:
+        avail = False
+    return BackendInfo(
+        name="xpu", available=avail, supports_training=avail,
+        supports_amp=avail, torch_device="xpu" if avail else None,
+    )
+
+
+def _probe_mlx() -> BackendInfo:
+    spec = importlib.util.find_spec("mlx")
+    if spec is None:
+        return BackendInfo("mlx", False, False, False, None,
+                           notes="mlx not installed")
+    try:
+        import mlx.core as mx
+        avail = bool(mx.metal.is_available())
+    except Exception:
+        avail = False
+    return BackendInfo(
+        name="mlx", available=avail, supports_training=False,
+        supports_amp=avail, torch_device=None,
+        notes="non-torch backend; Chronos uses chronos.mlx.* inference paths",
+    )
+
+
+def _probe_vulkan() -> BackendInfo:
+    """PyTorch Vulkan is only present in custom builds (USE_VULKAN=ON).
+    The stock wheel does not include it. When present, training autograd
+    is NOT supported upstream — only inference forward."""
+    try:
+        import torch
+        check = getattr(torch, "is_vulkan_available", None)
+        avail = bool(check()) if check is not None else False
+    except Exception:
+        avail = False
+    return BackendInfo(
+        name="vulkan", available=avail,
+        supports_training=False,  # upstream torch.vulkan has no autograd
+        supports_amp=False,
+        torch_device="vulkan" if avail else None,
+        notes="inference-only; requires custom torch build",
+    )
+
+
+def _probe_opencl() -> BackendInfo:
+    """Extension hook. A third-party module can register an OpenCL backend by
+    placing ``chronos.backend.ext.opencl.PROBE()`` that returns a BackendInfo
+    with available=True. The bundled stub returns False."""
+    try:
+        from chronos.backend.ext import opencl as opencl_ext
+        info = opencl_ext.PROBE()
+        if not isinstance(info, BackendInfo):
+            info = BackendInfo("opencl", False, False, False, None,
+                               notes="ext.opencl.PROBE() malformed")
+        return info
+    except Exception as e:
+        return BackendInfo("opencl", False, False, False, None,
+                           notes=f"ext not available: {e}")
+
+
+_PROBES = {
+    "cuda": _probe_cuda, "mps": _probe_mps, "xpu": _probe_xpu,
+    "mlx": _probe_mlx, "vulkan": _probe_vulkan, "opencl": _probe_opencl,
+    "cpu": lambda: BackendInfo("cpu", True, True, False, "cpu"),
+}
+
+
+class BackendDispatcher:
+    """Central entry point for backend selection across Chronos."""
+
+    def __init__(self):
+        self._cache: dict[str, BackendInfo] = {}
+
+    def info(self, name: str) -> BackendInfo:
+        if name not in self._cache:
+            probe = _PROBES.get(name)
+            if probe is None:
+                raise ValueError(f"unknown backend: {name}")
+            self._cache[name] = probe()
+        return self._cache[name]
+
+    def available(self) -> List[str]:
+        """Returns names in priority order, filtered to those actually usable."""
+        return [n for n in AUTO_PRIORITY if self.info(n).available]
+
+    def training_available(self) -> List[str]:
+        """Returns trainable backend names in training-priority order."""
+        return [
+            n for n in TRAINING_AUTO_PRIORITY
+            if self.info(n).available and self.info(n).supports_training
+        ]
+
+    def select(self, prefer: Optional[str] = None) -> str:
+        """Resolve a concrete backend name.
+
+        Resolution order:
+          1. ``CHRONOS_BACKEND`` env var if set.
+          2. ``prefer`` argument if usable.
+          3. First available in auto-priority order.
+          4. ``cpu`` (always available).
+        """
+        env = os.environ.get("CHRONOS_BACKEND")
+        if env:
+            if self.info(env).available:
+                return env
+            # Fall through with warning; don't crash.
+            print(f"[chronos.backend] CHRONOS_BACKEND={env} not available; auto-selecting.")
+        if prefer and self.info(prefer).available:
+            return prefer
+        for n in AUTO_PRIORITY:
+            if self.info(n).available:
+                return n
+        return "cpu"
+
+    def select_training(self, prefer: Optional[str] = None) -> str:
+        """Resolve a concrete training backend name.
+
+        Resolution order:
+          1. ``CHRONOS_TRAIN_BACKEND`` env var if set.
+          2. ``prefer`` argument when available and trainable.
+          3. First available backend in training priority order.
+          4. ``cpu``.
+
+        ``prefer`` may be ``None`` or ``"auto"`` to request automatic
+        selection. Non-trainable values such as ``mlx`` are ignored here.
+        """
+        env = os.environ.get("CHRONOS_TRAIN_BACKEND")
+        if env:
+            env = env.strip().lower()
+            if env == "auto":
+                env = ""
+            elif env in BACKENDS and self.info(env).available and self.info(env).supports_training:
+                return env
+            elif env:
+                print(f"[chronos.backend] CHRONOS_TRAIN_BACKEND={env} not available for training; auto-selecting.")
+
+        prefer = (prefer or "").strip().lower()
+        if prefer == "auto":
+            prefer = ""
+        if prefer and prefer in BACKENDS:
+            info = self.info(prefer)
+            if info.available and info.supports_training:
+                return prefer
+
+        for n in TRAINING_AUTO_PRIORITY:
+            info = self.info(n)
+            if info.available and info.supports_training:
+                return n
+        return "cpu"
+
+    def device_str(self, name: str) -> Optional[str]:
+        """PyTorch device string for a backend, or None for non-torch backends."""
+        return self.info(name).torch_device
+
+    def training_device_str(self, name: Optional[str] = None) -> str:
+        """PyTorch device string for a training backend."""
+        backend = self.select_training(name)
+        return self.info(backend).torch_device or "cpu"
+
+    def resolve_training_device(self, prefer: Optional[str] = None) -> tuple[str, str]:
+        """Resolve ``(backend, torch_device)`` for training.
+
+        Accepts backend-level requests such as ``auto`` / ``cuda`` / ``mps``
+        as well as explicit torch device strings like ``cuda:0``.
+        """
+        raw = (prefer or "").strip()
+        name = raw.lower()
+
+        explicit_map = {
+            "cuda:": "cuda",
+            "xpu:": "xpu",
+        }
+        for prefix, backend in explicit_map.items():
+            if name.startswith(prefix):
+                info = self.info(backend)
+                if info.available and info.supports_training:
+                    return backend, raw
+                print(f"[chronos.backend] requested training device {raw} is not available; auto-selecting.")
+                chosen = self.select_training()
+                return chosen, self.info(chosen).torch_device or "cpu"
+
+        if name in {"cpu", "mps", "cuda", "xpu"}:
+            chosen = self.select_training(name)
+            if chosen == name:
+                return chosen, self.info(chosen).torch_device or "cpu"
+
+        chosen = self.select_training(name or None)
+        return chosen, self.info(chosen).torch_device or "cpu"
+
+    def supports_training(self, name: str) -> bool:
+        return self.info(name).supports_training
+
+    def supports_amp(self, name: str) -> bool:
+        return self.info(name).supports_amp
+
+    def notes(self, name: str) -> str:
+        return self.info(name).notes
+
+    def describe(self) -> str:
+        """Human-readable summary, useful for --help and the WebUI."""
+        lines = ["Chronos compute backends:"]
+        for n in AUTO_PRIORITY:
+            i = self.info(n)
+            marker = "✓" if i.available else "·"
+            tr = "train+inf" if i.supports_training else ("infer" if i.available else "n/a")
+            extra = f"  — {i.notes}" if i.notes else ""
+            lines.append(f"  {marker} {n:<8} {tr:<10} dev={i.torch_device or '-':<6}{extra}")
+        return "\n".join(lines)
+
+    def describe_training(self) -> str:
+        """Human-readable summary for trainable backends only."""
+        lines = ["Chronos training backends:"]
+        for n in TRAINING_AUTO_PRIORITY:
+            i = self.info(n)
+            marker = "✓" if (i.available and i.supports_training) else "·"
+            extra = f"  — {i.notes}" if i.notes else ""
+            lines.append(f"  {marker} {n:<8} train      dev={i.torch_device or '-':<6}{extra}")
+        return "\n".join(lines)
+
+
+# Module-level convenience singleton
+_default = BackendDispatcher()
+
+
+def available() -> List[str]:
+    return _default.available()
+
+
+def select(prefer: Optional[str] = None) -> str:
+    return _default.select(prefer)
+
+
+def training_available() -> List[str]:
+    return _default.training_available()
+
+
+def select_training(prefer: Optional[str] = None) -> str:
+    return _default.select_training(prefer)
+
+
+def device_str(name: Optional[str] = None) -> Optional[str]:
+    return _default.device_str(name or select())
+
+
+def training_device_str(name: Optional[str] = None) -> str:
+    return _default.training_device_str(name)
+
+
+def resolve_training_device(prefer: Optional[str] = None) -> tuple[str, str]:
+    return _default.resolve_training_device(prefer)
+
+
+def describe() -> str:
+    return _default.describe()
+
+
+def describe_training() -> str:
+    return _default.describe_training()

chronos/backend/ext/__init__.py

@@ -0,0 +1,2 @@
+"""Backend extension plugin namespace. Third parties can add modules here
+(e.g. chronos/backend/ext/opencl.py with a PROBE() function)."""

chronos/backend/ext/opencl.py

@@ -0,0 +1,22 @@
+"""Opt-in OpenCL backend plugin stub.
+
+A real implementation should define::
+
+    from chronos.backend.dispatcher import BackendInfo
+    def PROBE() -> BackendInfo: ...
+    def to_device(tensor, ...): ...
+
+The bundled stub always reports 'unavailable'. Third parties can replace
+this file (or monkey-patch the module) to plug in their own OpenCL backend
+without modifying Chronos core.
+"""
+from chronos.backend.dispatcher import BackendInfo
+
+
+def PROBE() -> BackendInfo:
+    return BackendInfo(
+        name="opencl", available=False,
+        supports_training=False, supports_amp=False,
+        torch_device=None,
+        notes="stub; install chronos-opencl-plugin or provide a custom ext module",
+    )

chronos/cli.py

@@ -0,0 +1,62 @@
+"""
+chronos/cli.py — Project Chronos unified CLI
+
+Commands:
+    chronos train     — pretrain / fine-tune a Chronos model
+    chronos eval      — Phase 1 lookahead accuracy validation
+    chronos benchmark — Phase 3 end-to-end benchmark
+    chronos export    — export expert weights to SSD cluster layout
+"""
+import sys
+
+import argparse
+
+
+def cmd_train(args):
+    import train_chronos
+    sys.argv = ['train_chronos'] + args
+    train_chronos.main()
+
+
+def cmd_eval(args):
+    from chronos.eval.io_profiler import main
+    sys.argv = ['io_profiler'] + args
+    main()
+
+
+def cmd_benchmark(args):
+    from chronos.eval.benchmark import main
+    sys.argv = ['benchmark'] + args
+    main()
+
+
+def cmd_export(args):
+    from chronos.io.cluster_layout import main
+    sys.argv = ['cluster_layout'] + args
+    main()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog='chronos',
+        description='Project Chronos — On-device low-latency lookahead dual-layer MoE inference',
+    )
+    sub = parser.add_subparsers(dest='command', required=True)
+
+    sub.add_parser('train',     help='Train a Chronos model',           add_help=False)
+    sub.add_parser('eval',      help='Phase 1 lookahead accuracy eval',  add_help=False)
+    sub.add_parser('benchmark', help='Phase 3 end-to-end benchmark',     add_help=False)
+    sub.add_parser('export',    help='Export expert cluster layout',     add_help=False)
+
+    known, rest = parser.parse_known_args()
+    dispatch = {
+        'train':     cmd_train,
+        'eval':      cmd_eval,
+        'benchmark': cmd_benchmark,
+        'export':    cmd_export,
+    }
+    dispatch[known.command](rest)
+
+
+if __name__ == '__main__':
+    main()