mithril-client 0.1.0a1__cp314-cp314-macosx_11_0_arm64.whl

mithril/cli/utils/skypilot_passthrough.py

@@ -0,0 +1,38 @@
+"""SkyPilot passthrough helpers for the Mithril CLI."""
+
+from __future__ import annotations
+
+import importlib
+import sys
+from contextlib import suppress
+
+SKY_ALIAS_COMMANDS: tuple[str, ...] = ("exec", "status", "start", "down", "stop")
+
+
+def run_sky(*args: str) -> int:
+    """Run sky CLI with given arguments.
+
+    Note: sky.cli is not part of skypilot's public API, but it's the stable
+    entry point for the `sky` command (defined in skypilot's setup.py). We
+    invoke it directly to avoid requiring `sky` to be on PATH (e.g., when
+    installed via `uv tool`).
+    """
+    try:
+        sky_cli = importlib.import_module("sky.client.cli.command").cli
+        sky_cli.main(args=list(args), standalone_mode=False)
+    except SystemExit as e:
+        return e.code if isinstance(e.code, int) else 1
+    except Exception as e:  # noqa: BLE001
+        # SkyPilot is Click-based and may raise ClickException; avoid importing click
+        # directly and rely on duck-typing to preserve decent error output.
+        show = getattr(e, "show", None)
+        if callable(show):
+            with suppress(Exception):
+                show()
+            exit_code = getattr(e, "exit_code", 1)
+            return exit_code if isinstance(exit_code, int) else 1
+
+        print(f"Error: {e}", file=sys.stderr)
+        return 1
+    else:
+        return 0

mithril/cli/utils/streaming.py

