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

lightning_sdk/__init__.py

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

lightning_sdk/api/cloud_account_api.py

@@ -181,6 +181,11 @@ class CloudAccountApi:
         cloud_provider: Optional[Union["CloudProvider", str]],
         default_cloud_account: Optional[str],
     ) -> Optional[str]:
+        from lightning_sdk.machine import CloudProvider
+
+        if cloud_provider and not isinstance(cloud_provider, CloudProvider):
+            cloud_provider = CloudProvider(cloud_provider)
+
         if cloud_account:
             if cloud_provider:
                 cloud_account_resp = self.get_cloud_account_non_org(teamspace_id, cloud_account)

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,
@@ -569,46 +569,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/config/get.py

@@ -39,3 +39,19 @@ def get_studio() -> None:
     config = Config()
     studio = config.get_value(DefaultConfigKeys.studio)
     click.echo(studio)
+
+
+@get.command("cloud-account")
+def get_cloud_account() -> None:
+    """Get the default cloud account name from the config."""
+    config = Config()
+    cloud_account = config.get_value(DefaultConfigKeys.cloud_account)
+    click.echo(cloud_account)
+
+
+@get.command("cloud-provider")
+def get_cloud_provider() -> None:
+    """Get the default cloud provider name from the config."""
+    config = Config()
+    cloud_provider = config.get_value(DefaultConfigKeys.cloud_provider)
+    click.echo(cloud_provider)

lightning_sdk/cli/config/set.py

@@ -1,6 +1,7 @@
 import click
 
 from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.machine import CloudProvider
 from lightning_sdk.organization import Organization
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.config import Config, DefaultConfigKeys
@@ -75,3 +76,29 @@ def set_teamspace(teamspace_name: str) -> None:
         setattr(config, DefaultConfigKeys.teamspace_owner_type, "organization")
     else:
         setattr(config, DefaultConfigKeys.teamspace_owner_type, "user")
+
+
+@set_value.command("cloud-account")
+@click.argument("cloud_account_name")
+def set_cloud_account(cloud_account_name: str) -> None:
+    """Set the default cloud account name in the config."""
+    config = Config()
+    setattr(config, DefaultConfigKeys.cloud_account, cloud_account_name)
+
+
+@set_value.command("cloud-provider")
+@click.argument("cloud_provider_name")
+def set_cloud_provider(cloud_provider_name: str) -> None:
+    """Set the default cloud provider name in the config."""
+    config = Config()
+
+    try:
+        cloud_provider = CloudProvider(cloud_provider_name)
+    except ValueError:
+        # TODO: make this a generic CLI error
+        raise ValueError(
+            f"Could not resolve cloud provider: '{cloud_provider_name}'. "
+            f"Supported values are: {', '.join(m.name for m in list(CloudProvider))}"
+        ) from None
+
+    setattr(config, DefaultConfigKeys.cloud_provider, cloud_provider.name)

lightning_sdk/cli/studio/create.py

@@ -18,10 +18,16 @@ from lightning_sdk.studio import Studio
     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,
+)
 def create_studio(
     studio_name: Optional[str] = None,
     teamspace: Optional[str] = None,
     cloud_provider: Optional[str] = None,
+    cloud_account: Optional[str] = None,
 ) -> None:
     """Create a new Studio.
 
@@ -42,8 +48,17 @@ def create_studio(
     else:
         resolved_teamspace = None
 
+    if cloud_provider is not None:
+        cloud_provider = CloudProvider(cloud_provider)
+
     try:
-        studio = Studio(studio_name, teamspace=resolved_teamspace, create_ok=True, cloud_provider=cloud_provider)
+        studio = Studio(
+            studio_name,
+            teamspace=resolved_teamspace,
+            create_ok=True,
+            cloud_provider=cloud_provider,
+            cloud_account=cloud_account,
+        )
     except (RuntimeError, ValueError, ApiException) as e:
         print(e)
         if studio_name:

lightning_sdk/cli/studio/start.py

@@ -22,9 +22,17 @@ from lightning_sdk.studio import Studio
 @click.option("--interruptible", is_flag=True, help="Start the studio on an interruptible instance.")
 @click.option(
     "--cloud-provider",
-    help="The cloud provider to start the studio on. Defaults to teamspace default.",
+    help=(
+        "The cloud provider to start the studio on. Defaults to teamspace default. "
+        "Only used if --create is specified."
+    ),
     type=click.Choice(m.name for m in list(CloudProvider)),
 )
+@click.option(
+    "--cloud-account",
+    help="The cloud account to start the studio on. Defaults to teamspace default. Only used if --create is specified.",
+    type=click.STRING,
+)
 def start_studio(
     studio_name: Optional[str] = None,
     teamspace: Optional[str] = None,
@@ -32,6 +40,7 @@ def start_studio(
     machine: Optional[str] = None,
     interruptible: bool = False,
     cloud_provider: Optional[str] = None,
+    cloud_account: Optional[str] = None,
 ) -> None:
     """Start a Studio.
 
