lightning-sdk 0.2.14__py3-none-any.whl → 0.2.16__py3-none-any.whl

lightning_sdk/__init__.py

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

lightning_sdk/api/base_studio_api.py

@@ -0,0 +1,79 @@
+from typing import Any, List, Optional
+
+from lightning_sdk.lightning_cloud.openapi.models.update import Update as BaseStudioUpdateBody
+from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_environment_template import (
+    V1CloudSpaceEnvironmentTemplate,
+)
+from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_environment_type import V1CloudSpaceEnvironmentType
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+
+class BaseStudioApi:
+    def __init__(self) -> None:
+        self._client = LightningClient(retry=False, max_tries=0)
+
+    def get_base_studio(self, base_studio_id: str, org_id: str) -> V1CloudSpaceEnvironmentTemplate:
+        """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=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,
+        org_id: str,
+        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,
+    ) -> V1CloudSpaceEnvironmentTemplate:
+        base_studio = self.get_base_studio(base_studio_id, org_id)
+
+        # Get the current configuration for the base studio
+        update_body = BaseStudioUpdateBody(
+            org_id=base_studio.org_id,
+            name=base_studio.name,
+            allowed_machines=base_studio.config.allowed_machines,
+            default_machine=base_studio.config.default_machine,
+            environment_type=base_studio.config.environment_type,
+            machine_image_version=base_studio.config.machine_image_version,
+            setup_script_text=base_studio.config.setup_script_text,
+            disabled=base_studio.disabled,
+        )
+
+        # Apply changes only if the new value is not None
+        apply_change(update_body, "name", name)
+        apply_change(update_body, "allowed_machines", allowed_machines)
+        apply_change(update_body, "default_machine", default_machine)
+        apply_change(update_body, "environment_type", environment_type)
+        apply_change(update_body, "machine_image_version", machine_image_version)
+        apply_change(update_body, "setup_script_text", setup_script_text)
+        apply_change(update_body, "disabled", disabled)
+
+        return self._client.cloud_space_environment_template_service_update_cloud_space_environment_template(
+            id=base_studio_id,
+            body=update_body,
+        )
+
+
+def apply_change(spec: Any, key: str, value: Any) -> bool:
+    if value is None:
+        return False
+
+    if getattr(spec, key) != value:
+        setattr(spec, key, value)
+        return True
+
+    return False

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/license_api.py

@@ -0,0 +1,48 @@
+from typing import Optional
+
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+
+class LicenseApi:
+    def __init__(self) -> None:
+        self._client = LightningClient(retry=False, max_tries=0)
+
+    def valid_license(
+        self,
+        license_key: str,
+        product_name: str,
+        product_version: Optional[str] = None,
+        product_type: str = "package",
+    ) -> bool:
+        """Check if the license key is valid.
+
+        Args:
+            license_key: The license key to check.
+            product_name: The name of the product.
+            product_version: The version of the product.
+            product_type: The type of the product. Default is "package".
+
+        Returns:
+            True if the license key is valid, False otherwise.
+        """
+        response, code, _ = self._client.product_license_service_validate_product_license_with_http_info(
+            license_key=license_key,
+            product_name=product_name,
+            product_version=product_version,
+            product_type=product_type,
+        )
+        if code != 200:
+            raise ConnectionError(f"Failed to validate license key: {code} - {response}")
+        return response.valid
+
+    def list_user_licenses(self, user_id: str) -> list:
+        """List all licenses for a user.
+
+        Args:
+            user_id: The ID of the user.
+
+        Returns:
+            A list of licenses for the user.
+        """
+        response = self._client.product_license_service_list_user_licenses(user_id=user_id)
+        return response.licenses

lightning_sdk/api/llm_api.py

@@ -1,6 +1,12 @@
-from typing import List, Optional
+import base64
+import json
+from typing import Dict, Generator, List, Optional, Union
+
+from pip._vendor.urllib3 import HTTPResponse
 
 from lightning_sdk.lightning_cloud.openapi.models.v1_conversation_response_chunk import V1ConversationResponseChunk
+from lightning_sdk.lightning_cloud.openapi.models.v1_response_choice import V1ResponseChoice
+from lightning_sdk.lightning_cloud.openapi.models.v1_response_choice_delta import V1ResponseChoiceDelta
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
 
@@ -20,30 +26,85 @@ class LLMApi:
         result = self._client.assistants_service_list_assistants(user_id=user_id)
         return result.assistants
 
+    def _stream_chat_response(self, result: HTTPResponse) -> Generator[V1ConversationResponseChunk, None, None]:
+        for line in result.stream():
+            decoded_lines = line.decode("utf-8").strip()
+            for decoded_line in decoded_lines.splitlines():
+                try:
+                    payload = json.loads(decoded_line)
+                    result_data = payload.get("result", {})
+
+                    choices = []
+                    for choice in result_data.get("choices", []):
+                        delta = choice.get("delta", {})
+                        choices.append(
+                            V1ResponseChoice(
+                                delta=V1ResponseChoiceDelta(**delta),
+                                finish_reason=choice.get("finishReason"),
+                                index=choice.get("index"),
+                            )
+                        )
+
+                    yield V1ConversationResponseChunk(
+                        choices=choices,
+                        conversation_id=result_data.get("conversationId"),
+                        executable=result_data.get("executable"),
+                        id=result_data.get("id"),
+                        throughput=result_data.get("throughput"),
+                    )
+
+                except json.JSONDecodeError:
+                    print("Error decoding JSON:", decoded_line)
+
+    def _encode_image_bytes_to_data_url(self, image: str, mime_type: str = "image/jpeg") -> str:
+        with open(image, "rb") as image_file:
+            b64 = base64.b64encode(image_file.read()).decode("utf-8")
+            return f"data:{mime_type};base64,{b64}"
+
     def start_conversation(
         self,
         prompt: str,
         system_prompt: Optional[str],
-        max_completion_tokens: Optional[int],
+        max_completion_tokens: int,
         assistant_id: str,
-        conversation_id: Optional[str],
-    ) -> V1ConversationResponseChunk:
+        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,
+    ) -> 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,
+            "conversation_id": conversation_id,
+            "billing_project_id": billing_project_id,
+            "name": name,
+            "stream": stream,
+            "metadata": metadata or {},
         }
-        if conversation_id:
-            body["conversation_id"] = conversation_id
-        result = self._client.assistants_service_start_conversation(body, assistant_id)
-        return result.result
+        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
+        return self._stream_chat_response(result)
 
     def list_conversations(self, assistant_id: str) -> List[str]:
         result = self._client.assistants_service_list_conversations(assistant_id)

lightning_sdk/api/studio_api.py

@@ -5,7 +5,7 @@ import time
 import warnings
 import zipfile
 from threading import Event, Thread
-from typing import Any, Dict, Mapping, Optional, Tuple, Union
+from typing import Any, Dict, Generator, Mapping, Optional, Tuple, Union
 
 import backoff
 import requests
@@ -35,6 +35,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1CloudSpace,
     V1CloudSpaceInstanceConfig,
     V1CloudSpaceSeedFile,
+    V1CloudSpaceSourceType,
     V1CloudSpaceState,
     V1EndpointType,
     V1GetCloudSpaceInstanceStatusResponse,
@@ -110,6 +111,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 +119,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)
 
@@ -285,6 +288,52 @@ class StudioApi:
         for response in responses:
             yield response.result
 
+    def run_studio_commands_and_yield(
+        self, studio_id: str, teamspace_id: str, *commands: str, timeout: float, check_interval: float
+    ) -> Generator[Tuple[str, int], None, None]:
+        """Run given commands in a given Studio and yield the output and exit code for the given timeout.
+
+        Args:
+            timeout: wait for this many seconds for the command to finish.
+        """
+        response_submit = self._client.cloud_space_service_execute_command_in_cloud_space(
+            IdExecuteBody1("; ".join(commands), detached=True),
+            project_id=teamspace_id,
+            id=studio_id,
+        )
+
+        if not response_submit:
+            raise RuntimeError("Unable to submit command")
+
+        if response_submit.session_name == "":
+            raise RuntimeError("The session name should be defined.")
+
+        start_time = time.time()
+        exit_code = None
+        while True:
+            for resp in self._get_detached_command_status(
+                studio_id=studio_id,
+                teamspace_id=teamspace_id,
+                session_id=response_submit.session_name,
+            ):
+                if time.time() - start_time >= timeout:
+                    return
+
+                if resp.exit_code == -1:
+                    break
+
+                if exit_code is None:
+                    exit_code = resp.exit_code
+
+                elif exit_code != resp.exit_code:
+                    raise RuntimeError("Cannot determine exit code")
+
+                if resp.exit_code is not None and resp.exit_code != 0:
+                    raise RuntimeError(f"Command failed with exit code {resp.exit_code}. Output: {resp.output}")
+
+                yield resp.output, exit_code
+                time.sleep(check_interval)
+
     def run_studio_commands(self, studio_id: str, teamspace_id: str, *commands: str) -> Tuple[str, int]:
         """Run given commands in a given Studio."""
         response_submit = self._client.cloud_space_service_execute_command_in_cloud_space(

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

@@ -0,0 +1,83 @@
+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_template import (
+    V1CloudSpaceEnvironmentTemplate,
+)
+from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_environment_type import V1CloudSpaceEnvironmentType
+from lightning_sdk.organization import Organization
+from lightning_sdk.user import User
+from lightning_sdk.utils.resolve import _resolve_org, _resolve_user
+
+
+class BaseStudio:
+    def __init__(
+        self,
+        name: Optional[str] = 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._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)
+
+            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,
+        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,
+    ) -> V1CloudSpaceEnvironmentTemplate:
+        self._base_studio = self._base_studio_api.update_base_studio(
+            self._base_studio.id,
+            self._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) -> 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)