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

lightning_sdk/lightning_cloud/rest_client.py

@@ -35,6 +35,8 @@ from lightning_sdk.lightning_cloud.openapi import (
     LitRegistryServiceApi,
     PipelinesServiceApi,
     SchedulesServiceApi,
+    ProductLicenseServiceApi,
+    CloudSpaceEnvironmentTemplateServiceApi
 )
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lightning_cloud.source_code.logs_socket_api import LightningLogsSocketAPI
@@ -97,6 +99,8 @@ class GridRestClient(
         LitRegistryServiceApi,
         PipelinesServiceApi,
         SchedulesServiceApi,
+        ProductLicenseServiceApi,
+        CloudSpaceEnvironmentTemplateServiceApi
 ):
 
     def __init__(self, api_client: Optional[ApiClient] = None):

lightning_sdk/llm/llm.py

@@ -1,44 +1,63 @@
-from typing import Dict, List, Optional, Set, Tuple, Union
+from typing import Dict, Generator, List, Optional, Set, Tuple, Union
 
-from lightning_sdk.api import UserApi
 from lightning_sdk.api.llm_api import LLMApi
-from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
 from lightning_sdk.lightning_cloud.openapi import V1Assistant
+from lightning_sdk.lightning_cloud.openapi.models.v1_conversation_response_chunk import V1ConversationResponseChunk
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.organization import Organization
-from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _resolve_org, _resolve_user
+from lightning_sdk.owner import Owner
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.utils.resolve import _get_authed_user, _resolve_org
 
 
 class LLM:
     def __init__(
         self,
         name: str,
-        user: Union[str, "User", None] = None,
-        org: Union[str, "Organization", None] = None,
+        teamspace: Optional[str] = None,
     ) -> None:
-        self._auth = Auth()
-        self._user = None
-
-        try:
-            self._auth.authenticate()
-            self._user = User(name=UserApi()._get_user_by_id(self._auth.user_id).username)
-        except ConnectionError as e:
-            raise e
-
-        self._name = name
-        try:
-            self._user = _resolve_user(self._user or user)
-        except ValueError:
-            self._user = None
-
-        self._name = name
-        self._org, self._model_name = self._parse_model_name(name)
+        """Initializes the LLM instance with teamspace information, which is required for billing purposes.
+
+        Teamspace information is resolved through the following methods:
+        1. `.lightning/credentials.json` - Attempts to retrieve the teamspace from the local credentials file.
+        2. Environment Variables - Checks for `LIGHTNING_*` environment variables.
+        3. User Authentication - Redirects the user to the login page if teamspace information is not found.
+
+        Args:
+            name (str): The name of the model or resource.
+            teamspace (Optional[str]): The specified teamspace for billing. If not provided, it will be resolved
+                                       through the above methods.
+
+        Raises:
+            ValueError: If teamspace information cannot be resolved.
+        """
+        menu = _TeamspacesMenu()
+        user = _get_authed_user()
+        possible_teamspaces = menu._get_possible_teamspaces(user)
+        if teamspace is None:
+            if len(possible_teamspaces) == 1:
+                teamspace_name = next(iter(possible_teamspaces.values()))["name"]
+                self._teamspace = Teamspace(name=teamspace_name, org=None, user=user)
+            else:
+                self._teamspace = menu._resolve_teamspace(teamspace)
+        else:
+            self._teamspace = Teamspace(**menu._get_teamspace_from_name(teamspace, possible_teamspaces))
+
+        if self._teamspace is None:
+            raise ValueError("Teamspace is required for billing but could not be resolved. ")
+
+        self._user = user
+
+        self._model_provider, self._model_name = self._parse_model_name(name)
         try:
             # check if it is a org model
-            self._org = _resolve_org(self._org or org)
+            self._org = _resolve_org(self._model_provider)
         except ApiException:
-            self._org = None
+            if isinstance(self._teamspace.owner, Organization):
+                self._org = self._teamspace.owner
+            else:
+                self._org = None
 
         self._llm_api = LLMApi()
         self._public_models = self._build_model_lookup(self._get_public_models())
@@ -47,6 +66,18 @@ class LLM:
         self._model = self._get_model()
         self._conversations = {}
 
+    @property
+    def name(self) -> str:
+        return self._model_name
+
+    @property
+    def provider(self) -> str:
+        return self._model_provider
+
+    @property
+    def owner(self) -> Optional[Owner]:
+        return self._teamspace.owner
+
     def _parse_model_name(self, name: str) -> Tuple[str, str]:
         parts = name.split("/")
         if len(parts) == 1:
@@ -95,13 +126,23 @@ class LLM:
         available_models_str = "\n".join(available_models)
         raise ValueError(f"Model '{self._model_name}' not found. \nAvailable models: \n{available_models_str}")
 
-    def _get_conversations(self) -> Dict[str, str]:
-        # TODO: after updating backend, this will fetch conversations from backend
-        # conversations = self._llm_api.list_conversations(assistant_id=self._model.id)
-        return self._conversations
+    def _get_conversations(self) -> None:
+        conversations = self._llm_api.list_conversations(assistant_id=self._model.id)
+        for conversation in conversations:
+            if conversation.name and conversation.name not in self._conversations:
+                self._conversations[conversation.name] = conversation.id
 
-    def _fetch_conversations(self) -> None:
-        self._conversations = self._get_conversations()
+    def _stream_chat_response(
+        self, result: Generator[V1ConversationResponseChunk, None, None], conversation: Optional[str] = None
+    ) -> Generator[str, None, None]:
+        first_line = next(result, None)
+        if first_line:
+            if conversation and first_line.conversation_id:
+                self._conversations[conversation] = first_line.conversation_id
+            yield first_line.choices[0].delta.content
+
+        for line in result:
+            yield line.choices[0].delta.content
 
     def chat(
         self,
@@ -109,9 +150,10 @@ class LLM:
         system_prompt: Optional[str] = None,
         max_completion_tokens: Optional[int] = 500,
         conversation: Optional[str] = None,
-    ) -> str:
+        stream: bool = False,
+    ) -> Union[str, Generator[str, None, None]]:
         if conversation and conversation not in self._conversations:
-            self._fetch_conversations()
+            self._get_conversations()
 
         conversation_id = self._conversations.get(conversation) if conversation else None
         output = self._llm_api.start_conversation(
@@ -120,22 +162,26 @@ class LLM:
             max_completion_tokens=max_completion_tokens,
             assistant_id=self._model.id,
             conversation_id=conversation_id,
+            billing_project_id=self._teamspace.id,
+            name=conversation,
+            stream=stream,
         )
-        if conversation and not conversation_id:
-            self._conversations[conversation] = output.conversation_id
-        return output.choices[0].delta.content
+        if not stream:
+            if conversation and not conversation_id:
+                self._conversations[conversation] = output.conversation_id
+            return output.choices[0].delta.content
+        return self._stream_chat_response(output, conversation=conversation)
 
     def list_conversations(self) -> List[Dict]:
-        self._fetch_conversations()
+        self._get_conversations()
         return list(self._conversations.keys())
 
     def _get_conversation_messages(self, conversation_id: str) -> Optional[str]:
         return self._llm_api.get_conversation(assistant_id=self._model.id, conversation_id=conversation_id)
 
     def get_history(self, conversation: str) -> Optional[List[Dict]]:
-        # TODO: after updating backend, this will fetch conversation from backend
         if conversation not in self._conversations:
-            self._fetch_conversations()
+            self._get_conversations()
 
         if conversation not in self._conversations:
             raise ValueError(
@@ -152,6 +198,8 @@ class LLM:
         return history
 
     def reset_conversation(self, conversation: str) -> None:
+        if conversation not in self._conversations:
+            self._get_conversations()
         if conversation in self._conversations:
             self._llm_api.reset_conversation(
                 assistant_id=self._model.id,

lightning_sdk/services/__init__.py

@@ -1,5 +1,5 @@
 from lightning_sdk.services.file_endpoint import Client
-from lightning_sdk.services.finetune import LLMFinetune
+from lightning_sdk.services.finetune_llm import LLMFinetune
 from lightning_sdk.services.utilities import download_file
 
 __all__ = ["LLMFinetune", "Client", "download_file"]

lightning_sdk/services/license.py

@@ -0,0 +1,236 @@
+import importlib
+import json
+import os
+import socket
+import threading
+from functools import partial
+from pathlib import Path
+from typing import Optional
+
+from lightning_sdk.api.license_api import LicenseApi
+
+
+class LightningLicense:
+    """This class is used to manage the license for the Lightning SDK."""
+
+    _is_valid: Optional[bool] = None
+    _license_api: Optional[LicenseApi] = None
+    _stream_messages: Optional[callable] = None
+
+    def __init__(
+        self,
+        name: str,
+        license_key: Optional[str] = None,
+        product_version: Optional[str] = None,
+        product_type: str = "package",
+        stream_messages: callable = print,
+    ) -> None:
+        self._product_name = name
+        self._license_key = license_key
+        self._product_version = product_version
+        self.product_type = product_type
+        self._is_valid = None
+        self._license_api = None
+        self._stream_messages = stream_messages
+
+    def validate_license(self) -> bool:
+        """Validate the license key."""
+        if not self.is_online():
+            raise ConnectionError("No internet connection.")
+
+        self._license_api = LicenseApi()
+        return self._license_api.valid_license(
+            license_key=self.license_key,
+            product_name=self.product_name,
+            product_version=self.product_version,
+            product_type=self.product_type,
+        )
+
+    @staticmethod
+    def is_online(timeout: float = 2.0) -> bool:
+        """Check if the system is online by attempting to connect to a public DNS server (Google's).
+
+        This is a simple way to check for internet connectivity.
+
+        Args:
+            timeout: The timeout for the connection attempt.
+        """
+        try:
+            socket.create_connection(("8.8.8.8", 53), timeout=timeout)
+            return True
+        except OSError:
+            return False
+
+    @property
+    def is_valid(self) -> Optional[bool]:
+        """Check if the license key is valid.
+
+        license validation within package:
+          - user online with valid key -> everything as now
+          - user online with invalid key -> warning using wrong key + instructions
+          - user online with no key -> warning for missing license approval + instructions
+          - user offline with a key -> small warning  that key could not be verified
+          - user offline with no key -> warning for missing license approval + instructions
+        """
+        if isinstance(self._is_valid, bool):
+            # if the license key is already validated, return the cached value
+            return self._is_valid
+        if not self.product_version:
+            self._stream_messages("Product version is not set correctly, consider leave it empty for auto-determine.")
+        if not self.license_key:
+            self._stream_messages(
+                "License key is not set neither cannot be found in the package root or user home."
+                " Please make sure you have signed the license agreement and set the license key."
+                " For more information, please refer to the documentation.",
+            )
+        is_online = self.is_online()
+        if self.license_key and is_online:
+            self._is_valid = self.validate_license()
+        elif not is_online:
+            self._stream_messages(
+                "License key is set but the system is offline. "
+                "Please make sure you have a valid license key and the system is online."
+            )
+        return self._is_valid
+
+    @property
+    def has_required_details(self) -> bool:
+        """Check if the license key and product name are set."""
+        return bool(self.license_key and self.product_name and self.product_type)
+
+    @staticmethod
+    def _find_package_license_key(package_name: str) -> Optional[str]:
+        """Find the license key in the package root as .license_key or in user home as .lightning/licenses.json.
+
+        Args:
+            package_name: The name of the package. If not provided, it will be determined from the current module.
+        """
+        if not package_name:
+            return None
+        try:
+            pkg_locations = importlib.util.find_spec(package_name).submodule_search_locations
+            if not pkg_locations:
+                return None
+            license_file = os.path.join(pkg_locations[0], ".license_key")
+            with open(license_file) as fp:
+                return fp.read().strip()
+        except (FileNotFoundError, ModuleNotFoundError):
+            return None
+
+    @staticmethod
+    def _find_user_license_key(package_name: str) -> Optional[str]:
+        """Find the license key in the user home as .lightning/licenses.json.
+
+        Args:
+            package_name: The name of the package.
+        """
+        home = str(Path.home())
+        package_name = package_name.lower()
+        license_file = os.path.join(home, ".lightning", "licenses.json")
+        try:
+            with open(license_file) as fp:
+                licenses = json.load(fp)
+            # Check for the license key in the licenses.json file
+            for name in (package_name, package_name.replace("-", "_"), package_name.replace("_", "-")):
+                if name in licenses:
+                    return licenses[name]
+            return None
+        except (FileNotFoundError, json.JSONDecodeError):
+            return None
+
+    @staticmethod
+    def _determine_package_version(package_name: str) -> Optional[str]:
+        """Determine the product version based on the instantiation of the class.
+
+        Args:
+            package_name: The name of the package. If not provided, it will be determined from the current module.
+        """
+        try:
+            pkg = importlib.import_module(package_name)
+            return getattr(pkg, "__version__", None)
+        except ImportError:
+            return None
+
+    @property
+    def license_key(self) -> Optional[str]:
+        """Get the license key."""
+        if not self._license_key:
+            # If the license key is not set, fist try to find it in the package root
+            self._license_key = self._find_package_license_key(self.product_name.replace("-", "_"))
+            # If not found, try to find it in the user home
+            if not self._license_key:
+                self._license_key = self._find_user_license_key(self.product_name)
+        return self._license_key
+
+    @property
+    def product_name(self) -> str:
+        """Get the product name."""
+        return self._product_name
+
+    @property
+    def product_version(self) -> Optional[str]:
+        """Get the product version."""
+        if not self._product_version and self.product_type == "package":
+            self._product_version = self._determine_package_version(self.product_name.replace("-", "_"))
+        return self._product_version
+
+
+def check_license(
+    name: str,
+    license_key: Optional[str] = None,
+    product_version: Optional[str] = None,
+    product_type: str = "package",
+    stream_messages: callable = print,
+) -> None:
+    """Run the license check and stream outputs.
+
+    Args:
+        name: The name of the product.
+        license_key: The license key to check.
+        product_version: The version of the product.
+        product_type: The type of the product.
+        stream_messages: A callable to stream messages.
+    """
+    lit_license = LightningLicense(
+        name=name,
+        license_key=license_key,
+        product_version=product_version,
+        product_type=product_type,
+        stream_messages=stream_messages,
+    )
+    if lit_license.is_valid is False:
+        stream_messages(
+            "License key is not valid.\n"
+            f" Key: {lit_license.license_key}\n"
+            " Please make sure you have a valid license key."
+        )
+
+
+def check_license_in_background(
+    name: str,
+    license_key: Optional[str] = None,
+    product_version: Optional[str] = None,
+    product_type: str = "package",
+    stream_messages: callable = print,
+) -> threading.Thread:
+    """Run the license check in a background thread and stream outputs.
+
+    Args:
+        name: The name of the product.
+        license_key: The license key to check.
+        product_version: The version of the product.
+        product_type: The type of the product.
+        stream_messages: A callable to stream messages.
+    """
+    check_license_local = partial(
+        check_license,
+        name=name,
+        license_key=license_key,
+        product_version=product_version,
+        product_type=product_type,
+        stream_messages=stream_messages,
+    )
+
+    thread = threading.Thread(target=check_license_local, daemon=True)
+    thread.start()
+    return thread

lightning_sdk/studio.py

@@ -221,6 +221,36 @@ class Studio:
             self._studio.id, self._teamspace.id, machine, interruptible=interruptible
         )
 
+    def run_and_detach(self, *commands: str, timeout: float = 10, check_interval: float = 1) -> str:
+        """Runs given commands on the Studio and returns immediately.
+
+        The command will continue to run in the background.
+
+        Args:
+            timeout: wait for this many seconds for the command to finish.
+            check_interval: check the status of the command every this many seconds.
+        """
+        if check_interval > timeout:
+            raise ValueError("check_interval must be less than timeout")
+
+        if _LIGHTNING_DEBUG:
+            print(f"Running {commands=}")
+        status = self.status
+        if status != Status.Running:
+            raise RuntimeError(f"Cannot run a command in a studio that is not running. Studio {self.name} is {status}.")
+
+        iter_output = self._studio_api.run_studio_commands_and_yield(
+            self._studio.id, self._teamspace.id, *commands, timeout=timeout, check_interval=check_interval
+        )
+
+        output = ""
+        code = None
+        for line, exit_code in iter_output:
+            print(line)
+            output += line
+            code = exit_code
+        return output, code
+
     def run_with_exit_code(self, *commands: str) -> Tuple[str, int]:
         """Runs given commands on the Studio while returning output and exit code.
 

{lightning_sdk-0.2.14.dist-info → lightning_sdk-0.2.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 0.2.14
+Version: 0.2.15
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License