lightning-sdk 2025.9.16__py3-none-any.whl → 2025.9.29__py3-none-any.whl

lightning_sdk/cli/utils/studio_selection.py

@@ -1,12 +1,12 @@
 import os
 from contextlib import suppress
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 import click
 from simple_term_menu import TerminalMenu
 
 from lightning_sdk.cli.legacy.exceptions import StudioCliError
-from lightning_sdk.studio import Studio
+from lightning_sdk.studio import VM, Studio
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.utils.resolve import _get_authed_user
 
@@ -17,54 +17,58 @@ class StudiosMenu:
     It can be used to select a studio from a list of possible studios, or to resolve a studio from a name.
     """
 
-    def __init__(self, teamspace: Teamspace) -> None:
+    def __init__(self, teamspace: Teamspace, vm: bool = False) -> None:
         """Initialize the StudiosMenu with a teamspace.
 
         Args:
             teamspace: The teamspace to list studios from
         """
         self.teamspace = teamspace
+        self.vm = vm
 
-    def _get_studio_from_interactive_menu(self, possible_studios: Dict[str, Studio]) -> Studio:
+    def _get_studio_from_interactive_menu(self, possible_studios: Dict[str, Union[Studio, VM]]) -> Union[Studio, VM]:
         studio_names = sorted(possible_studios.keys())
-        terminal_menu = self._prepare_terminal_menu_studios(studio_names)
+        terminal_menu = self._prepare_terminal_menu_studios(studio_names, vm=self.vm)
         terminal_menu.show()
 
         selected_name = studio_names[terminal_menu.chosen_menu_index]
         return possible_studios[selected_name]
 
-    def _get_studio_from_name(self, studio: str, possible_studios: Dict[str, Studio]) -> Studio:
+    def _get_studio_from_name(self, studio: str, possible_studios: Dict[str, Union[Studio, VM]]) -> Union[Studio, VM]:
         if studio in possible_studios:
             return possible_studios[studio]
 
-        click.echo(f"Could not find Studio {studio}, please select it from the list:")
+        click.echo(f"Could not find {'VM' if self.vm else 'Studio'} {studio}, please select it from the list:")
         return self._get_studio_from_interactive_menu(possible_studios)
 
     @staticmethod
-    def _prepare_terminal_menu_studios(studio_names: List[str], title: Optional[str] = None) -> TerminalMenu:
+    def _prepare_terminal_menu_studios(
+        studio_names: List[str], title: Optional[str] = None, vm: bool = False
+    ) -> TerminalMenu:
         if title is None:
-            title = "Please select a Studio out of the following:"
+            title = f"Please select a {'VM' if vm else 'Studio'} out of the following:"
 
         return TerminalMenu(studio_names, title=title, clear_menu_on_exit=True)
 
-    def _get_possible_studios(self) -> Dict[str, Studio]:
+    def _get_possible_studios(self) -> Dict[str, Union[Studio, VM]]:
         """Get all available studios in the teamspace."""
         studios = {}
 
         user = _get_authed_user()
+        studios = self.teamspace.vms if self.vm else self.teamspace.studios
         for studio in self.teamspace.studios:
             if studio._studio.user_id == user.id:
                 studios[studio.name] = studio
         return studios
 
-    def __call__(self, studio: Optional[str] = None) -> Studio:
+    def __call__(self, studio: Optional[str] = None) -> Union[Studio, VM]:
         """Select a studio from the teamspace.
 
         Args:
             studio: Optional studio name to select. If not provided, will show interactive menu.
 
         Returns:
-            Selected Studio object
+            Selected Studio/VM object
 
         Raises:
             StudioCliError: If studio selection fails
@@ -73,15 +77,18 @@ class StudiosMenu:
             # try to resolve the studio from the name, environment or config
             resolved_studio = None
 
+            selected_cls = VM if self.vm else Studio
+
             with suppress(Exception):
