lightning-sdk 2025.8.18.post0__py3-none-any.whl → 2025.8.21__py3-none-any.whl

lightning_sdk/__init__.py

@@ -32,6 +32,6 @@ __all__ = [
     "User",
 ]
 
-__version__ = "2025.08.18.post0"
+__version__ = "2025.08.21"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/llm_api.py

@@ -146,7 +146,6 @@ class LLMApi:
                     {"contentType": "text", "parts": [prompt]},
                 ],
             },
-            "max_tokens": max_completion_tokens,
             "conversation_id": conversation_id,
             "billing_project_id": billing_project_id,
             "name": name,
@@ -159,6 +158,9 @@ class LLMApi:
             "parent_message_id": kwargs.get("parent_message_id", ""),
             "tools": tools,
         }
+        if max_completion_tokens is not None:
+            body["max_completion_tokens"] = max_completion_tokens
+
         if images:
             for image in images:
                 url = image
@@ -203,7 +205,6 @@ class LLMApi:
                     {"contentType": "text", "parts": [prompt]},
                 ],
             },
-            "max_completion_tokens": max_completion_tokens,
             "conversation_id": conversation_id,
             "billing_project_id": billing_project_id,
             "name": name,
@@ -216,6 +217,9 @@ class LLMApi:
             "parent_message_id": kwargs.get("parent_message_id", ""),
             "sent_at": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="microseconds"),
         }
+        if max_completion_tokens is not None:
+            body["max_completion_tokens"] = max_completion_tokens
+
         if images:
             for image in images:
                 url = image

lightning_sdk/api/studio_api.py

@@ -1,8 +1,7 @@
 import json
 import os
-import tempfile
 import time
-import zipfile
+from pathlib import Path
 from threading import Event, Thread
 from typing import Any, Dict, Generator, List, Mapping, Optional, Tuple, Union
 
@@ -12,6 +11,7 @@ from tqdm import tqdm
 
 from lightning_sdk.api.utils import (
     _create_app,
+    _download_studio_files,
     _DummyBody,
     _DummyResponse,
     _FileUploader,
@@ -205,6 +205,32 @@ class StudioApi:
             instance_id = code_status.in_use.cloud_space_instance_id
             print(f"Studio started | {teamspace_id=} {studio_id=} {instance_id=}")
 
+    def start_studio_async(
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        machine: Union[Machine, str],
+        interruptible: bool = False,
+        max_runtime: Optional[int] = None,
+    ) -> None:
+        """Start an existing Studio without blocking."""
+        # 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,
+                    **optional_kwargs_compute_body,
+                )
+            ),
+            teamspace_id,
+            studio_id,
+        )
+
     def stop_studio(self, studio_id: str, teamspace_id: str) -> None:
         """Stop an existing Studio."""
         self.stop_keeping_alive(teamspace_id=teamspace_id, studio_id=studio_id)
@@ -289,6 +315,79 @@ class StudioApi:
                 break
             time.sleep(1)
 