@@ -0,0 +1,235 @@
+"""Custom streaming utilities for Mithril CLI.
+
+This module provides a polling-based approach to monitoring SkyPilot requests,
+giving full control over the UI rather than parsing SkyPilot's log output.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, TypeVar
+
+from rich.live import Live
+from rich.spinner import Spinner
+from rich.text import Text
+
+from mithril.cli.utils.skypilot_passthrough import SKY_ALIAS_COMMANDS
+
+T = TypeVar("T")
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from rich.console import Console
+    from sky.server.common import RequestId
+
+    from mithril.sky import SkyClient
+
+TERMINAL_STATUSES = {"SUCCEEDED", "FAILED", "CANCELLED"}
+
+
+def stream_job_logs(
+    cluster_name: str,
+    job_id: int,
+    console: Console,
+    *,
+    sky: SkyClient,
+    follow: bool = True,
+    filter_fn: Callable[[str], str | None] | None = None,
+) -> int:
+    """Stream job logs with optional filtering.
+
+    Args:
+        cluster_name: The cluster name
+        job_id: The job ID to stream logs for
+        console: Rich console for output
+        sky: SkyClient instance
+        follow: Whether to follow the logs (like tail -f)
+        filter_fn: Optional function to filter/transform each line.
+                   Return the line to print, or None to skip it.
+
+    Returns:
+        Exit code (0 for success, non-zero for failure)
+    """
+    # Use preload_content=False to get an iterator
+    log_iter = sky.tail_logs(
+        cluster_name=cluster_name,
+        job_id=job_id,
+        follow=follow,
+        preload_content=False,
+    )
+
+    for line in log_iter:
+        if line is None:
+            # End of stream
+            break
+
+        if filter_fn:
+            filtered = filter_fn(line)
+            if filtered is not None:
+                console.print(filtered, end="", markup=False, highlight=False)
+        else:
+            console.print(line, end="", markup=False, highlight=False)
+
+    # Get the exit code by checking job status
+    job_statuses = sky.get(sky.job_status(cluster_name, job_ids=[job_id]))
+    status = job_statuses.get(job_id)
+
+    if status is None:
+        return 1  # Job not found
+
+    # JobStatus enum - check if succeeded
+    return 0 if str(status) == "SUCCEEDED" else 1
+
+
+def poll_launch(  # noqa: UP047 - PEP 695 syntax requires 3.12+
+    request_id: RequestId[T],
+    console: Console,
+    *,
+    sky: SkyClient,
+    poll_interval: float = 1.0,
+    quiet: bool = False,
+    initial_cluster_name: str | None = None,
+) -> tuple[T, str | None]:
+    """Poll a launch request until completion, showing custom progress.
+
+    Args:
+        request_id: The SkyPilot request ID from sky.launch()
+        console: Rich console for output
+        sky: SkyClient instance
+        poll_interval: How often to poll for status updates (seconds)
+        quiet: If True, suppress progress output
+        initial_cluster_name: Optional cluster name to seed status messages
+
+    Returns:
+        The result from the launch request (job_id, handle)
+    """
+    cluster_name = initial_cluster_name
+    initial_renderable, last_display_text = _render_status_message(
+        None,
+        cluster_name,
+    )
+
+    spinner = Spinner("dots", text="Launching...")
+    live = Live(spinner, console=console, refresh_per_second=10, transient=True)
+
+    try:
+        if not quiet:
+            spinner.update(text=initial_renderable)
+            live.start()
+        last_display_text, cluster_name = _run_poll_loop(
+            request_id=request_id,
+            poll_interval=poll_interval,
+            spinner=spinner,
+            live=live,
+            last_display_text=last_display_text,
+            cluster_name=cluster_name,
+            quiet=quiet,
+            sky=sky,
+        )
+    except KeyboardInterrupt:
+        if not quiet and cluster_name:
+            _print_provision_hint(console, cluster_name)
+        raise
+    finally:
+        if not quiet:
+            live.stop()
+
+    # Get the final result (this will raise if there was an error)
+    result = sky.get(request_id)
+
+    return result, cluster_name
+
+
+def format_status_message(status_msg: str, cluster_name: str | None) -> Text:
+    """Format a status message for display.
+
+    Appends the cluster name to status messages so users know which cluster
+    is being worked on. For example: "Launching" becomes "Launching (my-cluster)".
+
+    The cluster name is optional because it isn't available at the start of a
+    launch—SkyPilot assigns it during provisioning. Early status messages won't
+    have a cluster name yet.
+    """
+    formatted = status_msg
+    if cluster_name:
+        formatted = f"{formatted} ({cluster_name})"
+        formatted = (
+            f"{formatted} [dim]View logs: {_cli_command('logs')} --provision "
+            f"{cluster_name}[/dim]"
+        )
+    return Text.from_markup(formatted)
+
+
+def _default_launch_message(cluster_name: str | None) -> str:
+    if cluster_name:
+        return (
+            "Launching [dim]View logs: "
+            f"{_cli_command('logs')} --provision {cluster_name}[/dim]"
+        )
+    return "Launching..."
+
+
+def _render_status_message(
+    status_msg: str | None,
+    cluster_name: str | None,
+) -> tuple[Text, str]:
+    if status_msg:
+        renderable = format_status_message(status_msg, cluster_name)
+        return renderable, renderable.plain
+    message = _default_launch_message(cluster_name)
+    renderable = Text.from_markup(message)
+    return renderable, renderable.plain
+
+
+def _run_poll_loop(
+    *,
+    request_id: str,
+    poll_interval: float,
+    spinner: Spinner,
+    live: Live,
+    last_display_text: str,
+    cluster_name: str | None,
+    quiet: bool,
+    sky: SkyClient,
+) -> tuple[str, str | None]:
+    while True:
+        statuses = sky.api_status(request_ids=[request_id])
+
+        if not statuses:
+            time.sleep(poll_interval)
+            continue
+
+        status = statuses[0]
+        current_status = status.status
+        status_msg = status.status_msg
+        cluster_name = status.cluster_name or cluster_name
+
+        if not quiet:
+            renderable, display_text = _render_status_message(
+                status_msg,
+                cluster_name,
+            )
+            if display_text != last_display_text:
+                last_display_text = display_text
+                spinner.update(text=renderable)
+                live.update(spinner)
+
+        if current_status in TERMINAL_STATUSES:
+            break
+
+        time.sleep(poll_interval)
+
+    return last_display_text, cluster_name
+
+
+def _print_provision_hint(console: Console, cluster_name: str) -> None:
+    console.print(
+        f"[dim]View logs: {_cli_command('logs')} --provision {cluster_name}[/dim]"
+    )
+
+
+def _cli_command(command: str) -> str:
+    if command in SKY_ALIAS_COMMANDS:
+        return f"ml {command}"
+    return f"ml sky {command}"

mithril/cli/utils/volumes.py

@@ -0,0 +1,110 @@
+"""Volume synchronization between Mithril and SkyPilot.
+
+This module ensures that volumes referenced in a SkyPilot task are registered
+with SkyPilot before launch. It checks volumes against the Mithril API and
+automatically registers any that exist in Mithril but aren't yet known to SkyPilot.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from mithril.api.bindings.models.volume_model_interface import VolumeModelInterface
+
+# SkyPilot volume type constants from sky.utils.volume.VolumeType enum
+MITHRIL_FILE_SHARE_TYPE = "mithril-file-share"
+MITHRIL_BLOCK_TYPE = "mithril-block"
+
+if TYPE_CHECKING:
+    from mithril.api.bindings.models import VolumeModel
+    from mithril.api.client import MithrilClient
+    from mithril.sky import SkyClient
+
+
+def _get_skypilot_volumes(sky: SkyClient) -> set[str]:
+    """Get volume names registered with SkyPilot.
+
+    Args:
+        sky: SkyClient instance.
+
+    Returns:
+        Set of volume names currently registered in SkyPilot.
+    """
+    request_id = sky.volumes.ls()
+    volumes = sky.get(request_id)
+    return {vol.name for vol in volumes if vol.name}
+
+
+def _register_volume_with_skypilot(vol: VolumeModel, sky: SkyClient) -> None:
+    """Register a Mithril volume with SkyPilot.
+
+    Args:
+        vol: VolumeModel from Mithril API.
+        sky: SkyClient instance.
+
+    Raises:
+        Exception: If registration fails. sky.get() will raise RuntimeError,
+            RequestCancelled, or the original exception from the request.
+    """
+    # Map Mithril volume to SkyPilot volume spec
+    match vol.interface:
+        case VolumeModelInterface.FILE:
+            volume_type = MITHRIL_FILE_SHARE_TYPE
+        case VolumeModelInterface.BLOCK:
+            volume_type = MITHRIL_BLOCK_TYPE
+
+    # Note: SkyPilot treats GB and Gi identically (both map to 2**30 bytes)
+    sky_volume = sky.volumes.Volume(
+        name=vol.name,
+        type=volume_type,
+        infra=f"mithril/{vol.region}",
+        size=f"{vol.capacity_gb}GB",
+    )
+    request_id = sky.volumes.apply(sky_volume)
+    sky.get(request_id)
+
+
+def ensure_volumes_registered(
+    volume_names: list[str],
+    *,
+    mithril: MithrilClient,
+    sky: SkyClient,
+) -> list[str]:
+    """Ensure task volumes are registered with SkyPilot.
+
+    For each volume name:
+    1. Check if already registered with SkyPilot
+    2. If not, check if it exists in Mithril
+    3. If found in Mithril, register it with SkyPilot
+
+    Args:
+        volume_names: List of volume names to ensure are registered.
+        mithril: MithrilClient instance for Mithril API operations.
+        sky: SkyClient instance for SkyPilot operations.
+
+    Returns:
+        List of volume names that were newly registered.
+    """
+    if not volume_names:
+        return []
+
+    volume_names_set = set(volume_names)
+
+    # Check what's already in SkyPilot
+    sky_volumes = _get_skypilot_volumes(sky)
+    missing = volume_names_set - sky_volumes
+
+    if not missing:
+        return []
+
+    mithril_volumes = mithril.list_volumes()
+    mithril_by_name = {v.name: v for v in mithril_volumes}
+
+    registered: list[str] = []
+    for name in missing:
+        if name in mithril_by_name:
+            vol = mithril_by_name[name]
+            _register_volume_with_skypilot(vol, sky)
+            registered.append(name)
+
+    return registered