-                resolved_studio = Studio(name=studio, teamspace=self.teamspace, create_ok=False)
+                resolved_studio = selected_cls(name=studio, teamspace=self.teamspace, create_ok=False)
 
             if resolved_studio is not None:
                 return resolved_studio
 
             if os.environ.get("LIGHTNING_NON_INTERACTIVE", "0") == "1" and studio is None:
                 raise ValueError(
-                    "Studio selection is not supported in non-interactive mode. Please provide a studio name."
+                    f"{'VM' if self.vm else 'Studio'} selection is not supported in non-interactive mode. "
+                    "Please provide a studio name."
                 )
 
             click.echo(f"Listing studios in teamspace {self.teamspace.owner.name}/{self.teamspace.name}...")
@@ -101,6 +108,6 @@ class StudiosMenu:
 
         except Exception as e:
             raise StudioCliError(
-                "Could not resolve a Studio. "
+                f"Could not resolve a {'VM' if self.vm else 'Studio'}. "
                 "Please pass it as an argument or contact Lightning AI directly to resolve this issue."
             ) from e

lightning_sdk/cli/utils/teamspace_selection.py

@@ -1,15 +1,15 @@
 import os
+from contextlib import suppress
 from typing import Dict, List, Optional
 
 import click
 from simple_term_menu import TerminalMenu
 
-from lightning_sdk.api import OrgApi
 from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
-from lightning_sdk.teamspace import Teamspace
-from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _get_authed_user
+from lightning_sdk.teamspace import Organization, Teamspace
+from lightning_sdk.user import Owner, User
+from lightning_sdk.utils.resolve import ApiException, _get_authed_user, _resolve_teamspace
 
 
 class TeamspacesMenu:
@@ -18,51 +18,45 @@ class TeamspacesMenu:
     It can be used to select a teamspace from a list of possible teamspaces, or to resolve a teamspace from a name.
     """
 
-    def _get_teamspace_from_interactive_menu(self, possible_teamspaces: Dict[str, Dict[str, str]]) -> Dict[str, str]:
+    def __init__(self, owner: Optional[Owner] = None) -> None:
+        self._owner: Owner = owner
+
+    def _get_teamspace_from_interactive_menu(self, possible_teamspaces: Dict[str, str]) -> str:
         teamspace_ids = sorted(possible_teamspaces.keys())
-        terminal_menu = self._prepare_terminal_menu_teamspaces([possible_teamspaces[k] for k in teamspace_ids])
+        terminal_menu = self._prepare_terminal_menu_teamspaces(
+            [possible_teamspaces[k] for k in teamspace_ids], self._owner
+        )
         terminal_menu.show()
 
         selected_id = teamspace_ids[terminal_menu.chosen_menu_index]
         return possible_teamspaces[selected_id]
 
-    def _get_teamspace_from_name(
-        self, teamspace: str, possible_teamspaces: Dict[str, Dict[str, str]]
-    ) -> Dict[str, str]:
-        try:
-            owner, name = teamspace.split("/", maxsplit=1)
-        except ValueError as e:
-            raise ValueError(
-                f"Invalid teamspace format: '{teamspace}'. "
-                "Teamspace should be specified as '{teamspace_owner}/{teamspace_name}' "
-                "(e.g., 'my-org/my-teamspace')."
-            ) from e
-
+    def _get_teamspace_from_name(self, teamspace: str, possible_teamspaces: Dict[str, str]) -> str:
         for _, ts in possible_teamspaces.items():
-            if ts["name"] == name and (ts["user"] == owner or ts["org"] == owner):
+            if ts == teamspace:
                 return ts
 
-        click.echo(f"Could not find Teamspace {teamspace}, please select it from the list:")
+        click.echo(f"Could not find Teamspace {self._owner.name}/{teamspace}, please select it from the list:")
         return self._get_teamspace_from_interactive_menu(possible_teamspaces)
 
     @staticmethod
     def _prepare_terminal_menu_teamspaces(
-        possible_teamspaces: List[Dict[str, str]], title: Optional[str] = None
+        possible_teamspaces: List[str],
+        owner: Owner,
+        title: Optional[str] = None,
     ) -> TerminalMenu:
         if title is None:
-            title = "Please select a Teamspace out of the following:"
+            title = f"Please select a Teamspace (owned by {owner.name}) out of the following:"
 
-        return TerminalMenu(
-            [f"{t['user'] or t['org']}/{t['name']}" for t in possible_teamspaces], title=title, clear_menu_on_exit=True
-        )
+        return TerminalMenu(possible_teamspaces, title=title, clear_menu_on_exit=True)
 
     @staticmethod
-    def _get_possible_teamspaces(user: User) -> Dict[str, Dict[str, str]]:
-        org_api = OrgApi()
+    def _get_possible_teamspaces(user: User, owner: Owner) -> Dict[str, str]:
         user_api = user._user_api
 
-        user_api._get_organizations_for_authed_user()
-        memberships = user_api._get_all_teamspace_memberships(user_id=user.id)
+        memberships = user_api._get_all_teamspace_memberships(
+            user_id=user.id, org_id=owner.id if isinstance(owner, Organization) else None
+        )
 
         teamspaces = {}
         # get all teamspace memberships
@@ -70,56 +64,63 @@ class TeamspacesMenu:
             teamspace_id = membership.project_id
             teamspace_name = membership.name
 
-            # get organization if necessary
-            if membership.owner_type == "organization":
-                org_name = org_api._get_org_by_id(membership.owner_id).name
-                user_name = None
-            else:
-                org_name = None
-
-                # don't do a request if not necessary
-                if membership.owner_id == user.id:
-                    user_name = user.name
-                else:
-                    user_name = user_api._get_user_by_id(membership.owner_id).username
-
-            teamspaces[teamspace_id] = {"user": user_name, "org": org_name, "name": teamspace_name}
+            if membership.owner_id == owner.id:
+                teamspaces[teamspace_id] = teamspace_name
 
         return teamspaces
 
     def __call__(self, teamspace: Optional[str] = None) -> Teamspace:
-        try:
-            # try to resolve the teamspace from the name, environment or config
-            resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
+        resolved_teamspace = None
+        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
+        if resolved_teamspace is not None:
+            return resolved_teamspace
 
-            if resolved_teamspace is not None:
-                return resolved_teamspace
+        if self._owner is None:
+            # resolve owner if required
+            from lightning_sdk.cli.utils.owner_selection import OwnerMenu
 
-            if os.environ.get("LIGHTNING_NON_INTERACTIVE", "0") == "1":
-                raise ValueError(
-                    "Teamspace selection is not supported in non-interactive mode. "
-                    "Please provide a teamspace name in the format of 'owner/teamspace'."
-                )
+            menu = OwnerMenu()
+            self._owner = menu()
+
+        if isinstance(self._owner, Organization):
+            org = self._owner
+            user = None
+        elif isinstance(self._owner, User):
+            org = None
+            user = self._owner
+
+        with suppress(ApiException, ValueError, RuntimeError):
+            resolved_teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
 
-            # if the teamspace is not resolved, try to get the teamspace from the interactive menu
-            # this could mean that either no teamspace was provided or the provided teamspace is not valid
-            user = _get_authed_user()
+        if resolved_teamspace is not None:
+            return resolved_teamspace
 
-            possible_teamspaces = self._get_possible_teamspaces(user)
+        if os.environ.get("LIGHTNING_NON_INTERACTIVE", "0") == "1":
+            raise ValueError(
+                "Teamspace selection is not supported in non-interactive mode. "
+                "Please provide a teamspace name in the format of 'owner/teamspace'."
+            )
+
+        # if the teamspace is not resolved, try to get the teamspace from the interactive menu
+        # this could mean that either no teamspace was provided or the provided teamspace is not valid
+        try:
+            auth_user = _get_authed_user()
+
+            possible_teamspaces = self._get_possible_teamspaces(auth_user, self._owner)
             if teamspace is None:
-                teamspace = resolve_teamspace_owner_name_format(teamspace)
-                teamspace_dict = self._get_teamspace_from_interactive_menu(possible_teamspaces=possible_teamspaces)
+                teamspace_name = self._get_teamspace_from_interactive_menu(possible_teamspaces=possible_teamspaces)
             else:
-                teamspace_dict = self._get_teamspace_from_name(
+                teamspace_name = self._get_teamspace_from_name(
                     teamspace=teamspace, possible_teamspaces=possible_teamspaces
                 )
 
-            return Teamspace(**teamspace_dict)
+            return Teamspace(teamspace_name, org=org, user=user)
+
         except KeyboardInterrupt:
             raise KeyboardInterrupt from None
 
         except Exception as e:
             raise StudioCliError(
-                f"Could not find the given Teamspace {teamspace}. "
+                f"Could not find the given Teamspace owned by {self._owner.name}. "
                 "Please contact Lightning AI directly to resolve this issue."
             ) from e

lightning_sdk/cli/vm/__init__.py

@@ -0,0 +1,20 @@
+import click
+
+
+def register_commands(group: click.Group) -> None:
+    """Register studio commands with the given group."""
+    from lightning_sdk.cli.vm.create import create_vm
+    from lightning_sdk.cli.vm.delete import delete_vm
+    from lightning_sdk.cli.vm.list import list_vms
+    from lightning_sdk.cli.vm.ssh import ssh_vm
+    from lightning_sdk.cli.vm.start import start_vm
+    from lightning_sdk.cli.vm.stop import stop_vm
+    from lightning_sdk.cli.vm.switch import switch_vm
+
+    group.add_command(create_vm)
+    group.add_command(delete_vm)
+    group.add_command(list_vms)
+    group.add_command(ssh_vm)
+    group.add_command(start_vm)
+    group.add_command(stop_vm)
+    group.add_command(switch_vm)

lightning_sdk/cli/vm/create.py

@@ -0,0 +1,33 @@
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.studio.create import create_impl
+from lightning_sdk.machine import CloudProvider
+
+
+@click.command("create")
+@click.option("--name", help="The name of the VM to create. If not provided, a random name will be generated.")
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option(
+    "--cloud-provider",
+    help="The cloud provider to start the VM on. Defaults to teamspace default.",
+    type=click.Choice(m.name for m in list(CloudProvider)),
+)
+@click.option(
+    "--cloud-account",
+    help="The cloud account to create the VM on. Defaults to teamspace default.",
+    type=click.STRING,
+)
+def create_vm(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    cloud_provider: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+) -> None:
+    """Create a new VM.
+
+    Example:
+        lightning vm create
+    """
+    create_impl(name=name, teamspace=teamspace, cloud_provider=cloud_provider, cloud_account=cloud_account, vm=True)

lightning_sdk/cli/vm/delete.py

@@ -0,0 +1,25 @@
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.studio.delete import delete_impl
+
+
+@click.command("delete")
+@click.option(
+    "--name",
+    help=(
+        "The name of the VM to delete. "
+        "If not provided, will try to infer from environment, "
+        "use the default value from the config or prompt for interactive selection."
+    ),
+)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+def delete_vm(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Delete a VM.
+
+    Example:
+      lightning vm delete --name my-vm
+
+    """
+    return delete_impl(name=name, teamspace=teamspace, vm=True)

lightning_sdk/cli/vm/list.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.studio.list import list_impl
+
+
+@click.command("list")
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option(
+    "--all",
+    is_flag=True,
+    flag_value=True,
+    default=False,
+    help="List all VMs, not just the ones belonging to the authed user",
+)
+@click.option(
+    "--sort-by",
+    default=None,
+    type=click.Choice(["name", "teamspace", "status", "machine", "cloud-account"], case_sensitive=False),
+    help="the attribute to sort the VMs by.",
+)
+def list_vms(teamspace: Optional[str] = None, all: bool = False, sort_by: Optional[str] = None) -> None:  # noqa: A002
+    """List VMs in a teamspace.
+
+    Example:
+        lightning vm list --teamspace owner/teamspace
+
+    """
+    return list_impl(teamspace=teamspace, all=all, sort_by=sort_by, vm=True)

lightning_sdk/cli/vm/ssh.py

@@ -0,0 +1,31 @@
+from typing import List, Optional
+
+import click
+
+from lightning_sdk.cli.studio.ssh import ssh_impl
+
+
+@click.command("ssh")
+@click.option(
+    "--name",
+    help=(
+        "The name of the VM to ssh into. "
+        "If not provided, will try to infer from environment, "
+        "use the default value from the config or prompt for interactive selection."
+    ),
+)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)", type=click.STRING)
+@click.option(
+    "--option",
+    "-o",
+    help="Additional options to pass to the SSH command. Can be specified multiple times.",
+    multiple=True,
+    type=click.STRING,
+)
+def ssh_vm(name: Optional[str] = None, teamspace: Optional[str] = None, option: Optional[List[str]] = None) -> None:
+    """SSH into a VM.
+
+    Example:
+        lightning vm ssh --name my-vm
+    """
+    return ssh_impl(name=name, teamspace=teamspace, option=option, vm=False)

