lightning-sdk 2025.8.21__py3-none-any.whl → 2025.8.26__py3-none-any.whl

lightning_sdk/cli/studio/create.py

@@ -4,14 +4,15 @@ from typing import Optional
 
 import click
 
-from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.cli.utils.save_to_config import save_teamspace_to_config
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.machine import CloudProvider
 from lightning_sdk.studio import Studio
 
 
 @click.command("create")
-@click.argument("studio_name", required=False)
+@click.option("--name", help="The name of the studio 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",
@@ -24,7 +25,7 @@ from lightning_sdk.studio import Studio
     type=click.STRING,
 )
 def create_studio(
-    studio_name: Optional[str] = None,
+    name: Optional[str] = None,
     teamspace: Optional[str] = None,
     cloud_provider: Optional[str] = None,
     cloud_account: Optional[str] = None,
@@ -32,37 +33,27 @@ def create_studio(
     """Create a new Studio.
 
     Example:
-        lightning studio create [STUDIO_NAME]
-
-    STUDIO_NAME: the name of the studio to create.
-
-    If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
+        lightning studio create
     """
-    if teamspace is not None:
-        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
-        if resolved_teamspace is None:
-            raise ValueError(
-                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
-                "Does the teamspace exist?"
-            )
-    else:
-        resolved_teamspace = None
+    menu = TeamspacesMenu()
+
+    resolved_teamspace = menu(teamspace)
+    save_teamspace_to_config(resolved_teamspace, overwrite=False)
 
     if cloud_provider is not None:
         cloud_provider = CloudProvider(cloud_provider)
 
     try:
         studio = Studio(
-            studio_name,
+            name=name,
             teamspace=resolved_teamspace,
             create_ok=True,
             cloud_provider=cloud_provider,
             cloud_account=cloud_account,
         )
-    except (RuntimeError, ValueError, ApiException) as e:
-        print(e)
-        if studio_name:
-            raise ValueError(f"Could not create Studio: '{studio_name}'. Does the Studio exist?") from None
-        raise ValueError(f"Could not create Studio: '{studio_name}'. Please provide a Studio name") from None
+    except (RuntimeError, ValueError, ApiException):
+        if name:
+            raise ValueError(f"Could not create Studio: '{name}'. Does the Studio exist?") from None
+        raise ValueError(f"Could not create Studio: '{name}'. Please provide a Studio name") from None
 
     click.echo(f"Studio '{studio.name}' created successfully")

lightning_sdk/cli/studio/delete.py

@@ -4,41 +4,42 @@ from typing import Optional
 
 import click
 
-from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
-from lightning_sdk.studio import Studio
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 
 
 @click.command("delete")
-@click.argument("studio_name", required=False)
+@click.option(
+    "--name",
+    help=(
+        "The name of the studio 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)")
-def delete_studio(studio_name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+def delete_studio(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
     """Delete a Studio.
 
     Example:
-      lightning studio delete [STUDIO_NAME]
+      lightning studio delete --name my-studio
 
-    STUDIO_NAME: the name of the studio to delete.
-
-    If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
     """
-    # missing studio_name and teamspace are handled by the studio class
-    if teamspace is not None:
-        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
-        if resolved_teamspace is None:
-            raise ValueError(
-                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
-                "Does the teamspace exist?"
-            )
-    else:
-        resolved_teamspace = None
-
-    try:
-        studio = Studio(studio_name, teamspace=resolved_teamspace, create_ok=False)
-        studio.delete()
-    except Exception:
-        # TODO: make this a generic CLI error
-        if studio_name:
-            raise ValueError(f"Could not delete Studio: '{studio_name}'. Does the Studio exist?") from None
-        raise ValueError("No studio name provided. Use 'lightning studio delete <name>' to delete a studio.") from None
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
+
+    menu = StudiosMenu(resolved_teamspace)
+    studio = menu(studio=name)
+
+    studio_name = f"{studio.teamspace.owner.name}/{studio.teamspace.name}/{studio.name}"
+    confirmed = click.confirm(
+        f"Are you sure you want to delete studio '{studio_name}'?",
+        abort=True,
+    )
+    if not confirmed:
+        click.echo("Studio deletion cancelled")
+        return
+
+    studio.delete()
 
     click.echo(f"Studio '{studio.name}' deleted successfully")

lightning_sdk/cli/studio/list.py

@@ -6,8 +6,9 @@ import click
 from rich.table import Table
 
 from lightning_sdk.cli.utils.cloud_account_map import cloud_account_to_display_name
-from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
 from lightning_sdk.cli.utils.richt_print import rich_to_str
+from lightning_sdk.cli.utils.save_to_config import save_teamspace_to_config
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.studio import Studio
 
 
@@ -26,11 +27,9 @@ def list_studios(teamspace: Optional[str] = None, sort_by: Optional[str] = None)
         lightning studio list --teamspace owner/teamspace
 
     """
-    teamspace_resolved = resolve_teamspace_owner_name_format(teamspace)
-
-    if teamspace_resolved is None:
-        # TODO: make this a generic CLI error
-        raise ValueError(f"Could not resolve teamspace: {teamspace}")
+    menu = TeamspacesMenu()
+    teamspace_resolved = menu(teamspace=teamspace)
+    save_teamspace_to_config(teamspace_resolved, overwrite=False)
 
     studios = teamspace_resolved.studios
 

lightning_sdk/cli/studio/ssh.py

@@ -9,14 +9,22 @@ from typing import List, Optional
 
 import click
 
-from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud.login import Auth
-from lightning_sdk.studio import Studio
 from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH
 
 
 @click.command("ssh")
-@click.argument("studio_name", required=False)
+@click.option(
+    "--name",
+    help=(
+        "The name of the studio 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)", type=click.STRING)
 @click.option(
     "--option",
@@ -25,33 +33,22 @@ from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH
     multiple=True,
     type=click.STRING,
 )
-def ssh_studio(
-    studio_name: Optional[str] = None, teamspace: Optional[str] = None, option: Optional[List[str]] = None
-) -> None:
+def ssh_studio(name: Optional[str] = None, teamspace: Optional[str] = None, option: Optional[List[str]] = None) -> None:
     """SSH into a Studio.
 
     Example:
-        lightning studio ssh [STUDIO_NAME]
-
-    STUDIO_NAME: the name of the studio to SSH into.
-
-    If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
+        lightning studio ssh --name my-studio
     """
     auth = Auth()
     auth.authenticate()
     ssh_private_key_path = _download_ssh_keys(auth.api_key, force_download=False)
 
-    if teamspace is not None:
-        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
-        if resolved_teamspace is None:
-            raise ValueError(
-                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
-                "Does the teamspace exist?"
-            )
-    else:
-        resolved_teamspace = None
-
-    studio = Studio(studio_name, teamspace=resolved_teamspace)
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
+
+    menu = StudiosMenu(resolved_teamspace)
+    studio = menu(studio=name)
+    save_studio_to_config(studio)
 
     ssh_options = " -o " + " -o ".join(option) if option else ""
     ssh_command = f"ssh -i {ssh_private_key_path}{ssh_options} s_{studio._studio.id}@ssh.lightning.ai"

lightning_sdk/cli/studio/start.py

@@ -4,14 +4,22 @@ from typing import Optional
 
 import click
 
-from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
-from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.machine import CloudProvider, Machine
 from lightning_sdk.studio import Studio
 
 
 @click.command("start")
-@click.argument("studio_name", required=False)
+@click.option(
+    "--name",
+    help=(
+        "The name of the studio 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 studio if it doesn't exist")
 @click.option(
@@ -34,7 +42,7 @@ from lightning_sdk.studio import Studio
     type=click.STRING,
 )
 def start_studio(
-    studio_name: Optional[str] = None,
+    name: Optional[str] = None,
     teamspace: Optional[str] = None,
     create: bool = False,
     machine: Optional[str] = None,
@@ -45,37 +53,28 @@ def start_studio(
     """Start a Studio.
 
     Example:
-        lightning studio start [STUDIO_NAME]
+        lightning studio start --name my-studio
 
-    STUDIO_NAME: the name of the studio to start.
-
-    If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
     """
-    if teamspace is not None:
-        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
-        if resolved_teamspace is None:
-            raise ValueError(
-                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
-                "Does the teamspace exist?"
-            )
-    else:
-        resolved_teamspace = None
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
 
     if cloud_provider is not None:
         cloud_provider = CloudProvider(cloud_provider)
 
-    try:
+    if not create:
+        menu = StudiosMenu(resolved_teamspace)
+        studio = menu(studio=name)
+    else:
         studio = Studio(
-            studio_name,
+            name=name,
             teamspace=resolved_teamspace,
             create_ok=create,
             cloud_provider=cloud_provider,
             cloud_account=cloud_account,
         )
-    except (RuntimeError, ValueError, ApiException):
-        if studio_name:
-            raise ValueError(f"Could not start Studio: '{studio_name}'. Does the Studio exist?") from None
-        raise ValueError(f"Could not start Studio: '{studio_name}'. Please provide a Studio name") from None
+
+    save_studio_to_config(studio)
 
     Studio.show_progress = True
     studio.start(machine, interruptible=interruptible)

lightning_sdk/cli/studio/stop.py

@@ -4,41 +4,37 @@ from typing import Optional
 
 import click
 
-from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
-from lightning_sdk.studio import Studio
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 
 
 @click.command("stop")
-@click.argument("studio_name", required=False)
+@click.option(
+    "--name",
+    help=(
+        "The name of the studio 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)")
-def stop_studio(studio_name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+def stop_studio(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
     """Stop a Studio.
 
     Example:
-        lightning studio stop [STUDIO_NAME]
+        lightning studio stop --name my-studio
 
-    STUDIO_NAME: the name of the studio to stop.
-
-    If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
     """
     # missing studio_name and teamspace are handled by the studio class
-    if teamspace is not None:
-        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
-        if resolved_teamspace is None:
-            raise ValueError(
-                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
-                "Does the teamspace exist?"
-            )
-    else:
-        resolved_teamspace = None
-
-    try:
-        studio = Studio(studio_name, teamspace=resolved_teamspace)
-        studio.stop()
-    except Exception:
-        # TODO: make this a generic CLI error
-        if studio_name:
-            raise ValueError(f"Could not stop studio: '{studio_name}'. Does the studio exist?") from None
-        raise ValueError("No studio name provided. Use 'lightning studio stop <name>' to stop a studio.") from None
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
+
+    menu = StudiosMenu(resolved_teamspace)
+    studio = menu(studio=name)
+
+    studio.stop()
+
+    save_studio_to_config(studio)
 
     click.echo(f"Studio '{studio.name}' stopped successfully")

lightning_sdk/cli/studio/switch.py

@@ -4,14 +4,22 @@ from typing import Optional
 
 import click
 
-from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
-from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.machine import Machine
 from lightning_sdk.studio import Studio
 
 
 @click.command("switch")
-@click.argument("studio_name", required=False)
+@click.option(
+    "--name",
+    help=(
+        "The name of the studio 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(
     "--machine",
@@ -20,34 +28,22 @@ from lightning_sdk.studio import Studio
 )
 @click.option("--interruptible", is_flag=True, help="Switch the studio to an interruptible instance.")
 def switch_studio(
-    studio_name: Optional[str] = None,
+    name: Optional[str] = None,
     teamspace: Optional[str] = None,
     machine: Optional[str] = None,
     interruptible: bool = False,
 ) -> None:
     """Switch a Studio to a different machine type."""
-    if teamspace is not None:
-        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
-        if resolved_teamspace is None:
-            raise ValueError(
-                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
-                "Does the teamspace exist?"
-            )
-    else:
-        resolved_teamspace = None
-
-    try:
-        studio = Studio(
-            studio_name,
-            teamspace=resolved_teamspace,
-        )
-    except (RuntimeError, ValueError, ApiException):
-        if studio_name:
-            raise ValueError(f"Could not switch Studio: '{studio_name}'. Does the Studio exist?") from None
-        raise ValueError(f"Could not switch Studio: '{studio_name}'. Please provide a Studio name") from None
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
+
+    menu = StudiosMenu(resolved_teamspace)
+    studio = menu(studio=name)
 
     resolved_machine = Machine.from_str(machine)
     Studio.show_progress = True
     studio.switch_machine(resolved_machine, interruptible=interruptible)
 
+    save_studio_to_config(studio)
+
     click.echo(f"Studio '{studio.name}' switched to machine '{resolved_machine}' successfully")

lightning_sdk/cli/utils/resolve.py

@@ -4,7 +4,7 @@ from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.utils.resolve import _resolve_teamspace
 
 
-def resolve_teamspace_owner_name_format(teamspace_name: str) -> Optional[Teamspace]:
+def resolve_teamspace_owner_name_format(teamspace_name: Optional[str]) -> Optional[Teamspace]:
     teamspace_resolved = None
     if teamspace_name is None:
         return _resolve_teamspace(None, None, None)

lightning_sdk/cli/utils/save_to_config.py

@@ -0,0 +1,27 @@
+from lightning_sdk.organization import Organization
+from lightning_sdk.studio import Studio
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+
+def save_teamspace_to_config(teamspace: Teamspace, overwrite: bool = False) -> None:
+    saved = _save_to_config_if_not_exists(DefaultConfigKeys.teamspace_name, teamspace.name, overwrite)
+    saved = _save_to_config_if_not_exists(DefaultConfigKeys.teamspace_owner, teamspace.owner.name, overwrite=saved)
+    saved = _save_to_config_if_not_exists(
+        DefaultConfigKeys.teamspace_owner_type,
+        "organization" if isinstance(teamspace.owner, Organization) else "user",
+        overwrite=saved,
+    )
+
+
+def save_studio_to_config(studio: Studio, overwrite: bool = False) -> None:
+    saved = _save_to_config_if_not_exists(DefaultConfigKeys.studio, studio.name, overwrite)
+    save_teamspace_to_config(studio.teamspace, overwrite=saved)
+
+
+def _save_to_config_if_not_exists(key: str, value: str, overwrite: bool = False) -> bool:
+    cfg = Config()
+    if overwrite or cfg.get(key) is None:
+        cfg.set(key, value)
+        return True
+    return False

lightning_sdk/cli/utils/studio_selection.py

@@ -0,0 +1,106 @@
+import os
+from contextlib import suppress
+from typing import Dict, List, Optional
+
+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.teamspace import Teamspace
+from lightning_sdk.utils.resolve import _get_authed_user
+
+
+class StudiosMenu:
+    """This class is used to select a studio from a list of possible studios within a teamspace.
+
+    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:
+        """Initialize the StudiosMenu with a teamspace.
+
+        Args:
+            teamspace: The teamspace to list studios from
+        """
+        self.teamspace = teamspace
+
+    def _get_studio_from_interactive_menu(self, possible_studios: Dict[str, Studio]) -> Studio:
+        studio_names = sorted(possible_studios.keys())
+        terminal_menu = self._prepare_terminal_menu_studios(studio_names)
+        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:
+        if studio in possible_studios:
+            return possible_studios[studio]
+
+        click.echo(f"Could not find 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:
+        if title is None:
+            title = "Please select a Studio out of the following:"
+
+        return TerminalMenu(studio_names, title=title, clear_menu_on_exit=True)
+
+    def _get_possible_studios(self) -> Dict[str, Studio]:
+        """Get all available studios in the teamspace."""
+        studios = {}
+
+        user = _get_authed_user()
+        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:
+        """Select a studio from the teamspace.
+
+        Args:
+            studio: Optional studio name to select. If not provided, will show interactive menu.
+
+        Returns:
+            Selected Studio object
+
+        Raises:
+            StudioCliError: If studio selection fails
+        """
+        try:
+            # try to resolve the studio from the name, environment or config
+            resolved_studio = None
+
+            with suppress(Exception):
+                resolved_studio = Studio(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."
+                )
+
+            click.echo(f"Listing studios in teamspace {self.teamspace.owner.name}/{self.teamspace.name}...")
+
+            possible_studios = self._get_possible_studios()
+
+            if not possible_studios:
+                raise ValueError(f"No studios found in teamspace {self.teamspace.name}")
+
+            if studio is None:
+                return self._get_studio_from_interactive_menu(possible_studios)
+
+            return self._get_studio_from_name(studio, possible_studios)
+
+        except KeyboardInterrupt:
+            raise KeyboardInterrupt from None
+
+        except Exception as e:
+            raise StudioCliError(
+                "Could not resolve a Studio. "
+                "Please pass it as an argument or contact Lightning AI directly to resolve this issue."
+            ) from e