lightning-sdk 2025.7.10__py3-none-any.whl → 2025.7.22__py3-none-any.whl

lightning_sdk/__init__.py

@@ -4,7 +4,7 @@ from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__  # noq
 from lightning_sdk.deployment import Deployment
 from lightning_sdk.helpers import _check_version_and_prompt_upgrade, _set_tqdm_envvars_noninteractive
 from lightning_sdk.job import Job
-from lightning_sdk.machine import Machine
+from lightning_sdk.machine import CloudProvider, Machine
 from lightning_sdk.mmt import MMT
 from lightning_sdk.organization import Organization
 from lightning_sdk.plugin import JobsPlugin, MultiMachineTrainingPlugin, Plugin, SlurmJobsPlugin
@@ -16,6 +16,7 @@ from lightning_sdk.user import User
 __all__ = [
     "AIHub",
     "Agent",
+    "CloudProvider",
     "Deployment",
     "Job",
     "JobsPlugin",
@@ -31,6 +32,6 @@ __all__ = [
     "User",
 ]
 
-__version__ = "2025.07.10"
+__version__ = "2025.07.22"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/cloud_account_api.py

@@ -0,0 +1,154 @@
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+from lightning_sdk.lightning_cloud.openapi import (
+    Externalv1Cluster,
+    V1CloudProvider,
+    V1ClusterType,
+    V1ExternalCluster,
+    V1ListClusterAcceleratorsResponse,
+)
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+if TYPE_CHECKING:
+    from lightning_sdk.machine import CloudProvider
+
+
+class CloudAccountApi:
+    """Internal API client for API requests to cluster endpoints."""
+
+    def __init__(self) -> None:
+        self._client = LightningClient(max_tries=7)
+
+    def get_cloud_account(self, cloud_account_id: str, teamspace_id: str, org_id: str) -> Externalv1Cluster:
+        """Gets the cluster from given params cluster_id, project_id and owner.
+
+        Args:
+            cloud_account_id: the cloud account to get
+            teamspace_id: the teamspace the cloud_account is supposed to be associated with
+            org_id: The owning org of this teamspace
+
+        """
+        res = self._client.cluster_service_get_cluster(id=cloud_account_id, org_id=org_id, project_id=teamspace_id)
+        if not res:
+            raise ValueError(f"CloudAccount {cloud_account_id} does not exist")
+        return res
+
+    def list_cloud_accounts(self, teamspace_id: str) -> List[V1ExternalCluster]:
+        """Lists the cloud accounts for a given teamspace.
+
+        Args:
+            teamspace_id: The teamspace to list cloud accounts for
+
+        Returns:
+            A list of cloud accounts
+        """
+        res = self._client.cluster_service_list_project_clusters(
+            project_id=teamspace_id,
+        )
+        return res.clusters
+
+    def get_cloud_account_non_org(self, teamspace_id: str, cloud_account_id: str) -> Optional[V1ExternalCluster]:
+        for cluster in self.list_cloud_accounts(teamspace_id=teamspace_id):
+            if cluster.id == cloud_account_id:
+                return cluster
+
+        return None
+
+    def list_cloud_account_accelerators(self, cloud_account_id: str, org_id: str) -> V1ListClusterAcceleratorsResponse:
+        """Lists the accelerators for a given cloud account.
+
+        Args:
+            cloud_account_id: cluster ID to list accelerators for
+            org_id: The owning org of this project
+        """
+        res = self._client.cluster_service_list_cluster_accelerators(
+            id=cloud_account_id,
+            org_id=org_id,
+        )
+        if not res:
+            raise ValueError(f"CloudAccount {cloud_account_id} does not exist")
+        return res
+
+    def list_global_cloud_accounts(self, teamspace_id: str) -> List[V1ExternalCluster]:
+        """Lists the accelerators for a given teamspace.
+
+        Args:
+            teamspace_id: id of the teamspace to get the associated cloud_accounts for
+        """
+        cloud_accounts = self.list_cloud_accounts(teamspace_id=teamspace_id)
+        if not cloud_accounts:
+            raise ValueError(f"Teamspace {teamspace_id} does not exist")
+        filtered_cloud_accounts = filter(lambda x: x.spec.cluster_type == V1ClusterType.GLOBAL, cloud_accounts)
+        return list(filtered_cloud_accounts)
+
+    def get_cloud_account_provider_mapping(self, teamspace_id: str) -> Dict["CloudProvider", str]:
+        """Gets the cloud account <-> provider mapping."""
+        res = self.list_global_cloud_accounts(teamspace_id=teamspace_id)
+        return {self._get_cloud_account_provider(cloud_account): cloud_account.id for cloud_account in res}
+
+    @staticmethod
+    def _get_cloud_account_provider(cloud_account: Optional[V1ExternalCluster]) -> "CloudProvider":
+        """Determines the cloud provider based on the cloud_account configuration.
+
+        Args:
+            cloud_account: An optional Externalv1Cluster object containing cluster specifications
+
+        Returns:
+            CloudProvider: The determined cloud provider, defaults to AWS if no match is found
+        """
+        from lightning_sdk.machine import CloudProvider
+
+        if not cloud_account:
+            return CloudProvider.AWS
+
+        if cloud_account.spec and cloud_account.spec.driver:
+            if cloud_account.spec.driver == V1CloudProvider.LIGHTNING:
+                return CloudProvider.LIGHTNING
+
+            if cloud_account.spec.driver == V1CloudProvider.DGX:
+                return CloudProvider.DGX
+
+        if cloud_account.spec:
+            if cloud_account.spec.aws_v1:
+                return CloudProvider.AWS
+            if cloud_account.spec.google_cloud_v1:
+                return CloudProvider.GCP
+            if cloud_account.spec.lambda_labs_v1:
+                return CloudProvider.LAMBDA_LABS
+            if cloud_account.spec.vultr_v1:
+                return CloudProvider.VULTR
+            if cloud_account.spec.voltage_park_v1:
+                return CloudProvider.VOLTAGE_PARK
+            if cloud_account.spec.nebius_v1:
+                return CloudProvider.NEBIUS
+
+        return CloudProvider.AWS
+
+    def resolve_cloud_account(
+        self,
+        teamspace_id: str,
+        cloud_account: Optional[str],
+        cloud_provider: Optional[Union["CloudProvider", str]],
+        default_cloud_account: Optional[str],
+    ) -> Optional[str]:
+        if cloud_account:
+            if cloud_provider:
+                cloud_account_resp = self.get_cloud_account_non_org(teamspace_id, cloud_account)
+                cloud_provider_resp = self._get_cloud_account_provider(cloud_account_resp)
+                if cloud_provider_resp != cloud_provider:
+                    raise RuntimeError(
+                        f"Specified both cloud_provider ({cloud_provider}) and "
+                        "cloud_account ({cloud_account} has cloud provider {cloud_provider_resp}) which don't match!"
+                    )
+
+            return cloud_account
+
+        if cloud_provider:
+            cloud_account_mapping = self.get_cloud_account_provider_mapping(teamspace_id=teamspace_id)
+            if cloud_provider and cloud_provider in cloud_account_mapping:
+                return cloud_account_mapping[cloud_provider]
+
+        if default_cloud_account:
+            return default_cloud_account
+
+        return None

lightning_sdk/api/deployment_api.py

@@ -264,6 +264,7 @@ class DeploymentApi:
         custom_domain: Optional[str] = None,
         quantity: Optional[int] = None,
         include_credentials: Optional[bool] = None,
+        max_runtime: Optional[int] = None,
     ) -> V1Deployment:
         # Update the deployment in place
 