lightning_sdk/cli/vm/start.py

@@ -0,0 +1,60 @@
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.studio.start import start_impl
+from lightning_sdk.machine import CloudProvider, Machine
+
+
+@click.command("start")
+@click.option(
+    "--name",
+    help=(
+        "The name of the VM to start. "
+        "If not provided, will try to infer from environment, "
+        "use the default value from the config or prompt for interactive selection."
+    ),
+)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option("--create", is_flag=True, help="Create the VM if it doesn't exist")
+@click.option(
+    "--machine",
+    help="The machine type to start the VM on. Defaults to CPU-4",
+    type=click.Choice(m.name for m in Machine.__dict__.values() if isinstance(m, Machine) and m._include_in_cli),
+)
+@click.option("--interruptible", is_flag=True, help="Start the VM on an interruptible instance.")
+@click.option(
+    "--cloud-provider",
+    help=("The cloud provider to start the VM on. Defaults to teamspace default. Only used if --create is specified."),
+    type=click.Choice(m.name for m in list(CloudProvider)),
+)
+@click.option(
+    "--cloud-account",
+    help="The cloud account to start the VM on. Defaults to teamspace default. Only used if --create is specified.",
+    type=click.STRING,
+)
+def start_vm(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    create: bool = False,
+    machine: str = "CPU",
+    interruptible: bool = False,
+    cloud_provider: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+) -> None:
+    """Start a VM.
+
+    Example:
+        lightning vm start --name my-vm
+
+    """
+    return start_impl(
+        name=name,
+        teamspace=teamspace,
+        create=create,
+        machine=machine,
+        interruptible=interruptible,
+        cloud_provider=cloud_provider,
+        cloud_account=cloud_account,
+        vm=True,
+    )

lightning_sdk/cli/vm/stop.py

@@ -0,0 +1,25 @@
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.studio.stop import stop_impl
+
+
+@click.command("stop")
+@click.option(
+    "--name",
+    help=(
+        "The name of the VM to stop. "
+        "If not provided, will try to infer from environment, "
+        "use the default value from the config or prompt for interactive selection."
+    ),
+)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+def stop_vm(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Stop a VM.
+
+    Example:
+        lightning vm stop --name my-vm
+
+    """
+    return stop_impl(name=name, teamspace=teamspace, vm=True)

lightning_sdk/cli/vm/switch.py

@@ -0,0 +1,38 @@
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.studio.switch import switch_impl
+from lightning_sdk.machine import Machine
+
+
+@click.command("switch")
+@click.option(
+    "--name",
+    help=(
+        "The name of the VM to switch to a different machine. "
+        "If not provided, will try to infer from environment, "
+        "use the default value from the config or prompt for interactive selection."
+    ),
+)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option(
+    "--machine",
+    help="The machine type to switch the studio to.",
+    type=click.Choice(m.name for m in Machine.__dict__.values() if isinstance(m, Machine)),
+)
+@click.option("--interruptible", is_flag=True, help="Switch the studio to an interruptible instance.")
+def switch_vm(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    machine: Optional[str] = None,
+    interruptible: bool = False,
+) -> None:
+    """Switch a VM to a different machine type."""
+    return switch_impl(
+        name=name,
+        teamspace=teamspace,
+        machine=machine,
+        interruptible=interruptible,
+        vm=True,
+    )