lightning-sdk 0.1.49__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/api/lit_container_api.py

@@ -1,42 +1,254 @@
-from typing import Generator, List
+import inspect
+import time
+from typing import Any, Callable, Dict, Generator, Iterator, List, Optional
+
+import docker
+import requests
+from rich.console import Console
 
 from lightning_sdk.api.utils import _get_registry_url
+from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
 from lightning_sdk.lightning_cloud.openapi.models import V1DeleteLitRepositoryResponse
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.teamspace import Teamspace
 
 
+class LCRAuthFailedError(Exception):
+    def __init__(self) -> None:
+        super().__init__("Failed to authenticate with Lightning Container Registry. Please try again.")
+
+
+class DockerPushError(Exception):
+    pass
+
+
+class DockerNotRunningError(Exception):
+    def __init__(self, message: str = "Failed to connect to Docker") -> None:
+        self.message = message
+        super().__init__(self.message)
+
+    def print_help(self) -> None:
+        console = Console()
+        console.print("[bold red]Docker Error: Failed to connect to Docker. Is it running?[/bold red]")
+        console.print("[yellow]Troubleshooting:[/yellow]")
+        console.print("1. Check if Docker daemon is running")
+        console.print("2. Verify your user has proper permissions")
+        console.print("3. Try restarting Docker service")
+        console.print("4. Read more here: https://docs.docker.com/engine/daemon/start/")
+
+
+def retry_on_lcr_auth_failure(func: Callable) -> Callable:
+    def generator_wrapper(self: "LitContainerApi", *args: Any, **kwargs: Any) -> Callable:
+        try:
+            gen = func(self, *args, **kwargs)
+            yield from gen
+        except LCRAuthFailedError:
+            self.authenticate(reauth=True)
+            gen = func(self, *args, **kwargs)
+            yield from gen
+        return
+
+    def wrapper(self: "LitContainerApi", *args: Any, **kwargs: Any) -> Callable:
+        try:
+            return func(self, *args, **kwargs)
+        except LCRAuthFailedError:
+            self.authenticate(reauth=True)
+            return func(self, *args, **kwargs)
+
+    if inspect.isgeneratorfunction(func):
+        return generator_wrapper
+
+    return wrapper
+
+
 class LitContainerApi:
     def __init__(self) -> None:
         self._client = LightningClient(max_tries=3)
 
-    def list_containers(self, project_id: str) -> List:
-        project = self._client.lit_registry_service_get_lit_project_registry(project_id)
+        try:
+            self._docker_client = docker.from_env()
+            self._docker_client.ping()
+            self._docker_auth_config = {}
+        except docker.errors.DockerException:
+            raise DockerNotRunningError() from None
+
+    def authenticate(self, reauth: bool = False) -> bool:
+        resp = None
+        try:
+            authed_user = self._client.auth_service_get_user()
+            username = authed_user.username
+            api_key = authed_user.api_key
+            registry = _get_registry_url()
+            resp = self._docker_client.login(username, password=api_key, registry=registry, reauth=reauth)
+
+            if (
+                resp.get("username", None) == username
+                and resp.get("password", None) == api_key
+                and resp.get("serveraddress", None) == registry
+            ):
+                self._docker_auth_config = {"username": username, "password": api_key}
+                return True
+
+            # This is a new 200 response auth attempt from the client.
+            if "Status" in resp and resp["Status"] == "Login Succeeded":
+                self._docker_auth_config = {"username": username, "password": api_key}
+                return True
+
+            return False
+        except Exception as e:
+            print(f"Authentication error: {e} resp: {resp}")
+            return False
+
+    def list_containers(self, project_id: str, cloud_account: Optional[str] = None) -> List:
+        """Lists containers of the project ID.
+
+        :param project_id: The non-human readable project ID used internally to identify projects.
+        :param cloud_account: The cluster ID of the cloud account. If None, will use the default cluster.
+        :return:
+        """
+        if cloud_account is None:
+            project = self._client.lit_registry_service_get_lit_project_registry(project_id)
+        else:
+            project = self._client.lit_registry_service_get_lit_project_registry(project_id, cluster_id=cloud_account)
+
         return project.repositories
 
     def delete_container(self, project_id: str, container: str) -> V1DeleteLitRepositoryResponse:
+        """Deletes the container from LitCR. Garbage collection will come in and do the final sweep of the blob.
+
+        :param project_id: The non-human readable project ID used internally to identify projects.
+        :param container: The name of the docker container a user wants to push up, ie, nginx, vllm, etc
+        :return:
+        """
         try:
             return self._client.lit_registry_service_delete_lit_repository(project_id, container)
-        except Exception as ex:
-            raise ValueError(f"Could not delete container {container} from project {project_id}") from ex
+        except Exception as e:
+            raise ValueError(f"Could not delete container {container} from project {project_id}: {e!s}") from e
 
-    def upload_container(self, container: str, teamspace: Teamspace, tag: str) -> Generator[str, None, None]:
-        import docker
+    def get_container_url(
+        self, repository: str, tag: str, teamspace: Teamspace, cloud_account: Optional[str] = None
+    ) -> str:
+        """Docker container will be pushed to the URL returned from this function."""
+        registry_url = _get_registry_url()
+        container_basename = repository.split("/")[-1]
+        return (
+            f"{registry_url}/lit-container{f'-{cloud_account}' if cloud_account is not None else ''}/"
+            f"{teamspace.owner.name}/{teamspace.name}/{container_basename}"
+        )
 
-        try:
-            client = docker.from_env()
-            client.ping()
-        except docker.errors.DockerException as e:
-            raise RuntimeError(f"Failed to connect to Docker daemon: {e!s}. Is Docker running?") from None
+    @retry_on_lcr_auth_failure
+    def upload_container(
+        self,
+        container: str,
+        teamspace: Teamspace,
+        tag: str,
+        cloud_account: str,
+        platform: str,
+        return_final_dict: bool = False,
+    ) -> Generator[dict, None, Dict]:
+        """Upload container will push the container to LitCR.
+
+        It uses docker push API to interact with docker daemon which will then push the container to a storage
+        location defined by teamspace + cloud_account. Will retry pushes if not authenticated or if push errors happen
 
+        :param container: The name of the docker container a user wants to push up, ie, nginx, vllm, etc
+        :param teamspace: The teamspace he container is going to appear in. This is <OWNER_ID>/<TEAMSPACE_NAME>
+        :param tag: The container tag, default will latest.
+        :param cloud_account: If empty will be pushed to Lightning SaaS storage. Otherwise, this will be cluster_id.
+            Named cloud-account in the CLI options.
+        :param platform: If empty will be linux/amd64. This is important because our entire deployment infra runs on
+            linux/amd64. Will show user a warning otherwise.
+        :return: Generator[dict, None, dict]
+        """
         try:
-            client.images.get(container)
+            self._docker_client.images.get(f"{container}:{tag}")
         except docker.errors.ImageNotFound:
-            raise ValueError(f"Container {container} does not exist") from None
+            try:
+                self._docker_client.images.pull(repository=container, tag=tag, platform=platform)
+                self._docker_client.images.get(f"{container}:{tag}")
+            except docker.errors.ImageNotFound as e:
+                raise ValueError(f"Container {container}:{tag} does not exist") from e
+            except docker.errors.APIError as e:
+                raise ValueError(f"Could not pull container {container}") from e
+            except Exception as e:
+                raise ValueError(f"Unable to upload {container}:{tag}") from e
 
-        registry_url = _get_registry_url()
-        repository = f"{registry_url}/lit-container/{teamspace.owner.name}/{teamspace.name}/{container}"
-        tagged = client.api.tag(container, repository, tag)
+        repository = self.get_container_url(container, tag, teamspace, cloud_account)
+        tagged = self._docker_client.api.tag(f"{container}:{tag}", repository, tag)
         if not tagged:
-            raise ValueError(f"Could not tag container {container} with {repository}:{tag}")
-        return client.api.push(repository, stream=True, decode=True)
+            raise ValueError(f"Could not tag container {container}:{tag} with {repository}:{tag}")
+        yield from self._push_with_retry(repository, tag=tag)
+
+        if return_final_dict:
+            container_basename = repository.split("/")[-1]
+            yield {
+                "finish": True,
+                "url": f"{LIGHTNING_CLOUD_URL}/{teamspace.owner.name}/{teamspace.name}/containers/"
+                f"{container_basename}?section=tags"
+                f"{f'?clusterId={cloud_account}' if cloud_account is not None else ''}",
+                "repository": repository,
+            }
+
+    def _push_with_retry(self, repository: str, tag: str, max_retries: int = 3) -> Iterator[Dict[str, Any]]:
+        def is_auth_error(error_msg: str) -> bool:
+            auth_errors = ["unauthorized", "authentication required", "unauth"]
+            return any(err in error_msg.lower() for err in auth_errors)
+
+        def is_timeout_error(error_msg: str) -> bool:
+            timeout_errors = ["proxyconnect tcp", "i/o timeout"]
+            return any(err in error_msg.lower() for err in timeout_errors)
+
+        for attempt in range(max_retries):
+            try:
+                if attempt > 0:
+                    # This is important, if we don't set reauth here then we just keep using the
+                    # same authentication context that we know just failed.
+                    self.authenticate(reauth=True)
+                    time.sleep(2)
+
+                lines = self._docker_client.api.push(
+                    repository, tag=tag, stream=True, decode=True, auth_config=self._docker_auth_config
+                )
+                for line in lines:
+                    if isinstance(line, dict) and "error" in line:
+                        error = line["error"]
+                        if is_auth_error(error) or is_timeout_error(error):
+                            if attempt < max_retries - 1:
+                                break
+                            raise DockerPushError(f"Max retries reached while pushing: {error}")
+                        raise DockerPushError(f"Push error: {error}")
+                    yield line
+                else:
+                    return
+
+            except docker.errors.APIError as e:
+                if (is_auth_error(str(e)) or is_timeout_error(str(e))) and attempt < max_retries - 1:
+                    continue
+                raise DockerPushError("Pushing the container failed") from e
+
+        raise DockerPushError("Max push retries reached")
+
+    @retry_on_lcr_auth_failure
+    def download_container(
+        self, container: str, teamspace: Teamspace, tag: str, cloud_account: Optional[str] = None
+    ) -> Generator[str, None, None]:
+        """Will download container from LitCR. Optionally from a BYOC cluster.
+
+        :param container: The name of the container to download
+        :param teamspace: The teamspace containing the container
+        :param tag: The tag of the container to download
+        :return: Generator yielding download status
+        """
+        registry_url = _get_registry_url()
+        repository = (
+            f"{registry_url}/lit-container{f'-{cloud_account}' if cloud_account is not None else ''}/"
+            f"{teamspace.owner.name}/{teamspace.name}/{container}"
+        )
+        try:
+            self._docker_client.images.pull(repository, tag=tag, auth_config=self._docker_auth_config)
+        except requests.exceptions.HTTPError as e:
+            if "unauthorized" in e.response.text:
+                raise LCRAuthFailedError() from e
+        except docker.errors.APIError as e:
+            raise ValueError(f"Could not pull container {container} from {repository}:{tag}") from e
+        return self._docker_client.api.tag(repository, container, tag)

lightning_sdk/api/llm_api.py

