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

lightning_sdk/lightning_cloud/openapi/models/v1_user_requested_compute_config.py

@@ -42,6 +42,7 @@ class V1UserRequestedComputeConfig(object):
     """
     swagger_types = {
         'affinity_identifier': 'str',
+        'cluster_override': 'str',
         'count': 'int',
         'cpu_image_override': 'str',
         'disk_size': 'int',
@@ -56,6 +57,7 @@ class V1UserRequestedComputeConfig(object):
 
     attribute_map = {
         'affinity_identifier': 'affinityIdentifier',
+        'cluster_override': 'clusterOverride',
         'count': 'count',
         'cpu_image_override': 'cpuImageOverride',
         'disk_size': 'diskSize',
@@ -68,9 +70,10 @@ class V1UserRequestedComputeConfig(object):
         'spot': 'spot'
     }
 
-    def __init__(self, affinity_identifier: 'str' =None, count: 'int' =None, cpu_image_override: 'str' =None, disk_size: 'int' =None, gpu_image_override: 'str' =None, id: 'str' =None, name: 'str' =None, requested_run_duration_seconds: 'str' =None, same_compute_on_resume: 'bool' =None, shm_size: 'int' =None, spot: 'bool' =None):  # noqa: E501
+    def __init__(self, affinity_identifier: 'str' =None, cluster_override: 'str' =None, count: 'int' =None, cpu_image_override: 'str' =None, disk_size: 'int' =None, gpu_image_override: 'str' =None, id: 'str' =None, name: 'str' =None, requested_run_duration_seconds: 'str' =None, same_compute_on_resume: 'bool' =None, shm_size: 'int' =None, spot: 'bool' =None):  # noqa: E501
         """V1UserRequestedComputeConfig - a model defined in Swagger"""  # noqa: E501
         self._affinity_identifier = None
+        self._cluster_override = None
         self._count = None
         self._cpu_image_override = None
         self._disk_size = None
@@ -84,6 +87,8 @@ class V1UserRequestedComputeConfig(object):
         self.discriminator = None
         if affinity_identifier is not None:
             self.affinity_identifier = affinity_identifier
+        if cluster_override is not None:
+            self.cluster_override = cluster_override
         if count is not None:
             self.count = count
         if cpu_image_override is not None:
@@ -128,6 +133,27 @@ class V1UserRequestedComputeConfig(object):
 
         self._affinity_identifier = affinity_identifier
 
+    @property
+    def cluster_override(self) -> 'str':
+        """Gets the cluster_override of this V1UserRequestedComputeConfig.  # noqa: E501
+
+
+        :return: The cluster_override of this V1UserRequestedComputeConfig.  # noqa: E501
+        :rtype: str
+        """
+        return self._cluster_override
+
+    @cluster_override.setter
+    def cluster_override(self, cluster_override: 'str'):
+        """Sets the cluster_override of this V1UserRequestedComputeConfig.
+
+
+        :param cluster_override: The cluster_override of this V1UserRequestedComputeConfig.  # noqa: E501
+        :type: str
+        """
+
+        self._cluster_override = cluster_override
+
     @property
     def count(self) -> 'int':
         """Gets the count of this V1UserRequestedComputeConfig.  # noqa: E501

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,52 +1,113 @@
-from typing import Dict, List, Optional, Set, Tuple, Union
+import os
+import warnings
+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, _resolve_teamspace
 
 
 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
