lightning-sdk 2025.8.6rc2__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/api/user_api.py

@@ -1,15 +1,20 @@
-from typing import List, Union
+import re
+from typing import Dict, List, Optional, Union
 
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
+    SecretsIdBody1,
     V1CloudSpace,
+    V1CreateUserSecretRequest,
     V1GetUserResponse,
     V1ListCloudSpacesResponse,
     V1Membership,
     V1Organization,
     V1SearchUser,
+    V1Secret,
     V1UserFeatures,
 )
+from lightning_sdk.lightning_cloud.openapi.models.v1_secret_type import V1SecretType
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
 
@@ -56,8 +61,14 @@ class UserApi:
     def _get_all_teamspace_memberships(
         self,
         user_id: str,  # todo: this is unused, but still required
+        org_id: Optional[str] = None,
     ) -> List[V1Membership]:
-        return self._client.projects_service_list_memberships(filter_by_user_id=True).memberships
+        kwargs: Dict[str, Union[bool, str]] = {"filter_by_user_id": True}
+
+        if org_id is not None:
+            kwargs["organization_id"] = org_id
+
+        return self._client.projects_service_list_memberships(**kwargs).memberships
 
     def _get_authed_user_name(self) -> str:
         """Gets the currently logged-in user."""
@@ -69,3 +80,46 @@ class UserApi:
     def _get_feature_flags(self) -> V1UserFeatures:
         resp: V1GetUserResponse = self._client.auth_service_get_user()
         return resp.features
+
+    def get_secrets(self) -> Dict[str, str]:
+        """Get all secrets for the current user."""
+        secrets = self._get_secrets()
+        # this returns encrypted values for security. It doesn't make sense to show them,
+        # so we just return a placeholder
+        # not a security issue to replace in the client as we get the encrypted values from the server.
+        return {secret.name: "***REDACTED***" for secret in secrets if secret.type == V1SecretType.UNSPECIFIED}
+
+    def set_secret(self, key: str, value: str) -> None:
+        """Set a secret for the current user.
+
+        This will replace the existing secret if it exists and create a new one if it doesn't.
+        """
+        secrets = self._get_secrets()
+        for secret in secrets:
+            if secret.name == key:
+                return self._update_secret(secret.id, value)
+        return self._create_secret(key, value)
+
+    def _get_secrets(self) -> List[V1Secret]:
+        return self._client.secret_service_list_user_secrets().secrets
+
+    def _update_secret(self, secret_id: str, value: str) -> None:
+        self._client.secret_service_update_user_secret(
+            body=SecretsIdBody1(value=value),
+            id=secret_id,
+        )
+
+    def _create_secret(
+        self,
+        key: str,
+        value: str,
+    ) -> None:
+        self._client.secret_service_create_user_secret(body=V1CreateUserSecretRequest(name=key, value=value))
+
+    def verify_secret_name(self, name: str) -> bool:
+        """Verify if a secret name is valid.
+
+        A valid secret name starts with a letter or underscore, followed by letters, digits, or underscores.
+        """
+        pattern = r"^[A-Za-z_][A-Za-z0-9_]*$"
+        return re.match(pattern, name) is not None

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,6 +19,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     ModelsStoreApi,
     ProjectIdStorageBody,
     StorageCompleteBody,
+    StorageServiceApi,
     UploadIdCompleteBody,
     UploadIdPartsBody,
     V1CompletedPart,
@@ -367,38 +368,40 @@ _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:
@@ -413,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
@@ -447,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)
 
@@ -464,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"
 
@@ -539,6 +550,15 @@ def _download_model_files(
             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
     ) as part_executor:
@@ -549,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))
@@ -568,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,

lightning_sdk/base_studio.py

@@ -3,11 +3,11 @@ 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 import login
 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_org, _resolve_user
+from lightning_sdk.utils.resolve import _resolve_teamspace
 
 
 @dataclass
@@ -16,12 +16,15 @@ class BaseStudioInfo:
     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:
@@ -36,26 +39,35 @@ class BaseStudio:
         Raises:
             ConnectionError: If there is an issue with the authentication process.
         """
-        self._auth = login.Auth()
-        self._user = None
+        self._teamspace = 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
+        _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._user = _resolve_user(self._user or user)
-        self._org = _resolve_org(org)
+        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:
-            base_studio = self._base_studio_api.get_base_studio(name, self._org.id)
+            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 in organization {self._org.name}")
+                raise ValueError(f"Base studio with name {name} does not exist")
             self._base_studio = base_studio
 
     def update(
@@ -68,9 +80,11 @@ class BaseStudio:
         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,
-            self._org.id,
+            org_id,
             name=name,
             allowed_machines=allowed_machines,
             default_machine=default_machine,
@@ -80,20 +94,30 @@ class BaseStudio:
             disabled=disabled,
         )
 
-    def list(self, managed: bool = True) -> List[BaseStudioInfo]:
+    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[V1CloudSpaceEnvironmentTemplate]: A list of base studio templates.
+            List[BaseStudioInfo]: A list of base studio templates.
         """
-        result = []
-        for template in self._base_studio_api.get_all_base_studios(self._org.id, managed).templates:
-            result.append(
-                BaseStudioInfo(
-                    id=template.id,
-                    name=template.name,
-                    managed_id=template.managed_id,
-                    description=template.description,
-                ),
+        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,
             )
-        return result
+            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)

lightning_sdk/cli/config/set.py

@@ -0,0 +1,92 @@
+import click
+
+from lightning_sdk.cli.utils.save_to_config import save_teamspace_to_config
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.machine import CloudProvider
+from lightning_sdk.studio import Studio
+from lightning_sdk.utils.config import Config, DefaultConfigKeys
+from lightning_sdk.utils.resolve import _resolve_org, _resolve_user
+
+
+@click.group("set")
+def set_value() -> None:
+    """Set configuration values."""
+
+
+@set_value.command("user")
+@click.argument("user_name")
+def set_user(user_name: str) -> None:
+    """Set the default user name in the config."""
+    try:
+        _resolve_user(user_name)
+    except Exception:
+        # TODO: make this a generic CLI error
+        raise ValueError(f"Could not resolve user: '{user_name}'. Does the user exist?") from None
+
+    config = Config()
+    setattr(config, DefaultConfigKeys.user, user_name)
+
+
+@set_value.command("org")
+@click.argument("org_name")
+def set_org(org_name: str) -> None:
+    """Set the default organization name in the config."""
+    try:
+        _resolve_org(org_name)
+    except Exception:
+        # TODO: make this a generic CLI error
+        raise ValueError(f"Could not resolve organization: '{org_name}'. Does the organization exist?") from None
+
+    config = Config()
+    setattr(config, DefaultConfigKeys.organization, org_name)
+
+
+@set_value.command("studio")
+@click.argument("studio_name")
+def set_studio(studio_name: str) -> None:
+    """Set the default studio name in the config."""
+    try:
+        studio = Studio(studio_name)
+    except Exception:
+        # TODO: make this a generic CLI error
+        raise ValueError(f"Could not resolve studio: '{studio_name}'. Does the studio exist?") from None
+
+    config = Config()
+    setattr(config, DefaultConfigKeys.studio, studio.name)
+
+
+@set_value.command("teamspace")
+@click.argument("teamspace_name")
+def set_teamspace(teamspace_name: str) -> None:
+    """Set the default teamspace name in the config."""
+    menu = TeamspacesMenu()
+    teamspace_resolved = menu(teamspace=teamspace_name)
+
+    # explicit user action, so overwrite the config
+    save_teamspace_to_config(teamspace_resolved, overwrite=True)
+
+
+@set_value.command("cloud-account")
+@click.argument("cloud_account_name")
+def set_cloud_account(cloud_account_name: str) -> None:
+    """Set the default cloud account name in the config."""
+    config = Config()
+    setattr(config, DefaultConfigKeys.cloud_account, cloud_account_name)
+
+
+@set_value.command("cloud-provider")
+@click.argument("cloud_provider_name")
+def set_cloud_provider(cloud_provider_name: str) -> None:
+    """Set the default cloud provider name in the config."""
+    config = Config()
+
+    try:
+        cloud_provider = CloudProvider(cloud_provider_name)
+    except ValueError:
+        # TODO: make this a generic CLI error
+        raise ValueError(
+            f"Could not resolve cloud provider: '{cloud_provider_name}'. "
+            f"Supported values are: {', '.join(m.name for m in list(CloudProvider))}"
+        ) from None
+
+    setattr(config, DefaultConfigKeys.cloud_provider, cloud_provider.name)

lightning_sdk/cli/config/show.py

@@ -0,0 +1,9 @@
+import click
+
+from lightning_sdk.utils.config import Config
+
+
+@click.command("show")
+def show() -> None:
+    """Show configuration values."""
+    click.echo(Config())