lightning-sdk 0.1.49__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/api/utils.py

@@ -6,7 +6,7 @@ import re
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, TypedDict, Union
 
 import backoff
 import requests
@@ -19,11 +19,12 @@ from lightning_sdk.lightning_cloud.openapi import (
     ModelsStoreApi,
     ProjectIdStorageBody,
     StorageCompleteBody,
+    StorageServiceApi,
     UploadIdCompleteBody,
     UploadIdPartsBody,
-    UploadsUploadIdBody,
     V1CompletedPart,
     V1CompleteUpload,
+    V1PathMapping,
     V1PresignedUrl,
     V1SignedUrl,
     V1UploadProjectArtifactPartsResponse,
@@ -36,6 +37,12 @@ try:
     from lightning_sdk.lightning_cloud.openapi import AppsIdBody1 as AppsIdBody
 except ImportError:
     from lightning_sdk.lightning_cloud.openapi import AppsIdBody
+
+try:
+    from lightning_sdk.lightning_cloud.openapi import UploadsUploadIdBody1 as UploadsUploadIdBody
+except ImportError:
+    from lightning_sdk.lightning_cloud.openapi import UploadsUploadIdBody
+
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.machine import Machine
@@ -89,6 +96,8 @@ class _FileUploader:
                 unit="B",
                 unit_scale=True,
                 unit_divisor=1000,
+                position=-1,
+                mininterval=1,
             )
         else:
             self.progress_bar = None
@@ -137,7 +146,7 @@ class _FileUploader:
 
     def _request_urls(self, parts: List[int], upload_id: str) -> List[V1PresignedUrl]:
         """Requests urls for a batch of parts."""
-        body = UploadsUploadIdBody(cluster_id=self.cloud_account, filename=self.remote_path, parts=parts)
+        body = UploadsUploadIdBody(filename=self.remote_path, parts=parts, cluster_id=self.cloud_account)
         resp: V1UploadProjectArtifactPartsResponse = self.client.storage_service_upload_project_artifact_parts(
             body, self.teamspace_id, upload_id
         )
@@ -194,7 +203,6 @@ class _ModelFileUploader:
         model_id: str,
         version: str,
         teamspace_id: str,
-        cloud_account: str,
         file_path: str,
         remote_path: str,
         progress_bar: bool,
@@ -216,8 +224,9 @@ class _ModelFileUploader:
                 unit="B",
                 unit_scale=True,
                 unit_divisor=1000,
-                position=1,
                 leave=False,
+                position=-1,
+                mininterval=1,
             )
         else:
             self.progress_bar = None
@@ -321,38 +330,14 @@ class _DummyResponse:
         self.data = data
 
 
-# TODO: This should really come from some kind of metadata service
-_MACHINE_TO_COMPUTE_NAME: Dict[Machine, str] = {
-    Machine.CPU_SMALL: "m3.medium",
-    Machine.CPU: "cpu-4",
-    Machine.DATA_PREP: "data-large",
-    Machine.DATA_PREP_MAX: "data-max",
-    Machine.DATA_PREP_ULTRA: "data-ultra",
-    Machine.T4: "g4dn.2xlarge",
-    Machine.T4_X_4: "g4dn.12xlarge",
-    Machine.L4: "g6.4xlarge",
-    Machine.L4_X_4: "g6.12xlarge",
-    Machine.L4_X_8: "g6.48xlarge",
-    Machine.A10G: "g5.8xlarge",
-    Machine.A10G_X_4: "g5.12xlarge",
-    Machine.A10G_X_8: "g5.48xlarge",
-    Machine.L40S: "g6e.4xlarge",
-    Machine.L40S_X_4: "g6e.12xlarge",
-    Machine.L40S_X_8: "g6e.48xlarge",
-    Machine.A100_X_8: "p4d.24xlarge",
-    Machine.H100_X_8: "p5.48xlarge",
-    Machine.H200_X_8: "p5e.48xlarge",
-}
-
-
 def _machine_to_compute_name(machine: Union[Machine, str]) -> str:
     if isinstance(machine, Machine):
-        return _MACHINE_TO_COMPUTE_NAME[machine]
+        if machine.instance_type is not None:
+            return machine.instance_type
+        return machine.slug
     return machine
 
 
-_COMPUTE_NAME_TO_MACHINE: Dict[str, Machine] = {v: k for k, v in _MACHINE_TO_COMPUTE_NAME.items()}
-
 _DEFAULT_CLOUD_URL = "https://lightning.ai"
 _DEFAULT_REGISTRY_URL = "litcr.io"
 
