feat(cli): add lightning vm commands and a VM SDK class

python/lightning_sdk/__init__.py

@@ -14,6 +14,7 @@
 from lightning_sdk.studio import Studio
 from lightning_sdk.teamspace import ConnectionType, FolderLocation, Teamspace
 from lightning_sdk.user import User
+from lightning_sdk.vm import VM
 
 __all__ = [
     "MMT",
@@ -33,6 +34,7 @@
     "Studio",
     "Teamspace",
     "User",
+    "VM",
     "__version__",
 ]
 

python/lightning_sdk/api/vm_api.py

@@ -0,0 +1,124 @@
+"""Internal API client for VM (cloud instance) requests."""
+
+import time
+from typing import Callable, List, Optional
+
+from lightning_sdk.lightning_cloud.openapi import (
+    V1ClusterState,
+    V1CreateInstanceRequest,
+    V1ExternalCluster,
+    V1Instance,
+)
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+TERMINAL_STATUSES = frozenset({"failed", "reclaimed", "deleting"})
+
+
+class VMFailedError(RuntimeError):
+    """Raised when a VM reaches a terminal failure status while waiting."""
+
+
+class VMNotFoundError(RuntimeError):
+    """Raised when a VM disappears while waiting."""
+
+
+class VMApi:
+    """Internal API client for VM requests (mainly http requests)."""
+
+    def __init__(self) -> None:
+        self._client = LightningClient(max_tries=7)
+
+    def create_vm(
+        self,
+        name: str,
+        org_id: str,
+        teamspace_id: str,
+        cluster_id: str,
+        instance_type: str,
+        volume_size: Optional[int] = None,
+        spot: bool = False,
+    ) -> V1Instance:
+        """Create a VM (cloud instance)."""
+        body = V1CreateInstanceRequest(
+            name=name,
+            organization_id=org_id,
+            project_id=teamspace_id,
+            cluster_id=cluster_id,
+            instance_type=instance_type,
+            spot=spot,
+        )
+        if volume_size is not None:
+            body.volume_size = str(volume_size)
+        return self._client.cloud_instances_service_create_instance(body=body)
+
+    def get_vm(self, vm_id: str, org_id: str) -> Optional[V1Instance]:
+        """Get a VM by id, returning None if it does not exist."""
+        try:
+            return self._client.cloud_instances_service_get_instance(id=vm_id, organization_id=org_id)
+        except ApiException as ex:
+            if "Reason: Not Found" in str(ex):
+                return None
+            raise ex
+
+    def get_vm_by_name(self, name: str, teamspace_id: str) -> Optional[V1Instance]:
+        """Get a VM by name within a teamspace, returning None if not found."""
+        matches = [vm for vm in self.list_vms(teamspace_id) if vm.name == name]
+        if len(matches) > 1:
+            raise ValueError(f"Multiple VMs named {name!r}; use the VM id instead.")
+        return matches[0] if matches else None
+
+    def list_vms(self, teamspace_id: str) -> List[V1Instance]:
+        """List all VMs in a teamspace, paginating through all pages."""
+        vms: List[V1Instance] = []
+        page_token = None
+        while True:
+            kwargs = {"project_id": teamspace_id, "limit": "100"}
+            if page_token:
+                kwargs["page_token"] = page_token
+            response = self._client.cloud_instances_service_list_instances(**kwargs)
+            vms.extend(response.instances or [])
+            page_token = response.next_page_token
+            if not page_token:
+                break
+        return vms
+
+    def list_machine_clusters(self, org_id: str) -> List[V1ExternalCluster]:
+        """List the organization's running machine clusters."""
+        response = self._client.cluster_service_list_clusters(org_id=org_id)
+        return [
+            cluster
+            for cluster in (response.clusters or [])
+            if cluster.spec.machine_v1 is not None and cluster.status.phase == V1ClusterState.RUNNING
+        ]
+
+    def delete_vm(self, vm_id: str, org_id: str) -> None:
+        """Delete a VM by id."""
+        self._client.cloud_instances_service_delete_instance(id=vm_id, organization_id=org_id)
+
+    def wait_for_status(
+        self,
+        vm_id: str,
+        org_id: str,
+        target: str = "running",
+        timeout: float = 600.0,
+        poll_interval: float = 5.0,
+        sleep: Callable[[float], None] = time.sleep,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> V1Instance:
+        """Poll a VM until it reaches the target status, fails, disappears, or times out."""
+        deadline = clock() + timeout
+        while True:
+            vm = self.get_vm(vm_id, org_id)
+            if vm is None:
+                raise VMNotFoundError(f"VM {vm_id} no longer exists.")
+            if vm.status == target:
+                return vm
+            if vm.status in TERMINAL_STATUSES:
+                reason = f": {vm.status_reason}" if vm.status_reason else ""
+                raise VMFailedError(f"VM {vm_id} entered status {vm.status!r}{reason}")
+            if clock() >= deadline:
+                raise TimeoutError(
+                    f"VM {vm_id} did not reach status {target!r} within {timeout:.0f}s (last: {vm.status!r})"
+                )
+            sleep(poll_interval)

python/lightning_sdk/cli/entrypoint.py

@@ -35,6 +35,7 @@
     studio,
     teamspace,
     user,
+    vm,
 )
 from lightning_sdk.cli.legacy_redirects import (
     build_hidden_alias_group,
@@ -51,7 +52,7 @@
 click.rich_click.COMMAND_GROUPS = {
     "lightning": [
         {"name": "GET STARTED", "commands": ["login", "logout", "config", "completion"]},
-        {"name": "COMPUTE", "commands": ["studio", "base-studio", "machine", "container", "sandbox"]},
+        {"name": "COMPUTE", "commands": ["studio", "base-studio", "machine", "vm", "container", "sandbox"]},
         {"name": "TRAIN & DEPLOY", "commands": ["job", "mmt", "model", "deployment", "pipeline"]},
         {"name": "ACCESS", "commands": ["user", "teamspace", "auth", "api-key", "ssh", "license"]},
         {"name": "DATA & FILES", "commands": ["cp", "ls", "rm", "edit", "connection"]},
@@ -133,6 +134,7 @@ def logout() -> None:
 main_cli.add_command(machine)
 main_cli.add_command(api)
 main_cli.add_command(deployment)
+main_cli.add_command(vm)
 main_cli.add_command(container)
 main_cli.add_command(model)
 main_cli.add_command(pipeline)

python/lightning_sdk/cli/groups.py

@@ -30,6 +30,7 @@
 from lightning_sdk.cli.teamspace import register_commands as register_teamspace_commands
 from lightning_sdk.cli.user import register_commands as register_user_commands
 from lightning_sdk.cli.utils.logging import LightningCommand, LightningGroup
+from lightning_sdk.cli.vm import register_commands as register_vm_commands
 
 
 @click.group(name="studio", cls=LightningGroup)
@@ -91,6 +92,11 @@ def deployment() -> None:
     """Deploy autoscaling inference APIs."""
 
 
+@click.group(name="vm", cls=LightningGroup)
+def vm() -> None:
+    """Create and manage virtual machines."""
+
+
 @click.group(name="sandbox", cls=LightningGroup)
 def sandbox() -> None:
     """Ephemeral sandboxes for agents.
@@ -266,6 +272,7 @@ def edit(_ctx: click.Context) -> None:
 register_config_commands(config)
 register_api_commands(api)
 register_deployment_commands(deployment)
+register_vm_commands(vm)
 register_sandbox_commands(sandbox)
 register_container_commands(container)
 register_model_commands(model)

python/lightning_sdk/cli/vm/__init__.py

@@ -0,0 +1,25 @@
+"""VM CLI commands."""
+
+import rich_click as click
+
+
+def register_commands(group: click.Group) -> None:
+    """Register VM commands with the given group."""
+    from lightning_sdk.cli.utils.delete import register_delete_command
+    from lightning_sdk.cli.vm.create import create_vm
+    from lightning_sdk.cli.vm.delete import resolve_vm_delete
+    from lightning_sdk.cli.vm.inspect import inspect_vm
+    from lightning_sdk.cli.vm.list import list_vms
+    from lightning_sdk.cli.vm.ssh import ssh_vm
+
+    group.add_command(create_vm, name="create")
+    group.add_command(list_vms, name="list")
+    group.add_command(inspect_vm, name="inspect")
+    register_delete_command(
+        group,
+        label="VM",
+        help="Delete a virtual machine. This discards its disk.",
+        context_help="Override default teamspace (format: owner/teamspace).",
+        resolve_delete=resolve_vm_delete,
+    )
+    group.add_command(ssh_vm, name="ssh")

python/lightning_sdk/cli/vm/common.py

@@ -0,0 +1,76 @@
+"""Shared helpers for VM CLI commands."""
+
+import json
+from typing import Iterable, Optional
+
+import rich_click as click
+
+from lightning_sdk.api.vm_api import VMApi
+from lightning_sdk.cli.utils.teamspace_option import resolve_teamspace
+from lightning_sdk.lightning_cloud.openapi import V1Instance
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import Machine
+from lightning_sdk.models import _list_teamspaces
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.user import User
+
+MACHINE_VALUES = tuple(
+    [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli]
+)
+
+_SSH_KEY_HINT = " Generate one with 'lightning ssh generate'."
+
+
+def iter_teamspaces(teamspace: Optional[str], all_teamspaces: bool) -> Iterable[Teamspace]:
+    if not all_teamspaces or teamspace:
+        yield resolve_teamspace(teamspace)
+        return
+
+    for teamspace_slug in _list_teamspaces():
+        yield resolve_teamspace(teamspace_slug)
+
+
+def org_id_for(teamspace: Teamspace) -> str:
+    owner = teamspace.owner
+    if isinstance(owner, User):
+        raise click.ClickException(
+            f"VMs require a teamspace owned by an organization; '{teamspace.name}' is owned by a user."
+        )
+    return owner.id
+
+
+def resolve_vm(api: VMApi, teamspace: Teamspace, name_or_id: str) -> V1Instance:
+    org_id = org_id_for(teamspace)
+    try:
+        vm = api.get_vm_by_name(name_or_id, teamspace.id)
+    except ValueError as ex:
+        raise click.ClickException(str(ex)) from ex
+    if vm is None:
+        try:
+            vm = api.get_vm(name_or_id, org_id)
+        except ApiException as ex:
+            raise friendly_error(ex) from ex
+    if vm is None:
+        raise click.ClickException(
+            f"VM {name_or_id!r} was not found in teamspace '{teamspace.owner.name}/{teamspace.name}'."
+        )
+    return vm
+
+
+def server_message(ex: ApiException) -> str:
+    body = getattr(ex, "body", None)
+    if body:
+        try:
+            parsed = json.loads(body)
+            if isinstance(parsed, dict) and parsed.get("message"):
+                return str(parsed["message"])
+        except (TypeError, ValueError):
+            pass
+    return str(ex.reason or ex)
+
+
+def friendly_error(ex: ApiException) -> click.ClickException:
+    message = server_message(ex)
+    if "add an SSH key" in message:
+        message += _SSH_KEY_HINT
+    return click.ClickException(message)

python/lightning_sdk/cli/vm/create.py

@@ -0,0 +1,66 @@
+"""VM create command."""
+
+from typing import Optional
+
+import rich_click as click
+
+from lightning_sdk.cli.utils.logging import LightningCommand
+from lightning_sdk.cli.vm.common import MACHINE_VALUES, friendly_error, resolve_teamspace
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import Machine
+from lightning_sdk.vm import VM
+
+
+@click.command("create", cls=LightningCommand)
+@click.argument("name")
+@click.option("--machine", required=True, type=click.Choice(MACHINE_VALUES, case_sensitive=False), help="Machine type.")
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace).")
+@click.option(
+    "--cloud",
+    "cloud_account",
+    help=(
+        "Cloud account (cluster id) to create the VM in. "
+        "Defaults to the organization's machine cluster when there is exactly one."
+    ),
+)
+@click.option("--volume-size", type=click.IntRange(400, 800), help="Root disk size in GB (400-800).")
+@click.option("--spot", is_flag=True, default=False, help="Use an interruptible machine.")
+@click.option("--wait", is_flag=True, default=False, help="Block until the VM is running, then print the SSH command.")
+@click.option("--timeout", type=float, default=600.0, show_default=True, help="Seconds to wait with --wait.")
+def create_vm(
+    name: str,
+    machine: str,
+    teamspace: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    volume_size: Optional[int] = None,
+    spot: bool = False,
+    wait: bool = False,
+    timeout: float = 600.0,
+) -> None:
+    """Create a virtual machine."""
+    resolved_teamspace = resolve_teamspace(teamspace)
+    try:
+        vm = VM.create(
+            name=name,
+            machine=Machine.from_str(machine),
+            teamspace=resolved_teamspace,
+            cloud_account=cloud_account,
+            volume_size=volume_size,
+            spot=spot,
+            wait=wait,
+            timeout=timeout,
+        )
+    except ApiException as ex:
+        raise friendly_error(ex) from ex
+    except TimeoutError as ex:
+        raise click.ClickException(
+            f"{ex} The VM still exists; delete it with 'lightning vm delete <name>' if you no longer need it."
+        ) from ex
+    except (RuntimeError, ValueError) as ex:
+        raise click.ClickException(str(ex)) from ex
+
+    click.echo(f"Created VM {vm.name} ({vm.id}), status: {vm.status}")
+    if vm.ssh_command:
+        click.echo(f"Connect with: {vm.ssh_command}")
+    elif not wait:
+        click.echo(f"Run 'lightning vm ssh {vm.name}' once it is running.")

python/lightning_sdk/cli/vm/delete.py

@@ -0,0 +1,24 @@
+"""VM deletion resolver."""
+
+from typing import Optional
+
+from lightning_sdk.api.vm_api import VMApi
+from lightning_sdk.cli.utils.delete import DeleteAction
+from lightning_sdk.cli.vm.common import friendly_error, org_id_for, resolve_teamspace, resolve_vm
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+
+
+def resolve_vm_delete(name: str, teamspace: Optional[str]) -> DeleteAction:
+    """Resolve a VM and return its bound deletion action."""
+    resolved_teamspace = resolve_teamspace(teamspace)
+    api = VMApi()
+    vm = resolve_vm(api, resolved_teamspace, name)
+    org_id = org_id_for(resolved_teamspace)
+
+    def delete() -> None:
+        try:
+            api.delete_vm(vm.id, org_id)
+        except ApiException as ex:
+            raise friendly_error(ex) from ex
+
+    return delete