@@ -52,8 +61,17 @@ def start_studio(
     else:
         resolved_teamspace = None
 
+    if cloud_provider is not None:
+        cloud_provider = CloudProvider(cloud_provider)
+
     try:
-        studio = Studio(studio_name, teamspace=resolved_teamspace, create_ok=create, cloud_provider=cloud_provider)
+        studio = Studio(
+            studio_name,
+            teamspace=resolved_teamspace,
+            create_ok=create,
+            cloud_provider=cloud_provider,
+            cloud_account=cloud_account,
+        )
     except (RuntimeError, ValueError, ApiException):
         if studio_name:
             raise ValueError(f"Could not start Studio: '{studio_name}'. Does the Studio exist?") from None

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

lightning_sdk/lightning_cloud/openapi/api/k8_s_cluster_service_api.py

@@ -160,6 +160,119 @@ class K8SClusterServiceApi(object):
             _request_timeout=params.get('_request_timeout'),
             collection_formats=collection_formats)
 
+    def k8_s_cluster_service_list_aggregated_pod_metrics(self, project_id: 'str', cluster_id: 'str', **kwargs) -> 'V1ListAggregatedPodMetricsResponse':  # noqa: E501
+        """k8_s_cluster_service_list_aggregated_pod_metrics  # noqa: E501
+
+        This method makes a synchronous HTTP request by default. To make an
+        asynchronous HTTP request, please pass async_req=True
+        >>> thread = api.k8_s_cluster_service_list_aggregated_pod_metrics(project_id, cluster_id, async_req=True)
+        >>> result = thread.get()
+
+        :param async_req bool
+        :param str project_id: (required)
+        :param str cluster_id: (required)
+        :param str namespace:
+        :param datetime start: Date range.
+        :param datetime end:
+        :return: V1ListAggregatedPodMetricsResponse
+                 If the method is called asynchronously,
+                 returns the request thread.
+        """
+        kwargs['_return_http_data_only'] = True
+        if kwargs.get('async_req'):
+            return self.k8_s_cluster_service_list_aggregated_pod_metrics_with_http_info(project_id, cluster_id, **kwargs)  # noqa: E501
+        else:
+            (data) = self.k8_s_cluster_service_list_aggregated_pod_metrics_with_http_info(project_id, cluster_id, **kwargs)  # noqa: E501
+            return data
+
+    def k8_s_cluster_service_list_aggregated_pod_metrics_with_http_info(self, project_id: 'str', cluster_id: 'str', **kwargs) -> 'V1ListAggregatedPodMetricsResponse':  # noqa: E501
+        """k8_s_cluster_service_list_aggregated_pod_metrics  # noqa: E501
+
+        This method makes a synchronous HTTP request by default. To make an
+        asynchronous HTTP request, please pass async_req=True
+        >>> thread = api.k8_s_cluster_service_list_aggregated_pod_metrics_with_http_info(project_id, cluster_id, async_req=True)
+        >>> result = thread.get()
+
+        :param async_req bool
+        :param str project_id: (required)
+        :param str cluster_id: (required)
+        :param str namespace:
+        :param datetime start: Date range.
+        :param datetime end:
+        :return: V1ListAggregatedPodMetricsResponse
+                 If the method is called asynchronously,
+                 returns the request thread.
+        """
+
+        all_params = ['project_id', 'cluster_id', 'namespace', 'start', 'end']  # 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 k8_s_cluster_service_list_aggregated_pod_metrics" % key
+                )
+            params[key] = val
+        del params['kwargs']
+        # verify the required parameter 'project_id' is set
+        if ('project_id' not in params or
+                params['project_id'] is None):
+            raise ValueError("Missing the required parameter `project_id` when calling `k8_s_cluster_service_list_aggregated_pod_metrics`")  # noqa: E501
+        # verify the required parameter 'cluster_id' is set
+        if ('cluster_id' not in params or
+                params['cluster_id'] is None):
+            raise ValueError("Missing the required parameter `cluster_id` when calling `k8_s_cluster_service_list_aggregated_pod_metrics`")  # noqa: E501
+
+        collection_formats = {}
+
+        path_params = {}
+        if 'project_id' in params:
+            path_params['projectId'] = params['project_id']  # noqa: E501
+        if 'cluster_id' in params:
+            path_params['clusterId'] = params['cluster_id']  # noqa: E501
+
+        query_params = []
+        if 'namespace' in params:
+            query_params.append(('namespace', params['namespace']))  # noqa: E501
+        if 'start' in params:
+            query_params.append(('start', params['start']))  # noqa: E501
+        if 'end' in params:
+            query_params.append(('end', params['end']))  # noqa: E501
+
+        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/projects/{projectId}/clusters/{clusterId}/aggregated-metrics/pods', 'GET',
+            path_params,
+            query_params,
+            header_params,
+            body=body_params,
+            post_params=form_params,
+            files=local_var_files,
+            response_type='V1ListAggregatedPodMetricsResponse',  # 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 k8_s_cluster_service_list_cluster_metrics(self, project_id: 'str', cluster_id: 'str', **kwargs) -> 'V1ListClusterMetricsResponse':  # noqa: E501
         """k8_s_cluster_service_list_cluster_metrics  # noqa: E501
 

lightning_sdk/lightning_cloud/openapi/models/__init__.py

@@ -590,6 +590,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