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

lightning_sdk/llm/llm.py

@@ -0,0 +1,118 @@
+from typing import Dict, 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.lightning_cloud.openapi import V1Assistant
+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
+
+
+class LLM:
+    def __init__(
+        self,
+        name: str,
+        user: Union[str, "User", None] = None,
+        org: Union[str, "Organization", None] = 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)
+        try:
+            # check if it is a org model
+            self._org = _resolve_org(self._org or org)
+        except ApiException:
+            self._org = None
+
+        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 = {}
+
+    def _parse_model_name(self, name: str) -> Tuple[str, str]:
+        parts = name.split("/")
+        if len(parts) == 1:
+            # a user model or a org model
+            return None, parts[0]
+        if len(parts) == 2:
+            return parts[0], parts[1]
+        raise ValueError(
+            f"Model name must be in the format `organization/model_name` or `model_name`, but got '{name}'."
+        )
+
+    def _build_model_lookup(self, endpoints: List[str]) -> Dict[str, Set[str]]:
+        result = {}
+        for endpoint in endpoints:
+            result.setdefault(endpoint.model, []).append(endpoint)
+        return result
+
+    def _get_public_models(self) -> List[str]:
+        return self._llm_api.get_public_models()
+
+    def _get_org_models(self) -> List[str]:
+        return self._llm_api.get_org_models(self._org.id) if self._org else []
+
+    def _get_user_models(self) -> List[str]:
+        return self._llm_api.get_user_models(self._user.id) if self._user else []
+
+    def _get_model(self) -> V1Assistant:
+        # TODO how to handle multiple models with same model type? For now, just use the first one
+        if self._model_name in self._public_models:
+            return self._public_models.get(self._model_name)[0]
+        if self._model_name in self._org_models:
+            return self._org_models.get(self._model_name)[0]
+        if self._model_name in self._user_models:
+            return self._user_models.get(self._model_name)[0]
+
+        available_models = []
+        if self._public_models:
+            available_models.append(f"Public Models: {', '.join(self._public_models.keys())}")
+
+        if self._org and self._org_models:
+            available_models.append(f"Org ({self._org.name}) Models: {', '.join(self._org_models.keys())}")
+
+        if self._user and self._user_models:
+            available_models.append(f"User ({self._user.name}) Models: {', '.join(self._user_models.keys())}")
+
+        available_models_str = "\n".join(available_models)
+        raise ValueError(f"Model '{self._model_name}' not found. \nAvailable models: \n{available_models_str}")
+
+    def chat(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        max_completion_tokens: Optional[int] = 500,
+        conversation: Optional[str] = None,
+    ) -> str:
+        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,
+            assistant_id=self._model.id,
+            conversation_id=conversation_id,
+        )
+        if conversation and not conversation_id:
+            self._conversations[conversation] = output.conversation_id
+        return output.choices[0].delta.content
+
+    def list_conversations(self) -> List[Dict]:
+        return list(self._conversations.keys())

lightning_sdk/plugin.py

@@ -395,6 +395,25 @@ class SlurmJobsPlugin(_Plugin):
         return resp
 
 
+class CustomPortPlugin(_Plugin):
+    """Plugin handling the port of a given service."""
+
+    _plugin_run_name = "Custom Port"
+    _slug_name = "custom-port"
+
+    def run(self, name: Optional[str] = None, port: int = 8000) -> str:
+        """Starts a new port to the given Studio."""
+        if name is None:
+            name = _run_name("port")
+
+        return self._studio._studio_api.start_new_port(
+            teamspace_id=self._studio._teamspace.id,
+            studio_id=self._studio._studio.id,
+            name=name,
+            port=port,
+        )
+
+
 @runtime_checkable
 class _RunnablePlugin(Protocol):
     _plugin_run_name: str

lightning_sdk/serve.py

