lightning-sdk 2025.8.19.post0__py3-none-any.whl → 2025.8.26__py3-none-any.whl

lightning_sdk/cli/legacy/job_and_mmt_action.py

@@ -2,12 +2,12 @@ from typing import Optional
 
 from lightning_sdk.cli.legacy.jobs_menu import _JobsMenu
 from lightning_sdk.cli.legacy.mmts_menu import _MMTsMenu
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.job import Job
 from lightning_sdk.mmt import MMT
 
 
-class _JobAndMMTAction(_TeamspacesMenu, _JobsMenu, _MMTsMenu):
+class _JobAndMMTAction(TeamspacesMenu, _JobsMenu, _MMTsMenu):
     """Inspect resources of the Lightning AI platform to get additional details as JSON."""
 
     def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> Job:
@@ -20,7 +20,7 @@ class _JobAndMMTAction(_TeamspacesMenu, _JobsMenu, _MMTsMenu):
                 If not specified can be selected interactively.
 
         """
-        resolved_teamspace = self._resolve_teamspace(teamspace)
+        resolved_teamspace = self(teamspace)
         return self._resolve_job(name, teamspace=resolved_teamspace)
 
     def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> MMT:
@@ -33,5 +33,5 @@ class _JobAndMMTAction(_TeamspacesMenu, _JobsMenu, _MMTsMenu):
                 If not specified can be selected interactively.
 
         """
-        resolved_teamspace = self._resolve_teamspace(teamspace)
+        resolved_teamspace = self(teamspace)
         return self._resolve_mmt(name, teamspace=resolved_teamspace)

lightning_sdk/cli/legacy/list.py

@@ -8,7 +8,7 @@ from typing_extensions import Literal
 
 from lightning_sdk import Job, Machine, Studio, Teamspace
 from lightning_sdk.cli.legacy.clusters_menu import _ClustersMenu
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud.openapi import V1MultiMachineJob
 from lightning_sdk.lit_container import LitContainer
 from lightning_sdk.utils.resolve import _get_authed_user
@@ -48,7 +48,7 @@ def studios(
 ) -> None:
     """List studios for a given teamspace."""
     studios = []
-    menu = _TeamspacesMenu()
+    menu = TeamspacesMenu()
     if all and not teamspace:
         user = _get_authed_user()
         possible_teamspaces = menu._get_possible_teamspaces(user)
@@ -56,7 +56,7 @@ def studios(
             teamspace = Teamspace(**ts)
             studios.extend(teamspace.studios)
     else:
-        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        resolved_teamspace = menu(teamspace=teamspace)
         studios = resolved_teamspace.studios
 
     table = Table(
@@ -111,7 +111,7 @@ def jobs(
 ) -> None:
     """List jobs for a given teamspace."""
     jobs = []
-    menu = _TeamspacesMenu()
+    menu = TeamspacesMenu()
     if all and not teamspace:
         user = _get_authed_user()
         possible_teamspaces = menu._get_possible_teamspaces(user)
@@ -119,7 +119,7 @@ def jobs(
             teamspace = Teamspace(**ts)
             jobs.extend(teamspace.jobs)
     else:
-        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        resolved_teamspace = menu(teamspace=teamspace)
         jobs = resolved_teamspace.jobs
 
     table = Table(pad_edge=True)
@@ -182,7 +182,7 @@ def mmts(
 ) -> None:
     """List multi-machine jobs for a given teamspace."""
     jobs = []
-    menu = _TeamspacesMenu()
+    menu = TeamspacesMenu()
     if all and not teamspace:
         user = _get_authed_user()
         possible_teamspaces = menu._get_possible_teamspaces(user)
@@ -190,7 +190,7 @@ def mmts(
             teamspace = Teamspace(**ts)
             jobs.extend(teamspace.multi_machine_jobs)
     else:
-        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        resolved_teamspace = menu(teamspace=teamspace)
         jobs = resolved_teamspace.multi_machine_jobs
 
     table = Table(pad_edge=True)
@@ -242,9 +242,9 @@ def mmts(
 def containers(teamspace: Optional[str] = None, cloud_account: Optional[str] = None) -> None:
     """Display the list of available containers."""
     api = LitContainer()
-    menu = _TeamspacesMenu()
+    menu = TeamspacesMenu()
     clusters_menu = _ClustersMenu()
-    resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+    resolved_teamspace = menu(teamspace=teamspace)
 
     if not cloud_account:
         cloud_account = clusters_menu._resolve_cluster(resolved_teamspace)

lightning_sdk/cli/legacy/open.py

@@ -6,8 +6,8 @@ from typing import Optional
 import click
 from rich.console import Console
 
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
 from lightning_sdk.cli.legacy.upload import _upload_folder
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.studio import Studio
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.utils.resolve import _get_studio_url
@@ -49,8 +49,8 @@ def open(path: str = ".", teamspace: Optional[str] = None, cloud_account: Option
     try:
         resolved_teamspace = Teamspace()
     except ValueError:
-        menu = _TeamspacesMenu()
-        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        menu = TeamspacesMenu()
+        resolved_teamspace = menu(teamspace=teamspace)
 
     # default cloud account to current studios cloud account if run from studio
     # else it will fall back to teamspace default in the backend

lightning_sdk/cli/legacy/start.py

@@ -103,4 +103,5 @@ def studio(
     except KeyError:
         resolved_machine = machine
 
+    Studio.show_progress = True
     studio.start(resolved_machine)

lightning_sdk/cli/legacy/switch.py

@@ -59,4 +59,5 @@ def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU") ->
     except KeyError:
         resolved_machine = machine
 
+    Studio.show_progress = True
     studio.switch_machine(resolved_machine)

lightning_sdk/cli/legacy/upload.py

@@ -16,7 +16,7 @@ from lightning_sdk.api.lit_container_api import DockerNotRunningError, LCRAuthFa
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.models import upload_model as _upload_model
 from lightning_sdk.studio import Studio
@@ -137,8 +137,8 @@ def upload_container(
     platform: Optional[str] = "linux/amd64",
 ) -> None:
     """Upload a container to Lightning AI's container registry."""
-    menu = _TeamspacesMenu()
-    teamspace = menu._resolve_teamspace(teamspace)
+    menu = TeamspacesMenu()
+    teamspace = menu(teamspace)
     console = Console()
     with Progress(
         SpinnerColumn(),

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,29 @@ 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)
     click.echo(f"Studio '{studio.name}' started successfully")

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,33 +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