@@ -370,45 +355,53 @@ def _get_registry_url() -> str:
 
 
 def _sanitize_studio_remote_path(path: str, studio_id: str) -> str:
-    return f"/cloudspaces/{studio_id}/code/content/{path.replace('/teamspace/studios/this_studio/', '')}"
+    path = path.replace("/teamspace/studios/this_studio/", "")
+    root = f"/cloudspaces/{studio_id}/code/content/"
+    return os.path.join(root, path)
+
+
+def _resolve_teamspace_remote_path(path: str) -> str:
+    return f"/Uploads/{path.replace('/teamspace/studios/this_studio/', '')}"
 
 
 _DOWNLOAD_REQUEST_CHUNK_SIZE = 10 * _BYTES_PER_MB
 _DOWNLOAD_MIN_CHUNK_SIZE = 100 * _BYTES_PER_KB
 
 
+class _RefreshResponse(TypedDict):
+    url: str
+    size: int
+
+
 class _FileDownloader:
     def __init__(
         self,
-        client: LightningClient,
-        model_id: str,
-        version: str,
         teamspace_id: str,
         remote_path: str,
         file_path: str,
         executor: ThreadPoolExecutor,
         num_workers: int = 20,
         progress_bar: Optional[tqdm] = None,
+        url: Optional[str] = None,
+        size: Optional[int] = None,
+        refresh_fn: Optional[Callable[[], _RefreshResponse]] = None,
     ) -> None:
-        self.api = ModelsStoreApi(client.api_client)
-        self.model_id = model_id
-        self.version = version
         self.teamspace_id = teamspace_id
         self.local_path = file_path
         self.remote_path = remote_path
         self.progress_bar = progress_bar
         self.num_workers = num_workers
-        self._url = ""
-        self._size = 0
+        self._url = url
+        self._size = size
         self.executor = executor
+        self.refresh_fn = refresh_fn
 
     @backoff.on_exception(backoff.expo, ApiException, max_tries=10)
     def refresh(self) -> None:
-        response = self.api.models_store_get_model_file_url(
-            project_id=self.teamspace_id, model_id=self.model_id, version=self.version, filepath=self.remote_path
-        )
-        self._url = response.url
-        self._size = int(response.size)
+        if self.refresh_fn is not None:
+            response = self.refresh_fn()
+            self._url = response["url"]
+            self._size = response["size"]
 
     @property
     def url(self) -> str:
@@ -423,6 +416,11 @@ class _FileDownloader:
             return
         self.progress_bar.update(n)
 
+    def update_filename(self, desc: str) -> None:
+        if self.progress_bar is None:
+            return
+        self.progress_bar.set_description(f"{(desc[:72] + '...') if len(desc) > 75 else desc:<75.75}")
+
     @backoff.on_exception(backoff.expo, (requests.exceptions.HTTPError), max_tries=10)
     def _download_chunk(self, filename: str, start_end: Tuple[int]) -> None:
         start, end = start_end
@@ -457,6 +455,8 @@ class _FileDownloader:
                     f.write(b"\x00" * remaining_size)
 
     def _multipart_download(self, filename: str, num_workers: int) -> None:
+        self.update_filename(f"Downloading {self.remote_path}")
+
         num_chunks = num_workers
         chunk_size = math.ceil(self.size / num_chunks)
 
@@ -474,7 +474,8 @@ class _FileDownloader:
         concurrent.futures.wait(futures)
 
     def download(self) -> None:
-        self.refresh()
+        if self.url is None:
+            self.refresh()
 
         tmp_filename = f"{self.local_path}.download"
 