@@ -0,0 +1,306 @@
+import asyncio
+import base64
+import datetime
+import json
+import os
+import threading
+import time
+import warnings
+from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union
+from urllib.parse import urlencode
+
+from rich.prompt import Confirm
+
+from lightning_sdk.lightning_cloud import env
+from lightning_sdk.lightning_cloud.login import Auth, AuthServer
+from lightning_sdk.lightning_cloud.openapi.models import (
+    StreamResultOfV1ConversationResponseChunk,
+    V1ConversationResponseChunk,
+    V1ManagedModel,
+    V1ResponseChoice,
+    V1ResponseChoiceDelta,
+)
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+LITAI_CODE = os.environ.get("LITAI_CODE", "x334uv8t7v")
+
+
+class _AuthServer(AuthServer):
+    def __init__(self, model: str, *args: Any, **kwargs: Any) -> None:
+        self._model = model
+        super().__init__(*args, **kwargs)
+
+    def get_auth_url(self, port: int) -> str:
+        redirect_uri = f"http://localhost:{port}/login-complete"
+        params = urlencode({"redirectTo": redirect_uri, "okbhrt": LITAI_CODE, "litaimodel": self._model})
+        return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
+
+
+class _AuthLitAI(Auth):
+    def __init__(self, model: str, shall_confirm: bool = False) -> None:
+        super().__init__()
+        self._model = model
+        self._shall_confirm = shall_confirm
+
+    def _run_server(self) -> None:
+        if self._shall_confirm:
+            proceed = Confirm.ask(
+                "[bold yellow]LitAI needs to authenticate with Lightning AI.[/bold yellow]\n"
+                "This will open a browser window for login.\n"
+                "Do you want to continue?",
+                default=True,
+            )
+            if not proceed:
+                raise RuntimeError(
+                    "Login cancelled. Please login at https://lightning.ai/sign-up. Or run `lightning login` to login."
+                )
+        print("Opening browser for authentication...")
+        print("Please come back to the terminal after logging in.")
+        time.sleep(1)
+        _AuthServer(self._model).login_with_browser(self)
+
+
+def authenticate(model: str, shall_confirm: bool = True) -> None:
+    auth = _AuthLitAI(model, shall_confirm)
+    auth.authenticate()
+
+
+class LLMApi:
+    def __init__(self) -> None:
+        self._client = LightningClient(retry=False, max_tries=0)
+        self._assistant = None
+        self._model = None
+
+    def get_assistant(self, model_provider: str, model_name: str, user_name: str, org_name: str) -> str:
+        result = self._client.assistants_service_get_managed_model_assistant(
+            model_provider=model_provider, model_name=model_name, user_name=user_name, org_name=org_name
+        )
+        self._assistant = result
+        return result.id
+
+    def get_model_metadata(self, teamspace_id: str, model_name: str) -> V1ManagedModel:
+        if self._assistant and not self._model:
+            self._model = self._client.assistants_service_get_managed_model_by_name(
+                teamspace_id, self._assistant.managed_endpoint_id, model_name
+            )
+        return self._model
+
+    def _parse_stream_line(self, decoded_line: str) -> Optional[V1ConversationResponseChunk]:
+        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"),
+                    )
+                )
+
+            return V1ConversationResponseChunk(
+                choices=choices,
+                conversation_id=result_data.get("conversationId"),
+                executable=result_data.get("executable"),
+                id=result_data.get("id"),
+                throughput=result_data.get("throughput"),
+                stats=result_data.get("stats"),
+                usage=result_data.get("usage"),
+            )
+        except json.JSONDecodeError:
+            warnings.warn("Error decoding JSON:", decoded_line)
+            return None
+
+    def _stream_chat_response(
+        self, result: StreamResultOfV1ConversationResponseChunk
+    ) -> Generator[V1ConversationResponseChunk, None, None]:
+        for line in result.stream():
+            decoded_lines = line.decode("utf-8").strip()
+            for decoded_line in decoded_lines.splitlines():
+                chunk = self._parse_stream_line(decoded_line)
+                if chunk:
+                    yield chunk
+
+    def _encode_image_bytes_to_data_url(self, image: str) -> str:
+        with open(image, "rb") as image_file:
+            b64 = base64.b64encode(image_file.read()).decode("utf-8")
+            extension = image.split(".")[-1]
+            return f"data:image/{extension};base64,{b64}"
+
+    def start_conversation(
+        self,
+        prompt: str,
+        system_prompt: Optional[str],
+        max_completion_tokens: Optional[int],
+        assistant_id: str,
+        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,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        reasoning_effort: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Union[V1ConversationResponseChunk, Generator[V1ConversationResponseChunk, None, None]]:
+        is_internal_conversation = os.getenv("LIGHTNING_INTERNAL_CONVERSATION", "false").lower() == "true"
+        ephemeral = os.getenv("LIGHTNING_EPHEMERAL", "false").lower() == "true"
+        if ephemeral:
+            conversation_id = None
+            name = None
+        body = {
+            "message": {
+                "author": {"role": "user"},
+                "content": [
+                    {"contentType": "text", "parts": [prompt]},
+                ],
+            },
+            "conversation_id": conversation_id,
+            "billing_project_id": billing_project_id,
+            "name": name,
+            "stream": stream,
+            "metadata": metadata or {},
+            "internal_conversation": is_internal_conversation,
+            "system_prompt": system_prompt,
+            "ephemeral": ephemeral,
+            "parent_conversation_id": kwargs.get("parent_conversation_id", ""),
+            "parent_message_id": kwargs.get("parent_message_id", ""),
+            "tools": tools,
+        }
+        if max_completion_tokens is not None:
+            body["max_completion_tokens"] = max_completion_tokens
+
+        if reasoning_effort:
+            body["reasoning_effort"] = reasoning_effort
+
+        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)
+
+    async def async_start_conversation(
+        self,
+        prompt: str,
+        system_prompt: Optional[str],
+        max_completion_tokens: Optional[int],
+        assistant_id: str,
+        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,
+        reasoning_effort: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Union[V1ConversationResponseChunk, AsyncGenerator[V1ConversationResponseChunk, None]]:
+        is_internal_conversation = os.getenv("LIGHTNING_INTERNAL_CONVERSATION", "false").lower() == "true"
+        ephemeral = os.getenv("LIGHTNING_EPHEMERAL", "false").lower() == "true"
+        if ephemeral:
+            conversation_id = None
+            name = None
+        body = {
+            "message": {
+                "author": {"role": "user"},
+                "content": [
+                    {"contentType": "text", "parts": [prompt]},
+                ],
+            },
+            "conversation_id": conversation_id,
+            "billing_project_id": billing_project_id,
+            "name": name,
+            "stream": stream,
+            "metadata": metadata or {},
+            "internal_conversation": is_internal_conversation,
+            "system_prompt": system_prompt,
+            "ephemeral": ephemeral,
+            "parent_conversation_id": kwargs.get("parent_conversation_id", ""),
+            "parent_message_id": kwargs.get("parent_message_id", ""),
+            "sent_at": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="microseconds"),
+        }
+        if max_completion_tokens is not None:
+            body["max_completion_tokens"] = max_completion_tokens
+
+        if reasoning_effort:
+            body["reasoning_effort"] = reasoning_effort
+
+        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],
+                    }
+                )
+
+        if not stream:
+            thread = await asyncio.to_thread(
+                self._client.assistants_service_start_conversation, body, assistant_id, async_req=True
+            )
+            result = await asyncio.to_thread(thread.get)
+            return result.result
+
+        conversation_thread = await asyncio.to_thread(
+            self._client.assistants_service_start_conversation,
+            body,
+            assistant_id,
+            async_req=True,
+            _preload_content=False,
+        )
+
+        return self.stream_response(conversation_thread)
+
+    async def stream_response(self, thread: Any) -> AsyncGenerator[V1ConversationResponseChunk, None]:
+        loop = asyncio.get_event_loop()
+        response = await asyncio.to_thread(thread.get)
+
+        queue = asyncio.Queue()
+
+        def enqueue() -> None:
+            try:
+                for line in response:
+                    decoded_lines = line.decode("utf-8").strip()
+                    for decoded_line in decoded_lines.splitlines():
+                        chunk = self._parse_stream_line(decoded_line)
+                        if chunk:
+                            asyncio.run_coroutine_threadsafe(queue.put(chunk), loop)
+            finally:
+                asyncio.run_coroutine_threadsafe(queue.put(None), loop)
+
+        threading.Thread(target=enqueue, daemon=True).start()
+
+        while True:
+            item = await queue.get()
+            if item is None:
+                break
+            yield item
+
+    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)