@@ -291,6 +292,9 @@ class DeploymentApi:
         requires_release |= apply_change(deployment.spec, "spot", spot)
         requires_release |= apply_change(deployment.spec, "quantity", quantity)
         requires_release |= apply_change(deployment.spec, "include_credentials", include_credentials)
+        requires_release |= apply_change(
+            deployment.spec, "requested_run_duration_seconds", str(max_runtime) if max_runtime is not None else None
+        )
 
         if requires_release:
             if deployment.strategy is None:
@@ -569,6 +573,7 @@ def to_spec(
     quantity: Optional[int] = None,
     include_credentials: Optional[bool] = None,
     cloudspace_id: Optional[None] = None,
+    max_runtime: Optional[int] = None,
 ) -> V1JobSpec:
     if cloud_account is None:
         raise ValueError("The cloud account should be defined.")
@@ -585,6 +590,11 @@ def to_spec(
     if command is None and cloudspace_id is not None:
         raise ValueError("The command should be defined.")
 
+    # need to go via kwargs for typing compatibility since autogenerated apis accept None but aren't typed with None
+    optional_spec_kwargs = {}
+    if max_runtime:
+        optional_spec_kwargs["requested_run_duration_seconds"] = str(max_runtime)
+
     return V1JobSpec(
         cluster_id=cloud_account,
         command=command,
@@ -597,6 +607,7 @@ def to_spec(
         quantity=quantity,
         include_credentials=include_credentials,
         cloudspace_id=cloudspace_id,
+        **optional_spec_kwargs,
     )
 
 

lightning_sdk/api/job_api.py

@@ -215,6 +215,7 @@ class JobApiV2:
         path_mappings: Optional[Dict[str, str]],
         artifacts_local: Optional[str],  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
+        max_runtime: Optional[str] = None,
     ) -> V1Job:
         body = self._create_job_body(
             name=name,
@@ -231,6 +232,7 @@ class JobApiV2:
             path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            max_runtime=max_runtime,
         )
 
         job: V1Job = self._client.jobs_service_create_job(project_id=teamspace_id, body=body)
@@ -252,6 +254,7 @@ class JobApiV2:
         path_mappings: Optional[Dict[str, str]],
         artifacts_local: Optional[str],  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str],  # deprecated in favor of path_mappings)
