lightning-sdk 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

lightning_sdk/__init__.py

@@ -31,6 +31,6 @@ __all__ = [
     "User",
 ]
 
-__version__ = "0.2.15"
+__version__ = "0.2.17"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/base_studio_api.py

@@ -16,11 +16,17 @@ class BaseStudioApi:
         """Retrieve the base studio by its ID."""
         try:
             return self._client.cloud_space_environment_template_service_get_cloud_space_environment_template(
-                base_studio_id, org_id
+                base_studio_id, org_id=org_id
             )
         except ValueError as e:
             raise ValueError(f"Base studio {base_studio_id} does not exist") from e
 
+    def get_all_base_studios(self, org_id: str) -> List[V1CloudSpaceEnvironmentTemplate]:
+        """Retrieve all base studios for a given organization."""
+        return self._client.cloud_space_environment_template_service_list_cloud_space_environment_templates(
+            org_id=org_id
+        )
+
     def update_base_studio(
         self,
         base_studio_id: str,

lightning_sdk/api/cluster_api.py

@@ -1,4 +1,11 @@
-from lightning_sdk.lightning_cloud.openapi import Externalv1Cluster
+from typing import Dict, List, Optional
+
+from lightning_sdk.lightning_cloud.openapi import (
+    Externalv1Cluster,
+    V1CloudProvider,
+    V1ClusterType,
+    V1ListClusterAcceleratorsResponse,
+)
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
 
@@ -20,3 +27,78 @@ class ClusterApi:
         if not res:
             raise ValueError(f"Cluster {cluster_id} does not exist")
         return res
+
+    def list_cluster_accelerators(self, cluster_id: str, org_id: str) -> V1ListClusterAcceleratorsResponse:
+        """Lists the accelerators for a given cluster.
+
+        :param cluster_id: cluster ID test
+        :param project_id: the project the cluster is supposed to be associated with
+        :param org_id: The owning org of this cluster
+        """
+        res = self._client.cluster_service_list_cluster_accelerators(
+            id=cluster_id,
+            org_id=org_id,
+        )
+        if not res:
+            raise ValueError(f"Cluster {cluster_id} does not exist")
+        return res
+
+    def list_global_clusters(self, project_id: str, org_id: str) -> List[Externalv1Cluster]:
+        """Lists the accelerators for a given project.
+
+        :param project_id: project ID test
+        :param org_id: The owning org of this project
+        """
+        res = self._client.cluster_service_list_clusters(
+            project_id=project_id,
+            org_id=org_id,
+        )
+        if not res:
+            raise ValueError(f"Project {project_id} does not exist")
+        filtered_clusters = filter(lambda x: x.spec.cluster_type == V1ClusterType.GLOBAL, res.clusters)
+        return list(filtered_clusters)
+
+    def get_cluster_provider_mapping(self, project_id: str, org_id: str) -> Dict[V1CloudProvider, str]:
+        """Gets the cluster provider mapping."""
+        res = self.list_global_clusters(
+            project_id=project_id,
+            org_id=org_id,
+        )
+        return {self._get_cluster_provider(cluster): cluster.id for cluster in res}
+
+    def _get_cluster_provider(self, cluster: Optional[Externalv1Cluster]) -> V1CloudProvider:
+        """Determines the cloud provider based on the cluster configuration.
+
+        Args:
+            cluster: An optional Externalv1Cluster object containing cluster specifications
+
+        Returns:
+            V1CloudProvider: The determined cloud provider, defaults to AWS if no match is found
+        """
+        if not cluster:
+            return V1CloudProvider.AWS
+
+        if (
+            cluster.spec
+            and cluster.spec.driver
+            and cluster.spec.driver in [V1CloudProvider.LIGHTNING, V1CloudProvider.DGX]
+        ):
+            return cluster.spec.driver
+
+        if cluster.spec:
+            if cluster.spec.aws_v1:
+                return V1CloudProvider.AWS
+            if cluster.spec.google_cloud_v1:
+                return V1CloudProvider.GCP
+            if cluster.spec.lambda_labs_v1:
+                return V1CloudProvider.LAMBDA_LABS
+            if cluster.spec.vultr_v1:
+                return V1CloudProvider.VULTR
+            if cluster.spec.slurm_v1:
+                return V1CloudProvider.SLURM
+            if cluster.spec.voltage_park_v1:
+                return V1CloudProvider.VOLTAGE_PARK
+            if cluster.spec.nebius_v1:
+                return V1CloudProvider.NEBIUS
+
+        return V1CloudProvider.AWS

lightning_sdk/api/llm_api.py

@@ -1,5 +1,6 @@
+import base64
 import json
-from typing import Generator, List, Optional, Union
+from typing import Dict, Generator, List, Optional, Union
 
 from pip._vendor.urllib3 import HTTPResponse
 
@@ -55,25 +56,31 @@ class LLMApi:
                 except json.JSONDecodeError:
                     print("Error decoding JSON:", decoded_line)
 
+    def _encode_image_bytes_to_data_url(self, image: str) -> str:
+        with open(image, "rb") as image_file:
+            b64 = base64.b64encode(image_file.read()).decode("utf-8")
+            extension = image.split(".")[-1]
+            return f"data:image/{extension};base64,{b64}"
+
     def start_conversation(
         self,
         prompt: str,
         system_prompt: Optional[str],
         max_completion_tokens: int,
         assistant_id: str,
+        images: Optional[List[str]] = None,
         conversation_id: Optional[str] = None,
         billing_project_id: Optional[str] = None,
         name: Optional[str] = None,
+        metadata: Optional[Dict[str, str]] = None,
         stream: bool = False,
+        internal_conversation: bool = False,
     ) -> Union[V1ConversationResponseChunk, Generator[V1ConversationResponseChunk, None, None]]:
         body = {
             "message": {
                 "author": {"role": "user"},
                 "content": [
-                    {
-                        "contentType": "text",
-                        "parts": [prompt],
-                    }
+                    {"contentType": "text", "parts": [prompt]},
                 ],
             },
             "max_completion_tokens": max_completion_tokens,
@@ -81,7 +88,22 @@ class LLMApi:
             "billing_project_id": billing_project_id,
             "name": name,
             "stream": stream,
+            "metadata": metadata or {},
+            "internal_conversation": internal_conversation,
         }
+        if images:
+            for image in images:
+                url = image
+                if not image.startswith("http"):
+                    url = self._encode_image_bytes_to_data_url(image)
+
+                body["message"]["content"].append(
+                    {
+                        "contentType": "image",
+                        "parts": [url],
+                    }
+                )
+
         result = self._client.assistants_service_start_conversation(body, assistant_id, _preload_content=not stream)
         if not stream:
             return result.result

lightning_sdk/api/studio_api.py

@@ -32,9 +32,11 @@ from lightning_sdk.lightning_cloud.openapi import (
     IdForkBody1,
     IdStartBody,
     ProjectIdCloudspacesBody,
+    V1Assistant,
     V1CloudSpace,
     V1CloudSpaceInstanceConfig,
     V1CloudSpaceSeedFile,
+    V1CloudSpaceSourceType,
     V1CloudSpaceState,
     V1EndpointType,
     V1GetCloudSpaceInstanceStatusResponse,
@@ -46,6 +48,16 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1UserRequestedComputeConfig,
 )
 from lightning_sdk.lightning_cloud.openapi.models import ProjectIdEndpointsBody
+from lightning_sdk.lightning_cloud.openapi.models.project_id_agentmanagedendpoints_body import (
+    ProjectIdAgentmanagedendpointsBody,
+)
+from lightning_sdk.lightning_cloud.openapi.models.project_id_agents_body import (
+    ProjectIdAgentsBody,
+)
+from lightning_sdk.lightning_cloud.openapi.models.v1_endpoint import V1Endpoint
+from lightning_sdk.lightning_cloud.openapi.models.v1_managed_endpoint import V1ManagedEndpoint
+from lightning_sdk.lightning_cloud.openapi.models.v1_managed_model import V1ManagedModel
+from lightning_sdk.lightning_cloud.openapi.models.v1_upstream_managed import V1UpstreamManaged
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.machine import Machine
 
@@ -110,6 +122,7 @@ class StudioApi:
         name: str,
         teamspace_id: str,
         cloud_account: Optional[str] = None,
+        source: Optional[V1CloudSpaceSourceType] = None,
     ) -> V1CloudSpace:
         """Create a Studio with a given name in a given Teamspace on a possibly given cloud_account."""
         body = ProjectIdCloudspacesBody(
@@ -117,6 +130,7 @@ class StudioApi:
             name=name,
             display_name=name,
             seed_files=[V1CloudSpaceSeedFile(path="main.py", contents="print('Hello, Lightning World!')\n")],
+            source=source,
         )
         studio = self._client.cloud_space_service_create_cloud_space(body, teamspace_id)
 