@@ -10,7 +10,7 @@ from rich.console import Console
 from lightning_sdk import Deployment, Machine, Teamspace
 from lightning_sdk.api.deployment_api import AutoScaleConfig, DeploymentApi, Env, Secret
 from lightning_sdk.api.lit_container_api import LitContainerApi
-from lightning_sdk.api.utils import _get_cloud_url, _get_registry_url
+from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
 
 _DOCKER_NOT_RUNNING_MSG = (
@@ -189,12 +189,9 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
             if "status" in line:
                 yield {"status": line["status"].strip()}
 
-        registry_url = _get_registry_url()
         container_basename = repository.split("/")[-1]
-        repository = (
-            f"{registry_url}/lit-container{f'-{cloud_account}' if cloud_account is not None else ''}/"
-            f"{teamspace.owner.name}/{teamspace.name}/{container_basename}"
-        )
+        repository = lit_cr.get_container_url(repository, tag, teamspace, cloud_account)
+
         yield {
             "finish": True,
             "status": "Container pushed successfully",
@@ -303,6 +300,7 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
             include_credentials=include_credentials,
             cloudspace_id=cloudspace_id,
             from_onboarding=from_onboarding,
+            command="",
         )
 
         return {"deployment": deployment, "url": url}

lightning_sdk/studio.py

@@ -1,7 +1,10 @@
+import glob
 import os
 import warnings
 from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
 
+from tqdm.auto import tqdm
+
 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
@@ -265,6 +268,34 @@ class Studio:
             progress_bar=progress_bar,
         )
 
+    def upload_folder(self, folder_path: str, remote_path: Optional[str] = None, progress_bar: bool = True) -> None:
+        """Uploads a given folder to a remote path on the Studio."""
+        if folder_path is None:
+            raise ValueError("Cannot upload a folder that is None.")
+        folder_path = os.path.normpath(folder_path)
+        if os.path.isfile(folder_path):
+            raise NotADirectoryError(f"Cannot upload a file as a folder. '{folder_path}' is a file.")
+        if not os.path.exists(folder_path):
+            raise NotADirectoryError(f"Cannot upload a folder that does not exist. '{folder_path}' is not a directory.")
+        all_files = []
+        for fp in glob.glob(os.path.join(folder_path, "**"), recursive=True):
+            if not os.path.isfile(fp):
+                continue
+            rel_path = os.path.relpath(fp, folder_path)
+            remote_file = os.path.join(remote_path, rel_path) if remote_path else rel_path
+            all_files.append((fp, remote_file))
+
+        if progress_bar:
+            progress_bar = tqdm(total=len(all_files), desc="Uploading files", unit="file")
+        for local_file, remote_path in sorted(all_files, key=lambda p: p[1]):
+            if progress_bar:
+                progress_bar.set_description(f"Uploading {local_file}")
+            self.upload_file(local_file, remote_path=remote_path, progress_bar=False)
+            if progress_bar:
+                progress_bar.update(1)
+        if progress_bar:
+            progress_bar.close()
+
     def download_file(self, remote_path: str, file_path: Optional[str] = None) -> None:
         """Downloads a file from the Studio to a given target path."""
         if file_path is None:
@@ -445,6 +476,7 @@ class Studio:
     def _add_plugin(self, plugin_name: str) -> None:
         """Adds the just installed plugin to the internal list of plugins."""
         from lightning_sdk.plugin import (
+            CustomPortPlugin,
             InferenceServerPlugin,
             JobsPlugin,
             MultiMachineTrainingPlugin,
@@ -458,6 +490,7 @@ class Studio:
             "jobs": JobsPlugin,
             "multi-machine-training": MultiMachineTrainingPlugin,
             "inference-server": InferenceServerPlugin,
+            "custom-port": CustomPortPlugin,
         }.get(plugin_name, Plugin)
 
         description = self._list_installed_plugins()[plugin_name]

{lightning_sdk-0.2.11.dist-info → lightning_sdk-0.2.13.dist-info}/METADATA

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