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

lightning_sdk/__init__.py

@@ -9,7 +9,7 @@ from lightning_sdk.mmt import MMT
 from lightning_sdk.organization import Organization
 from lightning_sdk.plugin import JobsPlugin, MultiMachineTrainingPlugin, Plugin, SlurmJobsPlugin
 from lightning_sdk.status import Status
-from lightning_sdk.studio import Studio
+from lightning_sdk.studio import VM, Studio
 from lightning_sdk.teamspace import ConnectionType, FolderLocation, Teamspace
 from lightning_sdk.user import User
 
@@ -32,8 +32,9 @@ __all__ = [
     "Studio",
     "Teamspace",
     "User",
+    "VM",
 ]
 
-__version__ = "2025.09.29"
+__version__ = "2025.10.08"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/base_studio.py

@@ -16,6 +16,8 @@ class BaseStudioInfo:
     name: str
     managed_id: str
     description: str
+    creator: str
+    enabled: bool
 
 
 class BaseStudio:
@@ -80,20 +82,29 @@ class BaseStudio:
             disabled=disabled,
         )
 
-    def list(self, managed: bool = True) -> List[BaseStudioInfo]:
+    def list(self, managed: bool = True, include_disabled: bool = False) -> List[BaseStudioInfo]:
         """List all base studios in the organization.
 
+        Args:
+            managed: Whether to filter for managed base studios.
+            include_disabled: Whether to include disabled base studios in the results.
+
         Returns:
-            List[V1CloudSpaceEnvironmentTemplate]: A list of base studio templates.
+            List[BaseStudioInfo]: A list of base studio templates.
         """
-        result = []
-        for template in self._base_studio_api.get_all_base_studios(self._org.id, managed).templates:
-            result.append(
-                BaseStudioInfo(
-                    id=template.id,
-                    name=template.name,
-                    managed_id=template.managed_id,
-                    description=template.description,
-                ),
+        templates = self._base_studio_api.get_all_base_studios(self._org.id, managed).templates
+
+        return [
+            BaseStudioInfo(
+                id=template.id,
+                name=template.name,
+                managed_id=template.managed_id,
+                description=template.description,
+                creator="⚡ Lightning AI"
+                if template.managed_id
+                else UserApi()._get_user_by_id(template.user_id).username,
+                enabled=not template.disabled,
             )
-        return result
+            for template in templates
+            if include_disabled or not template.disabled
+        ]

lightning_sdk/cli/base_studio/__init__.py

@@ -0,0 +1,10 @@
+"""Base Studio CLI commands."""
+
+import click
+
+
+def register_commands(group: click.Group) -> None:
+    """Register base studio commands with the given group."""
+    from lightning_sdk.cli.base_studio.list import list_base_studios
+
+    group.add_command(list_base_studios)

lightning_sdk/cli/base_studio/list.py

@@ -0,0 +1,45 @@
+"""Base Studio list command."""
+
+import click
+from rich.table import Table
+
+from lightning_sdk.base_studio import BaseStudio
+from lightning_sdk.cli.utils.richt_print import rich_to_str
+
+
+@click.command("list")
+@click.option("--include-disabled", help="Include disabled Base Studios in the list.", is_flag=True)
+def list_base_studios(include_disabled: bool) -> None:
+    """List Base Studios in an org.
+
+    Example:
+        lightning base-studio list
+
+    """
+    return list_impl(include_disabled=include_disabled)
+
+
+def list_impl(include_disabled: bool) -> None:
+    base_studio_cls = BaseStudio()
+    base_studios = base_studio_cls.list(include_disabled=include_disabled) + base_studio_cls.list(
+        managed=False, include_disabled=include_disabled
+    )
+
+    table = Table(
+        pad_edge=True,
+    )
+
+    table.add_column("Name")
+    table.add_column("Description")
+    table.add_column("Creator")
+    table.add_column("Enabled")
+
+    for base_studio in base_studios:
+        table.add_row(
+            base_studio.name.lower().replace(" ", "-"),
+            base_studio.description or "",
+            base_studio.creator,
+            "Yes" if base_studio.enabled else "No",
+        )
+
+    click.echo(rich_to_str(table), color=True)

lightning_sdk/cli/entrypoint.py

@@ -17,6 +17,7 @@ from lightning_sdk.api.studio_api import _cloud_url
 
 # Import legacy groups directly from groups.py
 from lightning_sdk.cli.groups import (
+    base_studio,
     config,
     # job,
     # mmt,
@@ -85,6 +86,7 @@ main_cli.add_command(config)
 # main_cli.add_command(mmt)
 main_cli.add_command(studio)
 main_cli.add_command(vm)
+main_cli.add_command(base_studio)
 if os.environ.get("LIGHTNING_EXPERIMENTAL_CLI_ONLY", "0") != "1":
     #### LEGACY COMMANDS ####
     # these commands are currently supported for backwards compatibility, but will potentially be removed in the future.

lightning_sdk/cli/groups.py

@@ -2,6 +2,7 @@
 
 import click
 
+from lightning_sdk.cli.base_studio import register_commands as register_base_studio_commands
 from lightning_sdk.cli.config import register_commands as register_config_commands
 from lightning_sdk.cli.job import register_commands as register_job_commands
 from lightning_sdk.cli.mmt import register_commands as register_mmt_commands
@@ -34,9 +35,15 @@ def vm() -> None:
     """Manage Lightning AI VMs."""
 
 
+@click.group(name="base-studio")
+def base_studio() -> None:
+    """Manage Lightning AI Base Studios."""
+
+
 # Register config commands with the main config group
 register_job_commands(job)
 register_mmt_commands(mmt)
 register_studio_commands(studio)
 register_config_commands(config)
 register_vm_commands(vm)
+register_base_studio_commands(base_studio)

lightning_sdk/cli/legacy/download.py

@@ -14,7 +14,7 @@ from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
 from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.models import download_model
 from lightning_sdk.studio import Studio
-from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
+from lightning_sdk.utils.resolve import _get_authed_user
 
 
 def _expand_remote_path(path: str) -> str:
@@ -242,7 +242,7 @@ def _resolve_studio(studio: Optional[str]) -> Studio:
                     and (re.match(studio_name, st["name"]) or studio_name in st["name"]),
                     possible_studios,
                 )
-                if not possible_studios:
+                if not len(possible_studios):
                     raise ValueError(
                         f"Could not find Studio like '{studio}', please consider update your filtering pattern."
                     )
@@ -260,8 +260,7 @@ def _resolve_studio(studio: Optional[str]) -> Studio:
             "Please contact Lightning AI directly to resolve this issue."
         ) from e
 
-    with skip_studio_init():
-        return Studio(**selected_studio)
+    return Studio(**selected_studio)
 
 
 @download.command(name="licenses")

lightning_sdk/cli/legacy/upload.py

@@ -20,7 +20,7 @@ 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
-from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
+from lightning_sdk.utils.resolve import _get_authed_user
 
 _STUDIO_UPLOAD_STATUS_PATH = "~/.lightning/studios/uploads"
 
@@ -285,8 +285,7 @@ def _resolve_studio(studio: Optional[str]) -> Studio:
             "Please contact Lightning AI directly to resolve this issue."
         ) from e
 
-    with skip_studio_init():
-        return Studio(**selected_studio)
+    return Studio(**selected_studio)
 
 
 def _print_docker_push(lines: Generator, console: Console, progress: Progress, push_task: rich.progress.TaskID) -> None:

lightning_sdk/cli/studio/__init__.py

@@ -5,6 +5,7 @@ import click
 
 def register_commands(group: click.Group) -> None:
     """Register studio commands with the given group."""
+    from lightning_sdk.cli.studio.connect import connect_studio
     from lightning_sdk.cli.studio.create import create_studio
     from lightning_sdk.cli.studio.delete import delete_studio
     from lightning_sdk.cli.studio.list import list_studios
@@ -20,3 +21,4 @@ def register_commands(group: click.Group) -> None:
     group.add_command(start_studio)
     group.add_command(stop_studio)
     group.add_command(switch_studio)
+    group.add_command(connect_studio)

lightning_sdk/cli/studio/connect.py

@@ -0,0 +1,189 @@
+"""Studio connect command."""
+
+import subprocess
+import sys
+from typing import Dict, Optional, Set
+
+import click
+
+from lightning_sdk.base_studio import BaseStudio
+from lightning_sdk.cli.utils.richt_print import studio_name_link
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config, save_teamspace_to_config
+from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import CloudProvider, Machine
+from lightning_sdk.studio import Studio
+from lightning_sdk.utils.names import random_unique_name
+
+DEFAULT_MACHINE = "CPU"
+
+
+def _split_gpus_spec(gpus: str) -> tuple[str, int]:
+    machine_name, machine_val = gpus.split(":", 1)
+    machine_name = machine_name.strip()
+    machine_val = machine_val.strip()
+
+    if not machine_val.isdigit() or int(machine_val) <= 0:
+        raise ValueError(f"Invalid GPU count '{machine_val}'. Must be a positive integer.")
+
+    machine_num = int(machine_val)
+    return machine_name, machine_num
+
+
+def _construct_available_gpus(machine_options: Dict[str, str]) -> Set[str]:
+    # returns available gpus:count
+    available_gpus = set()
+    for v in machine_options.values():
+        if "_X_" in v:
+            gpu_type_num = v.replace("_X_", ":")
+            available_gpus.add(gpu_type_num)
+        else:
+            available_gpus.add(v)
+    return available_gpus
+
+
+def _get_machine_from_gpus(gpus: str) -> Machine:
+    machine_name = gpus
+    machine_num = 1
+
+    if ":" in gpus:
+        machine_name, machine_num = _split_gpus_spec(gpus)
+
+    machine_options = {
+        m.name.lower(): m.name for m in Machine.__dict__.values() if isinstance(m, Machine) and m._include_in_cli
+    }
+
+    if machine_num == 1:
+        # e.g. gpus=L4 or gpus=L4:1
+        gpu_key = machine_name.lower()
+        try:
+            return machine_options[gpu_key]
+        except KeyError:
+            available = ", ".join(_construct_available_gpus(machine_options))
+            raise ValueError(f"Invalid GPU type '{machine_name}'. Available options: {available}") from None
+
+    # Else: e.g. gpus=L4:4
+    gpu_key = f"{machine_name.lower()}_x_{machine_num}"
+    try:
+        return machine_options[gpu_key]
+    except KeyError:
+        available = ", ".join(_construct_available_gpus(machine_options))
+        raise ValueError(f"Invalid GPU configuration '{gpus}'. Available options: {available}") from None
+
+
+def _get_base_studio_id(studio_type: Optional[str]) -> Optional[str]:
+    base_studios = BaseStudio()
+    base_studios = base_studios.list()
+    template_id = None
+
+    if base_studios and len(base_studios):
+        # if not specified by user, use the first existing template studio
+        template_id = base_studios[0].id
+        # else, try to match the provided studio_type to base studio name
+        if studio_type:
+            normalized_studio_type = studio_type.lower().replace(" ", "-")
+            match = next(
+                (s for s in base_studios if s.name.lower().replace(" ", "-") == normalized_studio_type),
+                None,
+            )
+            if match:
+                template_id = match.id
+
+    return template_id
+
+
+@click.command("connect")
+@click.argument("name", required=False)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option(
+    "--cloud-provider",
+    help="The cloud provider to start the studio 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 studio on. Defaults to teamspace default.",
+    type=click.STRING,
+)
+@click.option(
+    "--machine",
+    help="The machine type to start the studio 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(
+    "--gpus",
+    help="The number and type of GPUs to start the studio on (format: TYPE:COUNT, e.g. L4:4)",
+    type=click.STRING,
+)
+@click.option(
+    "--studio-type",
+    help="The base studio template name to use for creating the studio. "
+    "Must be lowercase and hyphenated (use '-' instead of spaces). "
+    "Run 'lightning base-studio list' to see all available templates. "
+    "Defaults to the first available template.",
+    type=click.STRING,
+)
+def connect_studio(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    cloud_provider: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    machine: Optional[str] = None,
+    gpus: Optional[str] = None,
+    studio_type: Optional[str] = None,
+) -> None:
+    """Connect to a Studio.
+
+    Example:
+        lightning studio connect
+    """
+    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)
+
+    name = name or random_unique_name()
+
+    # check for available base studios
+    template_id = _get_base_studio_id(studio_type)
+
+    try:
+        studio = Studio(
+            name=name,
+            teamspace=resolved_teamspace,
+            create_ok=True,
+            cloud_provider=cloud_provider,
+            cloud_account=cloud_account,
+            studio_type=template_id,
+        )
+    except (RuntimeError, ValueError, ApiException):
+        raise ValueError(f"Could not create Studio: '{name}'") from None
+
+    click.echo(f"Connecting to Studio '{studio_name_link(studio)}' ...")
+
+    Studio.show_progress = True
+
+    if machine and gpus:
+        raise click.UsageError("Options --machine and --gpu are mutually exclusive. Provide only one.")
+    elif gpus:
+        machine = _get_machine_from_gpus(gpus.strip())
+    elif not machine:
+        machine = DEFAULT_MACHINE
+
+    save_studio_to_config(studio)
+    # by default, interruptible is False
+    studio.start(machine=machine, interruptible=False)
+
+    ssh_private_key_path = configure_ssh_internal()
+
+    ssh_option = "-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o LogLevel=ERROR"
+    try:
+        ssh_command = f"ssh -i {ssh_private_key_path} {ssh_option} s_{studio._studio.id}@ssh.lightning.ai"
+        subprocess.run(ssh_command.split())
+    except Exception as ex:
+        print(f"Failed to establish SSH connection: {ex}")
+        sys.exit(1)

lightning_sdk/cli/studio/ssh.py

@@ -1,19 +1,14 @@
 """Studio SSH command."""
 
-import os
-import platform
 import subprocess
-import uuid
-from pathlib import Path
 from typing import List, Optional
 
 import click
 
 from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
+from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal
 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.utils.config import _DEFAULT_CONFIG_FILE_PATH
 
 
 @click.command("ssh")
@@ -43,9 +38,7 @@ def ssh_studio(name: Optional[str] = None, teamspace: Optional[str] = None, opti
 
 
 def ssh_impl(name: Optional[str], teamspace: Optional[str], option: Optional[List[str]], vm: bool) -> None:
-    auth = Auth()
-    auth.authenticate()
-    ssh_private_key_path = _download_ssh_keys(auth.api_key, force_download=False)
+    ssh_private_key_path = configure_ssh_internal()
 
     menu = TeamspacesMenu()
     resolved_teamspace = menu(teamspace=teamspace)
@@ -63,53 +56,9 @@ def ssh_impl(name: Optional[str], teamspace: Optional[str], option: Optional[Lis
         subprocess.run(ssh_command.split())
     except Exception:
         # redownload the keys to be sure they are up to date
-        _download_ssh_keys(auth.api_key, force_download=True)
+        ssh_private_key_path = configure_ssh_internal(force_download=True)
         try:
             subprocess.run(ssh_command.split())
         except Exception:
             # TODO: make this a generic CLI error
             raise RuntimeError("Failed to establish SSH connection") from None
-
-
-def _download_ssh_keys(
-    api_key: str,
-    force_download: bool = False,
-    ssh_key_name: str = "lightning_rsa",
-) -> None:
-    """Download the SSH key for a User."""
-    ssh_private_key_path = os.path.join(os.path.expanduser(os.path.dirname(_DEFAULT_CONFIG_FILE_PATH)), ssh_key_name)
-
-    os.makedirs(os.path.dirname(ssh_private_key_path), exist_ok=True)
-
-    if not os.path.isfile(ssh_private_key_path) or force_download:
-        key_id = str(uuid.uuid4())
-        _download_file(
-            f"https://lightning.ai/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
-            ssh_private_key_path,
-            overwrite=True,
-            chmod=0o600,
-        )
-        _download_file(
-            f"https://lightning.ai/setup/ssh-public?t={api_key}&id={key_id}",
-            ssh_private_key_path + ".pub",
-            overwrite=True,
-        )
-
-    return ssh_private_key_path
-
-
-def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
-    """Download a file from a URL."""
-    import requests
-
-    if os.path.isfile(local_path) and not overwrite:
-        raise FileExistsError(f"The file {local_path} already exists and overwrite is set to False.")
-
-    response = requests.get(url, stream=True)
-    response.raise_for_status()
-
-    with open(local_path, "wb") as file:
-        for chunk in response.iter_content(chunk_size=8192):
-            file.write(chunk)
-    if chmod is not None:
-        os.chmod(local_path, 0o600)

lightning_sdk/cli/utils/ssh_connection.py

@@ -0,0 +1,59 @@
+import os
+import platform
+import uuid
+from pathlib import Path
+from typing import Optional
+
+from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH
+
+
+def configure_ssh_internal(force_download: bool = False) -> str:
+    """Internal function to configure SSH without Click decorators."""
+    auth = Auth()
+    auth.authenticate()
+    return download_ssh_keys(auth.api_key, force_download=force_download)
+
+
+def download_ssh_keys(
+    api_key: str | None,
+    force_download: bool = False,
+    ssh_key_name: str = "lightning_rsa",
+) -> str:
+    """Download the SSH key for a User."""
+    ssh_private_key_path = os.path.join(os.path.expanduser(os.path.dirname(_DEFAULT_CONFIG_FILE_PATH)), ssh_key_name)
+
+    os.makedirs(os.path.dirname(ssh_private_key_path), exist_ok=True)
+
+    if not os.path.isfile(ssh_private_key_path) or force_download:
+        key_id = str(uuid.uuid4())
+        download_file(
+            f"https://lightning.ai/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
+            Path(ssh_private_key_path),
+            overwrite=True,
+            chmod=0o600,
+        )
+        download_file(
+            f"https://lightning.ai/setup/ssh-public?t={api_key}&id={key_id}",
+            Path(ssh_private_key_path + ".pub"),
+            overwrite=True,
+        )
+
+    return ssh_private_key_path
+
+
+def download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
+    """Download a file from a URL."""
+    import requests
+
+    if os.path.isfile(local_path) and not overwrite:
+        raise FileExistsError(f"The file {local_path} already exists and overwrite is set to False.")
+
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+
+    with open(local_path, "wb") as file:
+        for chunk in response.iter_content(chunk_size=8192):
+            file.write(chunk)
+    if chmod is not None:
+        os.chmod(local_path, 0o600)

lightning_sdk/studio.py

@@ -8,6 +8,7 @@ from tqdm.auto import tqdm
 
 from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.api.studio_api import StudioApi
+from lightning_sdk.base_studio import BaseStudio
 from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.lightning_cloud.openapi import V1ClusterType
 from lightning_sdk.machine import CloudProvider, Machine
@@ -51,6 +52,8 @@ class Studio:
             If not specified, falls backto the teamspace default cloud account.
         create_ok: whether the studio will be created if it does not yet exist. Defaults to True
         provider: the provider of the machine, the studio should be created on.
+        studio_type: Type of studio to create. Only effective during initial creation;
+            ignored for existing studios.
 
     Note:
         Since a teamspace can either be owned by an org or by a user directly,
@@ -78,6 +81,7 @@ class Studio:
         source: Optional[str] = None,
         disable_secrets: bool = False,
         provider: Optional[Union[CloudProvider, str]] = None,  # deprecated in favor of cloud_provider
+        studio_type: Optional[str] = None,  # for base studio templates
     ) -> None:
         self._studio_api = StudioApi()
         self._cloud_account_api = CloudAccountApi()
@@ -113,6 +117,25 @@ class Studio:
                 default_cloud_account=self._teamspace.default_cloud_account,
             )
 
+        self._studio_type = None
+        if studio_type:
+            self._base_studio = BaseStudio()
+            self._available_base_studios = self._base_studio.list()
+            for bst in self._available_base_studios:
+                if (
+                    bst.id == studio_type
+                    or bst.name == studio_type
+                    or bst.name.lower().replace(" ", "-") == studio_type
+                ):
+                    self._studio_type = bst.id
+
+            if not self._studio_type:
+                raise ValueError(
+                    f"Could not find studio type with ID or name '{studio_type}'. "
+                    f"Available studio types: "
+                    f"{[bst.name.lower().replace(' ', '-') for bst in self._available_base_studios]}"
+                )
+
         # Resolve studio name if not provided: explicit → env (LIGHTNING_CLOUD_SPACE_ID) → config defaults
         if name is None and not getattr(self._skip_init, "value", False):
             studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
@@ -150,6 +173,7 @@ class Studio:
                         cloud_account=_cloud_account,
                         source=source,
                         disable_secrets=self._disable_secrets,
+                        cloud_space_environment_template_id=self._studio_type,
                     )
                 else:
                     raise e
