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

lightning_sdk/api/teamspace_api.py

@@ -1,26 +1,49 @@
 import os
+import re
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
+import requests
 from tqdm.auto import tqdm
 
-from lightning_sdk.api.utils import _download_model_files, _DummyBody, _get_model_version, _ModelFileUploader
+from lightning_sdk.api.utils import (
+    _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 (
+    Create,
     Externalv1LightningappInstance,
     ModelIdVersionsBody,
     ModelsStoreApi,
     ProjectIdAgentsBody,
     ProjectIdModelsBody,
+    ProjectIdSecretsBody,
+    SecretsIdBody,
     V1Assistant,
     V1CloudSpace,
+    V1ClusterAccelerator,
+    V1EfsConfig,
     V1Endpoint,
+    V1ExternalCluster,
+    V1GCSFolderDataConnection,
     V1Job,
+    V1LoginRequest,
+    V1Model,
     V1ModelVersionArchive,
     V1MultiMachineJob,
     V1Project,
     V1ProjectClusterBinding,
     V1PromptSuggestion,
+    V1R2DataConnection,
+    V1S3FolderDataConnection,
+    V1Secret,
+    V1SecretType,
     V1UpstreamOpenAI,
 )
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
@@ -33,7 +56,7 @@ class TeamspaceApi:
 
     def __init__(self) -> None:
         self._client = LightningClient(max_tries=7)
-        self._models: Optional[ModelsStoreApi] = None
+        self._models_api: Optional[ModelsStoreApi] = None
 
     def get_teamspace(self, name: str, owner_id: str) -> V1Project:
         """Get the current teamspace from the owner."""
@@ -50,7 +73,7 @@ class TeamspaceApi:
     def _get_teamspace_by_id(self, teamspace_id: str) -> V1Project:
         return self._client.projects_service_get_project(teamspace_id)
 
-    def list_teamspaces(self, owner_id: str, name: Optional[str] = None) -> Optional[V1Project]:
+    def list_teamspaces(self, owner_id: str, name: Optional[str] = None) -> Optional[List[V1Project]]:
         """Lists teamspaces from owner.
 
         If name is passed only teamspaces matching that name will be returned
@@ -165,48 +188,49 @@ class TeamspaceApi:
 
     # lazy property which is only created when needed
     @property
-    def models(self) -> ModelsStoreApi:
-        if not self._models:
-            self._models = ModelsStoreApi(self._client.api_client)
-        return self._models
+    def models_api(self) -> ModelsStoreApi:
+        if not self._models_api:
+            self._models_api = ModelsStoreApi(self._client.api_client)
+        return self._models_api
 
-    def get_model_version(self, name: str, version: str, teamspace_id: str) -> V1ModelVersionArchive:
+    def get_model_version(self, name: str, version: Optional[str], teamspace_id: str) -> V1ModelVersionArchive:
         return _get_model_version(client=self._client, name=name, version=version, teamspace_id=teamspace_id)
 
     def create_model(
         self,
         name: str,
+        version: Optional[str],
         metadata: Dict[str, str],
         private: bool,
         teamspace_id: str,
         cloud_account: str,
     ) -> V1ModelVersionArchive:
         # ask if such model already exists by listing models with specific name
-        models = self.models.models_store_list_models(project_id=teamspace_id, name=name).models
+        models = self.models_api.models_store_list_models(project_id=teamspace_id, name=name).models
         if len(models) == 0:
-            return self.models.models_store_create_model(
+            return self.models_api.models_store_create_model(
                 body=ProjectIdModelsBody(cluster_id=cloud_account, metadata=metadata, name=name, private=private),
                 project_id=teamspace_id,
             )
         assert len(models) == 1, "Multiple models with the same name found"
-        return self.models.models_store_create_model_version(
-            body=ModelIdVersionsBody(cluster_id=cloud_account),
+        return self.models_api.models_store_create_model_version(
+            body=ModelIdVersionsBody(cluster_id=cloud_account, version=version),
             project_id=teamspace_id,
             model_id=models[0].id,
         )
 
     def delete_model(self, name: str, version: Optional[str], teamspace_id: str) -> None:
         """Delete a model or a version from the model store."""
-        models = self.models.models_store_list_models(project_id=teamspace_id, name=name).models
-        assert len(models) == 1, "Multiple models with the same name found"
-        model_id = models[0].id
+        model = self.get_model(teamspace_id=teamspace_id, model_name=name)
         # decide if delete only version of whole model
         if version:
-            if version == "latest":
-                version = models[0].latest_version
-            self.models.models_store_delete_model_version(project_id=teamspace_id, model_id=model_id, version=version)
+            if version == "default":
+                version = model.default_version
+            self.models_api.models_store_delete_model_version(
+                project_id=teamspace_id, model_id=model.id, version=version
+            )
         else:
-            self.models.models_store_delete_model(project_id=teamspace_id, model_id=model_id)
+            self.models_api.models_store_delete_model(project_id=teamspace_id, model_id=model.id)
 
     def upload_model_file(
         self,
@@ -214,16 +238,15 @@ class TeamspaceApi:
         version: str,
         local_path: Path,
         remote_path: str,
-        cloud_account: str,
         teamspace_id: str,
         progress_bar: bool = True,
     ) -> None:
+        """Upload a file to the model store."""
         uploader = _ModelFileUploader(
             client=self._client,
             model_id=model_id,
             version=version,
             teamspace_id=teamspace_id,
-            cloud_account=cloud_account,
             file_path=str(local_path),
             remote_path=str(remote_path),
             progress_bar=progress_bar,
@@ -234,28 +257,28 @@ class TeamspaceApi:
         self,
         model_id: str,
         version: str,
-        root_path: Path,
-        filepaths: List[Path],
-        cloud_account: str,
+        file_paths: List[Path],
+        remote_paths: List[str],
         teamspace_id: str,
         progress_bar: bool = True,
     ) -> None:
-        main_pbar = tqdm(total=len(filepaths), desc="Uploading files...", position=0) if progress_bar else None
-        for filepath in filepaths:
+        """Upload files to the model store."""
+        main_pbar = tqdm(total=len(file_paths), desc="Uploading files...", position=0) if progress_bar else None
+        assert len(file_paths) == len(remote_paths), "File paths and remote paths must have the same length"
+        for filepath, remote_path in zip(file_paths, remote_paths):
             self.upload_model_file(
                 model_id=model_id,
                 version=version,
                 local_path=filepath,
-                remote_path=str(filepath.relative_to(root_path)),
-                cloud_account=cloud_account,
+                remote_path=remote_path,
                 teamspace_id=teamspace_id,
                 progress_bar=progress_bar,  # TODO: Global progress bar
             )
             if main_pbar:
                 main_pbar.update(1)
 
-    def complete_model_upload(self, model_id: str, version: str, teamspace_id: str) -> None:
-        self.models.models_store_complete_model_upload(
+    def _complete_model_upload(self, model_id: str, version: str, teamspace_id: str) -> None:
+        self.models_api.models_store_complete_model_upload(
             body=_DummyBody(),
             project_id=teamspace_id,
             model_id=model_id,
@@ -265,12 +288,14 @@ class TeamspaceApi:
     def download_model_files(
         self,
         name: str,
-        version: str,
+        version: Optional[str],
         download_dir: Path,
         teamspace_name: str,
         teamspace_owner_name: str,
         progress_bar: bool = True,
     ) -> List[str]:
+        if version is None:
+            version = "default"
         return _download_model_files(
             client=self._client,
             teamspace_name=teamspace_name,
@@ -295,3 +320,216 @@ class TeamspaceApi:
         ).lightningapps
         jobs = self._client.jobs_service_list_multi_machine_jobs(project_id=teamspace_id).multi_machine_jobs
         return apps, jobs
+
+    def list_machines(self, teamspace_id: str, cloud_account: str) -> List[V1ClusterAccelerator]:
+        response = self._client.cluster_service_list_project_cluster_accelerators(
+            project_id=teamspace_id, id=cloud_account
+        )
+        return response.accelerator
+
+    def get_model(self, teamspace_id: str, model_id: Optional[str] = None, model_name: Optional[str] = None) -> V1Model:
+        if model_id:
+            return self.models_api.models_store_get_model(project_id=teamspace_id, model_id=model_id)
+        if not model_name:
+            raise ValueError("Either `model_id` or `model_name` must be provided.")
+        # list models with specific name
+        models = self.models_api.models_store_list_models(project_id=teamspace_id, name=model_name).models
+        if len(models) == 0:
+            raise ValueError(f"Model '{model_name}' does not exist.")
+        if len(models) > 1:
+            raise RuntimeError(f"Model name '{model_name}' is not a unique with this teamspace.")
+        # if there is only one model with the name, return it
+        return models[0]
+
+    def list_models(self, teamspace_id: str) -> List[V1Model]:
+        response = self.models_api.models_store_list_models(project_id=teamspace_id)
+        return response.models
+
+    def list_model_versions(
+        self, teamspace_id: str, model_id: Optional[str] = None, model_name: Optional[str] = None
+    ) -> List[V1ModelVersionArchive]:
+        if model_name and not model_id:
+            model_id = self.get_model(teamspace_id=teamspace_id, model_name=model_name).id
+        response = self.models_api.models_store_list_model_versions(project_id=teamspace_id, model_id=model_id)
+        return response.versions
+
+    def upload_file(
+        self,
+        teamspace_id: str,
+        cloud_account: str,
+        file_path: str,
+        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,
+        )()
+
+    def download_file(
+        self,
+        path: str,
+        target_path: str,
+        teamspace_id: str,
+        cloud_account: str,
+        progress_bar: bool = True,
+    ) -> None:
+        """Downloads a given file in Teamspace drive 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
+
+        query_params = {
+            "clusterId": cloud_account,
+            "key": _resolve_teamspace_remote_path(path),
+            "token": token,
+        }
+
+        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 == 404:
+            raise FileNotFoundError(f"File {path} not found")
+        total_length = int(r.headers.get("content-length"))
+
+        if progress_bar:
+            pbar = tqdm(
+                desc=f"Downloading {os.path.split(path)[1]}",
+                total=total_length,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1000,
+            )
+
+            pbar_update = pbar.update
+        else:
+            pbar_update = lambda x: None
+
+        target_dir = os.path.split(target_path)[0]
+        if target_dir:
+            os.makedirs(target_dir, exist_ok=True)
+        with open(target_path, "wb") as f:
+            for chunk in r.iter_content(chunk_size=4096 * 8):
+                f.write(chunk)
+                pbar_update(len(chunk))
+
+    def download_folder(
+        self,
+        path: str,
+        target_path: str,
+        teamspace_id: str,
+        cloud_account: str,
+        progress_bar: bool = True,
+    ) -> None:
+        """Downloads a given folder from Teamspace drive to a target location."""
+        # TODO: Update this endpoint to permit basic auth
+        auth = Auth()
+        auth.authenticate()
+
+        prefix = _resolve_teamspace_remote_path(path)
+
+        # ensure we only download as a directory and not the entire prefix
+        if prefix.endswith("/") is False:
+            prefix = prefix + "/"
+
+        _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,
+        )
+
+    def get_secrets(self, teamspace_id: str) -> Dict[str, str]:
+        """Get all secrets for a teamspace."""
+        secrets = self._get_secrets(teamspace_id)
+        # 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, teamspace_id: str, key: str, value: str) -> None:
+        """Set a secret for a teamspace.
+
+        This will replace the existing secret if it exists and create a new one if it doesn't.
+        """
+        secrets = self._get_secrets(teamspace_id)
+        for secret in secrets:
+            if secret.name == key:
+                return self._update_secret(teamspace_id, secret.id, value)
+        return self._create_secret(teamspace_id, key, value)
+
+    def _get_secrets(self, teamspace_id: str) -> List[V1Secret]:
+        return self._client.secret_service_list_secrets(project_id=teamspace_id).secrets
+
+    def _update_secret(self, teamspace_id: str, secret_id: str, value: str) -> None:
+        self._client.secret_service_update_secret(
+            body=SecretsIdBody(value=value),
+            project_id=teamspace_id,
+            id=secret_id,
+        )
+
+    def _create_secret(
+        self,
+        teamspace_id: str,
+        key: str,
+        value: str,
+    ) -> None:
+        self._client.secret_service_create_secret(
+            body=ProjectIdSecretsBody(name=key, value=value, type=V1SecretType.UNSPECIFIED), project_id=teamspace_id
+        )
+
+    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
+
+    def new_folder(self, teamspace_id: str, name: str, cluster: Optional[V1ExternalCluster]) -> None:
+        create_request = Create(
+            name=name,
+            create_resources=True,
+            force=True,
+            writable=True,
+        )
+
+        if cluster is None:
+            create_request.r2 = V1R2DataConnection(name=name)
+        else:
+            create_request.cluster_id = cluster.id
+            create_request.access_cluster_ids = [cluster.id]
+
+            if cluster.spec.aws_v1:
+                create_request.s3_folder = V1S3FolderDataConnection()
+            elif cluster.spec.google_cloud_v1:
+                create_request.gcs_folder = V1GCSFolderDataConnection()
+
+        self._client.data_connection_service_create_data_connection(create_request, teamspace_id)
+
+    def new_connection(
+        self, teamspace_id: str, name: str, source: str, cluster: V1ExternalCluster, writable: bool, region: str
+    ) -> None:
+        create_request = Create(
+            name=name,
+            create_resources=False,
+            force=True,
+            writable=writable,
+            cluster_id=cluster.id,
+            access_cluster_ids=[cluster.id],
+        )
+
+        # TODO: Add support for other connection types
+        create_request.efs = V1EfsConfig(file_system_id=source, region=region)
+
+        self._client.data_connection_service_create_data_connection(create_request, teamspace_id)

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