diff --git a/python/lightning_sdk/__init__.py b/python/lightning_sdk/__init__.py index 0fa877e3..537f1703 100644 --- a/python/lightning_sdk/__init__.py +++ b/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__", ] diff --git a/python/lightning_sdk/api/vm_api.py b/python/lightning_sdk/api/vm_api.py new file mode 100644 index 00000000..f5174532 --- /dev/null +++ b/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) diff --git a/python/lightning_sdk/cli/entrypoint.py b/python/lightning_sdk/cli/entrypoint.py index 358e2fb2..584038dc 100644 --- a/python/lightning_sdk/cli/entrypoint.py +++ b/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) diff --git a/python/lightning_sdk/cli/groups.py b/python/lightning_sdk/cli/groups.py index 24452f71..4c659efd 100644 --- a/python/lightning_sdk/cli/groups.py +++ b/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) diff --git a/python/lightning_sdk/cli/vm/__init__.py b/python/lightning_sdk/cli/vm/__init__.py new file mode 100644 index 00000000..14bbe164 --- /dev/null +++ b/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") diff --git a/python/lightning_sdk/cli/vm/common.py b/python/lightning_sdk/cli/vm/common.py new file mode 100644 index 00000000..5ed2f1ab --- /dev/null +++ b/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) diff --git a/python/lightning_sdk/cli/vm/create.py b/python/lightning_sdk/cli/vm/create.py new file mode 100644 index 00000000..72c10584 --- /dev/null +++ b/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 ' 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.") diff --git a/python/lightning_sdk/cli/vm/delete.py b/python/lightning_sdk/cli/vm/delete.py new file mode 100644 index 00000000..142e5db8 --- /dev/null +++ b/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 diff --git a/python/lightning_sdk/cli/vm/inspect.py b/python/lightning_sdk/cli/vm/inspect.py new file mode 100644 index 00000000..e539b585 --- /dev/null +++ b/python/lightning_sdk/cli/vm/inspect.py @@ -0,0 +1,31 @@ +"""VM inspect command.""" + +import json +from datetime import datetime +from typing import Any, Optional + +import rich_click as click + +from lightning_sdk.api.vm_api import VMApi +from lightning_sdk.cli.utils.logging import LightningCommand +from lightning_sdk.cli.vm.common import resolve_teamspace, resolve_vm + + +def _json_safe(value: Any) -> Any: + if isinstance(value, datetime): + return value.isoformat() + if isinstance(value, list): + return [_json_safe(item) for item in value] + if isinstance(value, dict): + return {key: _json_safe(item) for key, item in value.items()} + return value + + +@click.command("inspect", cls=LightningCommand) +@click.argument("name") +@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace).") +def inspect_vm(name: str, teamspace: Optional[str] = None) -> None: + """Inspect a virtual machine as JSON.""" + resolved_teamspace = resolve_teamspace(teamspace) + vm = resolve_vm(VMApi(), resolved_teamspace, name) + click.echo(json.dumps(_json_safe(vm.to_dict()), indent=2, sort_keys=True)) diff --git a/python/lightning_sdk/cli/vm/list.py b/python/lightning_sdk/cli/vm/list.py new file mode 100644 index 00000000..ba1a73c4 --- /dev/null +++ b/python/lightning_sdk/cli/vm/list.py @@ -0,0 +1,52 @@ +"""VM list command.""" + +from typing import Optional + +import rich_click as click +from rich.table import Table + +from lightning_sdk.api.vm_api import VMApi +from lightning_sdk.cli.utils.logging import LightningCommand +from lightning_sdk.cli.utils.richt_print import rich_to_str +from lightning_sdk.cli.vm.common import iter_teamspaces, org_id_for + + +@click.command("list", cls=LightningCommand) +@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace).") +@click.option( + "--all", + "all_teamspaces", + is_flag=True, + flag_value=True, + default=False, + help="List VMs in all teamspaces visible to the selected owner.", +) +def list_vms(teamspace: Optional[str] = None, all_teamspaces: bool = False) -> None: + """List virtual machines in a teamspace.""" + api = VMApi() + rows = [] + for resolved_teamspace in iter_teamspaces(teamspace, all_teamspaces): + org_id_for(resolved_teamspace) + for vm in api.list_vms(resolved_teamspace.id): + rows.append((resolved_teamspace, vm)) + + table = Table(pad_edge=True) + table.add_column("Name", no_wrap=True) + table.add_column("Teamspace", no_wrap=True) + table.add_column("Status", no_wrap=True) + table.add_column("Machine", no_wrap=True) + table.add_column("Created", no_wrap=True) + table.add_column("SSH host", no_wrap=True) + + for resolved_teamspace, vm in sorted(rows, key=lambda row: row[1].name or ""): + created = vm.created_at.strftime("%Y-%m-%d %H:%M") if vm.created_at else "" + table.add_row( + vm.name or "", + f"{resolved_teamspace.owner.name}/{resolved_teamspace.name}", + vm.status or "", + vm.instance_type or "", + created, + vm.ssh_host or "", + ) + + click.echo(rich_to_str(table), color=True) diff --git a/python/lightning_sdk/cli/vm/ssh.py b/python/lightning_sdk/cli/vm/ssh.py new file mode 100644 index 00000000..68c3c651 --- /dev/null +++ b/python/lightning_sdk/cli/vm/ssh.py @@ -0,0 +1,38 @@ +"""VM ssh command.""" + +import os +import shlex +from typing import Optional, Sequence + +import rich_click as click + +from lightning_sdk.api.vm_api import VMApi +from lightning_sdk.cli.utils.logging import LightningCommand +from lightning_sdk.cli.vm.common import org_id_for, resolve_teamspace, resolve_vm + + +@click.command("ssh", cls=LightningCommand, context_settings={"ignore_unknown_options": True}) +@click.argument("name") +@click.argument("ssh_args", nargs=-1, type=click.UNPROCESSED) +@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace).") +@click.option( + "--timeout", type=float, default=600.0, show_default=True, help="Seconds to wait for the VM to be running." +) +def ssh_vm(name: str, ssh_args: Sequence[str] = (), teamspace: Optional[str] = None, timeout: float = 600.0) -> None: + """SSH into a virtual machine. Arguments after -- are passed to ssh.""" + resolved_teamspace = resolve_teamspace(teamspace) + api = VMApi() + vm = resolve_vm(api, resolved_teamspace, name) + + if vm.status != "running": + click.echo(f"VM {vm.name} is {vm.status}; waiting for it to be running...", err=True) + try: + vm = api.wait_for_status(vm.id, org_id_for(resolved_teamspace), timeout=timeout) + except (RuntimeError, TimeoutError) as ex: + raise click.ClickException(str(ex)) from ex + + if not vm.ssh_command: + raise click.ClickException(f"VM {vm.name} has no SSH command yet; try again in a moment.") + + argv = shlex.split(vm.ssh_command) + list(ssh_args) + os.execvp(argv[0], argv) diff --git a/python/lightning_sdk/vm.py b/python/lightning_sdk/vm.py new file mode 100644 index 00000000..16c1dddd --- /dev/null +++ b/python/lightning_sdk/vm.py @@ -0,0 +1,214 @@ +"""High-level VM (cloud instance) object.""" + +import os +from datetime import datetime +from typing import List, Optional, Union + +from lightning_sdk.api.utils import _machine_to_compute_name +from lightning_sdk.api.vm_api import VMApi +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.organization import Organization +from lightning_sdk.teamspace import Teamspace +from lightning_sdk.user import User +from lightning_sdk.utils.logging import TrackCallsMeta +from lightning_sdk.utils.resolve import _resolve_teamspace + +_TEAMSPACE_HELP = ( + "Pass a teamspace as 'owner/teamspace', set one of LIGHTNING_TEAMSPACE / LIGHTNING_ORG, " + "or configure a default with 'lightning config set teamspace'." +) + + +def _require_org_id(teamspace: Teamspace) -> str: + owner = teamspace.owner + if not isinstance(owner, Organization): + raise ValueError( + f"VMs require a teamspace owned by an organization; '{teamspace.name}' is owned by a user. " + "Pick an organization teamspace with --teamspace owner/teamspace." + ) + return owner.id + + +def _default_machine_cluster(org_id: str) -> str: + from_env = os.getenv("LIGHTNING_CLUSTER_ID") + if from_env: + return from_env + + clusters = VMApi().list_machine_clusters(org_id) + if not clusters: + raise ValueError( + "No machine cluster is available to this organization; " + "pass cloud_account=... (or --cloud) with a cluster id." + ) + if len(clusters) > 1: + candidates = ", ".join(sorted(cluster.id for cluster in clusters)) + raise ValueError( + f"This organization has several machine clusters ({candidates}); " + "pass cloud_account=... (or --cloud) with the one to use." + ) + return clusters[0].id + + +def _resolve( + teamspace: Optional[Union[str, Teamspace]], + org: Optional[Union[str, Organization]], + user: Optional[Union[str, User]], +) -> Teamspace: + resolved = _resolve_teamspace(teamspace=teamspace, org=org, user=user) + if resolved is None: + raise ValueError("Could not determine the teamspace for your VM. " + _TEAMSPACE_HELP) + return resolved + + +class VM(metaclass=TrackCallsMeta): + """A virtual machine in a Lightning teamspace. + + Args: + name_or_id: The VM name or id to load. + teamspace: The teamspace the VM lives in ('owner/teamspace' or a Teamspace). + org: The organization owning the teamspace. + user: The user owning the teamspace (VMs require org-owned teamspaces, kept for API symmetry). + """ + + def __init__( + self, + name_or_id: str, + teamspace: Optional[Union[str, Teamspace]] = None, + org: Optional[Union[str, Organization]] = None, + user: Optional[Union[str, User]] = None, + ) -> None: + self._teamspace = _resolve(teamspace, org, user) + self._org_id = _require_org_id(self._teamspace) + self._api = VMApi() + + instance = self._api.get_vm_by_name(name_or_id, self._teamspace.id) + if instance is None: + try: + instance = self._api.get_vm(name_or_id, self._org_id) + except ApiException as ex: + raise ValueError(f"Failed to look up VM '{name_or_id}': {ex.reason or ex}") from ex + if instance is None: + raise ValueError( + f"VM '{name_or_id}' was not found in teamspace " + f"'{self._teamspace.owner.name}/{self._teamspace.name}'." + ) + self._instance = instance + + @classmethod + def _from_instance(cls, instance: V1Instance, teamspace: Teamspace, org_id: str, api: VMApi) -> "VM": + vm = cls.__new__(cls) + vm._teamspace = teamspace + vm._org_id = org_id + vm._api = api + vm._instance = instance + return vm + + @classmethod + def create( + cls, + name: str, + machine: Union[Machine, str], + teamspace: Optional[Union[str, Teamspace]] = None, + org: Optional[Union[str, Organization]] = None, + user: Optional[Union[str, User]] = None, + cloud_account: Optional[str] = None, + volume_size: Optional[int] = None, + spot: bool = False, + wait: bool = False, + timeout: float = 600.0, + ) -> "VM": + """Create a VM and return it. With ``wait=True`` block until it is running.""" + resolved = _resolve(teamspace, org, user) + org_id = _require_org_id(resolved) + api = VMApi() + if isinstance(machine, str): + machine = Machine.from_str(machine) + cluster_id = cloud_account or _default_machine_cluster(org_id) + instance = api.create_vm( + name=name, + org_id=org_id, + teamspace_id=resolved.id, + cluster_id=cluster_id, + instance_type=_machine_to_compute_name(machine), + volume_size=volume_size, + spot=spot, + ) + vm = cls._from_instance(instance, resolved, org_id, api) + if wait: + vm.wait(timeout=timeout) + return vm + + @staticmethod + def list( + teamspace: Optional[Union[str, Teamspace]] = None, + org: Optional[Union[str, Organization]] = None, + user: Optional[Union[str, User]] = None, + ) -> List["VM"]: + """List the VMs in a teamspace.""" + resolved = _resolve(teamspace, org, user) + org_id = _require_org_id(resolved) + api = VMApi() + return [VM._from_instance(i, resolved, org_id, api) for i in api.list_vms(resolved.id)] + + def refresh(self) -> "VM": + """Re-fetch the VM from the API.""" + instance = self._api.get_vm(self.id, self._org_id) + if instance is None: + raise ValueError(f"VM '{self.name}' no longer exists.") + self._instance = instance + return self + + def wait(self, timeout: float = 600.0) -> "VM": + """Block until the VM is running. Raises on failure or timeout.""" + self._instance = self._api.wait_for_status(self.id, self._org_id, timeout=timeout) + return self + + def delete(self) -> None: + """Delete the VM. This discards its disk.""" + self._api.delete_vm(self.id, self._org_id) + + @property + def id(self) -> str: + return self._instance.id + + @property + def name(self) -> str: + return self._instance.name + + @property + def status(self) -> Optional[str]: + return self._instance.status + + @property + def machine(self) -> Optional[str]: + return self._instance.instance_type + + @property + def ssh_command(self) -> Optional[str]: + return self._instance.ssh_command + + @property + def ssh_host(self) -> Optional[str]: + return self._instance.ssh_host + + @property + def ssh_port(self) -> Optional[int]: + return self._instance.ssh_port + + @property + def ssh_user(self) -> Optional[str]: + return self._instance.ssh_user + + @property + def created_at(self) -> Optional[datetime]: + return self._instance.created_at + + @property + def teamspace(self) -> Teamspace: + return self._teamspace + + def __repr__(self) -> str: + """Return a debug representation of the VM.""" + return f"VM(name={self.name!r}, status={self.status!r}, teamspace={self._teamspace.name!r})" diff --git a/python/tests/api/test_vm_api.py b/python/tests/api/test_vm_api.py new file mode 100644 index 00000000..8907d72d --- /dev/null +++ b/python/tests/api/test_vm_api.py @@ -0,0 +1,191 @@ +from unittest.mock import MagicMock + +import pytest + +from lightning_sdk.api.vm_api import VMApi, VMFailedError, VMNotFoundError +from lightning_sdk.lightning_cloud.openapi import ( + V1ClusterState, + V1ClusterStatus, + V1ExternalCluster, + V1ExternalClusterSpec, + V1Instance, + V1ListClustersResponse, + V1ListInstancesResponse, + V1MachineDirectV1, +) +from lightning_sdk.lightning_cloud.openapi.rest import ApiException + + +def _api(monkeypatch) -> VMApi: + monkeypatch.setattr("lightning_sdk.api.vm_api.LightningClient", MagicMock()) + return VMApi() + + +def _not_found() -> ApiException: + return ApiException(status=404, reason="Not Found") + + +def test_create_vm_builds_request(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_create_instance.return_value = V1Instance(id="vm-1") + + result = api.create_vm( + name="sim-1", + org_id="org-1", + teamspace_id="ts-1", + cluster_id="cl-1", + instance_type="lit-h100-1", + volume_size=500, + spot=True, + ) + + assert result.id == "vm-1" + body = api._client.cloud_instances_service_create_instance.call_args.kwargs["body"] + assert body.name == "sim-1" + assert body.organization_id == "org-1" + assert body.project_id == "ts-1" + assert body.cluster_id == "cl-1" + assert body.instance_type == "lit-h100-1" + assert body.volume_size == "500" + assert body.spot is True + + +def test_create_vm_omits_volume_size_when_unset(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_create_instance.return_value = V1Instance(id="vm-1") + + api.create_vm(name="n", org_id="o", teamspace_id="t", cluster_id="c", instance_type="i") + + body = api._client.cloud_instances_service_create_instance.call_args.kwargs["body"] + assert body.volume_size is None + + +def test_get_vm_returns_none_on_404(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_get_instance.side_effect = _not_found() + + assert api.get_vm("vm-1", "org-1") is None + api._client.cloud_instances_service_get_instance.assert_called_once_with(id="vm-1", organization_id="org-1") + + +def test_get_vm_reraises_other_errors(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_get_instance.side_effect = ApiException(status=500, reason="boom") + + with pytest.raises(ApiException): + api.get_vm("vm-1", "org-1") + + +def test_list_vms_follows_pagination(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_list_instances.side_effect = [ + V1ListInstancesResponse(instances=[V1Instance(id="a")], next_page_token="p2"), + V1ListInstancesResponse(instances=[V1Instance(id="b")], next_page_token=""), + ] + + result = api.list_vms("ts-1") + + assert [vm.id for vm in result] == ["a", "b"] + calls = api._client.cloud_instances_service_list_instances.call_args_list + assert calls[0].kwargs == {"project_id": "ts-1", "limit": "100"} + assert calls[1].kwargs["page_token"] == "p2" + + +def test_get_vm_by_name_filters_list(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_list_instances.return_value = V1ListInstancesResponse( + instances=[V1Instance(id="a", name="x"), V1Instance(id="b", name="sim-1")], next_page_token="" + ) + + assert api.get_vm_by_name("sim-1", "ts-1").id == "b" + assert api.get_vm_by_name("nope", "ts-1") is None + + +def test_get_vm_by_name_rejects_duplicates(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_list_instances.return_value = V1ListInstancesResponse( + instances=[V1Instance(id="a", name="sim-1"), V1Instance(id="b", name="sim-1")], next_page_token="" + ) + + with pytest.raises(ValueError, match="Multiple VMs named 'sim-1'"): + api.get_vm_by_name("sim-1", "ts-1") + + +def _cluster(cluster_id: str, machine: bool, phase: str) -> V1ExternalCluster: + spec = V1ExternalClusterSpec(machine_v1=V1MachineDirectV1() if machine else None) + return V1ExternalCluster(id=cluster_id, spec=spec, status=V1ClusterStatus(phase=phase)) + + +def test_list_machine_clusters_filters_to_running_machine_clusters(monkeypatch): + api = _api(monkeypatch) + api._client.cluster_service_list_clusters.return_value = V1ListClustersResponse( + clusters=[ + _cluster("cl-machine", machine=True, phase=V1ClusterState.RUNNING), + _cluster("cl-kubernetes", machine=False, phase=V1ClusterState.RUNNING), + _cluster("cl-stopped", machine=True, phase=V1ClusterState.FAILED), + ] + ) + + result = api.list_machine_clusters("org-1") + + assert [cluster.id for cluster in result] == ["cl-machine"] + api._client.cluster_service_list_clusters.assert_called_once_with(org_id="org-1") + + +def test_delete_vm(monkeypatch): + api = _api(monkeypatch) + + api.delete_vm("vm-1", "org-1") + + api._client.cloud_instances_service_delete_instance.assert_called_once_with(id="vm-1", organization_id="org-1") + + +def _fake_clock(step: float): + now = [0.0] + + def clock() -> float: + now[0] += step + return now[0] + + return clock + + +def test_wait_for_status_returns_on_target(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_get_instance.side_effect = [ + V1Instance(id="vm-1", status="pending"), + V1Instance(id="vm-1", status="provisioning"), + V1Instance(id="vm-1", status="running"), + ] + sleep = MagicMock() + + result = api.wait_for_status("vm-1", "org-1", timeout=60, poll_interval=1, sleep=sleep, clock=_fake_clock(1)) + + assert result.status == "running" + assert sleep.call_count == 2 + + +def test_wait_for_status_raises_on_failure(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_get_instance.return_value = V1Instance( + id="vm-1", status="failed", status_reason="no capacity" + ) + + with pytest.raises(VMFailedError, match="no capacity"): + api.wait_for_status("vm-1", "org-1", timeout=60, sleep=MagicMock(), clock=_fake_clock(1)) + + +def test_wait_for_status_raises_when_vm_disappears(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_get_instance.side_effect = _not_found() + + with pytest.raises(VMNotFoundError): + api.wait_for_status("vm-1", "org-1", timeout=60, sleep=MagicMock(), clock=_fake_clock(1)) + + +def test_wait_for_status_times_out(monkeypatch): + api = _api(monkeypatch) + api._client.cloud_instances_service_get_instance.return_value = V1Instance(id="vm-1", status="pending") + + with pytest.raises(TimeoutError): + api.wait_for_status("vm-1", "org-1", timeout=5, poll_interval=1, sleep=MagicMock(), clock=_fake_clock(2)) diff --git a/python/tests/cli/vm/__init__.py b/python/tests/cli/vm/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/python/tests/cli/vm/test_vm.py b/python/tests/cli/vm/test_vm.py new file mode 100644 index 00000000..0019eea6 --- /dev/null +++ b/python/tests/cli/vm/test_vm.py @@ -0,0 +1,267 @@ +import re +from types import SimpleNamespace +from unittest.mock import MagicMock + +import click +from click.testing import CliRunner + +from lightning_sdk.cli.vm import register_commands +from lightning_sdk.cli.vm.create import create_vm +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 +from lightning_sdk.lightning_cloud.openapi import V1Instance +from lightning_sdk.lightning_cloud.openapi.rest import ApiException +from lightning_sdk.user import User +from tests.cli.help import assert_help_contains, mock_command_logging + + +def _plain(output: str) -> str: + return " ".join(output.split()) + + +def _teamspace() -> SimpleNamespace: + return SimpleNamespace(id="ts-1", name="research", owner=SimpleNamespace(id="org-1", name="ecorp")) + + +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") +_BOX_RE = re.compile(r"[\u2500-\u257f]") + + +def _flat(output: str) -> str: + text = _ANSI_RE.sub("", output) + text = _BOX_RE.sub(" ", text) + return " ".join(text.split()) + + +@mock_command_logging +def test_vm_help() -> None: + assert_help_contains( + "lightning vm --help", + "Usage: lightning vm [OPTIONS] COMMAND [ARGS]...", + "Create and manage virtual machines.", + "create", + "delete", + "inspect", + "list", + "ssh", + ) + + +@mock_command_logging +def test_create_prints_id_and_status(monkeypatch) -> None: + vm = SimpleNamespace(id="vm-1", name="sim-1", status="pending", ssh_command=None) + create = MagicMock(return_value=vm) + monkeypatch.setattr("lightning_sdk.cli.vm.create.resolve_teamspace", lambda teamspace: _teamspace()) + monkeypatch.setattr("lightning_sdk.cli.vm.create.VM", SimpleNamespace(create=create)) + + result = CliRunner().invoke(create_vm, ["sim-1", "--machine", "H100", "--volume-size", "500"]) + + assert result.exit_code == 0, result.output + assert "vm-1" in result.output + assert "pending" in result.output + kwargs = create.call_args.kwargs + assert kwargs["name"] == "sim-1" + assert kwargs["machine"].name == "H100" + assert kwargs["volume_size"] == 500 + assert kwargs["wait"] is False + + +@mock_command_logging +def test_create_wait_prints_ssh_command(monkeypatch) -> None: + vm = SimpleNamespace(id="vm-1", name="sim-1", status="running", ssh_command="ssh -p 20032 ubuntu@1.2.3.4") + create = MagicMock(return_value=vm) + monkeypatch.setattr("lightning_sdk.cli.vm.create.resolve_teamspace", lambda teamspace: _teamspace()) + monkeypatch.setattr("lightning_sdk.cli.vm.create.VM", SimpleNamespace(create=create)) + + result = CliRunner().invoke(create_vm, ["sim-1", "--machine", "H100", "--wait", "--timeout", "30"]) + + assert result.exit_code == 0, result.output + assert "ssh -p 20032 ubuntu@1.2.3.4" in result.output + assert create.call_args.kwargs["wait"] is True + assert create.call_args.kwargs["timeout"] == 30 + + +@mock_command_logging +def test_create_surfaces_server_message(monkeypatch) -> None: + exc = ApiException(status=412, reason="Failed Precondition") + exc.body = '{"code": 9, "message": "add an SSH key to your account first (Settings → Keys)"}' + monkeypatch.setattr("lightning_sdk.cli.vm.create.resolve_teamspace", lambda teamspace: _teamspace()) + monkeypatch.setattr("lightning_sdk.cli.vm.create.VM", SimpleNamespace(create=MagicMock(side_effect=exc))) + + result = CliRunner().invoke(create_vm, ["sim-1", "--machine", "H100"]) + + assert result.exit_code != 0 + assert "add an SSH key to your account first" in _plain(result.output) + assert "lightning ssh generate" in _plain(result.output) + + +@mock_command_logging +def test_list_renders_table(monkeypatch) -> None: + api = MagicMock() + api.list_vms.return_value = [ + V1Instance(id="vm-1", name="sim-1", status="running", instance_type="lit-h100-1", ssh_host="1.2.3.4"), + V1Instance(id="vm-2", name="sim-2", status="pending", instance_type="lit-h100-8"), + ] + monkeypatch.setattr("lightning_sdk.cli.vm.list.iter_teamspaces", lambda teamspace, all_teamspaces: [_teamspace()]) + monkeypatch.setattr("lightning_sdk.cli.vm.list.VMApi", MagicMock(return_value=api)) + + monkeypatch.setenv("COLUMNS", "200") + result = CliRunner().invoke(list_vms, []) + + assert result.exit_code == 0, result.output + assert "sim-1" in result.output + assert "running" in result.output + assert "1.2.3.4" in result.output + assert "lit-h100-8" in result.output + api.list_vms.assert_called_once_with("ts-1") + + +def _patch_lookup(monkeypatch, module: str, vm: V1Instance) -> MagicMock: + api = MagicMock() + monkeypatch.setattr(f"lightning_sdk.cli.vm.{module}.resolve_teamspace", lambda teamspace: _teamspace()) + monkeypatch.setattr(f"lightning_sdk.cli.vm.{module}.VMApi", MagicMock(return_value=api)) + monkeypatch.setattr(f"lightning_sdk.cli.vm.{module}.resolve_vm", lambda api, teamspace, name: vm) + return api + + +@mock_command_logging +def test_inspect_prints_json(monkeypatch) -> None: + vm = V1Instance(id="vm-1", name="sim-1", status="running", ssh_command="ssh -p 1 u@h") + _patch_lookup(monkeypatch, "inspect", vm) + + result = CliRunner().invoke(inspect_vm, ["sim-1"]) + + assert result.exit_code == 0, result.output + assert '"id": "vm-1"' in result.output + assert '"ssh_command": "ssh -p 1 u@h"' in result.output + + +def _vm_group() -> click.Group: + group = click.Group(name="vm") + register_commands(group) + return group + + +@mock_command_logging +def test_delete_prompts_and_aborts(monkeypatch) -> None: + vm = V1Instance(id="vm-1", name="sim-1") + api = _patch_lookup(monkeypatch, "delete", vm) + + result = CliRunner().invoke(_vm_group(), ["delete", "sim-1"], input="n\n") + + assert result.exit_code != 0 + api.delete_vm.assert_not_called() + + +@mock_command_logging +def test_delete_with_yes(monkeypatch) -> None: + vm = V1Instance(id="vm-1", name="sim-1") + api = _patch_lookup(monkeypatch, "delete", vm) + + result = CliRunner().invoke(_vm_group(), ["delete", "sim-1", "--yes"]) + + assert result.exit_code == 0, result.output + assert "VM deleted" in result.output + api.delete_vm.assert_called_once_with("vm-1", "org-1") + + +@mock_command_logging +def test_ssh_execs_command_with_extra_args(monkeypatch) -> None: + vm = V1Instance(id="vm-1", name="sim-1", status="running", ssh_command="ssh -p 20032 ubuntu@1.2.3.4") + api = _patch_lookup(monkeypatch, "ssh", vm) + execvp = MagicMock() + monkeypatch.setattr("lightning_sdk.cli.vm.ssh.os.execvp", execvp) + + result = CliRunner().invoke(ssh_vm, ["sim-1", "--", "-L", "8888:localhost:8888"]) + + assert result.exit_code == 0, result.output + execvp.assert_called_once_with("ssh", ["ssh", "-p", "20032", "ubuntu@1.2.3.4", "-L", "8888:localhost:8888"]) + api.wait_for_status.assert_not_called() + + +@mock_command_logging +def test_ssh_waits_when_not_running(monkeypatch) -> None: + pending = V1Instance(id="vm-1", name="sim-1", status="provisioning") + api = _patch_lookup(monkeypatch, "ssh", pending) + api.wait_for_status.return_value = V1Instance( + id="vm-1", name="sim-1", status="running", ssh_command="ssh -p 20032 ubuntu@1.2.3.4" + ) + execvp = MagicMock() + monkeypatch.setattr("lightning_sdk.cli.vm.ssh.os.execvp", execvp) + + result = CliRunner().invoke(ssh_vm, ["sim-1", "--timeout", "30"]) + + assert result.exit_code == 0, result.output + api.wait_for_status.assert_called_once_with("vm-1", "org-1", timeout=30) + execvp.assert_called_once() + + +@mock_command_logging +def test_create_rejects_out_of_range_volume_size(monkeypatch) -> None: + create = MagicMock() + monkeypatch.setattr("lightning_sdk.cli.vm.create.resolve_teamspace", lambda teamspace: _teamspace()) + monkeypatch.setattr("lightning_sdk.cli.vm.create.VM", SimpleNamespace(create=create)) + + result = CliRunner().invoke(create_vm, ["sim-1", "--machine", "H100", "--volume-size", "100"]) + + assert result.exit_code != 0 + assert "400" in result.output + create.assert_not_called() + + +@mock_command_logging +def test_create_timeout_mentions_cleanup(monkeypatch) -> None: + timeout = TimeoutError("VM vm-1 did not reach status 'running' within 30s (last: 'pending')") + monkeypatch.setattr("lightning_sdk.cli.vm.create.resolve_teamspace", lambda teamspace: _teamspace()) + monkeypatch.setattr("lightning_sdk.cli.vm.create.VM", SimpleNamespace(create=MagicMock(side_effect=timeout))) + + result = CliRunner().invoke(create_vm, ["sim-1", "--machine", "H100", "--wait"]) + + assert result.exit_code != 0 + assert "did not reach status" in _flat(result.output) + assert "lightning vm delete" in _flat(result.output) + + +@mock_command_logging +def test_list_rejects_user_owned_teamspace(monkeypatch) -> None: + owner = MagicMock(spec=User) + owner.name = "alice" + user_teamspace = SimpleNamespace(id="ts-2", name="personal", owner=owner) + monkeypatch.setattr("lightning_sdk.cli.vm.list.iter_teamspaces", lambda teamspace, all_teamspaces: [user_teamspace]) + monkeypatch.setattr("lightning_sdk.cli.vm.list.VMApi", MagicMock()) + + result = CliRunner().invoke(list_vms, []) + + assert result.exit_code != 0 + assert "owned by a user" in _flat(result.output) + + +@mock_command_logging +def test_inspect_reports_not_found(monkeypatch) -> None: + api = MagicMock() + api.get_vm_by_name.return_value = None + api.get_vm.return_value = None + monkeypatch.setattr("lightning_sdk.cli.vm.inspect.resolve_teamspace", lambda teamspace: _teamspace()) + monkeypatch.setattr("lightning_sdk.cli.vm.inspect.VMApi", MagicMock(return_value=api)) + + result = CliRunner().invoke(inspect_vm, ["ghost"]) + + assert result.exit_code != 0 + assert "was not found in teamspace 'ecorp/research'" in _flat(result.output) + + +@mock_command_logging +def test_inspect_surfaces_server_error_on_lookup(monkeypatch) -> None: + api = MagicMock() + api.get_vm_by_name.return_value = None + api.get_vm.side_effect = ApiException(status=500, reason="Internal Server Error") + monkeypatch.setattr("lightning_sdk.cli.vm.inspect.resolve_teamspace", lambda teamspace: _teamspace()) + monkeypatch.setattr("lightning_sdk.cli.vm.inspect.VMApi", MagicMock(return_value=api)) + + result = CliRunner().invoke(inspect_vm, ["vm-1"]) + + assert result.exit_code != 0 + assert "Internal Server Error" in _flat(result.output) + assert "Traceback" not in result.output + assert not isinstance(result.exception, ApiException) diff --git a/python/tests/core/test_vm.py b/python/tests/core/test_vm.py new file mode 100644 index 00000000..fe0a1cb4 --- /dev/null +++ b/python/tests/core/test_vm.py @@ -0,0 +1,225 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from lightning_sdk.api.utils import _machine_to_compute_name +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.organization import Organization +from lightning_sdk.user import User +from lightning_sdk.vm import VM, _require_org_id + + +def _org_teamspace() -> SimpleNamespace: + org = MagicMock(spec=Organization) + org.id = "org-1" + org.name = "ecorp" + return SimpleNamespace(id="ts-1", name="research", owner=org) + + +def _user_teamspace() -> SimpleNamespace: + user = MagicMock(spec=User) + user.name = "alice" + return SimpleNamespace(id="ts-2", name="personal", owner=user) + + +def _patch_resolution(monkeypatch, teamspace, api: MagicMock) -> None: + monkeypatch.setattr( + "lightning_sdk.vm._resolve_teamspace", + lambda teamspace=None, org=None, user=None, _resolved=teamspace: _resolved, + ) + monkeypatch.setattr("lightning_sdk.vm.VMApi", MagicMock(return_value=api)) + + +def test_require_org_id_rejects_user_owned_teamspace(): + with pytest.raises(ValueError, match="organization"): + _require_org_id(_user_teamspace()) + + +def test_require_org_id_returns_org_id(): + assert _require_org_id(_org_teamspace()) == "org-1" + + +def test_init_loads_by_name(monkeypatch): + api = MagicMock() + api.get_vm_by_name.return_value = V1Instance(id="vm-1", name="sim-1", status="running") + _patch_resolution(monkeypatch, _org_teamspace(), api) + + vm = VM("sim-1", teamspace="ecorp/research") + + assert vm.id == "vm-1" + assert vm.status == "running" + api.get_vm_by_name.assert_called_once_with("sim-1", "ts-1") + + +def test_init_falls_back_to_id(monkeypatch): + api = MagicMock() + api.get_vm_by_name.return_value = None + api.get_vm.return_value = V1Instance(id="vm-1", name="sim-1") + _patch_resolution(monkeypatch, _org_teamspace(), api) + + vm = VM("vm-1") + + assert vm.name == "sim-1" + api.get_vm.assert_called_once_with("vm-1", "org-1") + + +def test_init_raises_when_missing(monkeypatch): + api = MagicMock() + api.get_vm_by_name.return_value = None + api.get_vm.return_value = None + _patch_resolution(monkeypatch, _org_teamspace(), api) + + with pytest.raises(ValueError, match="not found"): + VM("ghost") + + +def test_init_wraps_api_errors_from_the_id_fallback(monkeypatch): + api = MagicMock() + api.get_vm_by_name.return_value = None + api.get_vm.side_effect = ApiException(status=500, reason="Internal Server Error") + _patch_resolution(monkeypatch, _org_teamspace(), api) + + with pytest.raises(ValueError, match="Internal Server Error"): + VM("vm-1") + + +def test_create_resolves_machine_and_cloud_account(monkeypatch): + monkeypatch.delenv("LIGHTNING_CLUSTER_ID", raising=False) + api = MagicMock() + api.create_vm.return_value = V1Instance(id="vm-1", name="sim-1", status="pending") + api.list_machine_clusters.return_value = [SimpleNamespace(id="cl-default")] + _patch_resolution(monkeypatch, _org_teamspace(), api) + + vm = VM.create("sim-1", machine=Machine.H100, volume_size=500) + + assert vm.id == "vm-1" + kwargs = api.create_vm.call_args.kwargs + assert kwargs["name"] == "sim-1" + assert kwargs["org_id"] == "org-1" + assert kwargs["teamspace_id"] == "ts-1" + assert kwargs["cluster_id"] == "cl-default" + assert kwargs["instance_type"] == _machine_to_compute_name(Machine.H100) + assert kwargs["volume_size"] == 500 + api.list_machine_clusters.assert_called_once_with("org-1") + + +def test_create_uses_cluster_id_from_the_environment(monkeypatch): + monkeypatch.setenv("LIGHTNING_CLUSTER_ID", "cl-env") + api = MagicMock() + api.create_vm.return_value = V1Instance(id="vm-1", name="sim-1", status="pending") + _patch_resolution(monkeypatch, _org_teamspace(), api) + + VM.create("sim-1", machine=Machine.H100) + + assert api.create_vm.call_args.kwargs["cluster_id"] == "cl-env" + api.list_machine_clusters.assert_not_called() + + +def test_create_errors_when_no_machine_cluster(monkeypatch): + monkeypatch.delenv("LIGHTNING_CLUSTER_ID", raising=False) + api = MagicMock() + api.list_machine_clusters.return_value = [] + _patch_resolution(monkeypatch, _org_teamspace(), api) + + with pytest.raises(ValueError, match="No machine cluster is available"): + VM.create("sim-1", machine=Machine.H100) + api.create_vm.assert_not_called() + + +def test_create_errors_when_multiple_machine_clusters(monkeypatch): + monkeypatch.delenv("LIGHTNING_CLUSTER_ID", raising=False) + api = MagicMock() + api.list_machine_clusters.return_value = [SimpleNamespace(id="cl-a"), SimpleNamespace(id="cl-b")] + _patch_resolution(monkeypatch, _org_teamspace(), api) + + with pytest.raises(ValueError, match="cl-a, cl-b"): + VM.create("sim-1", machine=Machine.H100) + api.create_vm.assert_not_called() + + +def test_create_with_explicit_cloud_account_and_wait(monkeypatch): + api = MagicMock() + api.create_vm.return_value = V1Instance(id="vm-1", name="sim-1", status="pending") + api.wait_for_status.return_value = V1Instance( + id="vm-1", name="sim-1", status="running", ssh_command="ssh -p 20032 ubuntu@1.2.3.4" + ) + _patch_resolution(monkeypatch, _org_teamspace(), api) + + vm = VM.create("sim-1", machine="H100", cloud_account="cl-explicit", wait=True, timeout=30) + + assert api.create_vm.call_args.kwargs["cluster_id"] == "cl-explicit" + api.wait_for_status.assert_called_once_with("vm-1", "org-1", timeout=30) + assert vm.status == "running" + assert vm.ssh_command == "ssh -p 20032 ubuntu@1.2.3.4" + + +def test_create_rejects_user_owned_teamspace(monkeypatch): + api = MagicMock() + _patch_resolution(monkeypatch, _user_teamspace(), api) + + with pytest.raises(ValueError, match="organization"): + VM.create("sim-1", machine=Machine.H100) + api.create_vm.assert_not_called() + + +def test_list(monkeypatch): + api = MagicMock() + api.list_vms.return_value = [V1Instance(id="a", name="x"), V1Instance(id="b", name="y")] + _patch_resolution(monkeypatch, _org_teamspace(), api) + + vms = VM.list() + + assert [vm.name for vm in vms] == ["x", "y"] + api.list_vms.assert_called_once_with("ts-1") + + +def test_refresh_wait_delete(monkeypatch): + api = MagicMock() + api.get_vm_by_name.return_value = V1Instance(id="vm-1", name="sim-1", status="pending") + api.get_vm.return_value = V1Instance(id="vm-1", name="sim-1", status="provisioning") + api.wait_for_status.return_value = V1Instance(id="vm-1", name="sim-1", status="running") + _patch_resolution(monkeypatch, _org_teamspace(), api) + vm = VM("sim-1") + + vm.refresh() + assert vm.status == "provisioning" + + vm.wait(timeout=12) + assert vm.status == "running" + api.wait_for_status.assert_called_once_with("vm-1", "org-1", timeout=12) + + vm.delete() + api.delete_vm.assert_called_once_with("vm-1", "org-1") + + +def test_properties_map_instance_fields(monkeypatch): + api = MagicMock() + api.get_vm_by_name.return_value = V1Instance( + id="vm-1", + name="sim-1", + status="running", + instance_type="lit-h100-1", + ssh_command="ssh -p 20032 ubuntu@1.2.3.4", + ssh_host="1.2.3.4", + ssh_port=20032, + ssh_user="ubuntu", + ) + _patch_resolution(monkeypatch, _org_teamspace(), api) + + vm = VM("sim-1") + + assert vm.machine == "lit-h100-1" + assert vm.ssh_host == "1.2.3.4" + assert vm.ssh_port == 20032 + assert vm.ssh_user == "ubuntu" + assert vm.teamspace.id == "ts-1" + assert "sim-1" in repr(vm) + + +def test_vm_is_exported(): + import lightning_sdk + + assert lightning_sdk.VM is VM