Add files via upload

Author: FonaTech

chronos/backend/__init__.py

@@ -10,9 +10,15 @@
     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.
@@ -26,6 +32,9 @@ def get_backend():
 
 __all__ = [
     "BackendDispatcher", "BackendInfo",
-    "available", "select", "device_str", "describe", "AUTO_PRIORITY",
+    "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

@@ -3,23 +3,24 @@
 
 Unified compute-backend dispatcher.
 
-Chronos dispatches across six backend names:
+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, unified memory; non-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:                cpu, cuda, mps, mlx
-Inference support (stock torch): cpu, cuda, mps, mlx, (vulkan if built-in)
-Inference via ext plugin:        opencl  (requires chronos.backend.ext.opencl
-                                         implementation; stub returns False)
+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
 
-The Vulkan and OpenCL hooks exist so that someone running a custom PyTorch
-build (or a custom kernel) can plug in without modifying core Chronos code.
-On stock pip-installed PyTorch, they report "not available" and fall back.
+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
 
@@ -32,12 +33,16 @@
 
 BACKENDS = ("cuda", "mlx", "mps", "xpu", "vulkan", "opencl", "cpu")
 
-# Auto-detect priority. mlx > cuda > mps > xpu > vulkan > opencl > cpu.
-# The intent: prefer whichever backend has the most optimized training path
-# on the host. On Apple Silicon, mlx has unified memory advantages that
-# mps (via torch) cannot match; on NVIDIA, cuda is always best.
+# 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:
@@ -101,9 +106,9 @@ def _probe_mlx() -> BackendInfo:
     except Exception:
         avail = False
     return BackendInfo(
-        name="mlx", available=avail, supports_training=avail,
+        name="mlx", available=avail, supports_training=False,
         supports_amp=avail, torch_device=None,
-        notes="non-torch backend; Chronos uses chronos.mlx.* paths",
+        notes="non-torch backend; Chronos uses chronos.mlx.* inference paths",
     )
 
 
@@ -167,6 +172,13 @@ 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.
 
@@ -189,10 +201,81 @@ def select(self, prefer: Optional[str] = None) -> str:
                 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
 
@@ -213,6 +296,16 @@ def describe(self) -> str:
             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()
@@ -226,9 +319,29 @@ 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()