lightning-sdk 2025.7.17__py3-none-any.whl → 2025.7.30rc0__py3-none-any.whl

lightning_sdk/api/studio_api.py

@@ -4,7 +4,7 @@ import tempfile
 import time
 import zipfile
 from threading import Event, Thread
-from typing import Any, Dict, Generator, Mapping, Optional, Tuple, Union
+from typing import Any, Dict, Generator, List, Mapping, Optional, Tuple, Union
 
 import backoff
 import requests
@@ -37,6 +37,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1CloudSpaceSeedFile,
     V1CloudSpaceSourceType,
     V1CloudSpaceState,
+    V1ClusterAccelerator,
     V1EndpointType,
     V1GetCloudSpaceInstanceStatusResponse,
     V1GetLongRunningCommandInCloudSpaceResponse,
@@ -121,7 +122,7 @@ class StudioApi:
         name: str,
         teamspace_id: str,
         cloud_account: Optional[str] = None,
-        source: Optional[V1CloudSpaceSourceType] = None,
+        source: Optional[Union[V1CloudSpaceSourceType, str]] = None,
         disable_secrets: bool = False,
         sandbox: bool = False,
         cloud_space_environment_template_id: Optional[str] = None,
@@ -168,12 +169,26 @@ class StudioApi:
         return startup_status and startup_status.top_up_restore_finished
 
     def start_studio(
-        self, studio_id: str, teamspace_id: str, machine: Union[Machine, str], interruptible: False
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        machine: Union[Machine, str],
+        interruptible: bool = False,
+        max_runtime: Optional[int] = None,
     ) -> None:
         """Start an existing Studio."""
+        # need to go via kwargs for typing compatibility since autogenerated apis accept None but aren't typed with None
+        optional_kwargs_compute_body = {}
+
+        if max_runtime is not None:
+            optional_kwargs_compute_body["requested_run_duration_seconds"] = str(max_runtime)
         self._client.cloud_space_service_start_cloud_space_instance(
             IdStartBody(
-                compute_config=V1UserRequestedComputeConfig(name=_machine_to_compute_name(machine), spot=interruptible)
+                compute_config=V1UserRequestedComputeConfig(
+                    name=_machine_to_compute_name(machine),
+                    spot=interruptible,
+                    **optional_kwargs_compute_body,
+                )
             ),
             teamspace_id,
             studio_id,
@@ -253,12 +268,24 @@ class StudioApi:
                 break
             time.sleep(1)
 
-    def get_machine(self, studio_id: str, teamspace_id: str) -> Machine:
+    def get_machine(self, studio_id: str, teamspace_id: str, cloud_account_id: str, org_id: str) -> Machine:
         """Get the current machine type the given Studio is running on."""
         response: V1CloudSpaceInstanceConfig = self._client.cloud_space_service_get_cloud_space_instance_config(
             project_id=teamspace_id, id=studio_id
         )
-        return Machine(response.compute_config.name, response.compute_config.name)
+        accelerators = self._get_machines_for_cloud_account(
+            teamspace_id=teamspace_id, cloud_account_id=cloud_account_id, org_id=org_id
+        )
+
+        for accelerator in accelerators:
+            if response.compute_config.name in (
+                accelerator.slug,
+                accelerator.slug_multi_cloud,
+                accelerator.instance_id,
+            ):
+                return Machine.from_str(accelerator.slug_multi_cloud)
+
+        return Machine.from_str(response.compute_config.name)
 
     def get_interruptible(self, studio_id: str, teamspace_id: str) -> bool:
         """Get whether the Studio is running on a interruptible instance."""
@@ -268,6 +295,22 @@ class StudioApi:
 
         return response.compute_config.spot
 
+    def _get_machines_for_cloud_account(
+        self, teamspace_id: str, cloud_account_id: str, org_id: str
+    ) -> List[V1ClusterAccelerator]:
+        from lightning_sdk.api.cloud_account_api import CloudAccountApi
+
+        cloud_account_api = CloudAccountApi()
+        accelerators = cloud_account_api.list_cloud_account_accelerators(
+            teamspace_id=teamspace_id,
+            cloud_account_id=cloud_account_id,
+            org_id=org_id,
+        )
+        if not accelerators:
+            return []
+
+        return list(filter(lambda acc: acc.enabled, accelerators.accelerator))
+
     def _get_detached_command_status(
         self, studio_id: str, teamspace_id: str, session_id: str
     ) -> V1GetLongRunningCommandInCloudSpaceResponse:
@@ -398,7 +441,7 @@ class StudioApi:
             body=body,
         )
 
