lightning-sdk 2025.11.13.post0__py3-none-any.whl

lightning_sdk/serve.py

@@ -0,0 +1,311 @@
+import os
+import shlex
+import subprocess
+from pathlib import Path
+from typing import Generator, List, Optional, Union
+
+import docker
+from rich.console import Console
+
+from lightning_sdk import CloudProvider, 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
+from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
+
+_DOCKER_NOT_RUNNING_MSG = (
+    "Deploying LitServe requires Docker to be running on the machine. "
+    "If Docker is not installed, please install it from https://docs.docker.com/get-docker/ "
+    "and start the Docker daemon before running this command."
+)
+
+
+class _LitServeDeployer:
+    def __init__(self, name: Optional[str], teamspace: Optional[Teamspace]) -> None:
+        self.name = name
+        self.teamspace = teamspace
+        self._console = Console()
+        self._client = None
+
+    @property
+    def client(self) -> docker.DockerClient:
+        if self._client is None:
+            try:
+                os.environ["DOCKER_BUILDKIT"] = "1"
+                self._client = docker.from_env()
+                self._client.ping()
+            except docker.errors.DockerException:
+                raise RuntimeError(_DOCKER_NOT_RUNNING_MSG) from None
+        return self._client
+
+    @property
+    def created(self) -> bool:
+        return DeploymentApi().get_deployment_by_name(self.name, self.teamspace.id) is not None
+
+    def dockerize_api(
+        self,
+        server_filename: str,
+        port: int = 8000,
+        gpu: bool = False,
+        tag: str = "litserve-model",
+        print_success: bool = True,
+    ) -> str:
+        import litserve as ls
+        from litserve import docker_builder
+
+        console = self._console
+        requirements = ""
+        if os.path.exists("requirements.txt"):
+            requirements = "-r requirements.txt"
+        else:
+            console.print(
+                f"requirements.txt not found at {os.getcwd()}. "
+                f"Make sure to install the required packages in the Dockerfile.",
+                style="yellow",
+            )
+
+        if os.path.exists("Dockerfile"):
+            console.print(
+                "Dockerfile already exists in the current directory, we will use it for building the container."
+            )
+            return os.path.abspath("Dockerfile")
+        current_dir = Path.cwd()
+        if not (current_dir / server_filename).is_file():
+            raise FileNotFoundError(f"Server file `{server_filename}` must be in the current directory: {os.getcwd()}")
+
+        version = ls.__version__
+        if gpu:
+            run_cmd = f"docker run --gpus all -p {port}:{port} {tag}:latest"
+            docker_template = docker_builder.CUDA_DOCKER_TEMPLATE
+        else:
+            run_cmd = f"docker run -p {port}:{port} {tag}:latest"
+            docker_template = docker_builder.DOCKERFILE_TEMPLATE
+        dockerfile_content = docker_template.format(
+            server_filename=server_filename,
+            port=port,
+            version=version,
+            requirements=requirements,
+        )
+        with open("Dockerfile", "w") as f:
+            f.write(dockerfile_content)
+
+        if print_success:
+            success_msg = f"""[bold]Dockerfile created successfully[/bold]
+Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additional dependencies or commands.
+
+[bold]Build the container with:[/bold]
+> [underline]docker build -t {tag} .[/underline]
+
+[bold]To run the Docker container on the machine:[/bold]
+> [underline]{run_cmd}[/underline]
+
+[bold]To push the container to a registry:[/bold]
+> [underline]docker push {tag}[/underline]
+
+Check out [blue][link=https://lightning.ai/docs/litserve/features]the docs[/link][/blue] for more details.
+"""
+            console.print(success_msg)
+        return os.path.abspath("Dockerfile")
+
+    @staticmethod
+    def generate_client() -> None:
+        from rich.console import Console
+
+        console = Console()
+        try:
+            from litserve.python_client import client_template
+        except ImportError:
+            raise ImportError(
+                "litserve is not installed. Please install it with `pip install lightning_sdk[serve]`"
+            ) from None
+
+        client_path = Path("client.py")
+        if client_path.exists():
+            console.print("Skipping client generation: client.py already exists", style="blue")
+        else:
+            try:
+                client_path.write_text(client_template)
+                console.print("✅ Client generated at client.py", style="bold green")
+            except OSError as e:
+                raise OSError(f"Failed to generate client.py: {e!s}") from None
+
+    def _docker_build_with_logs(
+        self, path: str, repository: str, tag: str, platform: str = "linux/amd64"
+    ) -> Generator[str, None, None]:
+        """Build Docker image using CLI with real-time log streaming.
+
+        Returns:
+            Tuple: (image_id, logs generator)
+
+        Raises:
+            RuntimeError: On build failure
+        """
+        cmd = f"docker build --platform {platform} -t {repository}:{tag} ."
+        proc = subprocess.Popen(
+            shlex.split(cmd),
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            bufsize=1,  # Line buffered
+        )
+
+        def log_generator() -> Generator[str, None, None]:
+            while True:
+                line = proc.stdout.readline()
+                if not line and proc.poll() is not None:
+                    break
+                yield line.strip()
+                if "error" in line.lower():
+                    proc.terminate()
+                    raise RuntimeError(f"Build failed: {line.strip()}")
+
+            if proc.returncode != 0:
+                raise RuntimeError(f"Build failed with exit code {proc.returncode}")
+
+        return log_generator()
+
+    def build_container(self, path: str, repository: str, tag: str) -> Generator[str, None, None]:
+        build_logs = self._docker_build_with_logs(path, repository, tag=tag)
+
+        for line in build_logs:
+            if "error" in line:
+                raise RuntimeError(f"Failed to build image: {line}")
+            else:
+                yield line.strip()
+
+    def push_container(
+        self,
+        repository: str,
+        tag: str,
+        teamspace: Teamspace,
+        lit_cr: LitContainerApi,
+        cloud_account: str,
+    ) -> Generator[dict, None, dict]:
+        lit_cr.authenticate()
+        push_status = lit_cr.upload_container(
+            repository, teamspace, tag=tag, cloud_account=cloud_account, platform=None
+        )
+        for line in push_status:
+            if "error" in line:
+                raise RuntimeError(f"Failed to push image: {line}")
+            if "status" in line:
+                yield {"status": line["status"].strip()}
+
+        container_basename = repository.split("/")[-1]
+        repository = lit_cr.get_container_url(repository, tag, teamspace, cloud_account)
+
+        yield {
+            "finish": True,
+            "status": "Container pushed successfully",
+            "url": f"{LIGHTNING_CLOUD_URL}/{teamspace.owner.name}/{teamspace.name}/containers/{container_basename}"
+            f"{f'?clusterId={cloud_account}' if cloud_account is not None else ''}",
+            "image": repository,
+        }
+
+    def _update_deployment(
+        self,
+        deployment: Deployment,
+        machine: Optional[Machine] = None,
+        image: Optional[str] = None,
+        entrypoint: Optional[str] = None,
+        command: Optional[str] = None,
+        env: Optional[List[Union[Env, Secret]]] = None,
+        min_replica: Optional[int] = 0,
+        max_replica: Optional[int] = 1,
+        spot: Optional[bool] = None,
+        replicas: Optional[int] = 1,
+        cloud_account: Optional[str] = None,
+        port: Optional[int] = 8000,
+        include_credentials: Optional[bool] = True,
+    ) -> None:
+        return deployment.update(
+            machine=machine,
+            image=image,
+            entrypoint=entrypoint,
+            command=command,
+            env=env,
+            max_replicas=max_replica,
+            min_replicas=min_replica,
+            replicas=replicas,
+            spot=spot,
+            cloud_account=cloud_account,
+            ports=[port],
+            include_credentials=include_credentials,
+        )
+
+    def run_on_cloud(
+        self,
+        deployment_name: str,
+        teamspace: Teamspace,
+        image: str,
+        metric: Optional[str] = None,
+        machine: Optional[Machine] = None,
+        min_replica: Optional[int] = 0,
+        max_replica: Optional[int] = 1,
+        spot: Optional[bool] = None,
+        replicas: Optional[int] = 1,
+        cloud_account: Optional[str] = None,
+        port: Optional[int] = 8000,
+        include_credentials: Optional[bool] = True,
+        cloudspace_id: Optional[str] = None,
+        from_onboarding: bool = False,
+        cloud_provider: Optional[CloudProvider] = None,
+    ) -> dict:
+        """Run a deployment on the cloud. If the deployment already exists, it will be updated.
+
+        Args:
+            deployment_name: The name of the deployment.
+            teamspace: The teamspace to run the deployment on.
+            image: The image to run the deployment on.
+            metric: The metric to use for autoscaling. Defaults to None.
+            machine: The machine to run the deployment on. Defaults to None.
+            min_replica: The minimum number of replicas to run. Defaults to 0.
+            max_replica: The maximum number of replicas to run. Defaults to 1.
+            spot: Whether to run the deployment on spot instances. Defaults to None.
+            replicas: The number of replicas to run. Defaults to 1.
+            cloud_account: The cloud account to run the deployment on. Defaults to None.
+            port: The port to run the deployment on. Defaults to 8000.
+            include_credentials: Whether to include credentials in the deployment. Defaults to True.
+            cloudspace_id: Connect to a Studio.
+            from_onboarding: Deployment created during onboarding.
+
+        Returns:
+            dict: The deployment and the URL of the deployment.
+        """
+        url = f"{_get_cloud_url()}/{teamspace.owner.name}/{teamspace.name}/jobs/{deployment_name}"
+        machine = machine or Machine.CPU
+        metric = metric or ("CPU" if machine.is_cpu() else "GPU")
+        deployment = Deployment(deployment_name, teamspace)
+        if deployment.is_started:
+            self._console.print(f"Deployment with name {deployment_name} already running. Updating the deployment.")
+            self._update_deployment(
+                deployment,
+                machine,
+                image,
+                min_replica=min_replica,
+                max_replica=max_replica,
+                spot=spot,
+                replicas=replicas,
+                cloud_account=cloud_account,
+                port=port,
+                include_credentials=include_credentials,
+            )
+            return {"deployment": deployment, "url": url, "updated": True}
+        autoscale = AutoScaleConfig(min_replicas=min_replica, max_replicas=max_replica, metric=metric, threshold=0.95)
+        deployment.start(
+            machine=machine,
+            image=image,
+            autoscale=autoscale,
+            spot=spot,
+            replicas=replicas,
+            cloud_account=cloud_account,
+            ports=[port],
+            include_credentials=include_credentials,
+            cloudspace_id=cloudspace_id,
+            from_litserve=True,
+            from_onboarding=from_onboarding,
+            command="",
+            cloud_provider=cloud_provider,
+        )
+
+        return {"deployment": deployment, "url": url}

