mithril-client 0.1.0a1__cp314-cp314-macosx_11_0_arm64.whl

mithril/api/client.py

@@ -0,0 +1,138 @@
+"""Mithril API client.
+
+Provides a high-level interface for interacting with the Mithril API.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import httpx
+
+from mithril.api.bindings import AuthenticatedClient
+from mithril.api.bindings.api.volumes import get_volumes_v2_volumes_get
+from mithril.api.bindings.models import HTTPValidationError
+from mithril.config import load_config
+
+if TYPE_CHECKING:
+    from mithril.api.bindings.models import VolumeModel
+    from mithril.config import Config
+
+DEFAULT_TIMEOUT = httpx.Timeout(30.0)
+
+
+class MithrilAPIError(Exception):
+    """Raised when the API returns an unexpected or invalid response."""
+
+
+class MithrilClient:
+    """Client for interacting with the Mithril API.
+
+    Can be constructed with explicit credentials or loaded from config:
+
+        # Load from environment/config file
+        client = MithrilClient()
+
+        # Explicit credentials
+        client = MithrilClient(api_url="https://api.mithril.ai", api_key="...")
+
+        # From existing config
+        client = MithrilClient.from_config(config)
+    """
+
+    def __init__(
+        self,
+        api_url: str | None = None,
+        api_key: str | None = None,
+        *,
+        project_id: str | None = None,
+        timeout: httpx.Timeout | None = None,
+    ) -> None:
+        """Initialize the client.
+
+        Args:
+            api_url: Base URL of the API. If not provided, loaded from config.
+            api_key: API key for authentication. If not provided, loaded from config.
+            project_id: Default project ID. If not provided, loaded from config.
+            timeout: Request timeout configuration.
+
+        Raises:
+            ConfigError: If credentials not provided and config is missing/invalid.
+        """
+        if api_url is not None and api_key is not None:
+            self._api_url = api_url
+            self._api_key = api_key
+            self._project_id = project_id or ""
+        else:
+            config = load_config()
+            self._api_url = api_url or config.api_url
+            self._api_key = api_key or config.api_key
+            self._project_id = project_id or config.project_id
+
+        self._timeout = timeout or DEFAULT_TIMEOUT
+        self._http_client: AuthenticatedClient | None = None
+
+    @classmethod
+    def from_config(cls, config: Config) -> MithrilClient:
+        """Create a client from an existing Config object.
+
+        Args:
+            config: Configuration object with API credentials.
+
+        Returns:
+            Configured MithrilClient instance.
+        """
+        return cls(
+            api_url=config.api_url,
+            api_key=config.api_key,
+            project_id=config.project_id,
+        )
+
+    @property
+    def project_id(self) -> str:
+        """The default project ID for this client."""
+        return self._project_id
+
+    @property
+    def _client(self) -> AuthenticatedClient:
+        """Lazily create the underlying HTTP client."""
+        if self._http_client is None:
+            self._http_client = AuthenticatedClient(
+                base_url=self._api_url,
+                token=self._api_key,
+                timeout=self._timeout,
+            )
+        return self._http_client
+
+    def list_volumes(
+        self,
+        project_id: str | None = None,
+        *,
+        region: str | None = None,
+    ) -> list[VolumeModel]:
+        """List volumes for a project.
+
+        Args:
+            project_id: Project to list volumes for. Defaults to client's project_id.
+            region: Optional region filter.
+
+        Returns:
+            List of VolumeModel objects.
+
+        Raises:
+            MithrilAPIError: If the API returns a validation error or
+                unexpected response.
+        """
+        project = project_id or self._project_id
+        result = get_volumes_v2_volumes_get.sync(
+            client=self._client,
+            project=project,
+            region=region,
+        )
+        if isinstance(result, list):
+            return result
+        error_info = (
+            result.detail if isinstance(result, HTTPValidationError) else result
+        )
+        msg = f"Unexpected API response when listing volumes: {error_info}"
+        raise MithrilAPIError(msg)

mithril/cli/__init__.py

@@ -0,0 +1,7 @@
+"""Mithril CLI."""
+
+from __future__ import annotations
+
+from mithril.cli.main import cli, main
+
+__all__ = ["cli", "main"]

mithril/cli/commands/__init__.py

@@ -0,0 +1,15 @@
+"""CLI commands package.
+
+Do not re-export command objects from this package (e.g. `from .launch import launch`).
+
+Why: doing so binds `mithril.cli.commands.launch` to the exported object, which can
+shadow the `mithril.cli.commands.launch` *module* and break code/tests that import the
+module to access helpers. This is common in CLIs because command objects are often
+defined at import time, and it's common to name the command the same as the module
+(e.g. `launch.py` exporting `launch`), making re-exports especially tempting.
+
+Prefer importing from the submodule:
+    from mithril.cli.commands.launch import launch  # ✓
+Avoid:
+    from mithril.cli.commands import launch         # ✗
+"""

mithril/cli/commands/help.py

@@ -0,0 +1,88 @@
+"""Help command for the Mithril CLI."""
+
+from __future__ import annotations
+
+from importlib.metadata import version
+
+from cyclopts import App
+from rich.console import Console
+
+from mithril._mcli import run as run_rust
+from mithril.cli.commands.launch import launch
+from mithril.cli.utils.skypilot_passthrough import SKY_ALIAS_COMMANDS, run_sky
+
+RUST_COMMANDS = {"ssh", "instance", "k8s"}
+
+
+def print_help() -> None:
+    """Print top-level help.
+
+    We route some commands directly to SkyPilot for passthrough, so we keep the
+    top-level help text explicit and stable.
+    """
+    ver = version("mithril-client")
+    help_text = f"""[dim]Mithril CLI {ver}[/dim]
+
+[dim]❊ Quickstart:[/dim]
+  [bold]1. ml setup[/bold]                  First-time setup
+  [bold]2. ml launch --gpus H100:8[/bold]   Spin up a GPU instance
+  [bold]3. ml status[/bold]                 Check running clusters
+  [bold]4. ml down[/bold]                   Tear down when you're done
+
+[dim]❊ Workload:[/dim]
+  launch    Launch a cluster and run a task.
+  exec      Execute a command or open an interactive shell on a cluster.
+  status    Show cluster and job status.
+  start     Start a stopped cluster.
+  stop      Stop/pause a cluster.
+  down      Delete a cluster.
+
+[dim]❊ Infrastructure:[/dim]
+  instance  Manage compute instances.
+  k8s       Manage Kubernetes clusters.
+  ssh       SSH into an instance.
+
+[dim]❊ Utilities:[/dim]
+  sky       Direct SkyPilot access.
+  setup     Configure credentials.
+  help      Show help manual.
+"""
+    Console(highlight=False).print(help_text)
+
+
+def help_cmd(command: tuple[str, ...] = ()) -> None:
+    """Show help for `ml` or a specific command.
+
+    For SkyPilot passthrough commands, this forwards to SkyPilot's help output.
+    For Rust commands (ssh, instance, k8s), this delegates to the Rust CLI.
+    """
+    if not command:
+        print_help()
+        return
+
+    cmd_name, *rest = command
+
+    # Rust commands: delegate to Rust CLI for help
+    if cmd_name in RUST_COMMANDS:
+        args = ["ml", cmd_name, *rest, "--help"]
+        raise SystemExit(run_rust(args))
+
+    # SkyPilot passthrough commands
+    if cmd_name == "sky":
+        args = [*rest, "--help"] if rest else ["--help"]
+        raise SystemExit(run_sky(*args))
+
+    if cmd_name in SKY_ALIAS_COMMANDS:
+        args = [cmd_name, *rest, "--help"] if rest else [cmd_name, "--help"]
+        raise SystemExit(run_sky(*args))
+
+    # First-class Python commands: render help via Cyclopts.
+    app = App(
+        name="ml",
+        version=version("mithril-client"),
+        version_flags=["--version"],
+        help_formatter="plain",
+    )
+    app.command(launch)
+
+    app([cmd_name, *rest, "--help"])

mithril/cli/commands/launch.py

@@ -0,0 +1,353 @@
+"""Launch command for the Mithril CLI."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from importlib.metadata import version as package_version
+from typing import TYPE_CHECKING, Annotated, Any, cast
+
+import click
+from cyclopts import Parameter
+from rich.console import Console
+
+from mithril.api.client import MithrilClient
+from mithril.cli.utils.skypilot_passthrough import SKY_ALIAS_COMMANDS
+from mithril.cli.utils.streaming import poll_launch, stream_job_logs
+from mithril.cli.utils.volumes import ensure_volumes_registered
+from mithril.sky import SkyClient
+
+if TYPE_CHECKING:
+    from sky import Resources as SkyResources
+    from sky import Task as SkyTask
+
+console = Console()
+INDENT_SYMBOL = "├── "
+INDENT_LAST_SYMBOL = "└── "
+
+
+@dataclass(frozen=True, slots=True)
+class ResourceParams:
+    """Resource-related CLI options."""
+
+    gpus: str | None = None
+    cpus: str | None = None
+    memory: str | None = None
+    cloud: str | None = None
+    region: str | None = None
+    use_spot: bool | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class TaskParams:
+    """Task-related CLI options."""
+
+    entrypoint: str | None
+    name: str | None = None
+    workdir: str | None = None
+    num_nodes: int | None = None
+    env: tuple[str, ...] = ()
+    resources: ResourceParams = field(default_factory=ResourceParams)
+
+
+def _exit(code: int) -> None:
+    raise SystemExit(code)
+
+
+# `ml launch` is its own command (rather than only exposing `ml sky launch ...`).
+# Today it's a thin proxy over SkyPilot: we translate flags into a `sky.Task` /
+# `sky.Resources` and delegate the launch.
+#
+# Keeping `ml launch` first-class gives us control over the launch sequence (e.g. better
+# logging/UX), and a stable place to add Mithril features that SkyPilot may choose not
+# to support (or that we eventually implement without going through SkyPilot).
+def launch(
+    entrypoint: str | None = None,
+    *,
+    cluster_name: Annotated[
+        str | None,
+        Parameter(name=["-c", "--cluster"], help="Cluster name."),
+    ] = None,
+    gpus: Annotated[
+        str | None,
+        Parameter(name="--gpus", help="GPU type and count (e.g., 'A100:4')."),
+    ] = None,
+    cpus: Annotated[
+        str | None,
+        Parameter(name="--cpus", help="Number of vCPUs."),
+    ] = None,
+    memory: Annotated[
+        str | None,
+        Parameter(name="--memory", help="Memory in GB."),
+    ] = None,
+    cloud: Annotated[
+        str | None,
+        Parameter(name="--cloud", help="Cloud provider."),
+    ] = None,
+    region: Annotated[str | None, Parameter(name="--region", help="Region.")] = None,
+    use_spot: Annotated[
+        bool | None,
+        Parameter(name="use-spot", negative="no-spot", help="Use spot instances."),
+    ] = None,
+    num_nodes: Annotated[
+        int | None,
+        Parameter(name="--num-nodes", help="Number of nodes."),
+    ] = None,
+    idle_minutes_to_autostop: Annotated[
+        int | None,
+        Parameter(
+            name=["-i", "--idle-minutes-to-autostop"],
+            help="Auto-stop after idle.",
+        ),
+    ] = None,
+    down: Annotated[
+        bool,
+        Parameter(name="--down", help="Tear down after jobs finish."),
+    ] = False,
+    detach_run: Annotated[
+        bool,
+        Parameter(name=["-d", "--detach-run"], help="Don't stream logs."),
+    ] = False,
+    retry_until_up: Annotated[
+        bool,
+        Parameter(name=["-r", "--retry-until-up"], help="Retry until up."),
+    ] = False,
+    yes: Annotated[
+        bool,
+        Parameter(name=["-y", "--yes"], help="Skip confirmation prompts."),
+    ] = False,
+    dryrun: Annotated[
+        bool,
+        Parameter(name="--dryrun", help="Show what would be launched."),
+    ] = False,
+    name: Annotated[
+        str | None,
+        Parameter(name=["-n", "--name"], help="Task name."),
+    ] = None,
+    workdir: Annotated[
+        str | None,
+        Parameter(name="--workdir", help="Working directory to sync."),
+    ] = None,
+    env: Annotated[
+        tuple[str, ...],
+        Parameter(name=["-e", "--env"], help="Environment variable (KEY=VALUE)."),
+    ] = (),
+    sky: SkyClient | None = None,
+) -> None:
+    r"""Launch a cluster and run a task.
+
+    Examples
+
+        ml launch task.yaml -c mycluster
+        ml launch 'echo hello' --gpus A100:1
+        ml launch --gpus A100:4 -c dev
+    """
+    if sky is None:
+        sky = SkyClient()
+    resources = ResourceParams(
+        gpus=gpus,
+        cpus=cpus,
+        memory=memory,
+        cloud=cloud,
+        region=region,
+        use_spot=use_spot,
+    )
+    task = build_task(
+        TaskParams(
+            entrypoint=entrypoint,
+            name=name,
+            workdir=workdir,
+            num_nodes=num_nodes,
+            env=env,
+            resources=resources,
+        ),
+        sky,
+    )
+
+    if dryrun:
+        console.print("[cyan]Dry run mode:[/cyan]")
+        console.print(f"  Cluster: {cluster_name or '(auto-generated)'}")
+        console.print(f"  Task: {task}")
+        console.print(f"  Resources: {task.resources}")
+        return
+
+    _print_version_banner(console, sky)
+
+    # Sync Mithril volumes to SkyPilot before launch. Volume values that are
+    # strings are persistent volume names; dicts are ephemeral configs (no sync
+    # needed).
+    if task.volumes:
+        volume_names = [v for v in task.volumes.values() if isinstance(v, str)]
+        if volume_names:
+            ensure_volumes_registered(
+                volume_names,
+                mithril=MithrilClient(),
+                sky=sky,
+            )
+    try:
+        request_id = sky.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,
+            _need_confirmation=not yes,
+        )
+    except click.exceptions.Abort:
+        # User pressed Ctrl+C during confirmation prompt
+        return
+
+    result, cluster_name = poll_launch(
+        request_id,
+        console,
+        sky=sky,
+        initial_cluster_name=cluster_name,
+    )
+    job_id, handle = cast(tuple[int | None, object | None], result)
+    cluster_name = _get_cluster_name(handle, cluster_name)
+
+    exit_code = 0
+    if not detach_run and job_id is not None:
+        assert cluster_name is not None, "cluster_name should be set for jobs"
+        exit_code = stream_job_logs(cluster_name, job_id, console, sky=sky)
+    hints = _format_command_hints(job_id=job_id, cluster_name=cluster_name)
+    if hints:
+        console.print(hints, markup=False)
+    if exit_code != 0:
+        _exit(exit_code)
+
+
+def build_task(params: TaskParams, sky: SkyClient) -> SkyTask:
+    """Build a sky.Task from CLI parameters."""
+    envs = parse_env_vars(params.env)
+
+    if params.entrypoint and params.entrypoint.endswith((".yaml", ".yml")):
+        task = sky.Task.from_yaml(params.entrypoint)
+        if params.name:
+            task.name = params.name
+        if params.workdir:
+            task.workdir = params.workdir
+        if envs:
+            task.update_envs(envs)
+        if _has_resource_overrides(params.resources):
+            task.set_resources(build_resources(params.resources, sky))
+        if params.num_nodes:
+            task.num_nodes = params.num_nodes
+        return task
+
+    resources = build_resources(params.resources, sky)
+    task = sky.Task(
+        name=params.name,
+        run=params.entrypoint,
+        workdir=params.workdir,
+        envs=envs if envs else None,
+        num_nodes=params.num_nodes or 1,
+    )
+    task.set_resources(resources)
+    return task
+
+
+def build_resources(params: ResourceParams, sky: SkyClient) -> SkyResources:
+    """Build sky.Resources from CLI parameters."""
+    kwargs: dict[str, Any] = {}
+    clouds = sky.clouds
+    if params.gpus:
+        kwargs["accelerators"] = params.gpus
+    if params.cpus:
+        kwargs["cpus"] = params.cpus
+    if params.memory:
+        kwargs["memory"] = params.memory
+    if params.region:
+        kwargs["region"] = params.region
+    if params.use_spot is not None:
+        kwargs["use_spot"] = params.use_spot
+    if params.cloud:
+        cloud_cls = getattr(clouds, params.cloud.upper(), None)
+        if cloud_cls is not None:
+            kwargs["cloud"] = cloud_cls()
+    return sky.Resources(**kwargs)
+
+
+def _has_resource_overrides(params: ResourceParams) -> bool:
+    """Return True if any resource overrides were provided."""
+    return any(
+        [
+            params.gpus,
+            params.cpus,
+            params.memory,
+            params.cloud,
+            params.region,
+            params.use_spot is not None,
+        ]
+    )
+
+
+def parse_env_vars(env: tuple[str, ...]) -> dict[str, str]:
+    """Parse environment variables from CLI args."""
+    envs: dict[str, str] = {}
+    for item in env:
+        if "=" in item:
+            key, _, value = item.partition("=")
+            envs[key] = value
+    return envs
+
+
+def _print_version_banner(console: Console, sky: SkyClient) -> None:
+    client_version = package_version("mithril-client")
+    server_version = sky.api_info().version
+    console.print(f"[dim]mithril-client {client_version}[/dim]", highlight=False)
+    if server_version:
+        console.print(
+            f"[dim]SkyPilot API server {server_version}[/dim]", highlight=False
+        )
+    else:
+        console.print("[dim]SkyPilot API server (version unknown)[/dim]")
+
+
+def _get_cluster_name(handle: object | None, fallback: str | None) -> str | None:
+    if handle is None:
+        return fallback
+    getter = getattr(handle, "get_cluster_name", None)
+    if callable(getter):
+        return cast("str | None", getter())
+    return cast("str | None", getattr(handle, "cluster_name", fallback))
+
+
+def _format_command_hints(job_id: int | None, cluster_name: str | None) -> str:
+    if cluster_name is None:
+        return ""
+    lines = ["📋 Useful Commands"]
+    if job_id is not None:
+        lines.append(f"Job ID: {job_id}")
+        lines.append(
+            f"{INDENT_SYMBOL}To cancel the job:\t\t "
+            f"{_cli_command('cancel')} {cluster_name} {job_id}"
+        )
+        lines.append(
+            f"{INDENT_SYMBOL}To stream job logs:\t\t "
+            f"{_cli_command('logs')} {cluster_name} {job_id}"
+        )
+        lines.append(
+            f"{INDENT_LAST_SYMBOL}To view job queue:\t\t "
+            f"{_cli_command('queue')} {cluster_name}"
+        )
+    lines.append(f"Cluster name: {cluster_name}")
+    lines.append(f"{INDENT_SYMBOL}To log into the head VM:\t ssh {cluster_name}")
+    lines.append(
+        f"{INDENT_SYMBOL}To submit a job:\t\t "
+        f"{_cli_command('exec')} {cluster_name} yaml_file"
+    )
+    lines.append(
+        f"{INDENT_SYMBOL}To stop the cluster:\t {_cli_command('stop')} {cluster_name}"
+    )
+    lines.append(
+        f"{INDENT_LAST_SYMBOL}To teardown the cluster:\t {_cli_command('down')} "
+        f"{cluster_name}"
+    )
+    return "\n" + "\n".join(lines)
+
+
+def _cli_command(command: str) -> str:
+    if command in SKY_ALIAS_COMMANDS:
+        return f"ml {command}"
+    return f"ml sky {command}"