-    def duplicate_studio(self, studio_id: str, teamspace_id: str, target_teamspace_id: str) -> Dict[str, str]:
+    def duplicate_studio(self, studio_id: str, teamspace_id: str, target_teamspace_id: str) -> Dict[str, Any]:
         """Duplicates the given Studio from a given Teamspace into a given target Teamspace."""
         target_teamspace = self._client.projects_service_get_project(target_teamspace_id)
         init_kwargs = {}
@@ -421,7 +464,7 @@ class StudioApi:
         init_kwargs["name"] = new_cloudspace.name
         init_kwargs["teamspace"] = target_teamspace.name
 
-        self.start_studio(new_cloudspace.id, target_teamspace_id, Machine.CPU, False)
+        self.start_studio(new_cloudspace.id, target_teamspace_id, Machine.CPU, False, None)
         return init_kwargs
 
     def delete_studio(self, studio_id: str, teamspace_id: str) -> None:

lightning_sdk/api/utils.py

@@ -331,7 +331,9 @@ class _DummyResponse:
 
 def _machine_to_compute_name(machine: Union[Machine, str]) -> str:
     if isinstance(machine, Machine):
-        return machine.instance_type
+        if machine.instance_type is not None:
+            return machine.instance_type
+        return machine.slug
     return machine
 
 
@@ -352,7 +354,9 @@ def _get_registry_url() -> str:
 
 
 def _sanitize_studio_remote_path(path: str, studio_id: str) -> str:
-    return f"/cloudspaces/{studio_id}/code/content/{path.replace('/teamspace/studios/this_studio/', '')}"
+    path = path.replace("/teamspace/studios/this_studio/", "")
+    root = f"/cloudspaces/{studio_id}/code/content/"
+    return os.path.join(root, path)
 
 
 def _resolve_teamspace_remote_path(path: str) -> str:

lightning_sdk/cli/clusters_menu.py

@@ -5,7 +5,7 @@ from rich.console import Console
 from simple_term_menu import TerminalMenu
 
 from lightning_sdk import Teamspace
-from lightning_sdk.api.cluster_api import ClusterApi
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.lightning_cloud.openapi import Externalv1Cluster, V1ProjectClusterBinding
 
 
@@ -29,9 +29,9 @@ class _ClustersMenu:
             selected_cluster_id = self._get_cluster_from_interactive_menu(
                 possible_clusters=teamspace.cloud_account_objs
             )
-            cluster_api = ClusterApi()
+            cloud_account_api = CloudAccountApi()
 
-            return cluster_api.get_cluster(
+            return cloud_account_api.get_cloud_account(
                 cluster_id=selected_cluster_id, org_id=teamspace.owner.id, project_id=teamspace.id
             )
         except KeyboardInterrupt:

lightning_sdk/cli/create.py

@@ -7,12 +7,15 @@ import click
 from rich.console import Console
 
 from lightning_sdk import Machine, Studio
-from lightning_sdk.api.cluster_api import ClusterApi
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
-from lightning_sdk.studio import Provider
+from lightning_sdk.machine import CloudProvider
+from lightning_sdk.utils.resolve import _resolve_deprecated_provider
 
-_MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
-_PROVIDER_VALUES = tuple([provider.value for provider in Provider])
+_MACHINE_VALUES = tuple(
+    [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli]
+)
+_PROVIDER_VALUES = tuple([provider.value for provider in CloudProvider])
 
 
 @click.group("create")
@@ -48,16 +51,26 @@ def create() -> None:
     ),
 )
 @click.option(
-    "--provider",
+    "--cloud-provider",
     default=None,
     type=click.Choice(_PROVIDER_VALUES),
     help="The provider to create the studio on. If --cloud-account is specified, this option is prioritized.",
 )