+    def switch_studio_machine_with_progress(
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        machine: Union[Machine, str],
+        interruptible: bool,
+        progress: Any,  # StudioProgressTracker - avoid circular import
+    ) -> None:
+        """Switches given Studio to a new machine type with progress tracking."""
+        progress.update_progress(10, "Requesting machine switch...")
+
+        self._request_switch(
+            studio_id=studio_id, teamspace_id=teamspace_id, machine=machine, interruptible=interruptible
+        )
+
+        progress.update_progress(20, "Waiting for machine allocation...")
+
+        # Wait until it's time to switch
+        requested_was_found = False
+        startup_status = None
+        base_progress = 20
+        max_wait_progress = 60
+        wait_counter = 0
+
+        while True:
+            status = self.get_studio_status(studio_id, teamspace_id)
+            requested_machine = status.requested
+
+            if requested_machine is not None:
+                requested_was_found = True
+                startup_status = requested_machine.startup_status
+
+            # if the requested machine was found in the past, use the in_use status instead.
+            # it might be that it either was cancelled or it actually is ready.
+            # Either way, since we're actually blocking below for the in use startup status
+            # it's safe to switch at this point
+            elif requested_was_found:
+                in_use_machine = status.in_use
+                if in_use_machine is not None:
+                    startup_status = in_use_machine.startup_status
+
+            if startup_status and startup_status.initial_restore_finished:
+                break
+
+            # Update progress gradually while waiting
+            wait_counter += 1
+            current_progress = min(base_progress + (wait_counter * 2), max_wait_progress)
+            progress.update_progress(current_progress, "Allocating new machine...")
+            time.sleep(1)
+
+        progress.update_progress(70, "Starting machine switch...")
+        self._client.cloud_space_service_switch_cloud_space_instance(teamspace_id, studio_id)
+
+        progress.update_progress(80, "Configuring new machine...")
+
+        # Wait until the new machine is ready to use
+        switch_counter = 0
+        while True:
+            in_use = self.get_studio_status(studio_id, teamspace_id).in_use
+            if in_use is None:
+                continue
+            startup_status = in_use.startup_status
+            if startup_status and startup_status.top_up_restore_finished:
+                break
+
+            # Update progress while waiting for machine to be ready
+            switch_counter += 1
+            current_progress = min(80 + switch_counter, 95)
+            progress.update_progress(current_progress, "Finalizing machine setup...")
+            time.sleep(1)
+
+        progress.complete("Machine switch completed successfully")
+
     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(
@@ -569,46 +668,24 @@ class StudioApi:
         progress_bar: bool = True,
     ) -> None:
         """Downloads a given folder from a Studio to a target location."""
-        # TODO: Update this endpoint to permit basic auth
+        # TODO: implement resumable downloads
         auth = Auth()
         auth.authenticate()
-        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
 
-        query_params = {
-            "clusterId": cloud_account,
-            "prefix": _sanitize_studio_remote_path(path, studio_id),
-            "token": token,
-        }
+        prefix = _sanitize_studio_remote_path(path, studio_id)
+        # ensure we only download as a directory and not the entire prefix
+        if prefix.endswith("/") is False:
+            prefix = prefix + "/"
 
-        r = requests.get(
-            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
-            params=query_params,
-            stream=True,
+        _download_studio_files(
+            client=self._client,
+            teamspace_id=teamspace_id,
+            cluster_id=cloud_account,
+            prefix=prefix,
+            download_dir=Path(target_path),
+            progress_bar=progress_bar,
         )
 
-        if progress_bar:
-            pbar = tqdm(
-                desc=f"Downloading {os.path.split(path)[1]}",
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1000,
-            )
-
-            pbar_update = pbar.update
-        else:
-            pbar_update = lambda x: None
-
-        if target_path:
-            os.makedirs(target_path, exist_ok=True)
-
-        with tempfile.TemporaryFile() as f:
-            for chunk in r.iter_content(chunk_size=4096 * 8):
-                f.write(chunk)
-                pbar_update(len(chunk))
-
-            with zipfile.ZipFile(f) as z:
-                z.extractall(target_path)
-
     def install_plugin(self, studio_id: str, teamspace_id: str, plugin_name: str) -> str:
         """Installs the given plugin."""
         resp: V1Plugin = self._client.cloud_space_service_install_plugin(

lightning_sdk/api/utils.py

@@ -6,7 +6,7 @@ import re
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, TypedDict, Union
 
 import backoff
 import requests
@@ -19,6 +19,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     ModelsStoreApi,
     ProjectIdStorageBody,
     StorageCompleteBody,
+    StorageServiceApi,
     UploadIdCompleteBody,
     UploadIdPartsBody,
     V1CompletedPart,
@@ -367,38 +368,40 @@ _DOWNLOAD_REQUEST_CHUNK_SIZE = 10 * _BYTES_PER_MB
 _DOWNLOAD_MIN_CHUNK_SIZE = 100 * _BYTES_PER_KB
 
 
+class _RefreshResponse(TypedDict):
+    url: str
+    size: int
+
+
 class _FileDownloader:
     def __init__(
         self,
-        client: LightningClient,
-        model_id: str,
-        version: str,
         teamspace_id: str,
         remote_path: str,
         file_path: str,
         executor: ThreadPoolExecutor,
         num_workers: int = 20,
         progress_bar: Optional[tqdm] = None,
+        url: Optional[str] = None,
+        size: Optional[int] = None,
+        refresh_fn: Optional[Callable[[], _RefreshResponse]] = None,
     ) -> None:
-        self.api = ModelsStoreApi(client.api_client)
-        self.model_id = model_id
-        self.version = version
         self.teamspace_id = teamspace_id
         self.local_path = file_path
         self.remote_path = remote_path
         self.progress_bar = progress_bar
         self.num_workers = num_workers
-        self._url = ""
-        self._size = 0
+        self._url = url
+        self._size = size
         self.executor = executor
+        self.refresh_fn = refresh_fn
 
     @backoff.on_exception(backoff.expo, ApiException, max_tries=10)
     def refresh(self) -> None:
-        response = self.api.models_store_get_model_file_url(
-            project_id=self.teamspace_id, model_id=self.model_id, version=self.version, filepath=self.remote_path
-        )
-        self._url = response.url
-        self._size = int(response.size)
+        if self.refresh_fn is not None:
+            response = self.refresh_fn()
+            self._url = response["url"]
+            self._size = response["size"]
 
     @property
     def url(self) -> str:
@@ -413,6 +416,11 @@ class _FileDownloader:
             return
         self.progress_bar.update(n)
 
+    def update_filename(self, desc: str) -> None:
+        if self.progress_bar is None:
+            return
+        self.progress_bar.set_description(f"{(desc[:72] + '...') if len(desc) > 75 else desc:<75.75}")
+
     @backoff.on_exception(backoff.expo, (requests.exceptions.HTTPError), max_tries=10)
     def _download_chunk(self, filename: str, start_end: Tuple[int]) -> None:
         start, end = start_end
@@ -447,6 +455,8 @@ class _FileDownloader:
                     f.write(b"\x00" * remaining_size)
 
     def _multipart_download(self, filename: str, num_workers: int) -> None:
+        self.update_filename(f"Downloading {self.remote_path}")
+
         num_chunks = num_workers
         chunk_size = math.ceil(self.size / num_chunks)
 
@@ -464,7 +474,8 @@ class _FileDownloader:
         concurrent.futures.wait(futures)
 
     def download(self) -> None:
-        self.refresh()
+        if self.url is None:
+            self.refresh()
 
         tmp_filename = f"{self.local_path}.download"
 
@@ -539,6 +550,15 @@ def _download_model_files(
             mininterval=1,
         )
 
+    def refresh_fn(filename: str) -> _RefreshResponse:
+        resp = api.models_store_get_model_file_url(
+            project_id=response.project_id,
+            model_id=response.model_id,
+            version=response.version,
+            filepath=filename,
+        )
+        return {"url": resp.url, "size": int(resp.size)}
+
     with ThreadPoolExecutor(max_workers=min(num_workers, len(response.filepaths))) as file_executor, ThreadPoolExecutor(
         max_workers=num_workers
     ) as part_executor:
@@ -549,15 +569,13 @@ def _download_model_files(
             local_file.parent.mkdir(parents=True, exist_ok=True)
 
             file_downloader = _FileDownloader(
-                client=client,
-                model_id=response.model_id,
-                version=response.version,
                 teamspace_id=response.project_id,
                 remote_path=filepath,
                 file_path=str(local_file),
                 num_workers=num_workers,
                 progress_bar=pbar,
                 executor=part_executor,
+                refresh_fn=lambda f=filepath: refresh_fn(f),
             )
 
             futures.append(file_executor.submit(file_downloader.download))
@@ -568,6 +586,78 @@ def _download_model_files(
         return response.filepaths
 
 
+def _download_studio_files(
+    client: LightningClient,
+    teamspace_id: str,
+    cluster_id: str,
+    prefix: str,
+    download_dir: Path,
+    progress_bar: bool,
+    num_workers: int = os.cpu_count() * 4,
+) -> None:
+    api = StorageServiceApi(client.api_client)
+    response = None
+
+    pbar = None
+    if progress_bar:
+        pbar = tqdm(
+            desc="Downloading files",
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1000,
+            position=-1,
+            mininterval=1,
+        )
+
+    def refresh_fn(filename: str) -> _RefreshResponse:
+        resp = api.storage_service_list_project_artifacts(
+            project_id=teamspace_id,
+            cluster_id=cluster_id,
+            page_token="",
+            include_download_url=True,
+            prefix=prefix + filename,
+            page_size=1,
+        )
+        return {"url": resp.artifacts[0].url, "size": int(resp.artifacts[0].size_bytes)}
+
+    with ThreadPoolExecutor(max_workers=num_workers) as file_executor, ThreadPoolExecutor(
+        max_workers=num_workers
+    ) as part_executor:
+        while response is None or (response is not None and response.next_page_token != ""):
+            response = api.storage_service_list_project_artifacts(
+                project_id=teamspace_id,
+                cluster_id=cluster_id,
+                page_token=response.next_page_token if response is not None else "",
+                include_download_url=True,
+                prefix=prefix,
+                page_size=1000,
+            )
+
+            page_futures = []
+            for file in response.artifacts:
+                local_file = download_dir / file.filename
+                local_file.parent.mkdir(parents=True, exist_ok=True)
+
+                file_downloader = _FileDownloader(
+                    teamspace_id=teamspace_id,
+                    remote_path=file.filename,
+                    file_path=str(local_file),
+                    num_workers=num_workers,
+                    progress_bar=pbar,
+                    executor=part_executor,
+                    url=file.url,
+                    size=int(file.size_bytes),
+                    refresh_fn=lambda f=file: refresh_fn(f.filename),
+                )
+
+                page_futures.append(file_executor.submit(file_downloader.download))
+
+            if page_futures:
+                concurrent.futures.wait(page_futures)
+
+            pbar.set_description("Download complete")
+
+
 def _create_app(
     client: CloudSpaceServiceApi,
     studio_id: str,

lightning_sdk/cli/legacy/create.py

@@ -7,10 +7,8 @@ import click
 from rich.console import Console
 
 from lightning_sdk import Machine, Studio
-from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
 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) and machine._include_in_cli]
@@ -83,14 +81,6 @@ def studio(
     menu = _TeamspacesMenu()
     teamspace_resolved = menu._resolve_teamspace(teamspace)
 
-    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
     if cloud_account is None:
@@ -107,11 +97,19 @@ def studio(
         console.print(f"Studio with name {name} already exists. Using {new_name} instead.")
         name = new_name
 
-    studio = Studio(name=name, teamspace=teamspace_resolved, cloud_account=cloud_account, create_ok=True)
+    studio = Studio(
+        name=name,
+        teamspace=teamspace_resolved,
+        cloud_account=cloud_account,
+        create_ok=True,
+        cloud_provider=cloud_provider,
+        provider=provider,
+    )
 
     console.print(f"Created Studio {studio.name}.")
 
     if start is not None:
         start_machine = getattr(Machine, start, start)
+        Studio.show_progress = True
         studio.start(start_machine)
         console.print(f"Started Studio {studio.name} on machine {start}")

lightning_sdk/cli/legacy/start.py

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

lightning_sdk/cli/legacy/switch.py

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

lightning_sdk/cli/studio/start.py

@@ -77,5 +77,6 @@ def start_studio(
             raise ValueError(f"Could not start Studio: '{studio_name}'. Does the Studio exist?") from None
         raise ValueError(f"Could not start Studio: '{studio_name}'. Please provide a Studio name") from None
 
+    Studio.show_progress = True
     studio.start(machine, interruptible=interruptible)
     click.echo(f"Studio '{studio.name}' started successfully")

lightning_sdk/cli/studio/switch.py

@@ -47,6 +47,7 @@ def switch_studio(
         raise ValueError(f"Could not switch Studio: '{studio_name}'. Please provide a Studio name") from None
 
     resolved_machine = Machine.from_str(machine)
+    Studio.show_progress = True
     studio.switch_machine(resolved_machine, interruptible=interruptible)
 
     click.echo(f"Studio '{studio.name}' switched to machine '{resolved_machine}' successfully")

lightning_sdk/lightning_cloud/openapi/__init__.py

@@ -638,6 +638,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_like_status import V1LikeSt
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_affiliate_links_response import V1ListAffiliateLinksResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_agent_job_artifacts_response import V1ListAgentJobArtifactsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_agent_jobs_response import V1ListAgentJobsResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_list_aggregated_pod_metrics_response import V1ListAggregatedPodMetricsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_assistants_response import V1ListAssistantsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_blog_posts_response import V1ListBlogPostsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_cloud_space_apps_response import V1ListCloudSpaceAppsResponse
@@ -852,6 +853,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_quest import V1Quest
 from lightning_sdk.lightning_cloud.openapi.models.v1_quest_status import V1QuestStatus
 from lightning_sdk.lightning_cloud.openapi.models.v1_queue_server_type import V1QueueServerType
 from lightning_sdk.lightning_cloud.openapi.models.v1_quotas import V1Quotas
+from lightning_sdk.lightning_cloud.openapi.models.v1_quote_annual_upsell_response import V1QuoteAnnualUpsellResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_quote_subscription_response import V1QuoteSubscriptionResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_r2_data_connection import V1R2DataConnection
 from lightning_sdk.lightning_cloud.openapi.models.v1_refresh_index_response import V1RefreshIndexResponse

lightning_sdk/lightning_cloud/openapi/api/billing_service_api.py

@@ -1356,6 +1356,91 @@ class BillingServiceApi(object):
             _request_timeout=params.get('_request_timeout'),
             collection_formats=collection_formats)
 
+    def billing_service_quote_annual_upsell(self, **kwargs) -> 'V1QuoteAnnualUpsellResponse':  # noqa: E501
+        """billing_service_quote_annual_upsell  # noqa: E501
+
+        This method makes a synchronous HTTP request by default. To make an
+        asynchronous HTTP request, please pass async_req=True
+        >>> thread = api.billing_service_quote_annual_upsell(async_req=True)
+        >>> result = thread.get()
+
+        :param async_req bool
+        :return: V1QuoteAnnualUpsellResponse
+                 If the method is called asynchronously,
+                 returns the request thread.
+        """
+        kwargs['_return_http_data_only'] = True
+        if kwargs.get('async_req'):
+            return self.billing_service_quote_annual_upsell_with_http_info(**kwargs)  # noqa: E501
+        else:
+            (data) = self.billing_service_quote_annual_upsell_with_http_info(**kwargs)  # noqa: E501
+            return data
+
+    def billing_service_quote_annual_upsell_with_http_info(self, **kwargs) -> 'V1QuoteAnnualUpsellResponse':  # noqa: E501
+        """billing_service_quote_annual_upsell  # noqa: E501
+
+        This method makes a synchronous HTTP request by default. To make an
+        asynchronous HTTP request, please pass async_req=True
+        >>> thread = api.billing_service_quote_annual_upsell_with_http_info(async_req=True)
+        >>> result = thread.get()
+
+        :param async_req bool
+        :return: V1QuoteAnnualUpsellResponse
+                 If the method is called asynchronously,
+                 returns the request thread.
+        """
+
+        all_params = []  # noqa: E501
+        all_params.append('async_req')
+        all_params.append('_return_http_data_only')
+        all_params.append('_preload_content')
+        all_params.append('_request_timeout')
+
+        params = locals()
+        for key, val in six.iteritems(params['kwargs']):
+            if key not in all_params:
+                raise TypeError(
+                    "Got an unexpected keyword argument '%s'"
+                    " to method billing_service_quote_annual_upsell" % key
+                )
+            params[key] = val
+        del params['kwargs']
+
+        collection_formats = {}
+
+        path_params = {}
+
+        query_params = []
+
+        header_params = {}
+
+        form_params = []
+        local_var_files = {}
+
+        body_params = None
+        # HTTP header `Accept`
+        header_params['Accept'] = self.api_client.select_header_accept(
+            ['application/json'])  # noqa: E501
+
+        # Authentication setting
+        auth_settings = []  # noqa: E501
+
+        return self.api_client.call_api(
+            '/v1/billing/annual-upsell', 'GET',
+            path_params,
+            query_params,
+            header_params,
+            body=body_params,
+            post_params=form_params,
+            files=local_var_files,
+            response_type='V1QuoteAnnualUpsellResponse',  # noqa: E501
+            auth_settings=auth_settings,
+            async_req=params.get('async_req'),
+            _return_http_data_only=params.get('_return_http_data_only'),
+            _preload_content=params.get('_preload_content', True),
+            _request_timeout=params.get('_request_timeout'),
+            collection_formats=collection_formats)
+
     def billing_service_quote_subscription(self, **kwargs) -> 'V1QuoteSubscriptionResponse':  # noqa: E501
         """billing_service_quote_subscription  # noqa: E501