lightning-sdk 2026.1.22__py3-none-any.whl → 2026.1.30__py3-none-any.whl

lightning_sdk/__version__.py

@@ -1,3 +1,3 @@
 """Version information for lightning_sdk."""
 
-__version__ = "2026.01.22"
+__version__ = "2026.01.30"

lightning_sdk/api/studio_api.py

@@ -13,17 +13,20 @@ import requests
 from tqdm import tqdm
 
 from lightning_sdk.api.utils import (
+    _MAX_SIZE_MULTI_PART_CHUNK,
+    _authenticate_and_get_token,
     _create_app,
     _DummyBody,
     _DummyResponse,
+    _FileUploader,
     _machine_to_compute_name,
     _sanitize_studio_remote_path,
+    _SinglePartFileUploader,
 )
 from lightning_sdk.api.utils import (
     _get_cloud_url as _cloud_url,
 )
 from lightning_sdk.constants import _LIGHTNING_DEBUG
-from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     AssistantsServiceCreateAssistantBody,
     AssistantsServiceCreateAssistantManagedEndpointBody,
@@ -49,7 +52,6 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1EnvVar,
     V1GetCloudSpaceInstanceStatusResponse,
     V1GetLongRunningCommandInCloudSpaceResponse,
-    V1LoginRequest,
     V1ManagedEndpoint,
     V1ManagedModel,
     V1Plugin,
@@ -661,14 +663,8 @@ class StudioApi:
         self.stop_keeping_alive(teamspace_id=teamspace_id, studio_id=studio_id)
         self._client.cloud_space_service_delete_cloud_space(project_id=teamspace_id, id=studio_id)
 
-    def _authenticate_and_get_token(self) -> str:
-        """Authenticate and return a token for API requests."""
-        auth = Auth()
-        auth.authenticate()
-        return self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
-
     def get_tree(self, studio_id: str, teamspace_id: str, path: str, query_params: Optional[dict] = None) -> None:
-        token = self._authenticate_and_get_token()
+        token = _authenticate_and_get_token(self._client)
 
         if query_params is None:
             query_params = {
@@ -729,33 +725,36 @@ class StudioApi:
         remote_path: str,
         progress_bar: bool,
     ) -> None:
-        """Uploads file to given remote path on the studio."""
-        token = self._authenticate_and_get_token()
-
-        query_params = {"token": token}
-        client_host = self._client.api_client.configuration.host
-        url = f"{client_host}/v1/projects/{teamspace_id}/artifacts/cloudspaces/{studio_id}/blobs/{remote_path}"
-
-        filesize = os.path.getsize(file_path)
-        with open(file_path, "rb") as f:
-            if progress_bar:
-                filesize = os.path.getsize(file_path)
-                with tqdm.wrapattr(
-                    f,
-                    "read",
-                    desc=f"Uploading {os.path.split(file_path)[1]}",
-                    total=filesize,
-                    unit="B",
-                    unit_scale=True,
-                    unit_divisor=1000,
-                ) as wrapped_file:
-                    r = requests.put(url, data=wrapped_file, params=query_params, timeout=30)
-            else:
-                r = requests.put(url, data=f, params=query_params, timeout=30)
+        """Uploads file to given remote path in the studio.
 
-        if r.status_code == 200:
-            return
-        raise RuntimeError(f"Failed to upload file '{file_path}' to the Studio. Status code: {r.status_code}")
+        Uses single-part upload for files <= 5MB, multipart upload for larger files.
+        """
+        file_size = os.path.getsize(file_path)
+        multipart_threshold = int(os.environ.get("LIGHTNING_MULTIPART_THRESHOLD", _MAX_SIZE_MULTI_PART_CHUNK))
+
+        if file_size <= multipart_threshold:
+            token = _authenticate_and_get_token(self._client)
+
+            query_params = {"token": token}
+            client_host = self._client.api_client.configuration.host
+            url = f"{client_host}/v1/projects/{teamspace_id}/artifacts/cloudspaces/{studio_id}/blobs/{remote_path}"
+
+            _SinglePartFileUploader(
+                client=self._client,
+                file_path=file_path,
+                url=url,
+                query_params=query_params,
+                progress_bar=progress_bar,
+            )()
+        else:
+            _FileUploader(
+                client=self._client,
+                teamspace_id=teamspace_id,
+                cloud_account=cloud_account,
+                file_path=file_path,
+                remote_path=_sanitize_studio_remote_path(remote_path, studio_id),
+                progress_bar=progress_bar,
+            )()
 
     def download_file(
         self,
@@ -768,7 +767,7 @@ class StudioApi:
     ) -> None:
         """Downloads a given file from a Studio to a target location."""
         # TODO: Update this endpoint to permit basic auth
