mlx-stack 0.1.0__py3-none-any.whl

mlx_stack/core/deps.py

@@ -0,0 +1,323 @@
+"""Dependency management module for mlx-stack.
+
+Checks for and auto-installs pinned versions of vllm-mlx and litellm
+as uv tools. Performs PATH lookup to detect installed tools, auto-installs
+via ``uv tool install <tool>==<version>`` when missing, shows progress
+during installation, verifies post-install availability, detects version
+mismatches with warnings, and provides clear manual install instructions
+on failure.
+
+Read-only commands (profile, config, recommend, models) do NOT trigger
+dependency checks.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import shutil
+import subprocess
+from dataclasses import dataclass
+
+from rich.console import Console
+
+# --------------------------------------------------------------------------- #
+# Pinned versions
+# --------------------------------------------------------------------------- #
+
+PINNED_VERSIONS: dict[str, str] = {
+    "vllm-mlx": "0.2.6",
+    "litellm": "1.67.2",
+}
+
+# Map tool name to the CLI binary name used for PATH lookup
+_TOOL_BINARY_MAP: dict[str, str] = {
+    "vllm-mlx": "vllm-mlx",
+    "litellm": "litellm",
+}
+
+# --------------------------------------------------------------------------- #
+# Exceptions
+# --------------------------------------------------------------------------- #
+
+
+class DependencyError(Exception):
+    """Raised when a required dependency cannot be found or installed."""
+
+
+class DependencyInstallError(DependencyError):
+    """Raised when ``uv tool install`` fails."""
+
+
+class DependencyVersionMismatchWarning(Exception):
+    """Sentinel for version mismatch warnings (not raised, used for typing)."""
+
+
+# --------------------------------------------------------------------------- #
+# Data classes
+# --------------------------------------------------------------------------- #
+
+
+@dataclass(frozen=True)
+class DependencyStatus:
+    """Status of a single managed dependency.
+
+    Attributes:
+        name: Package name (e.g. ``vllm-mlx``).
+        pinned_version: The version we want installed.
+        installed: Whether the tool binary is available on PATH.
+        installed_version: The detected installed version, or ``None``.
+        version_match: ``True`` when installed version matches pinned.
+    """
+
+    name: str
+    pinned_version: str
+    installed: bool
+    installed_version: str | None
+    version_match: bool | None
+
+
+# --------------------------------------------------------------------------- #
+# Internal helpers
+# --------------------------------------------------------------------------- #
+
+_console = Console(stderr=True)
+
+
+def _find_binary(tool: str) -> str | None:
+    """Return the full path to a tool's binary, or ``None`` if not found.
+
+    Uses :func:`shutil.which` for PATH lookup.
+    """
+    binary = _TOOL_BINARY_MAP.get(tool, tool)
+    return shutil.which(binary)
+
+
+def _get_installed_version(tool: str) -> str | None:
+    """Attempt to determine the installed version of a uv tool.
+
+    Runs ``uv tool list`` and parses the output for the tool name and
+    version.  Returns ``None`` when the tool is not found or the output
+    cannot be parsed.
+    """
+    uv_path = shutil.which("uv")
+    if uv_path is None:
+        return None
+
+    try:
+        env = {**os.environ, "NO_COLOR": "1"}
+        result = subprocess.run(
+            [uv_path, "tool", "list"],
+            capture_output=True,
+            text=True,
+            timeout=30,
+            env=env,
+        )
+    except (subprocess.TimeoutExpired, OSError):
+        return None
+
+    if result.returncode != 0:
+        return None
+
+    # Parse lines like "vllm-mlx v0.2.6" or "litellm v1.67.2"
+    for line in result.stdout.splitlines():
+        # Match "<tool> v<version>" pattern
+        pattern = rf"^{re.escape(tool)}\s+v?(\S+)"
+        match = re.match(pattern, line.strip())
+        if match:
+            return match.group(1)
+
+    return None
+
+
+def _install_tool(tool: str, version: str) -> None:
+    """Install a tool at a pinned version via ``uv tool install``.
+
+    Args:
+        tool: Package name (e.g. ``vllm-mlx``).
+        version: Exact version to install.
+
+    Raises:
+        DependencyInstallError: When ``uv tool install`` fails.
+        DependencyError: When ``uv`` itself is not available.
+    """
+    uv_path = shutil.which("uv")
+    if uv_path is None:
+        msg = (
+            "uv is not available on PATH. "
+            "Install it from https://docs.astral.sh/uv/ and try again."
+        )
+        raise DependencyError(msg)
+
+    install_spec = f"{tool}=={version}"
+    cmd = [uv_path, "tool", "install", install_spec]
+    cmd_str = f"uv tool install {install_spec}"
+
+    _console.print(f"[cyan]Installing {tool} v{version}...[/cyan]")
+
+    try:
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            timeout=300,
+        )
+    except subprocess.TimeoutExpired:
+        msg = (
+            f"Installation timed out: {cmd_str}\n\n"
+            f"Install manually with: {cmd_str}"
+        )
+        raise DependencyInstallError(msg) from None
+    except OSError as exc:
+        msg = (
+            f"Failed to run: {cmd_str}\n"
+            f"Error: {exc}\n\n"
+            f"Install manually with: {cmd_str}"
+        )
+        raise DependencyInstallError(msg) from None
+
+    if result.returncode != 0:
+        stderr = result.stderr.strip()
+        msg = (
+            f"Failed to install {tool}:\n"
+            f"  Command: {cmd_str}\n"
+            f"  Error: {stderr}\n\n"
+            f"Install manually with: {cmd_str}"
+        )
+        raise DependencyInstallError(msg)
+
+    _console.print(f"[green]✓ {tool} v{version} installed successfully.[/green]")
+
+
+def _verify_post_install(tool: str) -> bool:
+    """Verify that a tool is available on PATH after installation.
+
+    Returns ``True`` if the binary is found.
+    """
+    return _find_binary(tool) is not None
+
+
+# --------------------------------------------------------------------------- #
+# Public API
+# --------------------------------------------------------------------------- #
+
+
+def check_dependency(tool: str) -> DependencyStatus:
+    """Check the installation status of a single dependency.
+
+    Performs PATH lookup and version detection without installing anything.
+
+    Args:
+        tool: Package name (must be a key in :data:`PINNED_VERSIONS`).
+
+    Returns:
+        A :class:`DependencyStatus` with installation details.
+
+    Raises:
+        ValueError: If *tool* is not a known dependency.
+    """
+    if tool not in PINNED_VERSIONS:
+        msg = f"Unknown dependency '{tool}'. Known: {', '.join(sorted(PINNED_VERSIONS))}"
+        raise ValueError(msg)
+
+    pinned = PINNED_VERSIONS[tool]
+    binary_path = _find_binary(tool)
+    installed = binary_path is not None
+
+    installed_version: str | None = None
+    version_match: bool | None = None
+
+    if installed:
+        installed_version = _get_installed_version(tool)
+        if installed_version is not None:
+            version_match = installed_version == pinned
+
+    return DependencyStatus(
+        name=tool,
+        pinned_version=pinned,
+        installed=installed,
+        installed_version=installed_version,
+        version_match=version_match,
+    )
+
+
+def ensure_dependency(tool: str) -> DependencyStatus:
+    """Ensure a dependency is installed at the pinned version.
+
+    1. Check if the tool is already available on PATH.
+    2. If missing, auto-install via ``uv tool install <tool>==<version>``.
+    3. Verify post-install availability.
+    4. Warn on version mismatch (but do not block).
+
+    Args:
+        tool: Package name (must be a key in :data:`PINNED_VERSIONS`).
+
+    Returns:
+        A :class:`DependencyStatus` reflecting the final state.
+
+    Raises:
+        ValueError: If *tool* is not a known dependency.
+        DependencyError: If ``uv`` is not available.
+        DependencyInstallError: If auto-install fails.
+    """
+    status = check_dependency(tool)
+
+    if not status.installed:
+        # Auto-install
+        _install_tool(tool, status.pinned_version)
+
+        # Verify post-install
+        if not _verify_post_install(tool):
+            cmd_str = f"uv tool install {tool}=={status.pinned_version}"
+            msg = (
+                f"{tool} was not found on PATH after installation.\n"
+                f"This may be because the uv tool bin directory is not in your PATH.\n\n"
+                f"Try running:\n"
+                f"  {cmd_str}\n"
+                f"  export PATH=\"$HOME/.local/bin:$PATH\""
+            )
+            raise DependencyInstallError(msg)
+
+        # Re-check status after install
+        status = check_dependency(tool)
+
+    # Warn on version mismatch
+    if status.installed and status.version_match is False:
+        _warn_version_mismatch(tool, status)
+
+    return status
+
+
+def ensure_all_dependencies() -> list[DependencyStatus]:
+    """Ensure all managed dependencies are installed.
+
+    Calls :func:`ensure_dependency` for each tool in :data:`PINNED_VERSIONS`.
+
+    Returns:
+        A list of :class:`DependencyStatus` for each dependency.
+
+    Raises:
+        DependencyError: If any dependency cannot be installed.
+        DependencyInstallError: If any auto-install fails.
+    """
+    results: list[DependencyStatus] = []
+    for tool in PINNED_VERSIONS:
+        results.append(ensure_dependency(tool))
+    return results
+
+
+def _warn_version_mismatch(tool: str, status: DependencyStatus) -> None:
+    """Display a Rich warning about a version mismatch.
+
+    Args:
+        tool: The package name.
+        status: The dependency status with version info.
+    """
+    pinned = status.pinned_version
+    installed = status.installed_version or "unknown"
+    cmd = f"uv tool install {tool}=={pinned}"
+    _console.print(
+        f"[yellow]⚠ {tool} version mismatch: "
+        f"installed v{installed}, expected v{pinned}.[/yellow]\n"
+        f"  Upgrade/downgrade with: {cmd}"
+    )

mlx_stack/core/hardware.py

@@ -0,0 +1,304 @@
+"""Hardware detection module for Apple Silicon Macs.
+
+Detects chip model, GPU core count, unified memory, and memory bandwidth.
+Uses a lookup table of 17 known M-series variants (M1 through M5) and
+estimates bandwidth for unknown chips.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import subprocess
+from dataclasses import dataclass
+from typing import Any
+
+from mlx_stack.core.paths import ensure_data_home, get_profile_path
+
+# --------------------------------------------------------------------------- #
+# Lookup table — 17 known Apple Silicon M-series variants
+# bandwidth_gbps values sourced from Apple published specs
+# --------------------------------------------------------------------------- #
+CHIP_SPECS: dict[str, dict[str, float | int]] = {
+    "Apple M1": {"bandwidth_gbps": 68.25},
+    "Apple M1 Pro": {"bandwidth_gbps": 200.0},
+    "Apple M1 Max": {"bandwidth_gbps": 400.0},
+    "Apple M1 Ultra": {"bandwidth_gbps": 800.0},
+    "Apple M2": {"bandwidth_gbps": 100.0},
+    "Apple M2 Pro": {"bandwidth_gbps": 200.0},
+    "Apple M2 Max": {"bandwidth_gbps": 400.0},
+    "Apple M2 Ultra": {"bandwidth_gbps": 800.0},
+    "Apple M3": {"bandwidth_gbps": 100.0},
+    "Apple M3 Pro": {"bandwidth_gbps": 150.0},
+    "Apple M3 Max": {"bandwidth_gbps": 400.0},
+    "Apple M3 Ultra": {"bandwidth_gbps": 800.0},
+    "Apple M4": {"bandwidth_gbps": 120.0},
+    "Apple M4 Pro": {"bandwidth_gbps": 273.0},
+    "Apple M4 Max": {"bandwidth_gbps": 546.0},
+    "Apple M4 Ultra": {"bandwidth_gbps": 819.2},
+    "Apple M5 Max": {"bandwidth_gbps": 546.0},
+}
+
+
+class HardwareError(Exception):
+    """Raised when hardware detection fails or hardware is unsupported."""
+
+
+@dataclass(frozen=True)
+class HardwareProfile:
+    """Detected hardware profile for the current machine."""
+
+    chip: str
+    gpu_cores: int
+    memory_gb: int
+    bandwidth_gbps: float
+    is_estimate: bool
+
+    @property
+    def profile_id(self) -> str:
+        """Generate a profile ID like 'm4-pro-64'."""
+        # Strip 'Apple ' prefix, lowercase, replace spaces with hyphens
+        chip_part = self.chip.removeprefix("Apple ").lower().replace(" ", "-")
+        return f"{chip_part}-{self.memory_gb}"
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to a dictionary suitable for JSON output."""
+        return {
+            "chip": self.chip,
+            "gpu_cores": self.gpu_cores,
+            "memory_gb": self.memory_gb,
+            "bandwidth_gbps": self.bandwidth_gbps,
+            "profile_id": self.profile_id,
+        }
+
+
+# --------------------------------------------------------------------------- #
+# System command wrappers (mockable for tests)
+# --------------------------------------------------------------------------- #
+
+
+def _run_sysctl(key: str) -> str:
+    """Run sysctl and return the output for a given key.
+
+    Raises:
+        HardwareError: If the sysctl command fails.
+    """
+    try:
+        result = subprocess.run(
+            ["sysctl", "-n", key],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        if result.returncode != 0:
+            msg = f"sysctl failed for key '{key}': {result.stderr.strip()}"
+            raise HardwareError(msg)
+        return result.stdout.strip()
+    except FileNotFoundError:
+        msg = "sysctl command not found — are you running on macOS?"
+        raise HardwareError(msg) from None
+    except subprocess.TimeoutExpired:
+        msg = f"sysctl timed out reading key '{key}'"
+        raise HardwareError(msg) from None
+
+
+def _run_system_profiler() -> str:
+    """Run system_profiler SPDisplaysDataType and return output.
+
+    Raises:
+        HardwareError: If the system_profiler command fails.
+    """
+    try:
+        result = subprocess.run(
+            ["system_profiler", "SPDisplaysDataType"],
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        if result.returncode != 0:
+            msg = f"system_profiler failed: {result.stderr.strip()}"
+            raise HardwareError(msg)
+        return result.stdout
+    except FileNotFoundError:
+        msg = "system_profiler command not found — are you running on macOS?"
+        raise HardwareError(msg) from None
+    except subprocess.TimeoutExpired:
+        msg = "system_profiler timed out"
+        raise HardwareError(msg) from None
+
+
+# --------------------------------------------------------------------------- #
+# Detection functions
+# --------------------------------------------------------------------------- #
+
+
+def detect_chip() -> str:
+    """Detect the chip model name from sysctl.
+
+    Returns:
+        Chip name string, e.g. 'Apple M4 Pro'.
+
+    Raises:
+        HardwareError: If the chip cannot be detected or is not Apple Silicon.
+    """
+    brand = _run_sysctl("machdep.cpu.brand_string")
+    if not brand:
+        msg = "Could not read CPU brand string from sysctl"
+        raise HardwareError(msg)
+
+    # Validate it's Apple Silicon (M-series)
+    if not re.match(r"Apple M\d", brand):
+        msg = "mlx-stack requires Apple Silicon (M1 or later)"
+        raise HardwareError(msg)
+
+    return brand
+
+
+def detect_memory_gb() -> int:
+    """Detect unified memory in GB from sysctl hw.memsize.
+
+    Returns:
+        Memory in GB as an integer.
+
+    Raises:
+        HardwareError: If memory size cannot be read.
+    """
+    raw = _run_sysctl("hw.memsize")
+    try:
+        memsize_bytes = int(raw)
+    except ValueError:
+        msg = f"Unexpected hw.memsize value: {raw!r}"
+        raise HardwareError(msg) from None
+
+    return memsize_bytes // (1024 ** 3)
+
+
+def detect_gpu_cores() -> int:
+    """Detect GPU core count from system_profiler.
+
+    Returns:
+        Number of GPU cores.
+
+    Raises:
+        HardwareError: If GPU core count cannot be determined.
+    """
+    output = _run_system_profiler()
+
+    # Parse "Total Number of Cores: 40" from system_profiler output
+    match = re.search(r"Total Number of Cores:\s*(\d+)", output)
+    if not match:
+        msg = "Could not determine GPU core count from system_profiler output"
+        raise HardwareError(msg)
+
+    return int(match.group(1))
+
+
+def estimate_bandwidth(memory_gb: int) -> float:
+    """Estimate memory bandwidth for unknown chips based on memory size.
+
+    Uses a simple heuristic: larger memory configurations tend to have
+    higher bandwidth. This is a rough estimate — users should run
+    `mlx-stack bench --save` for accurate numbers.
+
+    Args:
+        memory_gb: Total unified memory in GB.
+
+    Returns:
+        Estimated bandwidth in GB/s.
+    """
+    if memory_gb <= 8:
+        return 68.0
+    if memory_gb <= 16:
+        return 100.0
+    if memory_gb <= 32:
+        return 200.0
+    if memory_gb <= 64:
+        return 400.0
+    if memory_gb <= 128:
+        return 546.0
+    return 800.0
+
+
+def lookup_bandwidth(chip: str, memory_gb: int) -> tuple[float, bool]:
+    """Look up bandwidth from the chip specs table.
+
+    Args:
+        chip: Full chip name (e.g. 'Apple M4 Pro').
+        memory_gb: Total unified memory in GB.
+
+    Returns:
+        Tuple of (bandwidth_gbps, is_estimate). is_estimate is True when
+        the chip is not in the lookup table and bandwidth was estimated.
+    """
+    specs = CHIP_SPECS.get(chip)
+    if specs is not None:
+        return float(specs["bandwidth_gbps"]), False
+    return estimate_bandwidth(memory_gb), True
+
+
+# --------------------------------------------------------------------------- #
+# Main detection entry point
+# --------------------------------------------------------------------------- #
+
+
+def detect_hardware() -> HardwareProfile:
+    """Detect hardware and build a HardwareProfile.
+
+    Detects chip model, GPU core count, unified memory, and memory
+    bandwidth. Known chips use the lookup table; unknown chips use
+    a bandwidth estimation.
+
+    Returns:
+        A HardwareProfile with all detected values.
+
+    Raises:
+        HardwareError: If hardware cannot be detected or is unsupported.
+    """
+    chip = detect_chip()
+    memory_gb = detect_memory_gb()
+    gpu_cores = detect_gpu_cores()
+    bandwidth_gbps, is_estimate = lookup_bandwidth(chip, memory_gb)
+
+    return HardwareProfile(
+        chip=chip,
+        gpu_cores=gpu_cores,
+        memory_gb=memory_gb,
+        bandwidth_gbps=bandwidth_gbps,
+        is_estimate=is_estimate,
+    )
+
+
+def save_profile(profile: HardwareProfile) -> None:
+    """Write the hardware profile to ~/.mlx-stack/profile.json.
+
+    Overwrites any existing profile file.
+
+    Args:
+        profile: The hardware profile to save.
+    """
+    ensure_data_home()
+    profile_path = get_profile_path()
+    profile_path.write_text(json.dumps(profile.to_dict(), indent=2) + "\n")
+
+
+def load_profile() -> HardwareProfile | None:
+    """Load the hardware profile from disk, if it exists.
+
+    Returns:
+        A HardwareProfile if the file exists and is valid, None otherwise.
+    """
+    profile_path = get_profile_path()
+    if not profile_path.exists():
+        return None
+
+    try:
+        data = json.loads(profile_path.read_text())
+        return HardwareProfile(
+            chip=data["chip"],
+            gpu_cores=data["gpu_cores"],
+            memory_gb=data["memory_gb"],
+            bandwidth_gbps=data["bandwidth_gbps"],
+            is_estimate=False,  # saved profiles are considered authoritative
+        )
+    except (json.JSONDecodeError, KeyError, TypeError):
+        return None