lightning_sdk/services/__init__.py

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

lightning_sdk/services/file_endpoint.py

@@ -0,0 +1,222 @@
+import os
+from pathlib import Path
+from time import sleep
+from typing import Any, Dict, Optional
+
+import requests
+import urllib3
+
+from lightning_sdk.api.utils import _FileUploader
+from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.services.utilities import _get_cluster, _get_project, _get_service_url
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
+
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+_LIGHTNING_SERVICE_EXECUTION_ID_HEADER = "X-Lightning-Service-Execution-Id"
+_AUTHORIZATION_HEADER = "Authorization"
+
+
+class Client:
+    """This class is used to communicate with the File Endpoint."""
+
+    def __init__(
+        self,
+        name: str,
+        teamspace: Optional[str],
+        cloud_account: Optional[str] = None,
+        cluster_id: Optional[str] = None,  # deprecated in favor of cloud_account
+    ) -> None:
+        """Constructor of the Client.
+
+        Args:
+            name: The name of the Studio File Endpoint Service.
+            teamspace: The name of the teamspace you want to attach the upload data and artifacts to be.
+            cloud_account: The name of the cloud account on which to upload the data.
+
+        """
+        cloud_account = _resolve_deprecated_cluster(cloud_account, cluster_id)
+
+        self._auth = Auth()
+
+        try:
+            self._auth.authenticate()
+        except ConnectionError as e:
+            raise e
+
+        self._name = name
+        self._teamspace = teamspace
+        self._client = LightningClient()
+        self._project = _get_project(client=self._client, project_name=teamspace)
+        self._cloud_account = _get_cluster(
+            client=self._client, project_id=self._project.project_id, cluster_id=cloud_account
+        )
+        self._file_endpoint = self._client.endpoint_service_get_file_endpoint_by_name(
+            project_id=self._project.project_id, name=self._name
+        )
+
+        self.headers = {_AUTHORIZATION_HEADER: f"Bearer {self._auth.api_key}"}
+
+        self._arguments = []
+        for argument in self._file_endpoint.arguments:
+            self._arguments.append(Argument(**argument.to_dict()))
+
+        self.url = _get_service_url(self._file_endpoint.cloudspace_id, self._file_endpoint.id)
+
+    def run(
+        self,
+        **kwargs: Dict[str, str],
+    ) -> None:
+        """The run method executes the file endpoint.
+
+        Args:
+            kwargs: The keyword arguments associated to the service
+
+        """
+        for argument in self._arguments:
+            if argument.is_file:
+                if argument.name not in kwargs:
+                    raise ValueError(f"This endpoint expects a file for the argument `{argument.name}`.")
+                value = kwargs[argument.name]
+                if not os.path.isfile(value):
+                    raise ValueError(f"This endpoint expects a file for the argument `{argument.name}`.")
+            else:
+                if argument.name not in kwargs:
+                    raise ValueError(f"This endpoint expects a value for the argument `{argument.name}`.")
+                value = kwargs[argument.name]
+                if os.path.isfile(value):
+                    raise ValueError(f"This endpoint doesn't expect a file for `{argument.name}`.")
+
+            argument.value = str(Path(value).resolve())
+
+        missing_names = [v.name for v in self._arguments if v.value is None]
+        if missing_names:
+            raise ValueError(f"You are missing values for the following arguments: {missing_names}")
+
+        # Avoid uploading duplicated files
+        files_to_upload = {}
+        for argument in self._arguments:
+            if not argument.should_upload:
+                continue
+            files_to_upload[argument.value] = argument
+
+        # TODO: Verify if the file exists in Teampace and avoid uploading if already there.
+        for argument in files_to_upload.values():
+            _FileUploader(
+                client=self._client,
+                teamspace_id=self._project.project_id,
+                cloud_account=self._cloud_account.cluster_id,
+                file_path=argument.value,
+                progress_bar=True,
+                remote_path=_sanitize_uploads_remote_path(argument.value),
+            )()
+
+        json = {
+            "teamspace_id": self._project.project_id,
+            "cluster_id": self._cloud_account.cluster_id,
+            "input": {},
+        }
+        for argument in self._arguments:
+            json["input"].update(**argument.to_dict())
+
+        response = requests.post(self.url, json=json, headers=self.headers)
+        if response.status_code != 200:
+            raise Exception(f"The endpoint isn't reachable. Status code: {response.status_code}")
+
+        self.headers = {
+            **self.headers,
+            _LIGHTNING_SERVICE_EXECUTION_ID_HEADER: response.headers[_LIGHTNING_SERVICE_EXECUTION_ID_HEADER],
+        }
+
+        self._check_progress(response.json())
+
+    def _check_progress(self, data: Dict[str, str]) -> Dict[str, str]:
+        """Check the current Studio status."""
+        while True:
+            url = f"{self.url}?run_id={data['run_id']}"
+            response = requests.post(url, headers=self.headers)
+
+            if response.status_code != 200:
+                raise Exception(f"The file endpoint had an error. Status code: {response.status_code}")
+
+            data = response.json()
+
+            # Display the progress status to the user.
+            print(data)
+
+            if data["stage"] == "completed":
+                break
+
+            if data["stage"] == "failed":
+                # TODO: Add more information on why the execution failed.
+                raise RuntimeError("The Studio File Endpoint failed")
+
+            # Wait until making the next request
+            sleep(3)
+
+        return data
+
+
+class Argument:
+    """A holder for the service argument."""
+
+    def __init__(self, name: str, type: str, **kwargs: Any) -> None:  # noqa: A002
+        self._name = name
+        self._type = type
+        self._value = None
+        self._kwargs = kwargs
+
+    @property
+    def is_text(self) -> bool:
+        """Whether this argument is of type Text."""
+        return self._type == "Text"
+
+    @property
+    def is_file(self) -> bool:
+        """Whether this argument is of type File."""
+        return self._type == "File"
+
+    @property
+    def value(self) -> Any:
+        """Returns the value."""
+        return self._value
+
+    @value.setter
+    def value(self, value: Any) -> None:
+        """Store the value."""
+        if self.is_file and not os.path.exists(value):
+            raise ValueError(f"The argument {self._name} should be a valid file.")
+        self._value = value
+
+    @property
+    def name(self) -> str:
+        """Returns the name."""
+        return self._name
+
+    def to_dict(self) -> Dict[str, str]:
+        """Convert the argument into its OpenAPI dataclass counterpart."""
+        if self.is_text:
+            return {self._name: str(self._value)}
+
+        if self.should_upload:
+            return {self._name: f"/teamspace/Uploads/{os.path.basename(self._value)}"}
+
+        return {self._name: self._value}
+
+    @property
+    def should_upload(self) -> bool:
+        """Whether the file should be uploaded."""
+        if not self.is_file:
+            return False
+
+        value = str(Path(self.value).resolve())
+        if value.startswith("/teamspace/"):
+            return False
+
+        return True
+
+
+def _sanitize_uploads_remote_path(file_path: str) -> str:
+    remote_path = os.path.basename(file_path)
+    return f"/Uploads/{remote_path}"