+        max_runtime: Optional[int] = None,
     ) -> ProjectIdJobsBody:
         env_vars = []
         if env is not None:
@@ -268,6 +271,11 @@ class JobApiV2:
             artifacts_remote=artifacts_remote,
         )
 
+        # need to go via kwargs for typing compatibility since autogenerated apis accept None but aren't typed with None
+        optional_spec_kwargs = {}
+        if max_runtime:
+            optional_spec_kwargs["requested_run_duration_seconds"] = str(max_runtime)
+
         spec = V1JobSpec(
             cloudspace_id=studio_id or "",
             cluster_id=cloud_account or "",
@@ -281,6 +289,7 @@ class JobApiV2:
             image_cluster_credentials=cloud_account_auth,
             image_secret_ref=image_credentials or "",
             path_mappings=path_mappings_list,
+            **optional_spec_kwargs,
         )
         return ProjectIdJobsBody(name=name, spec=spec)
 

lightning_sdk/api/llm_api.py

@@ -7,11 +7,12 @@ import threading
 import warnings
 from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union
 
-from pip._vendor.urllib3 import HTTPResponse
-
-from lightning_sdk.lightning_cloud.openapi.models.v1_conversation_response_chunk import V1ConversationResponseChunk
-from lightning_sdk.lightning_cloud.openapi.models.v1_response_choice import V1ResponseChoice
-from lightning_sdk.lightning_cloud.openapi.models.v1_response_choice_delta import V1ResponseChoiceDelta
+from lightning_sdk.lightning_cloud.openapi.models import (
+    StreamResultOfV1ConversationResponseChunk,
+    V1ConversationResponseChunk,
+    V1ResponseChoice,
+    V1ResponseChoiceDelta,
+)
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
 
@@ -54,7 +55,9 @@ class LLMApi:
             warnings.warn("Error decoding JSON:", decoded_line)
             return None
 
-    def _stream_chat_response(self, result: HTTPResponse) -> Generator[V1ConversationResponseChunk, None, None]:
+    def _stream_chat_response(
+        self, result: StreamResultOfV1ConversationResponseChunk
+    ) -> Generator[V1ConversationResponseChunk, None, None]:
         for line in result.stream():
             decoded_lines = line.decode("utf-8").strip()
             for decoded_line in decoded_lines.splitlines():
@@ -80,6 +83,7 @@ class LLMApi:
         name: Optional[str] = None,
         metadata: Optional[Dict[str, str]] = None,
         stream: bool = False,
+        tools: Optional[List[Dict[str, Any]]] = None,
         **kwargs: Any,
     ) -> Union[V1ConversationResponseChunk, Generator[V1ConversationResponseChunk, None, None]]:
         is_internal_conversation = os.getenv("LIGHTNING_INTERNAL_CONVERSATION", "false").lower() == "true"