+@click.option(
+    "--provider",
+    default=None,
+    type=click.Choice(_PROVIDER_VALUES),
+    help=(
+        "Deprecated. Use --cloud-provider instead. The provider to create the studio on. "
+        "If --cloud-account is specified, this option is prioritized."
+    ),
+)
 def studio(
     name: str,
     teamspace: Optional[str] = None,
     start: Optional[str] = None,
     cloud_account: Optional[str] = None,
+    cloud_provider: Optional[str] = None,
     provider: Optional[str] = None,
 ) -> None:
     """Create a new studio on the Lightning AI platform.
@@ -70,12 +83,13 @@ def studio(
     menu = _TeamspacesMenu()
     teamspace_resolved = menu._resolve_teamspace(teamspace)
 
-    if provider is not None:
-        cluster_api = ClusterApi()
-        cloud_account = cluster_api.get_cluster_provider_mapping(
-            teamspace_resolved.id,
-            teamspace_resolved.owner.id,
-        )[provider]
+    cloud_provider = str(_resolve_deprecated_provider(cloud_provider, provider))
+
+    if cloud_provider is not None:
+        cloud_account_api = CloudAccountApi()
+        cloud_account = cloud_account_api.resolve_cloud_account(
+            teamspace_resolved.id, cloud_account, cloud_provider, teamspace_resolved.default_cloud_account
+        )
 
     # 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/deploy/_auth.py

@@ -2,7 +2,8 @@ import os
 import time
 from datetime import datetime
 from enum import Enum
-from typing import List, Optional, TypedDict
+from typing import Any, List, Optional, TypedDict
+from urllib.parse import urlencode
 
 from rich.console import Console
 from rich.prompt import Confirm
@@ -10,6 +11,7 @@ from rich.prompt import Confirm
 from lightning_sdk import Teamspace
 from lightning_sdk.api import UserApi
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.lightning_cloud import env
 from lightning_sdk.lightning_cloud.login import Auth, AuthServer
 from lightning_sdk.lightning_cloud.openapi import V1CloudSpace
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
@@ -24,6 +26,17 @@ class _AuthMode(Enum):
     DEPLOY = "deploy"
 
 
+class _AuthServer(AuthServer):
+    def __init__(self, mode: _AuthMode, *args: Any, **kwargs: Any) -> None:
+        self._mode = mode
+        super().__init__(*args, **kwargs)
+
+    def get_auth_url(self, port: int) -> str:
+        redirect_uri = f"http://localhost:{port}/login-complete"
+        params = urlencode({"redirectTo": redirect_uri, "mode": self._mode.value, "okbhrt": LITSERVE_CODE})
+        return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
+
+
 class _AuthLitServe(Auth):
     def __init__(self, mode: _AuthMode, shall_confirm: bool = False) -> None:
         super().__init__()
@@ -33,7 +46,10 @@ class _AuthLitServe(Auth):
     def _run_server(self) -> None:
         if self._shall_confirm:
             proceed = Confirm.ask(
-                "Authenticating with Lightning AI. This will open a browser window. Continue?", default=True
+                "[bold yellow]LitServe needs to authenticate with Lightning AI to deploy your server.[/bold yellow]\n"
+                "This will open a browser window for login.\n"
+                "Do you want to continue?",
+                default=True,
             )
             if not proceed:
                 raise RuntimeError(
@@ -42,7 +58,7 @@ class _AuthLitServe(Auth):
         print("Opening browser for authentication...")
         print("Please come back to the terminal after logging in.")
         time.sleep(3)
-        AuthServer({"mode": self._mode, "okbhrt": LITSERVE_CODE}).login_with_browser(self)
+        _AuthServer(self._mode).login_with_browser(self)
 
 
 def authenticate(mode: _AuthMode, shall_confirm: bool = True) -> None:

lightning_sdk/cli/deploy/serve.py

@@ -25,7 +25,9 @@ from lightning_sdk.cli.deploy._auth import (
 from lightning_sdk.cli.deploy.devbox import _handle_devbox
 from lightning_sdk.serve import _LitServeDeployer
 
-_MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
+_MACHINE_VALUES = tuple(
+    [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli]
+)
 
 
 class _ServeGroup(click.Group):
@@ -317,12 +319,20 @@ def _handle_cloud(
     ls_deployer = _LitServeDeployer(name=deployment_name, teamspace=None)
     path = ls_deployer.dockerize_api(script_path, port=port, gpu=not machine.is_cpu(), tag=tag, print_success=False)
 
-    console.print(f"\nPlease review the Dockerfile at [u]{path}[/u] and make sure it is correct.", style="bold")
-    correct_dockerfile = True if non_interactive else Confirm.ask("Is the Dockerfile correct?", default=True)
+    console.print(f"\n[bold]LitServe generated a Dockerfile at:[/bold]\n[u]{path}[/u]\n")
+    console.print("Please check that it matches your server setup.")
+    correct_dockerfile = (
+        True
+        if non_interactive
+        else Confirm.ask("Have you reviewed the Dockerfile and confirmed it's correct?", default=True)
+    )
     if not correct_dockerfile:
-        console.print("Please fix the Dockerfile and try again.", style="red")
+        console.print("[red]Dockerfile review canceled. Please update the Dockerfile and try again.[/red]")
         return
 
+    console.print(
+        "Building your container image now.\n[cyan bold]Make sure Docker is installed and running.[/cyan bold]\n"
+    )
     with Progress(
         SpinnerColumn(),
         TextColumn("[progress.description]{task.description}"),
@@ -340,7 +350,12 @@ def _handle_cloud(
             progress.remove_task(build_task)
 
         except Exception as e:
-            console.print(f"❌ Deployment failed: {e}", style="red")
+            console.print(
+                "❌ Failed to build a container for your server.\n"
+                "Make sure Docker is installed and running, then try again.",
+                f"\n\nError details: {e}",
+                style="red",
+            )
             return
 
     # Push the container to the registry
@@ -373,6 +388,7 @@ def _handle_cloud(
         f"{resolved_teamspace.owner.name}/{resolved_teamspace.name}/{container_basename}"
     )
 
+    cloud_account = cloud_account or resolved_teamspace.default_cloud_account
     if from_onboarding:
         thread = Thread(
             target=ls_deployer.run_on_cloud,

lightning_sdk/cli/download.py

@@ -17,6 +17,27 @@ from lightning_sdk.studio import Studio
 from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
 
 
+def _expand_remote_path(path: str) -> str:
+    """Expand and normalize remote CLI paths.
+
+    - Strips leading `~/` or `~`
+    - Expands `~` to the user's home but returns relative to it
+    - Returns an empty string if path is empty or `~`
+    """
+    if not path:
+        return ""
+
+    local_home = os.path.expanduser("~")
+
+    # Expand to absolute path and remove the local home prefix if present
+    path = os.path.expanduser(path)
+    if path.startswith(local_home):
+        path = path[len(local_home) :]
+
+    # Remove any leading "/" or "~" remnants
+    return path.lstrip("/~")
+
+
 @click.group(name="download")
 def download() -> None:
     """Download resources from Lightning AI."""
@@ -86,17 +107,20 @@ def folder(
         raise ValueError("Either --studio or --teamspace must be provided, not both")
 
     if studio:
+        path = _expand_remote_path(path)
         resolved_downloader = _resolve_studio(studio)
     elif teamspace:
         menu = _TeamspacesMenu()
         resolved_downloader = menu._resolve_teamspace(teamspace)
+    else:
+        raise ValueError("Either --studio or --teamspace must be provided")
 
     if not path:
         local_path /= resolved_downloader.name
         path = ""
 
     try:
-        if not path:
+        if not path and teamspace:
             raise FileNotFoundError()
         resolved_downloader.download_folder(remote_path=path, target_path=str(local_path))
     except Exception as e:

lightning_sdk/cli/entrypoint.py

@@ -21,7 +21,7 @@ from lightning_sdk.cli.deploy.serve import deploy
 from lightning_sdk.cli.docker_cli import dockerize
 from lightning_sdk.cli.download import download
 from lightning_sdk.cli.generate import generate
-from lightning_sdk.cli.inspect import inspect
+from lightning_sdk.cli.inspection import inspect
 from lightning_sdk.cli.list import list_cli
 from lightning_sdk.cli.open import open
 from lightning_sdk.cli.run import run
@@ -50,7 +50,9 @@ def _notify_exception(exception_type: Type[BaseException], value: BaseException,
         renderables.append(Text("\n\nFull traceback:\n", style="bold yellow"))
         renderables.append(Syntax(tb_text, "python", theme="monokai", line_numbers=False, word_wrap=True))
     else:
-        renderables.append(Text("\n\nTo see the full traceback, set the LIGHTNING_DEBUG environment variable to 1."))
+        renderables.append(Text("\n\n🐞 To view the full traceback, set: LIGHTNING_DEBUG=1"))
+
+    renderables.append(Text("\n📘 Need help? Run: lightning <command> --help", style="cyan"))
 
     console.print(Panel(Group(*renderables), title="⚡ Lightning CLI Error", border_style="red"))
 

lightning_sdk/cli/list.py

@@ -274,7 +274,11 @@ def machines() -> None:
     table.add_column("Name")
 
     # Get all machine types from the enum
-    machine_types = [name for name in dir(Machine) if isinstance(getattr(Machine, name), Machine)]
+    machine_types = [
+        name
+        for name in dir(Machine)
+        if isinstance(getattr(Machine, name), Machine) and getattr(Machine, name)._include_in_cli
+    ]
 
     # Add rows to table
     for name in sorted(machine_types):

lightning_sdk/cli/run.py

@@ -8,7 +8,9 @@ from lightning_sdk.machine import Machine
 from lightning_sdk.mmt import MMT
 from lightning_sdk.teamspace import Teamspace
 
-_MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
+_MACHINE_VALUES = tuple(
+    [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli]
+)
 
 
 @click.group(name="run")

lightning_sdk/cli/start.py

@@ -4,10 +4,12 @@ import click
 
 from lightning_sdk import Machine, Studio
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
-from lightning_sdk.studio import Provider
+from lightning_sdk.machine import CloudProvider
 
-_MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
-_PROVIDER_VALUES = tuple([provider.value for provider in Provider])
+_MACHINE_VALUES = tuple(
+    [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli]
+)
+_PROVIDER_VALUES = tuple([provider.value for provider in CloudProvider])
 
 
 @click.group("start")
@@ -35,14 +37,28 @@ def start() -> None:
     type=click.Choice(_MACHINE_VALUES),
     help="The machine type to start the studio on.",
 )
+@click.option(
+    "--cloud-provider",
+    default=None,
+    type=click.Choice(_PROVIDER_VALUES),
+    help=("The provider to create the studio on. If --cloud-account is specified, this option is prioritized."),
+)
 @click.option(
     "--provider",
     default=None,
-    show_default=True,
     type=click.Choice(_PROVIDER_VALUES),
-    help="The provider to start the studio on.",
+    help=(
+        "Deprecated. Use --cloud-provider instead. The provider to create the studio on. "
+        "If --cloud-account is specified, this option is prioritized."
+    ),
 )
-def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU", provider: Optional[str] = None) -> None:
+def studio(
+    name: str,
+    teamspace: Optional[str] = None,
+    machine: str = "CPU",
+    cloud_provider: Optional[str] = None,
+    provider: Optional[str] = None,
+) -> None:
     """Start a studio on a given machine.
 
     Example:
@@ -59,10 +75,26 @@ def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU", pro
         owner, teamspace = None, None
 
     try:
-        studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False, provider=provider)
+        studio = Studio(
+            name=name,
+            teamspace=teamspace,
+            org=owner,
+            user=None,
+            create_ok=False,
+            cloud_provider=cloud_provider,
+            provider=provider,
+        )
     except (RuntimeError, ValueError, ApiException) as first_error:
         try:
-            studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False, provider=provider)
+            studio = Studio(
+                name=name,
+                teamspace=teamspace,
+                org=None,
+                user=owner,
+                create_ok=False,
+                cloud_provider=cloud_provider,
+                provider=provider,
+            )
         except (RuntimeError, ValueError, ApiException) as second_error:
             raise first_error from second_error
 

lightning_sdk/cli/switch.py

@@ -5,7 +5,9 @@ import click
 from lightning_sdk import Machine, Studio
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 
-_MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
+_MACHINE_VALUES = tuple(
+    [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli]
+)
 
 
 @click.group("switch")

lightning_sdk/deployment/deployment.py

@@ -130,6 +130,7 @@ class Deployment:
         include_credentials: Optional[bool] = None,
         from_onboarding: Optional[bool] = None,
         from_litserve: Optional[bool] = None,
+        max_runtime: Optional[int] = None,
     ) -> None:
         """The Lightning AI Deployment.
 
@@ -157,6 +158,10 @@ class Deployment:
             quantity: The number of machines per replica to deploy.
             include_credentials: Whether to include the environment variables for the SDK to authenticate
             from_onboarding: Whether the deployment is from onboarding.
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
 
         Note:
             Since a teamspace can either be owned by an org or by a user directly,
@@ -227,6 +232,7 @@ class Deployment:
                     quantity=quantity,
                     include_credentials=include_credentials if include_credentials is not None else True,
                     cloudspace_id=cloudspace_id,
+                    max_runtime=max_runtime,
                 ),
                 strategy=to_strategy(release_strategy),
             ),
@@ -262,6 +268,7 @@ class Deployment:
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
         quantity: Optional[int] = None,
         include_credentials: Optional[bool] = None,
+        max_runtime: Optional[int] = None,
     ) -> None:
         cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
 
@@ -288,6 +295,7 @@ class Deployment:
             release_strategy=release_strategy,
             quantity=quantity,
             include_credentials=include_credentials,
+            max_runtime=max_runtime,
         )
 
     def stop(self) -> None: