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

lightning_sdk/__init__.py

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

lightning_sdk/api/base_studio_api.py

@@ -0,0 +1,73 @@
+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
+            )
+        except ValueError as e:
+            raise ValueError(f"Base studio {base_studio_id} does not exist") from e
+
+    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/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,11 @@
-from typing import List, Optional
+import json
+from typing import 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,14 +25,47 @@ 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 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:
+        conversation_id: Optional[str] = None,
+        billing_project_id: Optional[str] = None,
+        name: Optional[str] = None,
+        stream: bool = False,
+    ) -> Union[V1ConversationResponseChunk, Generator[V1ConversationResponseChunk, None, None]]:
         body = {
             "message": {
                 "author": {"role": "user"},
@@ -39,8 +77,23 @@ class LLMApi:
                 ],
             },
             "max_completion_tokens": max_completion_tokens,
+            "conversation_id": conversation_id,
+            "billing_project_id": billing_project_id,
+            "name": name,
+            "stream": stream,
         }
-        if conversation_id:
-            body["conversation_id"] = conversation_id
-        result = self._client.assistants_service_start_conversation(body, assistant_id)
-        return result.result
+        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)
+        return result.conversations
+
+    def get_conversation(self, assistant_id: str, conversation_id: str) -> V1ConversationResponseChunk:
+        result = self._client.assistants_service_get_conversation(assistant_id, conversation_id)
+        return result.messages
+
+    def reset_conversation(self, assistant_id: str, conversation_id: str) -> None:
+        self._client.assistants_service_delete_conversation(assistant_id, conversation_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
@@ -285,6 +285,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/base_studio.py

@@ -0,0 +1,70 @@
+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()
+
+        self._base_studio = self._base_studio_api.get_base_studio(name, self._org.id)
+
+    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,
+        )

lightning_sdk/cli/delete.py

@@ -17,7 +17,7 @@ def delete() -> None:
 
 
 @delete.command(name="container")
-@click.argument("container")
+@click.argument("name")
 @click.option(
     "--teamspace",
     default=None,
@@ -27,18 +27,16 @@ def delete() -> None:
         "If not provided, can be selected in an interactive menu."
     ),
 )
-def delete_container(container: str, teamspace: Optional[str] = None) -> None:
-    """Delete the docker container CONTAINER."""
+def container(name: str, teamspace: Optional[str] = None) -> None:
+    """Delete the docker container NAME."""
     api = LitContainer()
     menu = _TeamspacesMenu()
     resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
     try:
-        api.delete_container(container, resolved_teamspace.name, resolved_teamspace.owner.name)
-        Console().print(f"Container {container} deleted successfully.")
+        api.delete_container(name, resolved_teamspace.name, resolved_teamspace.owner.name)
+        Console().print(f"Container {name} deleted successfully.")
     except Exception as e:
-        raise StudioCliError(
-            f"Could not delete container {container} from project {resolved_teamspace.name}: {e}"
-        ) from None
+        raise StudioCliError(f"Could not delete container {name} from project {resolved_teamspace.name}: {e}") from None
 
 
 @delete.command(name="job")

lightning_sdk/cli/download.py

@@ -1,3 +1,4 @@
+import json
 import os
 import re
 from pathlib import Path
@@ -6,6 +7,7 @@ from typing import Optional
 import click
 from rich.console import Console
 
+from lightning_sdk.api.license_api import LicenseApi
 from lightning_sdk.api.lit_container_api import LitContainerApi
 from lightning_sdk.cli.exceptions import StudioCliError
 from lightning_sdk.cli.studios_menu import _StudiosMenu
@@ -208,3 +210,26 @@ def _resolve_studio(studio: Optional[str]) -> Studio:
 
     with skip_studio_init():
         return Studio(**selected_studio)
+
+
+@download.command(name="licenses")
+def licenses() -> None:
+    """Download licenses for all products/packages.
+
+    Example:
+      lightning download licenses
+
+    """
+    user = _get_authed_user()
+    api = LicenseApi()
+    licenses = api.list_user_licenses(user.id)
+
+    user_home = Path.home()
+    lit_dir = user_home / ".lightning"
+    lit_dir.mkdir(parents=True, exist_ok=True)
+    licenses_file = lit_dir / "licenses.json"
+
+    licenses_short = {ll.product_name: ll.license_key for ll in licenses if ll.is_valid}
+    with licenses_file.open("w") as fp:
+        json.dump(licenses_short, fp, indent=4)
+    Console().print(f"Licenses downloaded to {licenses_file}", style="green")