@@ -105,6 +109,7 @@ class LLMApi:
             "ephemeral": ephemeral,
             "parent_conversation_id": kwargs.get("parent_conversation_id", ""),
             "parent_message_id": kwargs.get("parent_message_id", ""),
+            "tools": tools,
         }
         if images:
             for image in images:

lightning_sdk/api/mmt_api.py

@@ -87,6 +87,7 @@ class MMTApiV2:
         path_mappings: Optional[Dict[str, str]],
         artifacts_local: Optional[str],  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
+        max_runtime: Optional[int],
     ) -> V1MultiMachineJob:
         body = self._create_mmt_body(
             name=name,
@@ -104,6 +105,7 @@ class MMTApiV2:
             path_mappings=path_mappings,
             artifacts_local=artifacts_local,  # deprecated in favor of path_mappings
             artifacts_remote=artifacts_remote,  # deprecated in favor of path_mappings
+            max_runtime=max_runtime,
         )
 
         job: V1MultiMachineJob = self._client.jobs_service_create_multi_machine_job(project_id=teamspace_id, body=body)
@@ -126,6 +128,7 @@ class MMTApiV2:
         path_mappings: Optional[Dict[str, str]],
         artifacts_local: Optional[str],  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
+        max_runtime: Optional[int] = None,
     ) -> ProjectIdMultimachinejobsBody:
         env_vars = []
         if env is not None:
@@ -142,6 +145,11 @@ class MMTApiV2:
             artifacts_remote=artifacts_remote,
         )
 
+        # need to go via kwargs for typing compatibility since autogenerated apis accept None but aren't typed with None
+        optional_spec_kwargs = {}
+        if max_runtime:
+            optional_spec_kwargs["requested_run_duration_seconds"] = str(max_runtime)
+
         spec = V1JobSpec(
             cloudspace_id=studio_id or "",
             cluster_id=cloud_account or "",
@@ -155,6 +163,7 @@ class MMTApiV2:
             image_cluster_credentials=cloud_account_auth,
             image_secret_ref=image_credentials or "",
             path_mappings=path_mappings_list,
+            **optional_spec_kwargs,
         )
         return ProjectIdMultimachinejobsBody(
             name=name, spec=spec, cluster_id=cloud_account or "", machines=num_machines

lightning_sdk/api/pipeline_api.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, List, Optional, Union
 
-from lightning_sdk.api.cluster_api import ClusterApi
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.lightning_cloud.openapi.models import (
     ProjectIdPipelinesBody,
     ProjectIdSchedulesBody,
@@ -23,7 +23,7 @@ class PipelineApi:
 
     def __init__(self) -> None:
         self._client = LightningClient(max_tries=0, retry=False)
-        self._cluster_api = ClusterApi()
+        self._cloud_account_api = CloudAccountApi()
 
     def get_pipeline_by_id(self, project_id: str, pipeline_id_or_name: str) -> Optional[V1Pipeline]:
         if pipeline_id_or_name.startswith("pip_"):
@@ -75,6 +75,7 @@ class PipelineApi:
                     resource_id=pipeline.id,
                     parent_resource_id=parent_pipeline_id or "",
                     resource_type=V1ScheduleResourceType.PIPELINE,
+                    timezone=schedule.timezone,
                 )
 
                 self._client.schedules_service_create_schedule(body, teamspace.id)
@@ -97,7 +98,7 @@ class PipelineApi:
 
         from lightning_sdk.pipeline.utils import _get_cloud_account
 
-        clusters = self._cluster_api.list_clusters(project_id=teamspace.id)
+        clusters = self._cloud_account_api.list_cloud_accounts(teamspace_id=teamspace.id)
 
         selected_cluster = None
         selected_cluster_id = _get_cloud_account(steps)

lightning_sdk/api/studio_api.py

@@ -121,7 +121,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 +168,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,
@@ -398,7 +412,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 +435,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/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,13 @@ 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])
+_PROVIDER_VALUES = tuple([provider.value for provider in CloudProvider])
 
 
 @click.group("create")
@@ -48,16 +49,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 +81,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

@@ -317,12 +317,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 +348,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 +386,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/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