@@ -264,10 +288,11 @@ class Studio:
             else:
                 interruptible = self.teamspace.start_studios_on_interruptible
 
+        new_machine = machine
+        if not isinstance(machine, Machine):
+            new_machine = Machine.from_str(machine)
+
         if status == Status.Running:
-            new_machine = machine
-            if not isinstance(machine, Machine):
-                new_machine = Machine.from_str(machine)
             if new_machine != self.machine:
                 raise RuntimeError(
                     f"Requested to start {self._cls_name} on {new_machine}, "
@@ -289,7 +314,11 @@ class Studio:
             with StudioProgressTracker("start", show_progress=True) as progress:
                 # Start the studio without blocking
                 self._studio_api.start_studio_async(
-                    self._studio.id, self._teamspace.id, machine, interruptible=interruptible, max_runtime=max_runtime
+                    self._studio.id,
+                    self._teamspace.id,
+                    new_machine,
+                    interruptible=interruptible,
+                    max_runtime=max_runtime,
                 )
 
                 # Track progress through completion
@@ -299,7 +328,7 @@ class Studio:
         else:
             # Use the blocking version if no progress is needed
             self._studio_api.start_studio(
-                self._studio.id, self._teamspace.id, machine, interruptible=interruptible, max_runtime=max_runtime
+                self._studio.id, self._teamspace.id, new_machine, interruptible=interruptible, max_runtime=max_runtime
             )
 
         self._setup()

lightning_sdk/utils/progress.py

@@ -17,30 +17,31 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_get_cloud_space_instance_st
 class StartupPhase(Enum):
     """Studio startup phase messages."""
 
-    ALLOCATING_MACHINE = "Allocating machine from cloud provider..."
-    STUDIO_STARTING = "Studio is starting up..."
-    SETTING_UP_ENVIRONMENT = "Setting up Studio environment..."
-    RESTORING_STATE = "Restoring Studio state..."
-    FINALIZING_SETUP = "Finalizing Studio setup..."
-    COMPLETED = "Studio started successfully"
+    STARTING_STUDIO = "Starting Studio..."
+    GETTING_MACHINE = "Getting a machine..."
 
+    SWITCHING_STUDIO = "Switching Studio..."
 
-def get_switching_progress_message(percentage: int, is_base_studio: bool, is_new_cloud_space: bool) -> str:
+    SETTING_UP_MACHINE = "Setting up machine..."
+    RESTORING_STUDIO = "Restoring Studio..."
+    PREPARING_STUDIO = "Preparing Studio..."
+    RESTORING_BASE_STUDIO = "Restoring Base Studio..."
+    SETTING_UP_BASE_STUDIO = "Setting up Base Studio..."
+    DONE = "Done"
+
+
+def get_switching_progress_message(percentage: int, is_base_studio: bool) -> str:
     """Get progress message for switching studios."""
     percentage = max(0, min(100, round(percentage)))
 
     if percentage > 98:
-        message = "Done"
+        message = StartupPhase.DONE.value
     elif percentage > 80:
-        if is_new_cloud_space:
-            message = "Setting up Base Studio..." if is_base_studio else "Preparing Studio..."
-        else:
-            message = "Restoring Studio..."
+        message = StartupPhase.RESTORING_BASE_STUDIO.value if is_base_studio else StartupPhase.RESTORING_STUDIO.value
     elif percentage > 60:
-        message = "Setting up machine from cloud provider"
+        message = StartupPhase.SETTING_UP_MACHINE.value
     else:
-        message = "Allocating machine from cloud provider"
-
+        message = StartupPhase.SWITCHING_STUDIO.value
     return f"({percentage}%) {message}"
 
 
@@ -110,15 +111,13 @@ class StudioProgressTracker:
         if self.progress:
             self.progress.stop()
 
-    def update_progress(
-        self, percentage: int, message: str = "", is_base_studio: bool = False, is_new_cloud_space: bool = False
-    ) -> None:
+    def update_progress(self, percentage: int, message: str = "", is_base_studio: bool = False) -> None:
         """Update progress bar with current percentage and message."""
         if not self.progress or self.task_id is None:
             return
 
         if self.operation_type == "switch":
-            display_message = get_switching_progress_message(percentage, is_base_studio, is_new_cloud_space)
+            display_message = get_switching_progress_message(percentage, is_base_studio)
         else:
             display_message = message or f"{self.operation_type.capitalize()}ing Studio..."
 
@@ -153,7 +152,7 @@ class StudioProgressTracker:
         message_stability_delay = 3.0  # Seconds to wait before changing message
 
         # Show initial progress immediately
-        self.update_progress(5, StartupPhase.ALLOCATING_MACHINE.value)
+        self.update_progress(5, StartupPhase.STARTING_STUDIO.value)
 
         while True:
             try:
@@ -163,7 +162,7 @@ class StudioProgressTracker:
                 # Default fallback progress based on time
                 time_based_progress = min(95, int((elapsed / timeout) * 100))
                 current_progress = max(last_progress, time_based_progress)
-                current_message = StartupPhase.ALLOCATING_MACHINE.value
+                current_message = StartupPhase.STARTING_STUDIO.value
 
                 # Check if we have detailed status information
                 if hasattr(status, "in_use") and status.in_use:
@@ -178,7 +177,7 @@ class StudioProgressTracker:
                             hasattr(startup_status, "top_up_restore_finished")
                             and startup_status.top_up_restore_finished
                         ):
-                            self.complete(StartupPhase.COMPLETED.value)
+                            self.complete(StartupPhase.DONE.value)
                             break
 
                         # Check other phases in descending priority
@@ -186,14 +185,14 @@ class StudioProgressTracker:
                             hasattr(startup_status, "initial_restore_finished")
                             and startup_status.initial_restore_finished
                         ):
-                            current_progress = max(current_progress, 85)
-                            current_message = StartupPhase.FINALIZING_SETUP.value
+                            current_progress = max(current_progress, 80)
+                            current_message = StartupPhase.PREPARING_STUDIO.value
                         elif hasattr(startup_status, "container_ready") and startup_status.container_ready:
-                            current_progress = max(current_progress, 70)
-                            current_message = StartupPhase.RESTORING_STATE.value
-                        elif hasattr(startup_status, "machine_ready") and startup_status.machine_ready:
                             current_progress = max(current_progress, 60)
-                            current_message = StartupPhase.SETTING_UP_ENVIRONMENT.value
+                            current_message = StartupPhase.SETTING_UP_MACHINE.value
+                        elif hasattr(startup_status, "machine_ready") and startup_status.machine_ready:
+                            current_progress = max(current_progress, 30)
+                            current_message = StartupPhase.GETTING_MACHINE.value
 
                     # Check general phase information
                     if hasattr(in_use, "phase") and in_use.phase:
@@ -201,7 +200,7 @@ class StudioProgressTracker:
 
                         if phase == "CLOUD_SPACE_INSTANCE_STATE_RUNNING":
                             current_progress = max(current_progress, 80)
-                            current_message = StartupPhase.STUDIO_STARTING.value
+                            current_message = StartupPhase.SETTING_UP_MACHINE.value
                         elif phase == "CLOUD_SPACE_INSTANCE_STATE_PENDING":
                             # Track time in pending phase for smoother progress
                             if "pending" not in phase_start_times:
@@ -211,7 +210,7 @@ class StudioProgressTracker:
                             # Progress more smoothly through pending phase (10-60%)
                             pending_progress = 10 + min(50, int((pending_elapsed / 60) * 50))
                             current_progress = max(current_progress, pending_progress)
-                            current_message = StartupPhase.ALLOCATING_MACHINE.value
+                            current_message = StartupPhase.GETTING_MACHINE.value
 
                 # Check for requested machine status (pre-allocation)
                 elif hasattr(status, "requested") and status.requested:
@@ -222,7 +221,7 @@ class StudioProgressTracker:
                     # Progress through allocation phase (5-30%)
                     allocation_progress = 5 + min(25, int((allocation_elapsed / 30) * 25))
                     current_progress = max(current_progress, allocation_progress)
-                    current_message = StartupPhase.ALLOCATING_MACHINE.value
+                    current_message = StartupPhase.GETTING_MACHINE.value
 
                 # Ensure progress never decreases and moves smoothly
                 if current_progress > last_progress:
@@ -262,12 +261,12 @@ class StudioProgressTracker:
 
                 # Only update message if enough time has passed
                 current_time = time.time()
-                should_update_message = StartupPhase.ALLOCATING_MACHINE.value != self._last_message and (
+                should_update_message = StartupPhase.GETTING_MACHINE.value != self._last_message and (
                     current_time - last_message_time >= message_stability_delay or last_message_time == 0
                 )
 
                 if should_update_message:
-                    self.update_progress(fallback_progress, StartupPhase.ALLOCATING_MACHINE.value)
+                    self.update_progress(fallback_progress, StartupPhase.GETTING_MACHINE.value)
                     last_message_time = current_time
                 else:
                     # Update progress but keep existing message

{lightning_sdk-2025.9.29.dist-info → lightning_sdk-2025.10.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 2025.9.29
+Version: 2025.10.8
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License

{lightning_sdk-2025.9.29.dist-info → lightning_sdk-2025.10.8.dist-info}/RECORD

@@ -1,8 +1,8 @@
 docs/source/conf.py,sha256=r8yX20eC-4mHhMTd0SbQb5TlSWHhO6wnJ0VJ_FBFpag,13249
-lightning_sdk/__init__.py,sha256=HBcv4oavnBLmLhOoGKn7B71XVE1xFmSPUw0GM-UjU4w,1221
+lightning_sdk/__init__.py,sha256=MLcIAWdrZoDrz0HJ7KhaLRbFrSg8Qs_XFS8p2vr2cwQ,1235
 lightning_sdk/agents.py,sha256=ly6Ma1j0ZgGPFyvPvMN28JWiB9dATIstFa5XM8pMi6I,1577
 lightning_sdk/ai_hub.py,sha256=iI1vNhgcz_Ff1c3rN1ogN7dK-r-HXRj6NMtS2cA14UA,6925
-lightning_sdk/base_studio.py,sha256=_Pwwl37R9GRd7t-f2kO5aQXiLNrP4sUtUNht2ZkP8LE,3678
+lightning_sdk/base_studio.py,sha256=FIjHUiZtRf42m5YeqYYaJxcdgOtJVAJLZU42gzmldO8,4109
 lightning_sdk/constants.py,sha256=ztl1PTUBULnqTf3DyKUSJaV_O20hNtUYT6XvAYIrmIk,749
 lightning_sdk/helpers.py,sha256=KWMWnORHItIIA3PGn71YPs-7RjzGi8IXa2kQ5Qo4U8M,2459
 lightning_sdk/lit_container.py,sha256=8ys49TXI9MO89jLTA7MwDrKrssTRARAIF9OVmolDLq0,5273
@@ -14,7 +14,7 @@ lightning_sdk/plugin.py,sha256=f3P5-xZY6x-MX0Fs2z_Q2erSxPSiHZARO0BVkCezHw4,15192
 lightning_sdk/sandbox.py,sha256=_NvnWotEXW2rBiVFZZ4krKXxVjuAqfNh04qELSM0-Pg,5786
 lightning_sdk/serve.py,sha256=uW7zLhQ3X90ifetpxzTb8FNxifv5vIs7qZlgfEjVKzk,11794
 lightning_sdk/status.py,sha256=lLGAuSvXBoXQFEEsEYwdCi0RcSNatUn5OPjJVjDtoM0,386
-lightning_sdk/studio.py,sha256=_mPNeH_IBoWC-bwEiGXCeWizl1D2YXdWMUVUH3QyFsw,33145
+lightning_sdk/studio.py,sha256=pij-pT-Cd11j5-LMAHxORLRP9a-Gcow87ZkTPcluu38,34350
 lightning_sdk/teamspace.py,sha256=iYaEaZH8N8Yj4XXj6M4zIH0tsjPVrAf_BzAAv6Jt9VI,26556
 lightning_sdk/user.py,sha256=TSYh38rxoi7qKOfrK2JYh_Nknya2Kbz2ngDIY85fFOY,1778
 lightning_sdk/api/__init__.py,sha256=xrp_RNECJGQtL5rZHF69WOzEuEIbWSLtjWAJAz4R5K4,500
@@ -35,8 +35,10 @@ lightning_sdk/api/teamspace_api.py,sha256=UZ6Ka6TL7zggs5qhQB4wl47VRusVTMpth1-toG
 lightning_sdk/api/user_api.py,sha256=o9gZmtvqNfj4kdSpo2fyyRuFAP7bM5w7mx0Oj9__Ads,4702
 lightning_sdk/api/utils.py,sha256=1NJgW4HCfzMqgnAqbMA7RRr2v3iM3KTuPIUQK5klDeQ,27127
 lightning_sdk/cli/__init__.py,sha256=lksw08t_ZIuHOH47LCIqSVHeZ8cUXI2aJWHYhyujYHM,32
-lightning_sdk/cli/entrypoint.py,sha256=kvzeMKNuMoGSfsvg8v9QsBuMHR4U3pcWqBEHfvQQzhM,4595
-lightning_sdk/cli/groups.py,sha256=OPH6PHN82w5ERX00wzmstqU0NjEHmViVLAmy8-y8jQk,1131
+lightning_sdk/cli/entrypoint.py,sha256=wMAyzfKRqZ6EZ73SxpuJDk3u8J3o-I5vF0VnESJh2QE,4646
+lightning_sdk/cli/groups.py,sha256=3l9_U-8T613RJ2fHZy_-iJbtpq4fH8auWH7eFGSdZmg,1373
+lightning_sdk/cli/base_studio/__init__.py,sha256=npBEmXWNVuiawm98TFVNHZf038kD4aW3gzVHsjbUPNI,272
+lightning_sdk/cli/base_studio/list.py,sha256=DCXzSIY4IB0lHEstP0PA-8e7VqQQSUXL6OVv8VvwBKM,1244
 lightning_sdk/cli/config/__init__.py,sha256=_aJ7uZZOdjomZ-ABmz8Cu3fLmnSujnQK4hKAEKM4aF4,395
 lightning_sdk/cli/config/get.py,sha256=JXGRNVGxELS6Es2RESgStzTKnjg5rr-M6yOmx2ty9A8,1606
 lightning_sdk/cli/config/set.py,sha256=9p8q8nUWGWNcAcflQ3hS_tj-59jY7cmkzccuiZl92h4,3117
@@ -50,7 +52,7 @@ lightning_sdk/cli/legacy/connect.py,sha256=0qJ0yvykoPBkmCvvEbWKppt7wlUmRq3LNNGEs
 lightning_sdk/cli/legacy/create.py,sha256=ftp7cP8dqtO3h35mZUkpUjFW1tHwKxPQoUOPsabcy0M,3682
 lightning_sdk/cli/legacy/delete.py,sha256=IC04CNqJFIEfgl6Tlt4uEl3t81yZCb8VYLRwy3B8xiA,3816
 lightning_sdk/cli/legacy/docker_cli.py,sha256=EozLC4qnvHhgukmnu35itfI5n7v-abCpwKPkPy3eOV4,997
-lightning_sdk/cli/legacy/download.py,sha256=7ZOkqnAIEww12LM_cDO8hBhVbMYkChl82-vXHixcNwg,11357
+lightning_sdk/cli/legacy/download.py,sha256=_Yaw1slwjE7ygDomjEJliXLTE-Gfj1d3zkH3daYVASk,11311
 lightning_sdk/cli/legacy/entrypoint.py,sha256=dk_BKvNzQMYREdAGTTdHF8nZWp_EuvCii-WOz4YFwrk,3784
 lightning_sdk/cli/legacy/exceptions.py,sha256=QUF3OMAMZwBikvlusimSHSBjb6ywvHpfAumJBEaodSw,169
 lightning_sdk/cli/legacy/generate.py,sha256=3tLaqALiuWG3AEIQtNHChMNpaJJU9zV5VYGshiEmKug,1582
@@ -66,17 +68,18 @@ lightning_sdk/cli/legacy/stop.py,sha256=jSiuxIunDnL3A3DQ1m8GZx01Qt4Zx5vIA3u7U81C
 lightning_sdk/cli/legacy/studios_menu.py,sha256=TA9rO6_fFHGMz0Nt4rJ6iV80X5pZE4xShrSiyXoU-oQ,4129
 lightning_sdk/cli/legacy/switch.py,sha256=G3y9BmP-0d_OKJub5fMZ0ddnGZkOFBtregMkx_KWxcw,1945
 lightning_sdk/cli/legacy/teamspace_menu.py,sha256=u5_lyB1aCD8KJIsOuxi0bvPi75McU5EbluMLyD_NBUw,4104
-lightning_sdk/cli/legacy/upload.py,sha256=7FTUnA4CvE7FaNTmNKhjmr-7Ly72QGvDQVqvp1EeQpY,13773
+lightning_sdk/cli/legacy/upload.py,sha256=fFJct9F_1vLqTMLtSjuv4BN3AbUpeGebwv3ZXH0NiqI,13722
 lightning_sdk/cli/legacy/deploy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 lightning_sdk/cli/legacy/deploy/_auth.py,sha256=wpwhkF2tmURlhEBaZ3QgOnHQELy3XLPt2d4LWb7JosY,7700
 lightning_sdk/cli/legacy/deploy/devbox.py,sha256=SdOtYdGefzldn-WCa6sz31zysyf2nP5D7O2x1DE9z_E,6839
 lightning_sdk/cli/legacy/deploy/serve.py,sha256=bxN04mHlArjHtxR_QkCM0S_RPmL3WIPsgnPgFadSXwI,15478
 lightning_sdk/cli/mmt/__init__.py,sha256=CLgr-ZHHLS6Db_JGqpxbn4G2pYrKi4Qn-uhi8e0kFNc,145
-lightning_sdk/cli/studio/__init__.py,sha256=Reb5a0iHJfNMxmnP30DAmYqpJ-YDQoJcgyVrD2cYDNo,823
+lightning_sdk/cli/studio/__init__.py,sha256=jqUuEniYlf3fwJni-MGLII3oprpsBSQEFqzBJwPIxMk,925
+lightning_sdk/cli/studio/connect.py,sha256=yb4Ga4lLjrg5C0sLPwAoEkkTpT7rv165AwqBOsYBIWc,6488
 lightning_sdk/cli/studio/create.py,sha256=kTxLtvTFjo-SgkdzOIngdXu_6yYLHwVbkfMhHjpIFkI,2505
 lightning_sdk/cli/studio/delete.py,sha256=Gi1KLCrSgjTiyoQe9NRvF4ZpsNyeM07KHnBxfdJg2FE,1483
 lightning_sdk/cli/studio/list.py,sha256=ieLiHRIfAb7E5e_r48BeTC03ib8RY6maeF7nkmrqaek,3301
-lightning_sdk/cli/studio/ssh.py,sha256=5CRw62p9rUb4BIL1WZP-Ehzcgfq_hVkY5ixTurxUKKo,3866
+lightning_sdk/cli/studio/ssh.py,sha256=zHpVUe2QiFvM91Pj5kzNCnacavwrptxNsPpxyxpKLno,2194
 lightning_sdk/cli/studio/start.py,sha256=GHfV_1I1XwV8UHOWBP2WVfd1B2BHyw1lQ9HMv6Eomes,3236
 lightning_sdk/cli/studio/stop.py,sha256=hVXIyrv68rvrwIQgyEPUKt3Th8S1iipOVflopPnoVOY,1398
 lightning_sdk/cli/studio/switch.py,sha256=x6F1biDtfR0dk1t7X3NSkmfj5jkzGJe6-nvj3SonQlQ,2069
@@ -87,6 +90,7 @@ lightning_sdk/cli/utils/owner_selection.py,sha256=PbKbFoGblxi60lF_kZxnlqxjCmI9r9
 lightning_sdk/cli/utils/resolve.py,sha256=RWQ8MgXEJCdsoY50y8Pa9B-A6NA_SvxOqL5pVKfvSr8,917
 lightning_sdk/cli/utils/richt_print.py,sha256=psSY65nLAQbxK_K0w93Qq9857aUUbm77FLS1sc3Oecg,1262
 lightning_sdk/cli/utils/save_to_config.py,sha256=mMjjfA4Yv1e7MtE03iADTXcTSAuUL0Ws9OZObx6ufo0,1164
+lightning_sdk/cli/utils/ssh_connection.py,sha256=cfppkWOLRLVSE-LhkENRANgdPNpT3U6beg8Lyx1xoxo,1984
 lightning_sdk/cli/utils/studio_selection.py,sha256=z0SMd82sOt_mfTiQZgVY27MxLkMAjM__SWhlOlfXJT0,4347
 lightning_sdk/cli/utils/teamspace_selection.py,sha256=n-FFfWLTV63SxoK6Fi2VQikBIfKYqMNTew9b9v-1UQk,4945
 lightning_sdk/cli/vm/__init__.py,sha256=tAJpluJVDOF3fUQyRwko5ZYIpqglczMyRv9Kz4OVaIk,711
@@ -1203,11 +1207,11 @@ lightning_sdk/utils/config.py,sha256=LfCAbko_f-1TgFoYD1KT68tZHEIgn3M2y1HUqTMruUo
 lightning_sdk/utils/dynamic.py,sha256=glUTO1JC9APtQ6Gr9SO02a3zr56-sPAXM5C3NrTpgyQ,1959
 lightning_sdk/utils/enum.py,sha256=h2JRzqoBcSlUdanFHmkj_j5DleBHAu1esQYUsdNI-hU,4106
 lightning_sdk/utils/names.py,sha256=1EuXbIh7wldkDp1FG10oz9vIOyWrpGWeFFVy-DQBgzA,18162
-lightning_sdk/utils/progress.py,sha256=IXfEcUF-rL5jIw0Hir6eSxN7VBZfR--1O2LaEhGAU70,12698
+lightning_sdk/utils/progress.py,sha256=bLWw39fzq29PMWoFXaPIVfoS3Ug245950oWOFJ2ZaiU,12596
 lightning_sdk/utils/resolve.py,sha256=4TyEnIgIrvkSvYk5i5PmcIogD_5Y9pBhiphRLfLMttc,10477
-lightning_sdk-2025.9.29.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
-lightning_sdk-2025.9.29.dist-info/METADATA,sha256=CYT1gpwGdPiXJFMDKm0cmhiilbuIEgxvrFT0ndOwhuc,4130
-lightning_sdk-2025.9.29.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-lightning_sdk-2025.9.29.dist-info/entry_points.txt,sha256=OoZa4Fc8NMs6GSN0cdA1J8e6couzAcL82CbM1yo4f_M,122
-lightning_sdk-2025.9.29.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
-lightning_sdk-2025.9.29.dist-info/RECORD,,
+lightning_sdk-2025.10.8.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
+lightning_sdk-2025.10.8.dist-info/METADATA,sha256=N8b-kMsJedFH9daVWxNta30bWOraumpFMRESf8W-Sh4,4130
+lightning_sdk-2025.10.8.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+lightning_sdk-2025.10.8.dist-info/entry_points.txt,sha256=OoZa4Fc8NMs6GSN0cdA1J8e6couzAcL82CbM1yo4f_M,122
+lightning_sdk-2025.10.8.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
+lightning_sdk-2025.10.8.dist-info/RECORD,,