mlx_routing.py

"""MLX backend trainer registry — routes task → MLX trainer class.

All trainer imports are deferred to function bodies so the MLX classes don't
load on non-Apple Silicon machines during CLI startup (project convention:
lazy heavy imports).
"""

from __future__ import annotations


class _MLXRegistry(dict):
    """Dict-like registry that lazy-imports MLX trainer classes."""

    def _resolve(self, task: str):
        if task == "sft":
            from soup_cli.trainer.mlx_sft import MLXSFTTrainerWrapper

            return MLXSFTTrainerWrapper
        if task == "dpo":
            from soup_cli.trainer.mlx_dpo import MLXDPOTrainerWrapper

            return MLXDPOTrainerWrapper
        if task == "grpo":
            from soup_cli.trainer.mlx_grpo import MLXGRPOTrainerWrapper

            return MLXGRPOTrainerWrapper
        raise KeyError(task)

    def __contains__(self, key: object) -> bool:
        return key in ("sft", "dpo", "grpo")

    def __getitem__(self, key: str):
        return self._resolve(key)

    def keys(self):  # noqa: D401
        return ("sft", "dpo", "grpo")


MLX_TRAINER_REGISTRY: _MLXRegistry = _MLXRegistry()


def get_mlx_trainer(task: str):
    """Return the MLX trainer class for a task, or raise ValueError."""
    try:
        return MLX_TRAINER_REGISTRY[task]
    except KeyError as exc:
        supported = ", ".join(MLX_TRAINER_REGISTRY.keys())
        raise ValueError(
            f"MLX backend does not support task '{task}'. "
            f"Supported: {supported}. Use backend=transformers for full task coverage."
        ) from exc


def resolve_trainer(cfg, trainer_kwargs: dict | None = None):
    """Backend-first trainer resolution used by the train CLI.

    Returns ``(trainer_cls, kwargs)`` — ``trainer_cls`` is the MLX trainer
    class when ``cfg.backend == "mlx"`` (raising ValueError for unsupported
    tasks), or ``None`` so the caller falls through to the transformers task
    chain. ``trainer_kwargs`` are forwarded unchanged in both cases.
    """
    kwargs = dict(trainer_kwargs or {})
    if getattr(cfg, "backend", None) == "mlx":
        return get_mlx_trainer(cfg.task), kwargs
    return None, kwargs