lightning_sdk/services/finetune_llm.py

@@ -0,0 +1,35 @@
+from typing import Literal
+
+from lightning_sdk.services.file_endpoint import Client
+
+
+class LLMFinetune(Client):
+    """The LLM Finetune is the client to the LLM Finetune Service Studio.
+
+    Learn more: https://lightning.ai/lightning-ai/studios/llm-finetune-service~01h5rahq6gbhw5m4bzyws0at5h.
+
+    """
+
+    def __init__(self, teamspace: str) -> None:
+        super().__init__(name="lightning-al/llm-finetunes", teamspace=teamspace)
+
+    def run(
+        self,
+        data_path: str,
+        model: Literal["llama-2-7b", "tiny-llama"] = "tiny-llama",
+        mode: Literal["lora", "full"] = "lora",
+        epochs: int = 3,
+        learning_rate: float = 0.0002,
+        micro_batch_size: int = 2,
+        global_batch_size: int = 8,
+    ) -> None:
+        """The run method executes the LLM Finetune Service."""
+        super().run(
+            data_path=data_path,
+            model=model,
+            mode=mode,
+            epochs=epochs,
+            learning_rate=learning_rate,
+            micro_batch_size=micro_batch_size,
+            global_batch_size=global_batch_size,
+        )

lightning_sdk/services/utilities.py