-        token = self._authenticate_and_get_token()
+        token = _authenticate_and_get_token(self._client)
 
         query_params = {
             "clusterId": cloud_account,
@@ -864,7 +863,7 @@ class StudioApi:
             print(f"No files found in {path}")
             return
 
-        token = self._authenticate_and_get_token()
+        token = _authenticate_and_get_token(self._client)
 
         total_size = sum(f.get("size", 0) for f in files)
 
@@ -910,7 +909,7 @@ class StudioApi:
         if info["type"] != "file":
             raise IsADirectoryError(f"The path '{path}' is a directory. Use 'remove_folder()' to remove directories.")
 
-        token = self._authenticate_and_get_token()
+        token = _authenticate_and_get_token(self._client)
 
         query_params = {"token": token}
         client_host = self._client.api_client.configuration.host
@@ -933,7 +932,7 @@ class StudioApi:
         if info["type"] == "file":
             raise ValueError(f"The path '{path}' is a file. Use 'remove_file()' to remove files.")
 
-        token = self._authenticate_and_get_token()
+        token = _authenticate_and_get_token(self._client)
 
         query_params = {"token": token}
         client_host = self._client.api_client.configuration.host

lightning_sdk/api/teamspace_api.py

@@ -1,5 +1,8 @@
+import concurrent
 import os
 import re
+import warnings
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
@@ -7,12 +10,15 @@ import requests
 from tqdm.auto import tqdm
 
 from lightning_sdk.api.utils import (
+    _MAX_SIZE_MULTI_PART_CHUNK,
+    _authenticate_and_get_token,
     _download_model_files,
-    _download_teamspace_files,
     _DummyBody,
+    _FileUploader,
     _get_model_version,
     _ModelFileUploader,
     _resolve_teamspace_remote_path,
+    _SinglePartFileUploader,
 )
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
@@ -32,7 +38,6 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1ExternalCluster,
     V1GCSFolderDataConnection,
     V1Job,
-    V1LoginRequest,
     V1Model,
     V1ModelVersionArchive,
     V1MultiMachineJob,
@@ -391,6 +396,58 @@ class TeamspaceApi:
         response = self.models_api.models_store_list_model_versions(project_id=teamspace_id, model_id=model_id)
         return response.versions
 
+    def get_uploads_tree(self, teamspace_id: str, path: str, query_params: Optional[dict] = None) -> None:
+        token = _authenticate_and_get_token(self._client)
+
+        if query_params is None:
+            query_params = {
+                "token": token,
+            }
+        else:
+            query_params["token"] = token
+        r = requests.get(
+            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/uploads/trees/{path}",
+            params=query_params,
+        )
+        return r.json()
+
+    def get_path_info(self, teamspace_id: str, path: str = "") -> dict:
+        path = path.strip("/")
+
+        if "/" in path:
+            parent_path = path.rsplit("/", 1)[0]
+            target_name = path.rsplit("/", 1)[1]
+        else:
+            if path == "":
+                # root directory
+                return {"exists": True, "type": "directory", "size": None}
+            parent_path = ""
+            target_name = path
+
+        tree = self.get_uploads_tree(teamspace_id, path=parent_path)
+        tree_items = tree.get("tree", [])
+        for item in tree_items:
+            item_name = item.get("path", "")
+            if item_name == target_name:
+                item_type = item.get("type")
+                # if type == "blob" it's a file, if "tree" it's a directory
+                return {
+                    "exists": True,
+                    "type": "file" if item_type == "blob" else "directory",
+                    "size": item.get("size", 0) if item_type == "blob" else None,
+                }
+        warnings.warn(f"If '{path}' is a directory, it may be empty and thus not detected.")
+        return {"exists": False, "type": None, "size": None}
+
+    def list_uploads_files(
+        self,
+        teamspace_id: str,
+        path: str = "",
+    ) -> List[Dict]:
+        """Recursively list all files in a /Uploads/ directory tree."""
+        path = path.strip("/")
+        return self.get_uploads_tree(teamspace_id, path, query_params={"recursive": "true"}).get("tree", [])
+
     def upload_file(
         self,
         teamspace_id: str,
@@ -398,73 +455,65 @@ class TeamspaceApi:
         file_path: str,
         remote_path: str,
         progress_bar: bool,
+        headers: Optional[Dict[str, str]] = None,
     ) -> None:
-        """Uploads file to given remote path in the Teamspace drive."""
-        auth = Auth()
-        auth.authenticate()
-        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
-
-        query_params = {"token": token, "clusterId": cloud_account}
-        client_host = self._client.api_client.configuration.host
-        url = f"{client_host}/v1/projects/{teamspace_id}/artifacts/uploads/blobs/{remote_path}"
-
-        filesize = os.path.getsize(file_path)
-        with open(file_path, "rb") as f:
-            if progress_bar:
-                filesize = os.path.getsize(file_path)
-                with tqdm.wrapattr(
-                    f,
-                    "read",
-                    desc=f"Uploading {os.path.split(file_path)[1]}",
-                    total=filesize,
-                    unit="B",
-                    unit_scale=True,
-                    unit_divisor=1000,
-                ) as wrapped_file:
-                    r = requests.put(url, data=wrapped_file, params=query_params, timeout=30)
-            else:
-                r = requests.put(url, data=f, params=query_params, timeout=30)
-
-        if r.status_code == 200:
-            return
-        raise RuntimeError(f"Failed to upload file '{file_path}' to the Teamspace drive. Status code: {r.status_code}")
+        """Uploads file to given remote path in the Teamspace drive.
+
+        Uses single-part upload for files <= 5MB, multipart upload for larger files.
+        """
+        file_size = os.path.getsize(file_path)
+
+        multipart_threshold = int(os.environ.get("LIGHTNING_MULTIPART_THRESHOLD", _MAX_SIZE_MULTI_PART_CHUNK))
+
+        if file_size <= multipart_threshold:
+            token = _authenticate_and_get_token(self._client)
+
+            query_params = {"token": token, "clusterId": cloud_account}
+            client_host = self._client.api_client.configuration.host
+            url = f"{client_host}/v1/projects/{teamspace_id}/artifacts/blobs/{remote_path}"
+
+            _SinglePartFileUploader(
+                client=self._client,
+                file_path=file_path,
+                url=url,
+                query_params=query_params,
+                progress_bar=progress_bar,
+                headers=headers,
+            )()
+        else:
+            _FileUploader(
+                client=self._client,
+                teamspace_id=teamspace_id,
+                cloud_account=cloud_account,
+                file_path=file_path,
+                remote_path=_resolve_teamspace_remote_path(remote_path),
+                progress_bar=progress_bar,
+            )()
 
     def download_file(
         self,
         path: str,
         target_path: str,
         teamspace_id: str,
+        cloud_account: Optional[str] = None,
         progress_bar: bool = True,
     ) -> None:
-        """Downloads a given file in Teamspace drive to a target location."""
+        """Downloads a given file in Teamspace drive /Uploads/ to a target location."""
         # TODO: Update this endpoint to permit basic auth
-        auth = Auth()
-        auth.authenticate()
-        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
+        token = _authenticate_and_get_token(self._client)
 
-        cluster_ids = [ca.cluster_id for ca in self.list_cloud_accounts(teamspace_id)]
+        query_params = {
+            "token": token,
+        }
 
-        found = False
-        for cluster_id in cluster_ids:
-            query_params = {
-                "clusterId": cluster_id,
-                "key": _resolve_teamspace_remote_path(path),
-                "token": token,
-            }
-
-            r = requests.get(
-                f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
-                params=query_params,
-                stream=True,
-            )
-
-            if r.status_code == 200:
-                found = True
-                break
-
-        if not found:
-            raise FileNotFoundError(f"The provided path does not exist in the teamspace: {path}")
+        if cloud_account:
+            query_params["clusterId"] = cloud_account
 
+        r = requests.get(
+            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/uploads/blobs/{path}",
+            params=query_params,
+            stream=True,
+        )
         total_length = int(r.headers.get("content-length"))
 
         if progress_bar:
@@ -488,33 +537,101 @@ class TeamspaceApi:
                 f.write(chunk)
                 pbar_update(len(chunk))
 
+    def _download_single_file(
+        self,
+        file_info: Dict,
+        base_path: str,
+        download_dir: Path,
+        teamspace_id: str,
+        token: str,
+        cloud_account: Optional[str] = None,
+        pbar: Optional[tqdm] = True,
+    ) -> None:
+        """Download a single file from Teamspace drive /Uploads/ with progress tracking."""
+        relative_path = file_info["path"].lstrip("/")
+        local_file = download_dir / relative_path
+        local_file.parent.mkdir(parents=True, exist_ok=True)
+
+        file_path = os.path.join(base_path, relative_path) if base_path else relative_path
+
+        query_params = {
+            "token": token,
+        }
+        if cloud_account:
+            query_params["clusterId"] = cloud_account
+
+        r = requests.get(
+            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/uploads/blobs/{file_path}",
+            params=query_params,
+            stream=True,
+        )
+
+        with open(str(local_file), "wb") as f:
+            for chunk in r.iter_content(chunk_size=4096 * 8):
+                f.write(chunk)
+                if pbar:
+                    pbar.update(len(chunk))
+
     def download_folder(
         self,
         path: str,
         target_path: str,
         teamspace_id: str,
-        cloud_account: str,
+        cloud_account: Optional[str] = None,
         progress_bar: bool = True,
+        num_workers: Optional[int] = None,
     ) -> None:
-        """Downloads a given folder from Teamspace drive to a target location."""
+        """Downloads a given folder from Teamspace drive /Uploads/ to a target location."""
         # TODO: Update this endpoint to permit basic auth
-        auth = Auth()
-        auth.authenticate()
+        if num_workers is None:
+            num_workers = os.cpu_count() * 4
 
-        prefix = _resolve_teamspace_remote_path(path)
+        # Normalize the path
+        path = path.strip("/")
+        download_dir = Path(target_path)
+        download_dir.mkdir(parents=True, exist_ok=True)
 
-        # ensure we only download as a directory and not the entire prefix
-        if prefix.endswith("/") is False:
-            prefix = prefix + "/"
+        files = self.list_uploads_files(teamspace_id, path)
 
-        _download_teamspace_files(
-            client=self._client,
-            teamspace_id=teamspace_id,
-            cluster_id=cloud_account,
-            prefix=prefix,
-            download_dir=Path(target_path),
-            progress_bar=progress_bar,
-        )
+        if not files:
+            print(f"No files found in {path}")
+            return
+
+        token = _authenticate_and_get_token(self._client)
+
+        total_size = sum(f.get("size", 0) for f in files)
+
+        pbar = None
+        if progress_bar:
+            pbar = tqdm(
+                desc="Downloading files",
+                total=total_size,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1000,
+                mininterval=1,
+            )
+
+        with ThreadPoolExecutor(max_workers=num_workers) as executor:
+            futures = [
+                executor.submit(
+                    self._download_single_file,
+                    file_info,
+                    path,
+                    download_dir,
+                    teamspace_id,
+                    token,
+                    cloud_account,
+                    pbar,
+                )
+                for file_info in files
+            ]
+            concurrent.futures.wait(futures)
+
+        if pbar:
+            pbar.set_description("Download complete")
+            pbar.refresh()
+            pbar.close()
 
     def get_secrets(self, teamspace_id: str) -> Dict[str, str]:
         """Get all secrets for a teamspace."""

lightning_sdk/api/utils.py

@@ -15,6 +15,7 @@ import requests
 from tqdm.auto import tqdm
 
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__, _LIGHTNING_DEBUG
+from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     CloudSpaceServiceApi,
     CloudSpaceServiceCreateCloudSpaceAppInstanceBody,
@@ -29,6 +30,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     StorageServiceUploadProjectArtifactPartsBody,
     V1CompletedPart,
     V1CompleteUpload,
+    V1LoginRequest,
     V1PathMapping,
     V1PresignedUrl,
     V1SignedUrl,
@@ -57,6 +59,66 @@ _MAX_BATCH_SIZE = 50
 _MAX_WORKERS = 10
 
 
+class _SinglePartFileUploader:
+    """A class handling upload files to studio and teamspace drive with new endpoint."""
+
+    def __init__(
+        self,
+        client: LightningClient,
+        file_path: str,
+        url: str,
+        query_params: Dict[str, str],
+        progress_bar: bool,
+        headers: Optional[Dict[str, str]] = None,
+    ) -> None:
+        self.client = client
+        self.local_path = file_path
+        self.url = url
+        self.query_params = query_params
+        self.headers = headers
+        self.filesize = os.path.getsize(file_path)
+
+        if progress_bar:
+            self.progress_bar = tqdm(
+                desc=f"Uploading {os.path.split(file_path)[1]}",
+                total=self.filesize,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1000,
+                position=-1,
+                mininterval=1,
+            )
+        else:
+            self.progress_bar = None
+
+    def __call__(self) -> None:
+        self._upload_with_retry()
+
+    @backoff.on_exception(
+        backoff.expo, (requests.exceptions.HTTPError, requests.exceptions.RequestException), max_tries=10
+    )
+    def _upload_with_retry(self) -> None:
+        with open(self.local_path, "rb") as f:
+            if self.progress_bar is not None:
+                with tqdm.wrapattr(
+                    f,
+                    "read",
+                    desc=f"Uploading {os.path.split(self.local_path)[1]}",
+                    total=self.filesize,
+                    unit="B",
+                    unit_scale=True,
+                    unit_divisor=1000,
+                ) as wrapped_file:
+                    r = requests.put(
+                        self.url, data=wrapped_file, params=self.query_params, timeout=30, headers=self.headers
+                    )
+            else:
+                r = requests.put(self.url, data=f, params=self.query_params, timeout=30, headers=self.headers)
+
+        if r.status_code != 200:
+            raise RuntimeError(f"Failed to upload file '{self.local_path}'. Status code: {r.status_code}")
+
+
 class _FileUploader:
     """A class handling the upload to studios.
 
@@ -356,7 +418,7 @@ def _sanitize_studio_remote_path(path: str, studio_id: str) -> str:
 
 
 def _resolve_teamspace_remote_path(path: str) -> str:
-    return f"/Uploads/{path.replace('/teamspace/uploads/', '')}"
+    return f"{path.replace('/teamspace/', '')}"
 
 
 _DOWNLOAD_REQUEST_CHUNK_SIZE = 10 * _BYTES_PER_MB
@@ -816,3 +878,9 @@ def to_iso_z(dt: datetime) -> str:
         dt = dt.replace(tzinfo=timezone.utc)
         return dt.astimezone(timezone.utc).isoformat(timespec="milliseconds")
     return dt.isoformat(timespec="milliseconds")
+
+
+def _authenticate_and_get_token(client: Any) -> str:
+    auth = Auth()
+    auth.authenticate()
+    return client.auth_service_login(V1LoginRequest(auth.api_key)).token

lightning_sdk/cli/cp/__init__.py

@@ -4,6 +4,8 @@ from typing import Any, Literal, Optional
 
 import click
 
+from lightning_sdk.cli.cp.teamspace_uploads import cp_download as teamspace_uploads_cp_download
+from lightning_sdk.cli.cp.teamspace_uploads import cp_upload as teamspace_uploads_cp_upload
 from lightning_sdk.cli.studio.cp import cp_download as studio_cp_download
 from lightning_sdk.cli.studio.cp import cp_upload as studio_cp_upload
 
@@ -17,7 +19,7 @@ def parse_lit_url(url: str) -> tuple[str, list[str], Literal["studios", "uploads
 
     if path[2] == "studios":
         resource_type = "studios"
-    elif "uploads" in path[3]:
+    elif path[2] == "uploads":
         resource_type = "uploads"
     else:
         raise ValueError("URL must contain either 'studios' or 'uploads'")
@@ -39,12 +41,16 @@ def route_cp_operation(source: str, destination: str, **options: Any) -> None:
     if source_is_lit:
         resource_type = parse_lit_url(source)
         if resource_type == "studios":
-            return studio_cp_download(source, destination, options)
+            return studio_cp_download(source, destination, options.get("recursive", False))
+        if resource_type == "uploads":
+            return teamspace_uploads_cp_download(source, destination, options)
         raise ValueError(f"Resource type: {resource_type} is not supported")
     else:
         resource_type = parse_lit_url(destination)
         if resource_type == "studios":
-            return studio_cp_upload(source, destination, options)
+            return studio_cp_upload(source, destination, options.get("recursive", False))
+        if resource_type == "uploads":
+            return teamspace_uploads_cp_upload(source, destination, options)
         raise ValueError(f"Resource type: {resource_type} is not supported")
 
 
@@ -52,13 +58,10 @@ def register_commands(command: click.Command) -> None:
     """Register cp command callback."""
 
     def new_callback(source: str, destination: Optional[str], recursive: bool, **kwargs: Any) -> None:
-        try:
-            route_cp_operation(
-                source=source,
-                destination=destination,
-                recursive=recursive,
-            )
-        except Exception:
-            raise
+        route_cp_operation(
+            source=source,
+            destination=destination,
+            recursive=recursive,
+        )
 
     command.callback = new_callback