lightning-sdk 0.1.2__py3-none-any.whl → 0.1.46__py3-none-any.whl

lightning_sdk/api/teamspace_api.py

@@ -1,7 +1,25 @@
-from typing import List, Optional
+import os
+from pathlib import Path
+from typing import Dict, List, Optional
 
+from tqdm.auto import tqdm
+
+from lightning_sdk.api.utils import _download_model_files, _DummyBody, _get_model_version, _ModelFileUploader
 from lightning_sdk.lightning_cloud.login import Auth
-from lightning_sdk.lightning_cloud.openapi import V1CloudSpace, V1Project, V1ProjectClusterBinding
+from lightning_sdk.lightning_cloud.openapi import (
+    ModelIdVersionsBody,
+    ModelsStoreApi,
+    ProjectIdAgentsBody,
+    ProjectIdModelsBody,
+    V1Assistant,
+    V1CloudSpace,
+    V1Endpoint,
+    V1ModelVersionArchive,
+    V1Project,
+    V1ProjectClusterBinding,
+    V1PromptSuggestion,
+    V1UpstreamOpenAI,
+)
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
 __all__ = ["TeamspaceApi"]
@@ -11,13 +29,14 @@ class TeamspaceApi:
     """Internal API client for Teamspace requests (mainly http requests)."""
 
     def __init__(self) -> None:
-        self._client = LightningClient(max_tries=3)
+        self._client = LightningClient(max_tries=7)
+        self._models: Optional[ModelsStoreApi] = None
 
     def get_teamspace(self, name: str, owner_id: str) -> V1Project:
         """Get the current teamspace from the owner."""
         teamspaces = self.list_teamspaces(name=name, owner_id=owner_id)
 
-        if len(teamspaces) == 0:
+        if not teamspaces:
             raise ValueError(f"Teamspace {name} does not exist")
 
         if len(teamspaces) > 1:
@@ -28,7 +47,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]) -> Optional[V1Project]:
+    def list_teamspaces(self, owner_id: str, name: Optional[str] = None) -> Optional[V1Project]:
         """Lists teamspaces from owner.
 
         If name is passed only teamspaces matching that name will be returned
@@ -38,22 +57,21 @@ class TeamspaceApi:
         # -> list projects authed users are members of + filter later on
         res = self._client.projects_service_list_memberships(filter_by_user_id=True)
 
-        return [
-            self._get_teamspace_by_id(m.project_id)
-            for m in filter(
-                # only return teamspaces actually owned by the id
-                lambda x: x.owner_id == owner_id,
-                # if name is provided, filter for teamspaces matching that name
-                filter(lambda x: name is None or x.name == name or x.display_name == name, res.memberships),
-            )
-        ]
+        teamspaces = []
+        for teamspace in res.memberships:
+            # if name is provided, filter for teamspaces matching that name
+            match_name = name is None or teamspace.name == name or teamspace.display_name == name
+            # and only return teamspaces actually owned by the id
+            if match_name and teamspace.owner_id == owner_id:
+                teamspaces.append(self._get_teamspace_by_id(teamspace.project_id))
+        return teamspaces
 
-    def list_studios(self, teamspace_id: str, cluster_id: str = "") -> List[V1CloudSpace]:
+    def list_studios(self, teamspace_id: str, cloud_account: str = "") -> List[V1CloudSpace]:
         """List studios in teamspace."""
         kwargs = {"project_id": teamspace_id, "user_id": self._get_authed_user_id()}
 
-        if cluster_id:
-            kwargs["cluster_id"] = cluster_id
+        if cloud_account:
+            kwargs["cluster_id"] = cloud_account
 
         cloudspaces = []
 
@@ -69,12 +87,193 @@ class TeamspaceApi:
 
         return cloudspaces
 
-    def list_clusters(self, teamspace_id: str) -> List[V1ProjectClusterBinding]:
-        """Lists clusters in a teamspace."""
+    def list_cloud_accounts(self, teamspace_id: str) -> List[V1ProjectClusterBinding]:
+        """Lists cloud_accounts in a teamspace."""
         return self._client.projects_service_list_project_cluster_bindings(project_id=teamspace_id).clusters
 
     def _get_authed_user_id(self) -> str:
-        """Gets the currently logged in user."""
+        """Gets the currently logged-in user."""
         auth = Auth()
         auth.authenticate()
         return auth.user_id
+
+    def get_default_cloud_account(self, teamspace_id: str) -> str:
+        """Get the default cloud account id of the teamspace."""
+        return self._client.projects_service_get_project(teamspace_id).project_settings.preferred_cluster
+
+    def _determine_cloud_account(self, teamspace_id: str) -> str:
+        """Attempts to determine the cloud account id of the teamspace.
+
+        Raises an error if it's ambiguous.
+
+        """
+        # when you run  from studio, the cloud account is with env. vars
+        cloud_account = os.getenv("LIGHTNING_CLUSTER_ID")
+        if cloud_account:
+            return cloud_account
+
+        # if there is only one cluster, use that and ignore default setting :D
+        cloud_accounts = [c.cluster_id for c in self.list_cloud_accounts(teamspace_id=teamspace_id)]
+        if len(cloud_accounts) == 1:
+            return cloud_accounts[0]
+        # otherwise, try to determine the default cloud_account, another API call but we do not care :(
+        default_cloud_account = self.get_default_cloud_account(teamspace_id=teamspace_id)
+        if default_cloud_account:
+            return default_cloud_account
+        raise RuntimeError(
+            "Could not determine the current cloud account. Please provide it manually as input."
+            f" Choices are: {', '.join(cloud_accounts)}"
+        )
+
+    def create_agent(
+        self,
+        name: str,
+        teamspace_id: str,
+        api_key: str,
+        base_url: str,
+        model: str,
+        org_id: Optional[str] = "",
+        prompt_template: Optional[str] = "",
+        description: Optional[str] = "",
+        prompt_suggestions: Optional[List[str]] = None,
+        file_uploads_enabled: Optional[bool] = None,
+    ) -> V1Assistant:
+        openai_endpoint = V1UpstreamOpenAI(api_key=api_key, base_url=base_url)
+
+        endpoint = V1Endpoint(
+            name=name,
+            openai=openai_endpoint,
+            project_id=teamspace_id,
+        )
+
+        ([V1PromptSuggestion(content=suggestion) for suggestion in prompt_suggestions] if prompt_suggestions else None)
+
+        body = ProjectIdAgentsBody(
+            endpoint=endpoint,
+            name=name,
+            model=model,
+            org_id=org_id,
+            prompt_template=prompt_template,
+            description=description,
+            file_uploads_enabled=file_uploads_enabled,
+        )
+
+        return self._client.assistants_service_create_assistant(body=body, project_id=teamspace_id)
+
+    # 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 get_model_version(self, name: str, version: 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,
+        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
+        if len(models) == 0:
+            return self.models.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),
+            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
+        # 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)
+        else:
+            self.models.models_store_delete_model(project_id=teamspace_id, model_id=model_id)
+
+    def upload_model_file(
+        self,
+        model_id: str,
+        version: str,
+        local_path: Path,
+        remote_path: str,
+        cloud_account: str,
+        teamspace_id: str,
+        progress_bar: bool = True,
+    ) -> None:
+        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,
+        )
+        uploader()
+
+    def upload_model_files(
+        self,
+        model_id: str,
+        version: str,
+        root_path: Path,
+        filepaths: List[Path],
+        cloud_account: 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:
+            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,
+                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(
+            body=_DummyBody(),
+            project_id=teamspace_id,
+            model_id=model_id,
+            version=version,
+        )
+
+    def download_model_files(
+        self,
+        name: str,
+        version: str,
+        download_dir: Path,
+        teamspace_name: str,
+        teamspace_owner_name: str,
+        progress_bar: bool = True,
+    ) -> List[str]:
+        return _download_model_files(
+            client=self._client,
+            teamspace_name=teamspace_name,
+            teamspace_owner_name=teamspace_owner_name,
+            name=name,
+            version=version,
+            download_dir=download_dir,
+            progress_bar=progress_bar,
+        )

lightning_sdk/api/user_api.py

@@ -1,12 +1,14 @@
-from typing import List
+from typing import List, Union
 
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     V1CloudSpace,
+    V1GetUserResponse,
     V1ListCloudSpacesResponse,
     V1Membership,
     V1Organization,
     V1SearchUser,
+    V1UserFeatures,
 )
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
@@ -15,22 +17,28 @@ class UserApi:
     """Internal API Client for user requests (mainly http requests)."""
 
     def __init__(self) -> None:
-        super().__init__()
+        self._client = LightningClient(max_tries=7)
 
-        self._client = LightningClient(max_tries=3)
-
-    def get_user(self, name: str) -> V1SearchUser:
+    def get_user(self, name: str) -> Union[V1SearchUser, V1GetUserResponse]:
         """Gets user by name."""
+        authed_user = self._client.auth_service_get_user()
+        if authed_user.username == name:
+            return authed_user
+
+        # if it's not the authed user, lookup by name
+        # TODO: This API won't necesarily return the correct thing
         response = self._client.user_service_search_users(query=name)
 
         users = [u for u in response.users if u.username == name]
-        if not len(users):
+        if not users:
             raise ValueError(f"User {name} does not exist.")
         return users[0]
 
     def _get_user_by_id(self, user_id: str) -> V1SearchUser:
         response = self._client.user_service_search_users(query=user_id)
         users = [u for u in response.users if u.id == user_id]
+        if not users:
+            raise ValueError(f"User {user_id} does not exist.")
         return users[0]
 
     def _get_organizations_for_authed_user(
@@ -39,18 +47,25 @@ class UserApi:
         """Returns Organizations for the current authed user."""
         return self._client.organizations_service_list_organizations().organizations
 
-    def _get_cloudspaces_for_user(self, user_id: str, project_id: str) -> List[V1CloudSpace]:
+    def _get_cloudspaces_for_user(self, project_id: str, user_id: str = "") -> List[V1CloudSpace]:
         resp: V1ListCloudSpacesResponse = self._client.cloud_space_service_list_cloud_spaces(
             project_id=project_id, user_id=user_id
         )
         return resp.cloudspaces
 
-    def _get_all_teamspace_memberships(self, user_id: str) -> List[V1Membership]:
+    def _get_all_teamspace_memberships(
+        self,
+        user_id: str,  # todo: this is unused, but still required
+    ) -> List[V1Membership]:
         return self._client.projects_service_list_memberships(filter_by_user_id=True).memberships
 
     def _get_authed_user_name(self) -> str:
-        """Gets the currently logged in user."""
+        """Gets the currently logged-in user."""
         auth = Auth()
         auth.authenticate()
         user = self._get_user_by_id(auth.user_id)
         return user.username
+
+    def _get_feature_flags(self) -> V1UserFeatures:
+        resp: V1GetUserResponse = self._client.auth_service_get_user()
+        return resp.features