mithril/cli/main.py

@@ -0,0 +1,68 @@
+"""Mithril CLI."""
+
+from __future__ import annotations
+
+import sys
+from importlib.metadata import version
+
+from cyclopts import App, Group
+
+from mithril.cli.commands.help import help_cmd, print_help
+from mithril.cli.commands.launch import launch
+from mithril.cli.utils.skypilot_passthrough import SKY_ALIAS_COMMANDS, run_sky
+
+_PACKAGE_NAME = "mithril-client"
+
+
+def cli(argv: list[str] | None = None) -> int:
+    """CLI router; returns an exit code."""
+    argv = list(sys.argv[1:] if argv is None else argv)
+
+    if not argv or argv[0] in {"-h", "--help"}:
+        print_help()
+        return 0
+
+    if argv[0] == "--version":
+        print(version(_PACKAGE_NAME))
+        return 0
+
+    cmd, rest = argv[0], argv[1:]
+
+    if cmd == "sky":
+        # Preserve Click-like behavior: `ml sky` shows SkyPilot help.
+        args = rest or ["--help"]
+        return run_sky(*args)
+
+    if cmd in SKY_ALIAS_COMMANDS:
+        return run_sky(cmd, *rest)
+
+    # Our first-class commands.
+    app = App(
+        name="ml",
+        version=version(_PACKAGE_NAME),
+        version_flags=["--version"],
+        # Match SkyPilot's Click-style help output more closely by disabling the
+        # rich "boxed table" help renderer.
+        help_formatter="plain",
+        help_format="plaintext",
+        # Use "Options" heading to match clap's style
+        group_parameters=Group("Options", sort_key=0),
+    )
+    app.command(help_cmd, name="help")
+    app.command(launch)
+
+    try:
+        app(argv)
+    except SystemExit as e:
+        return e.code if isinstance(e.code, int) else 1
+    else:
+        return 0
+
+
+def main() -> None:
+    """Entry point."""
+    raise SystemExit(cli())
+
+
+if __name__ == "__main__":
+    main()

mithril/cli/utils/__init__.py

@@ -0,0 +1 @@
+"""Internal CLI utilities (not user-facing commands)."""