mithril/config.py

@@ -0,0 +1,47 @@
+"""Mithril configuration management.
+
+Delegates to the Rust implementation for config loading.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from mithril._mcli import MithrilConfig as RustConfig
+from mithril._mcli import load_config as rust_load_config
+
+DEFAULT_API_URL = "https://api.mithril.ai"
+
+
+class ConfigError(Exception):
+    """Raised when configuration is missing or invalid."""
+
+
+@dataclass(frozen=True, slots=True)
+class Config:
+    """Resolved configuration for Mithril operations."""
+
+    api_key: str
+    project_id: str
+    api_url: str
+
+
+def load_config() -> Config:
+    """Load and validate configuration from environment and config file.
+
+    Returns:
+        Validated Config object.
+
+    Raises:
+        ConfigError: If required configuration (api_key, project_id) is missing.
+    """
+    try:
+        rust_cfg: RustConfig = rust_load_config()
+        return Config(
+            api_key=rust_cfg.api_key,
+            project_id=rust_cfg.project_id,
+            api_url=rust_cfg.api_url,
+        )
+    except RuntimeError as e:
+        # Convert Rust RuntimeError to Python ConfigError
+        raise ConfigError(str(e)) from e

mithril/sky/__init__.py

@@ -0,0 +1,141 @@
+"""SkyPilot wrapper for stable imports and testability.
+
+This module provides two interfaces:
+
+1. **Module-level API** for SDK developers::
+
+       from mithril import sky
+
+       task = sky.Task(run="echo hello")
+       resources = sky.Resources(accelerators="A100:4")
+       sky.launch(task=task, cluster_name="dev", ...)
+
+2. **SkyClient class** for CLI code and testing (supports parameter injection)::
+
+       from mithril.sky import SkyClient
+
+       client = SkyClient()
+       # or inject a mock in tests
+       build_task(params, sky=mock_client)
+
+The module lazily imports SkyPilot to keep `ml --help` fast.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeVar
+
+from mithril.sky.client import SkyClient
+
+T = TypeVar("T")
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from sky import Task as SkyTask
+    from sky import backends
+    from sky.server.common import RequestId
+    from sky.server.requests import payloads
+    from sky.skylet import job_lib
+
+__all__ = [
+    "SkyClient",
+    "api_info",
+    "api_status",
+    "get",
+    "job_status",
+    "launch",
+    "tail_logs",
+]
+
+# Default client instance for module-level API
+_client = SkyClient()
+
+# Module-level attributes (Task, Resources, clouds, volumes) are accessed via
+# __getattr__ to maintain lazy loading - SkyPilot is only imported when first accessed.
+_LAZY_ATTRS = {"Task", "Resources", "clouds", "volumes"}
+
+
+def __getattr__(name: str) -> object:
+    """Lazily expose SkyClient attributes at module level."""
+    if name in _LAZY_ATTRS:
+        return getattr(_client, name)
+    msg = f"module {__name__!r} has no attribute {name!r}"
+    raise AttributeError(msg)
+
+
+def __dir__() -> list[str]:
+    """Include lazy attributes in dir() for discoverability."""
+    return [*globals().keys(), *_LAZY_ATTRS]
+
+
+def launch(
+    *,
+    task: SkyTask,
+    cluster_name: str | None = None,
+    retry_until_up: bool = False,
+    idle_minutes_to_autostop: int | None = None,
+    dryrun: bool = False,
+    down: bool = False,
+) -> RequestId[tuple[int | None, backends.ResourceHandle | None]]:
+    """Launch a task on SkyPilot.
+
+    Args:
+        task: A sky.Task object
+        cluster_name: Name of the cluster (auto-generated if None)
+        retry_until_up: Keep retrying until the cluster is up
+        idle_minutes_to_autostop: Auto-stop after idle for this many minutes
+        dryrun: If True, don't actually launch
+        down: Tear down cluster after job completes
+
+    Returns:
+        Request ID for the launch operation
+    """
+    return _client.launch(
+        task=task,
+        cluster_name=cluster_name,
+        retry_until_up=retry_until_up,
+        idle_minutes_to_autostop=idle_minutes_to_autostop,
+        dryrun=dryrun,
+        down=down,
+    )
+
+
+def api_status(*, request_ids: list[str]) -> list[payloads.RequestPayload]:
+    """Get status of SkyPilot requests."""
+    return _client.api_status(request_ids=request_ids)
+
+
+def get(request_id: RequestId[T]) -> T:  # noqa: UP047 - PEP 695 syntax requires 3.12+
+    """Get result of a SkyPilot request."""
+    return _client.get(request_id)
+
+
+def tail_logs(
+    *,
+    cluster_name: str,
+    job_id: int,
+    follow: bool = True,
+    preload_content: bool = False,
+) -> Iterable[str | None]:
+    """Stream logs from a SkyPilot job."""
+    return _client.tail_logs(
+        cluster_name=cluster_name,
+        job_id=job_id,
+        follow=follow,
+        preload_content=preload_content,
+    )
+
+
+def job_status(
+    cluster_name: str,
+    *,
+    job_ids: list[int] | None = None,
+) -> RequestId[dict[int | None, job_lib.JobStatus | None]]:
+    """Get status of a SkyPilot job (returns a request ID)."""
+    return _client.job_status(cluster_name, job_ids=job_ids)
+
+
+def api_info() -> object:
+    """Get API server info (status, version, etc.)."""
+    return _client.api_info()