+        """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:
+            # get current teamspace
+            self._teamspace = _resolve_teamspace(teamspace=None, org=None, user=None)
+        else:
+            self._teamspace = Teamspace(**menu._get_teamspace_from_name(teamspace, possible_teamspaces))
+
+        if self._teamspace is None:
+            # select the first available teamspace
+            first_teamspace = next(iter(possible_teamspaces.values()), None)
+
+            if first_teamspace:
+                self._teamspace = Teamspace(
+                    name=first_teamspace["name"],
+                    org=first_teamspace["org"],
+                    user=first_teamspace["user"],
+                )
+                warnings.warn(
+                    f"No teamspace given. Using teamspace: {self._teamspace.name}.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
+        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:
-            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._llm_api = LLMApi()
 
-        self._name = name
-        self._org, 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)
+
+            try:
+                # check if user has access to the org
+                self._org_models = self._build_model_lookup(self._get_org_models())
+            except ApiException:
+                warnings.warn(
+                    f"User is not authenticated to access the model in organization: '{self._model_provider}'.\n"
+                    " Proceeding with appropriate org models, user models, or public models.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                self._model_provider = None
+                raise
         except ApiException:
-            self._org = None
+            if isinstance(self._teamspace.owner, Organization):
+                self._org = self._teamspace.owner
+            else:
+                self._org = None
+            self._org_models = self._build_model_lookup(self._get_org_models())
 
-        self._llm_api = LLMApi()
         self._public_models = self._build_model_lookup(self._get_public_models())
-        self._org_models = self._build_model_lookup(self._get_org_models())
         self._user_models = self._build_model_lookup(self._get_user_models())
         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,47 +156,76 @@ 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 _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
 
-    def _fetch_conversations(self) -> None:
-        self._conversations = self._get_conversations()
+        for line in result:
+            yield line.choices[0].delta.content
 
     def chat(
         self,
         prompt: str,
         system_prompt: Optional[str] = None,
         max_completion_tokens: Optional[int] = 500,
+        images: Optional[Union[List[str], str]] = None,
         conversation: Optional[str] = None,
-    ) -> str:
+        metadata: Optional[Dict[str, str]] = None,
+        stream: bool = False,
+        upload_local_images: bool = False,
+    ) -> Union[str, Generator[str, None, None]]:
         if conversation and conversation not in self._conversations:
-            self._fetch_conversations()
+            self._get_conversations()
+
+        if images:
+            if isinstance(images, str):
+                images = [images]
+            for image in images:
+                if not isinstance(image, str):
+                    raise NotImplementedError(f"Image type {type(image)} are not supported yet.")
+                if not image.startswith("http") and upload_local_images:
+                    self._teamspace.upload_file(file_path=image, remote_path=f"images/{os.path.basename(image)}")
 
         conversation_id = self._conversations.get(conversation) if conversation else None
         output = self._llm_api.start_conversation(
             prompt=prompt,
             system_prompt=system_prompt,
             max_completion_tokens=max_completion_tokens,
+            images=images,
             assistant_id=self._model.id,
             conversation_id=conversation_id,
+            billing_project_id=self._teamspace.id,
+            metadata=metadata,
+            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 +242,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

@@ -1,13 +1,16 @@
 import glob
 import os
 import warnings
+from enum import Enum
 from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
 
 from tqdm.auto import tqdm
 
+from lightning_sdk.api.cluster_api import ClusterApi
 from lightning_sdk.api.studio_api import StudioApi
 from lightning_sdk.api.utils import _machine_to_compute_name
 from lightning_sdk.constants import _LIGHTNING_DEBUG
+from lightning_sdk.lightning_cloud.openapi import V1CloudSpaceSourceType
 from lightning_sdk.machine import Machine
 from lightning_sdk.organization import Organization
 from lightning_sdk.owner import Owner
@@ -24,6 +27,19 @@ if TYPE_CHECKING:
 _logger = _setup_logger(__name__)
 
 
+class Provider(Enum):
+    # Machine providers based on v1CloudProvider
+    AWS = "AWS"
+    GCP = "GCP"
+    VULTR = "VULTR"
+    LAMBDA_LABS = "LAMBDA_LABS"
+    DGX = "DGX"
+    VOLTAGE_PARK = "VOLTAGE_PARK"
+    NEBIUS = "NEBIUS"
+    CLOUDFLARE = "CLOUDFLARE"
+    LIGHTNING = "LIGHTNING"
+
+
 class Studio:
     """A single Lightning AI Studio.
 
@@ -38,6 +54,8 @@ class Studio:
         cloud_account: the name of the cloud account, the studio should be created on.
             Doesn't matter when the studio already exists.
         create_ok: whether the studio will be created if it does not yet exist. Defaults to True
+        provider: the provider of the machine, the studio should be created on.
+
     Note:
         Since a teamspace can either be owned by an org or by a user directly,
         only one of the arguments can be provided.
@@ -56,8 +74,11 @@ class Studio:
         cloud_account: Optional[str] = None,
         create_ok: bool = True,
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
+        provider: Optional[str] = None,
+        source: Optional[V1CloudSpaceSourceType] = None,
     ) -> None:
         self._studio_api = StudioApi()
+        self._cluster_api = ClusterApi()
 
         self._teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
         self._cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
@@ -65,6 +86,16 @@ class Studio:
 
         self._plugins = {}
 
+        if provider is not None:
+            if isinstance(provider, str) and provider in Provider.__members__:
+                provider = Provider(provider)
+            else:
+                raise ValueError(f"Invalid provider: {provider}. Must be one of {Provider.__members__.keys()}.")
+            self._cloud_account = self._cluster_api.get_cluster_provider_mapping(
+                self._teamspace.id,
+                self._teamspace.owner.id,
+            )[provider.value]
+
         if name is None:
             studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
             if studio_id is None:
@@ -76,7 +107,7 @@ class Studio:
             except ValueError as e:
                 if create_ok:
                     self._studio = self._studio_api.create_studio(
-                        name, self._teamspace.id, cloud_account=self._cloud_account
+                        name, self._teamspace.id, cloud_account=self._cloud_account, source=source
                     )
                 else:
                     raise ValueError(f"Studio {name} does not exist.") from e
@@ -221,6 +252,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.