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

lightning_sdk/__version__.py

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

lightning_sdk/api/studio_api.py

@@ -13,6 +13,7 @@ import requests
 from tqdm import tqdm
 
 from lightning_sdk.api.utils import (
+    _authenticate_and_get_token,
     _create_app,
     _DummyBody,
     _DummyResponse,
@@ -23,7 +24,6 @@ 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 +49,6 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1EnvVar,
     V1GetCloudSpaceInstanceStatusResponse,
     V1GetLongRunningCommandInCloudSpaceResponse,
-    V1LoginRequest,
     V1ManagedEndpoint,
     V1ManagedModel,
     V1Plugin,
@@ -661,14 +660,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 = {
@@ -730,7 +723,7 @@ class StudioApi:
         progress_bar: bool,
     ) -> None:
         """Uploads file to given remote path on the studio."""
-        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
@@ -768,7 +761,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 +857,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 +903,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 +926,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,11 @@ import requests
 from tqdm.auto import tqdm
 
 from lightning_sdk.api.utils import (
+    _authenticate_and_get_token,
     _download_model_files,
-    _download_teamspace_files,
     _DummyBody,
     _get_model_version,
     _ModelFileUploader,
-    _resolve_teamspace_remote_path,
 )
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
@@ -32,7 +34,6 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1ExternalCluster,
     V1GCSFolderDataConnection,
     V1Job,
-    V1LoginRequest,
     V1Model,
     V1ModelVersionArchive,
     V1MultiMachineJob,
@@ -391,6 +392,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,
@@ -399,10 +452,8 @@ class TeamspaceApi:
         remote_path: str,
         progress_bar: bool,
     ) -> 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
+        """Uploads file to given remote path in the Teamspace drive /Uploads/."""
+        token = _authenticate_and_get_token(self._client)
 
         query_params = {"token": token, "clusterId": cloud_account}
         client_host = self._client.api_client.configuration.host
@@ -434,37 +485,25 @@ class TeamspaceApi:
         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
-
-        cluster_ids = [ca.cluster_id for ca in self.list_cloud_accounts(teamspace_id)]
-
-        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,
-            )
+        token = _authenticate_and_get_token(self._client)
 
-            if r.status_code == 200:
-                found = True
-                break
+        query_params = {
+            "token": token,
+        }
 
-        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 +527,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,
@@ -816,3 +818,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

lightning_sdk/cli/cp/teamspace_uploads.py

@@ -0,0 +1,93 @@
+import os
+from pathlib import Path
+
+from rich.console import Console
+
+from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.cli.utils.filesystem import parse_teamspace_uploads_path, resolve_teamspace
+
+
+def cp_upload(
+    local_file_path: str,
+    teamspace_path: str,
+    options: dict[str, any],
+) -> None:
+    console = Console()
+    recursive = options.get("recursive", False)
+    cloud_account = options.get("cloud_account", None)
+    if not Path(local_file_path).exists():
+        raise FileNotFoundError(f"The provided path does not exist: {local_file_path}")
+
+    teamspace_path_result = parse_teamspace_uploads_path(teamspace_path)
+
+    selected_teamspace = resolve_teamspace(teamspace_path_result["teamspace"], teamspace_path_result["owner"])
+    console.print(f"Uploading to {selected_teamspace.owner.name}/{selected_teamspace.name}")
+
+    if Path(local_file_path).is_dir():
+        if not recursive:
+            raise ValueError(f"'{local_file_path}' is a directory. Use -r flag to copy directories recursively.")
+        selected_teamspace.upload_folder(
+            local_file_path, teamspace_path_result["destination"], cloud_account=cloud_account
+        )
+    else:
+        if teamspace_path.endswith(("/", "\\")):
+            # if destination ends with / or \, treat it as a directory
+            file_name = os.path.basename(local_file_path)
+            teamspace_path_result["destination"] = os.path.join(teamspace_path_result["destination"], file_name)
+        selected_teamspace.upload_file(
+            local_file_path, teamspace_path_result["destination"], cloud_account=cloud_account
+        )
+
+    studio_url = (
+        _get_cloud_url().replace(":443", "") + "/" + selected_teamspace.owner.name + "/" + selected_teamspace.name
+    )
+    console.print(f"See your file at {studio_url}")
+
+
+def cp_download(
+    teamspace_path: str,
+    local_path: str,
+    options: dict[str, any],
+) -> None:
+    console = Console()
+    teamspace_path_result = parse_teamspace_uploads_path(teamspace_path)
+    recursive = options.get("recursive", False)
+
+    selected_teamspace = resolve_teamspace(teamspace_path_result["teamspace"], teamspace_path_result["owner"])
+
+    # check if file/folder exists
+    path_info = selected_teamspace._teamspace_api.get_path_info(
+        selected_teamspace._teamspace.id, path=teamspace_path_result["destination"]
+    )
+    if not path_info["exists"]:
+        raise FileNotFoundError(
+            f"The provided path does not exist in the teamspace drive: {teamspace_path_result['destination']} "
+            "Note that empty folders may not be detected as existing."
+        )
+
+    console.print(f"Downloading from {selected_teamspace.owner.name}/{selected_teamspace.name}")
+    if path_info["type"] == "directory":
+        if not recursive:
+            raise ValueError(
+                f"'{teamspace_path_result['destination']}' is a directory. Use -r flag to copy directories recursively."
+            )
+        folder_name = os.path.basename(teamspace_path_result["destination"].rstrip("/"))
+        if local_path in ("./", "."):
+            if folder_name == "":
+                folder_name = f"{selected_teamspace.name}_downloads"
+            target_path = os.path.join(local_path, folder_name)
+        else:
+            target_path = local_path
+
+        selected_teamspace.download_folder(teamspace_path_result["destination"], target_path)
+        console.print(f"See your folder at {target_path}")
+    else:
+        if os.path.isdir(local_path) or local_path.endswith(("/", "\\")):
+            # if local_path ends with / or \ or is a directory, treat it as a directory
+            file_name = os.path.basename(teamspace_path_result["destination"])
+            target_path = os.path.join(local_path, file_name)
+        else:
+            target_path = local_path
+        os.makedirs(os.path.dirname(target_path), exist_ok=True)
+        selected_teamspace.download_file(teamspace_path_result["destination"], target_path)
+        console.print(f"See your file at {target_path}")