@@ -731,6 +745,56 @@ class StudioApi:
         )
         return endpoint.urls[0]
 
+    def create_assistant(self, studio_id: str, teamspace_id: str, port: int, assistant_name: str) -> V1Assistant:
+        target_teamspace = self._client.projects_service_get_project(teamspace_id)
+        org_id = ""
+        if target_teamspace.owner_type == "ORGANIZATION":
+            org_id = target_teamspace.owner_id
+        endpoint = self._client.endpoint_service_create_endpoint(
+            project_id=teamspace_id,
+            body=ProjectIdEndpointsBody(
+                ports=[str(port)],
+                cloudspace=V1UpstreamCloudSpace(
+                    cloudspace_id=studio_id,
+                    port=str(port),
+                    type=V1EndpointType.PLUGIN_API,
+                ),
+            ),
+        )
+        valid_url = endpoint.urls[0]
+        managed_endpoint = self._client.assistants_service_create_assistant_managed_endpoint(
+            body=ProjectIdAgentmanagedendpointsBody(
+                endpoint=V1ManagedEndpoint(
+                    name=assistant_name,
+                    base_url=valid_url + "/v1",
+                    models_metadata=[
+                        V1ManagedModel(
+                            name=assistant_name,
+                        )
+                    ],
+                ),
+                org_id=org_id,
+            ),
+            project_id=teamspace_id,
+        )
+
+        body = ProjectIdAgentsBody(
+            endpoint=V1Endpoint(
+                cloudspace=V1UpstreamCloudSpace(cloudspace_id=studio_id),
+                name=assistant_name,
+                managed=V1UpstreamManaged(id=managed_endpoint.endpoint.id),
+            ),
+            name=assistant_name,
+            model=assistant_name,
+            cloudspace_id=studio_id,
+            model_provider="",
+        )
+
+        return self._client.assistants_service_create_assistant(
+            body=body,
+            project_id=teamspace_id,
+        )
+
     def _create_app(
         self, studio_id: str, teamspace_id: str, cloud_account: str, plugin_type: str, **other_arguments: Any
     ) -> Externalv1LightningappInstance:

lightning_sdk/api/teamspace_api.py

@@ -1,10 +1,20 @@
 import os
+import tempfile
+import zipfile
 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,
+    _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 (
     Externalv1LightningappInstance,
@@ -17,6 +27,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1ClusterAccelerator,
     V1Endpoint,
     V1Job,
+    V1LoginRequest,
     V1Model,
     V1ModelVersionArchive,
     V1MultiMachineJob,
@@ -331,3 +342,118 @@ class TeamspaceApi:
             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,
+        )
+        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()
+        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
+
+        query_params = {
+            "clusterId": cloud_account,
+            "prefix": _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 progress_bar:
+            pbar = tqdm(
+                desc=f"Downloading {os.path.split(path)[1]}",
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1000,
+            )
+
+            pbar_update = pbar.update
+        else:
+            pbar_update = lambda x: None
+
+        if target_path:
+            os.makedirs(target_path, exist_ok=True)
+
+        with tempfile.TemporaryFile() as f:
+            for chunk in r.iter_content(chunk_size=4096 * 8):
+                f.write(chunk)
+                pbar_update(len(chunk))
+
+            with zipfile.ZipFile(f) as z:
+                z.extractall(target_path)

lightning_sdk/api/utils.py

@@ -355,6 +355,10 @@ def _sanitize_studio_remote_path(path: str, studio_id: str) -> str:
     return f"/cloudspaces/{studio_id}/code/content/{path.replace('/teamspace/studios/this_studio/', '')}"
 
 
+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
 

lightning_sdk/base_studio.py

@@ -45,7 +45,12 @@ class BaseStudio:
 
         self._base_studio_api = BaseStudioApi()
 
-        self._base_studio = self._base_studio_api.get_base_studio(name, self._org.id)
+        if name is not None:
+            base_studio = self._base_studio_api.get_base_studio(name, self._org.id)
+
+            if base_studio is None:
+                raise ValueError(f"Base studio with name {name} does not exist in organization {self._org.name}")
+            self._base_studio = base_studio
 
     def update(
         self,
@@ -68,3 +73,11 @@ class BaseStudio:
             setup_script_text=setup_script_text,
             disabled=disabled,
         )
+
+    def list(self) -> List[V1CloudSpaceEnvironmentTemplate]:
+        """List all base studios in the organization.
+
+        Returns:
+            List[V1CloudSpaceEnvironmentTemplate]: A list of base studio templates.
+        """
+        return self._base_studio_api.get_all_base_studios(self._org.id)

lightning_sdk/cli/create.py

@@ -7,9 +7,12 @@ import click
 from rich.console import Console
 
 from lightning_sdk import Machine, Studio
+from lightning_sdk.api.cluster_api import ClusterApi
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.studio import Provider
 
 _MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
+_PROVIDER_VALUES = tuple([provider.value for provider in Provider])
 
 
 @click.group("create")
@@ -44,8 +47,18 @@ def create() -> None:
         "or fall back to the teamspace default."
     ),
 )
+@click.option(
+    "--provider",
+    default=None,
+    type=click.Choice(_PROVIDER_VALUES),
+    help="The provider to create the studio on. If --cloud-account is specified, this option is prioritized.",
+)
 def studio(
-    name: str, teamspace: Optional[str] = None, start: Optional[str] = None, cloud_account: Optional[str] = None
+    name: str,
+    teamspace: Optional[str] = None,
+    start: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    provider: Optional[str] = None,
 ) -> None:
     """Create a new studio on the Lightning AI platform.
 
@@ -57,6 +70,13 @@ def studio(
     menu = _TeamspacesMenu()
     teamspace_resolved = menu._resolve_teamspace(teamspace)
 
+    if provider is not None:
+        cluster_api = ClusterApi()
+        cloud_account = cluster_api.get_cluster_provider_mapping(
+            teamspace_resolved.id,
+            teamspace_resolved.owner.id,
+        )[provider]
+
     # default cloud account to current studios cloud account if run from studio
     # else it will fall back to teamspace default in the backend
     if cloud_account is None: