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

lightning_sdk/__version__.py

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

lightning_sdk/api/k8s_api.py

@@ -1,8 +1,8 @@
 import json
 import logging
-from typing import Any, Dict, TypedDict, Union
-
-import pandas as pd
+from collections import defaultdict
+from datetime import datetime
+from typing import Any, Dict, List, TypedDict
 
 from lightning_sdk.api.utils import ApiException
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
@@ -63,42 +63,88 @@ class K8sClusterApi:
         except Exception:
             return str(e.reason)
 
-    def get_billing_usage(self, print_data: bool = False, **kwargs: Dict[str, Any]) -> Union[pd.DataFrame, pd.Series]:
+    def get_billing_usage(self, print_data: bool = False, **kwargs: Dict[str, Any]) -> List[Dict[str, Any]]:
         """Gets the k8s usage metrics.
 
         Returns:
-            The k8s usage metrics as a DataFrame or Series.
+            The k8s usage metrics as a list of dictionaries.
         """
         try:
             response = self._client.k8_s_cluster_service_list_cluster_metrics(self.cloud_account, **kwargs)
             cluster_metrics = [entry.to_dict() for entry in response.cluster_metrics]
 
-            df = pd.DataFrame.from_records(cluster_metrics)
-            if df.empty:
-                return df
-
-            df["hour"] = pd.to_datetime(df["timestamp"]).dt.floor("h")
-
-            # Calculate the mean of num_allocated_gpus for each hour
-            aggregated = df.groupby("hour", as_index=False)["num_allocated_gpus"].mean()
-            # Merge the aggregated values back into the original DataFrame
-            df = df.merge(aggregated, on="hour", suffixes=("", "_mean"))
-
-            # Replace the original num_allocated_gpus with the mean values
-            df["num_allocated_gpus"] = df["num_allocated_gpus_mean"]
-
-            # We group the data by hour and take the first occurrence to avoid duplicates
-            df = df.drop_duplicates(subset="hour", keep="first")
-
-            # Convert timestamp to hourly floor and rename columns
-            df["billed_gpus"] = df.apply(_calculate_billed_k8s_gpus, axis=1)
+            if not cluster_metrics:
+                return []
+
+            # Parse timestamps and floor to hour, then group by hour
+            hourly_data = defaultdict(lambda: {"allocated_gpus": [], "first_entry": None})
+
+            for entry in cluster_metrics:
+                # Parse timestamp and floor to hour
+                timestamp = entry["timestamp"]
+                if isinstance(timestamp, str):
+                    timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
+                hour = timestamp.replace(minute=0, second=0, microsecond=0)
+
+                # Store allocated GPUs for averaging
+                hourly_data[hour]["allocated_gpus"].append(entry["num_allocated_gpus"])
+
+                # Keep first entry for each hour (for other fields)
+                if hourly_data[hour]["first_entry"] is None:
+                    hourly_data[hour]["first_entry"] = entry
+
+            # Build result list with aggregated data
+            result = []
+            for hour, data in sorted(hourly_data.items()):
+                entry = data["first_entry"]
+
+                # Calculate mean of allocated GPUs for this hour
+                mean_allocated_gpus = sum(data["allocated_gpus"]) / len(data["allocated_gpus"])
+
+                # Create row with mean allocated GPUs
+                row = {
+                    "hour": hour,
+                    "num_gpus": entry["num_gpus"],
+                    "num_requested_gpus": entry["num_requested_gpus"],
+                    "num_allocated_gpus": mean_allocated_gpus,
+                }
+
+                # Calculate billed GPUs
+                row["billed_gpus"] = _calculate_billed_k8s_gpus(
+                    {
+                        "num_allocated_gpus": mean_allocated_gpus,
+                        "num_requested_gpus": row["num_requested_gpus"],
+                        "num_gpus": row["num_gpus"],
+                    }
+                )
+
+                result.append(row)
 
-            # Keep only the required columns
-            df = df[["hour", "num_gpus", "num_requested_gpus", "num_allocated_gpus", "billed_gpus"]]
             if print_data:
-                with pd.option_context("display.max_rows", None, "display.max_columns", None):
-                    print(df)
-            return df
+                # Print using rich table (local import)
+                from rich.console import Console
+                from rich.table import Table
+
+                table = Table(title="K8s Billing Usage")
+                table.add_column("Hour", style="cyan")
+                table.add_column("Available GPUs", justify="right", style="green")
+                table.add_column("Requested GPUs", justify="right", style="yellow")
+                table.add_column("Allocated GPUs (Avg)", justify="right", style="magenta")
+                table.add_column("Billed GPUs", justify="right", style="red")
+
+                for row in result:
+                    table.add_row(
+                        str(row["hour"]),
+                        str(row["num_gpus"]),
+                        str(row["num_requested_gpus"]),
+                        f"{row['num_allocated_gpus']:.2f}",
+                        str(row["billed_gpus"]),
+                    )
+
+                console = Console()
+                console.print(table)
+
+            return result
         except ApiException as e:
             msg = self._parse_request_failure_body(e)
             logger.error(f"Failed to retrieve Kubernetes usage data: {msg}")

lightning_sdk/api/studio_api.py

@@ -1,7 +1,9 @@
+import concurrent
 import json
 import os
 import time
 import warnings
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from threading import Event, Thread
 from typing import Any, Dict, Generator, List, Mapping, Optional, Tuple, Union
@@ -12,10 +14,8 @@ from tqdm import tqdm
 
 from lightning_sdk.api.utils import (
     _create_app,
-    _download_teamspace_files,
     _DummyBody,
     _DummyResponse,
-    _FileUploader,
     _machine_to_compute_name,
     _sanitize_studio_remote_path,
 )
@@ -410,6 +410,21 @@ class StudioApi:
 
         progress.complete("Machine switch completed successfully")
 
+    def machine_is_supported(self, machine: Machine, teamspace_id: str, cloud_account_id: str, org_id: str) -> bool:
+        """Check if the machine is available in provided cloud_account."""
+        accelerators = self._get_machines_for_cloud_account(
+            teamspace_id=teamspace_id, cloud_account_id=cloud_account_id, org_id=org_id
+        )
+
+        for accelerator in accelerators:
+            if accelerator.accelerator_type == "GPU":
+                accelerator_resources_count = accelerator.resources.gpu
+            else:
+                accelerator_resources_count = accelerator.resources.cpu
+            if machine.accelerator_count == accelerator_resources_count and machine.family == accelerator.family:
+                return True
+        return False
+
     def machine_has_capacity(self, machine: Machine, teamspace_id: str, cloud_account_id: str, org_id: str) -> bool:
         """Check capacity of the requested machine."""
         accelerators = self._get_machines_for_cloud_account(
@@ -646,14 +661,21 @@ 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 get_tree(self, studio_id: str, teamspace_id: str, path: str) -> None:
+    def _authenticate_and_get_token(self) -> str:
+        """Authenticate and return a token for API requests."""
         auth = Auth()
         auth.authenticate()
-        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
+        return self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
 
-        query_params = {
-            "token": 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()
+
+        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/cloudspaces/{studio_id}/trees/{path}",
             params=query_params,
@@ -688,6 +710,16 @@ class StudioApi:
         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_files(
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        path: str = "",
+    ) -> List[Dict]:
+        """Recursively list all files in a directory tree."""
+        path = path.strip("/")
+        return self.get_tree(studio_id, teamspace_id, path, query_params={"recursive": "true"}).get("tree", [])
+
     def upload_file(
         self,
         studio_id: str,
@@ -698,14 +730,32 @@ class StudioApi:
         progress_bar: bool,
     ) -> None:
         """Uploads file to given remote path on the studio."""
-        _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,
-        )()
+        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)
+
+        if r.status_code == 200:
+            return
+        raise RuntimeError(f"Failed to upload file '{file_path}' to the Studio. Status code: {r.status_code}")
 
     def download_file(
         self,
@@ -718,9 +768,7 @@ class StudioApi:
     ) -> None:
         """Downloads a given file from a Studio 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 = self._authenticate_and_get_token()
 
         query_params = {
             "clusterId": cloud_account,
@@ -729,7 +777,7 @@ class StudioApi:
         }
 
         r = requests.get(
-            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
+            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/cloudspaces/{studio_id}/blobs/{path}",
             params=query_params,
             stream=True,
         )
@@ -756,6 +804,39 @@ class StudioApi:
                 f.write(chunk)
                 pbar_update(len(chunk))
 
+    def _download_single_studio_file(
+        self,
+        file_info: Dict,
+        base_path: str,
+        download_dir: Path,
+        studio_id: str,
+        teamspace_id: str,
+        token: str,
+        pbar: Optional[tqdm],
+    ) -> None:
+        """Download a single file from Studio 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,
+        }
+
+        r = requests.get(
+            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/cloudspaces/{studio_id}/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,
@@ -764,25 +845,106 @@ class StudioApi:
         teamspace_id: str,
         cloud_account: str,
         progress_bar: bool = True,
+        num_workers: Optional[int] = None,
     ) -> None:
         """Downloads a given folder from a Studio to a target location."""
         # TODO: implement resumable downloads
-        auth = Auth()
-        auth.authenticate()
 
-        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 + "/"
+        if num_workers is None:
+            num_workers = os.cpu_count() * 4
 
-        _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,
-        )
+        # Normalize the path
+        path = path.strip("/")
+        download_dir = Path(target_path)
+        download_dir.mkdir(parents=True, exist_ok=True)
+
+        files = self.list_files(studio_id, teamspace_id, path)
+
+        if not files:
+            print(f"No files found in {path}")
+            return
+
+        token = self._authenticate_and_get_token()
+
+        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_studio_file,
+                    file_info,
+                    path,
+                    download_dir,
+                    studio_id,
+                    teamspace_id,
+                    token,
+                    pbar,
+                )
+                for file_info in files
+            ]
+            concurrent.futures.wait(futures)
+
+        if pbar:
+            pbar.set_description("Download complete")
+            pbar.refresh()
+            pbar.close()
+
+    def remove_file(self, studio_id: str, teamspace_id: str, path: str) -> None:
+        """Removes a file from a Studio."""
+        info = self.get_path_info(studio_id, teamspace_id, path=path)
+
+        if not info["exists"]:
+            raise FileNotFoundError(f"The path '{path}' does not exist in the Studio.")
+
+        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()
+
+        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/{path}"
+
+        r = requests.delete(url, params=query_params, timeout=30)
+
+        if r.status_code == 204:
+            return
+
+        raise RuntimeError(f"Failed to remove file '{path}' from the Studio. Status code: {r.status_code}")
+
+    def remove_folder(self, studio_id: str, teamspace_id: str, path: str) -> None:
+        """Removes a folder (directory) from a Studio."""
+        info = self.get_path_info(studio_id, teamspace_id, path=path)
+
+        if not info["exists"]:
+            raise FileNotFoundError(f"The path '{path}' does not exist in the Studio.")
+
+        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()
+
+        query_params = {"token": token}
+        client_host = self._client.api_client.configuration.host
+        url = f"{client_host}/v1/projects/{teamspace_id}/artifacts/cloudspaces/{studio_id}/trees/{path}"
+
+        r = requests.delete(url, params=query_params, timeout=30)
+
+        if r.status_code == 204:
+            return
+
+        raise RuntimeError(f"Failed to remove folder '{path}' from the Studio. Status code: {r.status_code}")
 
     def install_plugin(self, studio_id: str, teamspace_id: str, plugin_name: str) -> str:
         """Installs the given plugin."""

lightning_sdk/api/teamspace_api.py

@@ -10,7 +10,6 @@ from lightning_sdk.api.utils import (
     _download_model_files,
     _download_teamspace_files,
     _DummyBody,
-    _FileUploader,
     _get_model_version,
     _ModelFileUploader,
     _resolve_teamspace_remote_path,
@@ -401,14 +400,34 @@ class TeamspaceApi:
         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,
-        )()
+        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}")
 
     def download_file(
         self,

lightning_sdk/cli/cp/__init__.py

@@ -0,0 +1,64 @@
+"""CP CLI commands."""
+
+from typing import Any, Literal, Optional
+
+import click
+
+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 "uploads" in path[3]:
+        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)
+        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)
+        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:
+        try:
+            route_cp_operation(
+                source=source,
+                destination=destination,
+                recursive=recursive,
+            )
+        except Exception:
+            raise
+
+    command.callback = new_callback

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)