@@ -511,8 +512,8 @@ def _get_model_version(client: LightningClient, teamspace_id: str, name: str, ve
         raise ValueError(f"Model `{name}` does not exist")
     elif len(models) > 1:
         raise ValueError("Multiple models with the same name found")
-    if version == "latest":
-        return models[0].latest_version
+    if version is None or version == "default":
+        return models[0].default_version
     versions = api.models_store_list_model_versions(project_id=teamspace_id, model_id=models[0].id).versions
     if not versions:
         raise ValueError(f"Model `{name}` does not have any versions")
@@ -545,7 +546,18 @@ def _download_model_files(
             total=float(response.size_bytes),
             unit_scale=True,
             unit_divisor=1000,
+            position=-1,
+            mininterval=1,
+        )
+
+    def refresh_fn(filename: str) -> _RefreshResponse:
+        resp = api.models_store_get_model_file_url(
+            project_id=response.project_id,
+            model_id=response.model_id,
+            version=response.version,
+            filepath=filename,
         )
+        return {"url": resp.url, "size": int(resp.size)}
 
     with ThreadPoolExecutor(max_workers=min(num_workers, len(response.filepaths))) as file_executor, ThreadPoolExecutor(
         max_workers=num_workers
@@ -557,15 +569,13 @@ def _download_model_files(
             local_file.parent.mkdir(parents=True, exist_ok=True)
 
             file_downloader = _FileDownloader(
-                client=client,
-                model_id=response.model_id,
-                version=response.version,
                 teamspace_id=response.project_id,
                 remote_path=filepath,
                 file_path=str(local_file),
                 num_workers=num_workers,
                 progress_bar=pbar,
                 executor=part_executor,
+                refresh_fn=lambda f=filepath: refresh_fn(f),
             )
 
             futures.append(file_executor.submit(file_downloader.download))
@@ -576,6 +586,78 @@ def _download_model_files(
         return response.filepaths
 
 
+def _download_teamspace_files(
+    client: LightningClient,
+    teamspace_id: str,
+    cluster_id: str,
+    prefix: str,
+    download_dir: Path,
+    progress_bar: bool,
+    num_workers: int = os.cpu_count() * 4,
+) -> None:
+    api = StorageServiceApi(client.api_client)
+    response = None
+
+    pbar = None
+    if progress_bar:
+        pbar = tqdm(
+            desc="Downloading files",
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1000,
+            position=-1,
+            mininterval=1,
+        )
+
+    def refresh_fn(filename: str) -> _RefreshResponse:
+        resp = api.storage_service_list_project_artifacts(
+            project_id=teamspace_id,
+            cluster_id=cluster_id,
+            page_token="",
+            include_download_url=True,
+            prefix=prefix + filename,
+            page_size=1,
+        )
+        return {"url": resp.artifacts[0].url, "size": int(resp.artifacts[0].size_bytes)}
+
+    with ThreadPoolExecutor(max_workers=num_workers) as file_executor, ThreadPoolExecutor(
+        max_workers=num_workers
+    ) as part_executor:
+        while response is None or (response is not None and response.next_page_token != ""):
+            response = api.storage_service_list_project_artifacts(
+                project_id=teamspace_id,
+                cluster_id=cluster_id,
+                page_token=response.next_page_token if response is not None else "",
+                include_download_url=True,
+                prefix=prefix,
+                page_size=1000,
+            )
+
+            page_futures = []
+            for file in response.artifacts:
+                local_file = download_dir / file.filename
+                local_file.parent.mkdir(parents=True, exist_ok=True)
+
+                file_downloader = _FileDownloader(
+                    teamspace_id=teamspace_id,
+                    remote_path=file.filename,
+                    file_path=str(local_file),
+                    num_workers=num_workers,
+                    progress_bar=pbar,
+                    executor=part_executor,
+                    url=file.url,
+                    size=int(file.size_bytes),
+                    refresh_fn=lambda f=file: refresh_fn(f.filename),
+                )
+
+                page_futures.append(file_executor.submit(file_downloader.download))
+
+            if page_futures:
+                concurrent.futures.wait(page_futures)
+
+            pbar.set_description("Download complete")
+
+
 def _create_app(
     client: CloudSpaceServiceApi,
     studio_id: str,
@@ -614,3 +696,54 @@ def remove_datetime_prefix(text: str) -> str:
     # lines looks something like
     # '[2025-01-08T14:15:03.797142418Z] ⚡  ~ echo Hello\n[2025-01-08T14:15:03.803077717Z] Hello\n'
     return re.sub(r"^\[.*?\] ", "", text, flags=re.MULTILINE)
+
+
+def resolve_path_mappings(
+    mappings: Dict[str, str],
+    artifacts_local: Optional[str],
+    artifacts_remote: Optional[str],
+) -> List[V1PathMapping]:
+    path_mappings_list = []
+    for k, v in mappings.items():
+        splitted = str(v).rsplit(":", 1)
+        connection_name: str
+        connection_path: str
+        if len(splitted) == 1:
+            connection_name = splitted[0]
+            connection_path = ""
+        else:
+            connection_name, connection_path = splitted
+
+        path_mappings_list.append(
+            V1PathMapping(
+                connection_name=connection_name,
+                connection_path=connection_path,
+                container_path=k,
+            )
+        )
+
+    if artifacts_remote:
+        splitted = str(artifacts_remote).rsplit(":", 2)
+        if len(splitted) not in (2, 3):
+            raise RuntimeError(
+                f"Artifacts remote need to be of format efs:connection_name[:path] but got {artifacts_remote}"
+            )
+        else:
+            if not artifacts_local:
+                raise RuntimeError("If Artifacts remote is specified, artifacts local should be specified as well")
+
+            if len(splitted) == 2:
+                _, connection_name = splitted
+                connection_path = ""
+            else:
+                _, connection_name, connection_path = splitted
+
+            path_mappings_list.append(
+                V1PathMapping(
+                    connection_name=connection_name,
+                    connection_path=connection_path,
+                    container_path=artifacts_local,
+                )
+            )
+
+    return path_mappings_list

lightning_sdk/base_studio.py

@@ -0,0 +1,123 @@
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+from lightning_sdk.api.base_studio_api import BaseStudioApi
+from lightning_sdk.api.user_api import UserApi
+from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_environment_type import V1CloudSpaceEnvironmentType
+from lightning_sdk.organization import Organization
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.user import User
+from lightning_sdk.utils.resolve import _resolve_teamspace
+
+
+@dataclass
+class BaseStudioInfo:
+    id: str
+    name: str
+    managed_id: str
+    description: str
+    creator: str
+    enabled: bool
+
+
+class BaseStudio:
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        teamspace: Optional[Union[str, Teamspace]] = None,
+        org: Optional[Union[str, Organization]] = None,
+        user: Optional[Union[str, User]] = None,
+    ) -> None:
+        """Initializes the BaseStudio instance with organization and user information.
+
+        Args:
+            org (Optional[Union[str, Organization]]): The organization for the base studio. If not provided,
+                                                      it will be resolved through the authentication process.
+            user (Optional[Union[str, User]]): The user for the base studio. If not provided, it will be resolved
+                                               through the authentication process.
+
+        Raises:
+            ConnectionError: If there is an issue with the authentication process.
+        """
+        self._teamspace = None
+
+        _teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+        if _teamspace is None:
+            raise ValueError("Couldn't resolve teamspace from the provided name, org, or user")
+
+        self._teamspace = _teamspace
+
+        # self._auth = login.Auth()
+        # self._user = None
+
+        # try:
+        #     self._auth.authenticate()
+        #     if user is None:
+        #         self._user = User(name=UserApi()._get_user_by_id(self._auth.user_id).username)
+        # except ConnectionError as e:
+        #     raise e
+
+        # self._user = _resolve_user(self._user or user)
+        # self._org = _resolve_org(org)
+
+        self._base_studio_api = BaseStudioApi()
+
+        if name is not None:
+            org_id = self._teamspace._org.id if self._teamspace._org is not None else None
+            base_studio = self._base_studio_api.get_base_studio(name, org_id)
+
+            if base_studio is None:
+                raise ValueError(f"Base studio with name {name} does not exist")
+            self._base_studio = base_studio
+
+    def update(
+        self,
+        name: Optional[str] = None,
+        allowed_machines: Optional[List[str]] = None,
+        default_machine: Optional[str] = None,
+        disabled: Optional[bool] = None,
+        environment_type: Optional[V1CloudSpaceEnvironmentType] = None,
+        machine_image_version: Optional[str] = None,
+        setup_script_text: Optional[str] = None,
+    ) -> None:
+        org_id = self._teamspace._org.id if self._teamspace._org is not None else None
+        # TODO: if not in an org, can't update them
+        self._base_studio = self._base_studio_api.update_base_studio(
+            self._base_studio.id,
+            org_id,
+            name=name,
+            allowed_machines=allowed_machines,
+            default_machine=default_machine,
+            environment_type=environment_type,
+            machine_image_version=machine_image_version,
+            setup_script_text=setup_script_text,
+            disabled=disabled,
+        )
+
+    def list(self, include_disabled: bool = False) -> List[BaseStudioInfo]:
+        """List all base studios in the organization.
+
+        Args:
+            managed: Whether to filter for managed base studios.
+            include_disabled: Whether to include disabled base studios in the results.
+
+        Returns:
+            List[BaseStudioInfo]: A list of base studio templates.
+        """
+        org_id = self._teamspace._org.id if self._teamspace._org is not None else None
+        templates = self._base_studio_api.get_all_base_studios(org_id).templates
+
+        return [
+            BaseStudioInfo(
+                id=template.id,
+                name=template.name,
+                managed_id=template.managed_id,
+                description=template.description,
+                creator="⚡ Lightning AI"
+                if template.managed_id
+                else UserApi()._get_user_by_id(template.user_id).username,
+                enabled=not template.disabled,
+            )
+            for template in templates
+            if include_disabled or not template.disabled
+        ]

lightning_sdk/cli/__init__.py

@@ -0,0 +1 @@
+"""Lightning SDK CLI module."""

lightning_sdk/cli/base_studio/__init__.py

@@ -0,0 +1,10 @@
+"""Base Studio CLI commands."""
+
+import click
+
+
+def register_commands(group: click.Group) -> None:
+    """Register base studio commands with the given group."""
+    from lightning_sdk.cli.base_studio.list import list_base_studios
+
+    group.add_command(list_base_studios)

lightning_sdk/cli/base_studio/list.py

@@ -0,0 +1,43 @@
+"""Base Studio list command."""
+
+import click
+from rich.table import Table
+
+from lightning_sdk.base_studio import BaseStudio
+from lightning_sdk.cli.utils.richt_print import rich_to_str
+
+
+@click.command("list")
+@click.option("--include-disabled", help="Include disabled Base Studios in the list.", is_flag=True)
+def list_base_studios(include_disabled: bool) -> None:
+    """List Base Studios in an org.
+
+    Example:
+        lightning base-studio list
+
+    """
+    return list_impl(include_disabled=include_disabled)
+
+
+def list_impl(include_disabled: bool) -> None:
+    base_studio_cls = BaseStudio()
+    base_studios = base_studio_cls.list(include_disabled=include_disabled)
+
+    table = Table(
+        pad_edge=True,
+    )
+
+    table.add_column("Name")
+    table.add_column("Description")
+    table.add_column("Creator")
+    table.add_column("Enabled")
+
+    for base_studio in base_studios:
+        table.add_row(
+            base_studio.name.lower().replace(" ", "-"),
+            base_studio.description or "",
+            base_studio.creator,
+            "Yes" if base_studio.enabled else "No",
+        )
+
+    click.echo(rich_to_str(table), color=True)

lightning_sdk/cli/config/__init__.py

@@ -0,0 +1,14 @@
+"""Config CLI commands."""
+
+import click
+
+
+def register_commands(group: click.Group) -> None:
+    """Register config commands with the given group."""
+    from lightning_sdk.cli.config.get import get
+    from lightning_sdk.cli.config.set import set_value
+    from lightning_sdk.cli.config.show import show
+
+    group.add_command(get)
+    group.add_command(set_value)
+    group.add_command(show)

lightning_sdk/cli/config/get.py

@@ -0,0 +1,57 @@
+import click
+
+from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+
+@click.group("get")
+def get() -> None:
+    """Get configuration values."""
+
+
+@get.command("user")
+def get_user() -> None:
+    """Get the default user name from the config."""
+    config = Config()
+    user = config.get_value(DefaultConfigKeys.user)
+    click.echo(user)
+
+
+@get.command("org")
+def get_org() -> None:
+    """Get the default organization name from the config."""
+    config = Config()
+    org = config.get_value(DefaultConfigKeys.organization)
+    click.echo(org)
+
+
+@get.command("teamspace")
+def get_teamspace() -> None:
+    """Get the default teamspace name from the config."""
+    config = Config()
+    teamspace_name = config.get_value(DefaultConfigKeys.teamspace_name)
+    teamspace_owner = config.get_value(DefaultConfigKeys.teamspace_owner)
+    click.echo(f"{teamspace_owner}/{teamspace_name}")
+
+
+@get.command("studio")
+def get_studio() -> None:
+    """Get the default sutdio name from the config."""
+    config = Config()
+    studio = config.get_value(DefaultConfigKeys.studio)
+    click.echo(studio)
+
+
+@get.command("cloud-account")
+def get_cloud_account() -> None:
+    """Get the default cloud account name from the config."""
+    config = Config()
+    cloud_account = config.get_value(DefaultConfigKeys.cloud_account)
+    click.echo(cloud_account)
+
+
+@get.command("cloud-provider")
+def get_cloud_provider() -> None:
+    """Get the default cloud provider name from the config."""
+    config = Config()
+    cloud_provider = config.get_value(DefaultConfigKeys.cloud_provider)
+    click.echo(cloud_provider)