@@ -0,0 +1,125 @@
+import contextlib
+import os
+from typing import Optional
+
+import requests
+import urllib3
+
+from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.lightning_cloud.openapi import V1Membership, V1ProjectClusterBinding
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+_CHUNK_SIZE = 1024 * 1024
+
+
+def _get_project(client: LightningClient, project_name: Optional[str] = None) -> V1Membership:
+    """Get a project membership for the user from the backend."""
+    projects = client.projects_service_list_memberships()
+    if len(projects.memberships) == 0:
+        raise ValueError("No valid projects found. Please reach out to lightning.ai team to create a project")
+
+    if project_name is None:
+        return projects.memberships[0]
+
+    matches = []
+    for membership in projects.memberships:
+        if membership.name == project_name or membership.display_name == project_name:
+            matches.append(membership)
+
+    if len(matches) == 1:
+        return matches[0]
+
+    if len(matches) >= 2:
+        raise ValueError(f"We found several teamspaces. Which one do you want to use {[m.name for m in matches]}")
+
+    raise ValueError("No valid projects found. Please reach out to lightning.ai team to create a project")
+
+
+def _get_cluster(
+    client: LightningClient, project_id: str, cluster_id: Optional[str] = None, allow_neoclouds: bool = False
+) -> V1ProjectClusterBinding:
+    """Get a project membership for the user from the backend."""
+    clusters = client.projects_service_list_project_cluster_bindings(project_id=project_id)
+    if cluster_id:
+        for cluster in clusters.clusters:
+            if cluster.cluster_id == cluster_id:
+                return cluster
+        raise ValueError(
+            f"No valid cluster found with the provided {cluster_id}."
+            f"Found {[c.cluster_id for c in clusters.clusters]}."
+        )
+
+    # filter neoclouds out
+    if not allow_neoclouds:
+        cluster_objs = client.cluster_service_list_clusters(project_id=project_id)
+        # filter for aws or gcp cluster
+        valid_clusters = filter(
+            lambda c: c.spec.aws_v1 is not None or c.spec.google_cloud_v1 is not None, cluster_objs.clusters
+        )
+        valid_clusters = {c.id for c in valid_clusters}
+
+        if len(valid_clusters):
+            clusters.clusters = list(filter(lambda c: c.cluster_id in valid_clusters, clusters.clusters))
+
+    clusters = sorted(clusters.clusters, key=lambda x: x.created_at)
+    if len(clusters):
+        return clusters[0]
+    return None
+
+
+def _get_service_url(cloud_space_id: str, file_endpoint_id: str) -> str:
+    url = _get_cloud_url()
+    domain = _get_domain(url)
+    protocol = _get_protocol(url)
+    return f"{protocol}//{file_endpoint_id}-{cloud_space_id}.cloudspaces.{domain}"
+
+
+def _get_domain(url: str) -> str:
+    base_url = url.split("//")[1].split("/")[0]
+    if "localhost:9800" in base_url:
+        return "local.litng.ai:8118"
+    if "lightning.ai" in base_url and "localhost:8888" in base_url:
+        return "litng.ai"
+    return base_url
+
+
+def _get_protocol(url: str) -> str:
+    return url.split("//")[0]
+
+
+def download_file(filepath: str, cache_dir: str = "/cache") -> str:
+    """Download the file passed to the service execution by the user."""
+    service_id = os.getenv("LIGHTNING_SERVICE_EXECUTION_ID")
+
+    if service_id is None:
+        raise RuntimeError("The ServiceId is required. Please reach out to lightning.ai team.")
+
+    # Note: Make the function idempotent
+    prefix = os.path.join(cache_dir, service_id)
+    if not filepath.startswith(prefix):
+        saved_filepath = filepath
+        if saved_filepath.startswith("/"):
+            saved_filepath = saved_filepath[1:]
+        saved_filepath = os.path.join(prefix, saved_filepath)
+    else:
+        saved_filepath = filepath
+        filepath = filepath.replace(prefix, "")
+
+    if os.path.exists(saved_filepath):
+        return saved_filepath
+
+    client = LightningClient(retry=False)
+    download_artifacts = client.endpoint_service_download_service_execution_artifact(
+        project_id=os.getenv("LIGHTNING_CLOUD_PROJECT_ID"), id=service_id, page_token=None, filepath=filepath
+    )
+    artifact = download_artifacts.artifacts[0]
+
+    os.makedirs(os.path.dirname(saved_filepath), exist_ok=True)
+
+    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+    with contextlib.suppress(ConnectionError):
+        request = requests.get(artifact.url, stream=True, verify=False)
+        with open(saved_filepath, "wb") as fp:
+            for chunk in request.iter_content(chunk_size=_CHUNK_SIZE):
+                fp.write(chunk)  # type: ignore
+    return saved_filepath