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

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,13 +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,
-    _FileUploader,
     _get_model_version,
     _ModelFileUploader,
-    _resolve_teamspace_remote_path,
 )
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
@@ -33,7 +34,6 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1ExternalCluster,
     V1GCSFolderDataConnection,
     V1Job,
-    V1LoginRequest,
     V1Model,
     V1ModelVersionArchive,
     V1MultiMachineJob,
@@ -392,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,
@@ -400,52 +452,58 @@ class TeamspaceApi:
         remote_path: str,
         progress_bar: bool,
     ) -> None:
-        """Uploads file to given remote path in the Teamspace drive."""
-        _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,
-        )()
+        """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
+        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}")
 
     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
-
-        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,
-            }
+        token = _authenticate_and_get_token(self._client)
 
-            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
+        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:
@@ -469,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

@@ -0,0 +1,67 @@
+"""CP CLI commands."""
+
+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
+
+
+def parse_lit_url(url: str) -> tuple[str, list[str], Literal["studios", "uploads"]]:
+    """Parse lit:// URL and extract resource type."""
+    if "://" not in url:
+        raise ValueError("URL must contain '://'")
+
+    path = url.split("://")[-1].split("/")
+
+    if path[2] == "studios":
+        resource_type = "studios"
+    elif path[2] == "uploads":
+        resource_type = "uploads"
+    else:
+        raise ValueError("URL must contain either 'studios' or 'uploads'")
+
+    return resource_type
+
+
+def route_cp_operation(source: str, destination: str, **options: Any) -> None:
+    """Route copy operation based on URL structure."""
+    source_is_lit = source.startswith("lit://")
+    dest_is_lit = destination.startswith("lit://")
+
+    if source_is_lit and dest_is_lit:
+        raise ValueError("Cannot copy between two remote URLs. One path must be local.")
+
+    if not source_is_lit and not dest_is_lit:
+        raise ValueError("At least one path must be a lit://")
+
+    if source_is_lit:
+        resource_type = parse_lit_url(source)
+        if resource_type == "studios":
+            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.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")
+
+
+def register_commands(command: click.Command) -> None:
+    """Register cp command callback."""
+
+    def new_callback(source: str, destination: Optional[str], recursive: bool, **kwargs: Any) -> None:
+        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}")

lightning_sdk/cli/entrypoint.py

@@ -12,6 +12,7 @@ from lightning_sdk.api.studio_api import _cloud_url
 from lightning_sdk.cli.groups import (
     base_studio,
     config,
+    cp,
     # job,
     license,
     # mmt,
@@ -77,6 +78,7 @@ main_cli.add_command(studio)
 main_cli.add_command(vm)
 main_cli.add_command(base_studio)
 main_cli.add_command(license)
+main_cli.add_command(cp)
 
 if os.environ.get("LIGHTNING_EXPERIMENTAL_CLI_ONLY", "0") != "1":
     #### LEGACY COMMANDS ####

lightning_sdk/cli/groups.py

@@ -4,6 +4,7 @@ import click
 
 from lightning_sdk.cli.base_studio import register_commands as register_base_studio_commands
 from lightning_sdk.cli.config import register_commands as register_config_commands
+from lightning_sdk.cli.cp import register_commands as register_cp_commands
 from lightning_sdk.cli.job import register_commands as register_job_commands
 from lightning_sdk.cli.license import register_commands as register_license_commands
 from lightning_sdk.cli.mmt import register_commands as register_mmt_commands
@@ -46,6 +47,26 @@ def license() -> None:  # noqa: A001
     """Manage Lightning AI Product Licenses."""
 
 
+@click.command(name="cp")
+@click.argument("source")
+@click.argument("destination", required=False)
+@click.option("--recursive", "-r", is_flag=True, help="Copy directories recursively")
+@click.pass_context
+def cp() -> None:
+    """Copy files between local filesystem, Studios, and teamspace drives.
+
+    \b
+    URL formats:
+      Studios:          lit://<owner>/<teamspace>/studios/<studio-name>/<path>
+      Teamspace drives: lit://<owner>/<teamspace>/uploads/<path>
+
+    \b
+    Examples:
+      lightning studio cp source.txt lit://<owner>/<my-teamspace>/studios/<my-studio>/destination.txt
+      lightning studio cp -r source_folder/ lit://<owner>/<my-teamspace>/studios/<my-studio>/destination_folder/
+    """
+
+
 # Register config commands with the main config group
 register_job_commands(job)
 register_mmt_commands(mmt)
@@ -54,3 +75,4 @@ register_config_commands(config)
 register_vm_commands(vm)
 register_base_studio_commands(base_studio)
 register_license_commands(license)
+register_cp_commands(cp)

lightning_sdk/cli/legacy/clusters_menu.py

@@ -32,8 +32,8 @@ class _ClustersMenu:
 
             cloud_account_api = CloudAccountApi()
 
-            resolved_cluster_obj = cloud_account_api.get_cloud_account(
-                cloud_account_id=selected_cluster_id, org_id=teamspace.owner.id, teamspace_id=teamspace.id
+            resolved_cluster_obj = cloud_account_api.get_cloud_account_non_org(
+                cloud_account_id=selected_cluster_id, teamspace_id=teamspace.id
             )
 
             return None if resolved_cluster_obj.spec.cluster_type == V1ClusterType.GLOBAL else resolved_cluster_obj.id

lightning_sdk/cli/legacy/deploy/_auth.py

@@ -8,7 +8,7 @@ from urllib.parse import urlencode
 from rich.console import Console
 from rich.prompt import Confirm
 
-from lightning_sdk import Teamspace
+from lightning_sdk import Organization, Teamspace
 from lightning_sdk.api import UserApi
 from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud import env
@@ -74,13 +74,14 @@ def authenticate(mode: _AuthMode, shall_confirm: bool = True) -> None:
 def select_teamspace(teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace:
     if teamspace is None:
         menu = TeamspacesMenu()
-        possible_teamspaces = menu._get_possible_teamspaces(_get_authed_user())
+        auth_user = _get_authed_user()
+        possible_teamspaces = menu._get_possible_teamspaces(auth_user)
         if len(possible_teamspaces) == 1:
-            name = next(iter(possible_teamspaces.values()))["name"]
-            return Teamspace(name=name, org=org, user=user)
-
+            teamspace_name = next(iter(possible_teamspaces.values()))
+            if isinstance(menu._owner, Organization):
+                return Teamspace(name=teamspace_name, org=menu._owner, user=None)
+            return Teamspace(name=teamspace_name, org=None, user=menu._owner)
         return menu(teamspace)
-
     return _resolve_teamspace(teamspace=teamspace, org=org, user=user)