lightning-sdk 2025.12.17__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__ = "2025.12.17"
+__version__ = "2026.01.27"

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
@@ -11,11 +13,10 @@ import requests
 from tqdm import tqdm
 
 from lightning_sdk.api.utils import (
+    _authenticate_and_get_token,
     _create_app,
-    _download_teamspace_files,
     _DummyBody,
     _DummyResponse,
-    _FileUploader,
     _machine_to_compute_name,
     _sanitize_studio_remote_path,
 )
@@ -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,
@@ -410,6 +409,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 +660,15 @@ 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:
-        auth = Auth()
-        auth.authenticate()
-        token = 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 = _authenticate_and_get_token(self._client)
 
-        query_params = {
-            "token": 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 +703,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 +723,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 = _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}"
+
+        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 +761,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 = _authenticate_and_get_token(self._client)
 
         query_params = {
             "clusterId": cloud_account,
@@ -729,7 +770,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 +797,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 +838,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 = _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_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 = _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/{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 = _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}/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."""