lightning-sdk 0.1.41__py3-none-any.whl → 0.1.42__py3-none-any.whl

lightning_sdk/api/utils.py

@@ -1,3 +1,4 @@
+import concurrent.futures
 import errno
 import math
 import os
@@ -8,7 +9,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 import backoff
 import requests
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__, _LIGHTNING_DEBUG
 from lightning_sdk.lightning_cloud.openapi import (
@@ -66,14 +67,14 @@ class _FileUploader:
         self,
         client: LightningClient,
         teamspace_id: str,
-        cluster_id: str,
+        cloud_account: str,
         file_path: str,
         remote_path: str,
         progress_bar: bool,
     ) -> None:
         self.client = client
         self.teamspace_id = teamspace_id
-        self.cluster_id = cluster_id
+        self.cloud_account = cloud_account
 
         self.local_path = file_path
 
@@ -107,7 +108,7 @@ class _FileUploader:
 
     def _multipart_upload(self, count: int) -> None:
         """Does a parallel multipart upload."""
-        body = ProjectIdStorageBody(cluster_id=self.cluster_id, filename=self.remote_path)
+        body = ProjectIdStorageBody(cluster_id=self.cloud_account, filename=self.remote_path)
         resp: V1UploadProjectArtifactResponse = self.client.storage_service_upload_project_artifact(
             body=body, project_id=self.teamspace_id
         )
@@ -123,7 +124,7 @@ class _FileUploader:
                 completed.extend(self._process_upload_batch(executor=p, batch=batch, upload_id=resp.upload_id))
 
         completed_body = StorageCompleteBody(
-            cluster_id=self.cluster_id, filename=self.remote_path, parts=completed, upload_id=resp.upload_id
+            cluster_id=self.cloud_account, filename=self.remote_path, parts=completed, upload_id=resp.upload_id
         )
         self.client.storage_service_complete_upload_project_artifact(body=completed_body, project_id=self.teamspace_id)
 
@@ -135,7 +136,7 @@ class _FileUploader:
 
     def _request_urls(self, parts: List[int], upload_id: str) -> List[V1PresignedUrl]:
         """Requests urls for a batch of parts."""
-        body = UploadsUploadIdBody(cluster_id=self.cluster_id, filename=self.remote_path, parts=parts)
+        body = UploadsUploadIdBody(cluster_id=self.cloud_account, filename=self.remote_path, parts=parts)
         resp: V1UploadProjectArtifactPartsResponse = self.client.storage_service_upload_project_artifact_parts(
             body, self.teamspace_id, upload_id
         )
@@ -192,7 +193,7 @@ class _ModelFileUploader:
         model_id: str,
         version: str,
         teamspace_id: str,
-        cluster_id: str,
+        cloud_account: str,
         file_path: str,
         remote_path: str,
         progress_bar: bool,
@@ -201,7 +202,6 @@ class _ModelFileUploader:
         self.model_id = model_id
         self.version = version
         self.teamspace_id = teamspace_id
-        self.cluster_id = cluster_id
         self.local_path = file_path
         self.remote_path = remote_path
 
@@ -215,6 +215,8 @@ class _ModelFileUploader:
                 unit="B",
                 unit_scale=True,
                 unit_divisor=1000,
+                position=1,
+                leave=False,
             )
         else:
             self.progress_bar = None
@@ -376,6 +378,7 @@ class _FileDownloader:
         teamspace_id: str,
         remote_path: str,
         file_path: str,
+        executor: ThreadPoolExecutor,
         num_workers: int = 20,
         progress_bar: Optional[tqdm] = None,
     ) -> None:
@@ -389,7 +392,7 @@ class _FileDownloader:
         self.num_workers = num_workers
         self._url = ""
         self._size = 0
-        self.refresh()
+        self.executor = executor
 
     @backoff.on_exception(backoff.expo, ApiException, max_tries=10)
     def refresh(self) -> None:
@@ -445,26 +448,26 @@ class _FileDownloader:
                 if remaining_size > 0:
                     f.write(b"\x00" * remaining_size)
 
-    def _multipart_download(self, filename: str, max_workers: int) -> None:
-        num_chunks = max_workers
+    def _multipart_download(self, filename: str, num_workers: int) -> None:
+        num_chunks = num_workers
         chunk_size = math.ceil(self.size / num_chunks)
 
         if chunk_size < _DOWNLOAD_MIN_CHUNK_SIZE:
             num_chunks = math.ceil(self.size / _DOWNLOAD_MIN_CHUNK_SIZE)
             chunk_size = _DOWNLOAD_MIN_CHUNK_SIZE
 
-        num_workers = min(max_workers, num_chunks)
-
         ranges = []
         for part_number in range(num_chunks):
             start = part_number * chunk_size
             end = min(start + chunk_size - 1, self.size - 1)
             ranges.append((start, end))
 
-        with ThreadPoolExecutor(max_workers=num_workers) as executor:
-            executor.map(partial(self._download_chunk, filename), ranges)
+        futures = [self.executor.submit(self._download_chunk, filename, r) for r in ranges]
+        concurrent.futures.wait(futures)
 
     def download(self) -> None:
+        self.refresh()
+
         tmp_filename = f"{self.local_path}.download"
 
         try:
@@ -536,31 +539,40 @@ def _download_model_files(
             unit_divisor=1000,
         )
 
-    for filepath in response.filepaths:
-        local_file = download_dir / filepath
-        local_file.parent.mkdir(parents=True, exist_ok=True)
-
-        file_downloader = _FileDownloader(
-            client=client,
-            model_id=response.model_id,
-            version=response.version,
-            teamspace_id=response.project_id,
-            remote_path=filepath,
-            file_path=str(local_file),
-            num_workers=num_workers,
-            progress_bar=pbar,
-        )
+    with ThreadPoolExecutor(max_workers=min(num_workers, len(response.filepaths))) as file_executor, ThreadPoolExecutor(
+        max_workers=num_workers
+    ) as part_executor:
+        futures = []
+
+        for filepath in response.filepaths:
+            local_file = download_dir / filepath
+            local_file.parent.mkdir(parents=True, exist_ok=True)
+
+            file_downloader = _FileDownloader(
+                client=client,
+                model_id=response.model_id,
+                version=response.version,
+                teamspace_id=response.project_id,
+                remote_path=filepath,
+                file_path=str(local_file),
+                num_workers=num_workers,
+                progress_bar=pbar,
+                executor=part_executor,
+            )
+
+            futures.append(file_executor.submit(file_downloader.download))
 
-        file_downloader.download()
+        # wait for all threads
+        concurrent.futures.wait(futures)
 
-    return response.filepaths
+        return response.filepaths
 
 
 def _create_app(
     client: CloudSpaceServiceApi,
     studio_id: str,
     teamspace_id: str,
-    cluster_id: str,
+    cloud_account: str,
     plugin_type: str,
     **other_arguments: Any,
 ) -> Externalv1LightningappInstance:
@@ -573,7 +585,7 @@ def _create_app(
         del other_arguments["interruptible"]
 
     body = AppsIdBody(
-        cluster_id=cluster_id,
+        cluster_id=cloud_account,
         plugin_arguments=other_arguments,
         service_id=os.getenv(_LIGHTNING_SERVICE_EXECUTION_ID_KEY),
         unique_id=__GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id],
@@ -584,6 +596,6 @@ def _create_app(
     ).lightningappinstance
 
     if _LIGHTNING_DEBUG:
-        print(f"Create App: {resp.id=} {teamspace_id=} {studio_id=} {cluster_id=}")
+        print(f"Create App: {resp.id=} {teamspace_id=} {studio_id=} {cloud_account=}")
 
     return resp

lightning_sdk/cli/ai_hub.py

@@ -32,7 +32,7 @@ class _AIHub(_StudiosMenu):
     def deploy(
         self,
         api_id: str,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         name: Optional[str] = None,
         teamspace: Optional[str] = None,
         org: Optional[str] = None,
@@ -41,9 +41,9 @@ class _AIHub(_StudiosMenu):
 
         Args:
           api_id: API template ID.
-          cluster: Cluster to deploy the API to. Defaults to user's default cluster.
+          cloud_account: Cloud Account to deploy the API to. Defaults to user's default cloud account.
           name: Name of the deployed API. Defaults to the name of the API template.
           teamspace: Teamspace to deploy the API to. Defaults to user's default teamspace.
           org: Organization to deploy the API to. Defaults to user's default organization.
         """
-        return self._hub.run(api_id, cluster=cluster, name=name, teamspace=teamspace, org=org)
+        return self._hub.run(api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)

lightning_sdk/cli/entrypoint.py

@@ -6,6 +6,7 @@ from lightning_sdk.cli.ai_hub import _AIHub
 from lightning_sdk.cli.download import _Downloads
 from lightning_sdk.cli.legacy import _LegacyLightningCLI
 from lightning_sdk.cli.run import _Run
+from lightning_sdk.cli.serve import _Docker, _LitServe
 from lightning_sdk.cli.upload import _Uploads
 from lightning_sdk.lightning_cloud.login import Auth
 
@@ -19,8 +20,9 @@ class StudioCLI:
         self.download = _Downloads()
         self.upload = _Uploads()
         self.aihub = _AIHub()
-
         self.run = _Run(legacy_run=_LegacyLightningCLI() if _LIGHTNING_AVAILABLE else None)
+        self.serve = _LitServe()
+        self.dockerize = _Docker()
 
     def login(self) -> None:
         """Login to Lightning AI Studios."""

lightning_sdk/cli/mmt.py

@@ -31,14 +31,15 @@ class MMTCLI:
             teamspace: The teamspace the job should be associated with. Defaults to the current teamspace.
             org: The organization owning the teamspace (if any). Defaults to the current organization.
             user: The user owning the teamspace (if any). Defaults to the current user.
-            cluster: The cluster to run the job on. Defaults to the studio cluster if running with studio compute env.
-                If not provided will fall back to the teamspaces default cluster.
+            cloud_account: The cloud account to run the job on.
+                Defaults to the studio cloud account if running with studio compute env.
+                If not provided will fall back to the teamspaces default cloud account.
             env: Environment variables to set inside the job.
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             image_credentials: The credentials used to pull the image. Required if the image is private.
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
-            cluster_auth: Whether to authenticate with the cluster to pull the image.
-                Required if the registry is part of a cluster provider (e.g. ECR).
+            cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
+                Required if the registry is part of a cloud provider (e.g. ECR).
             artifacts_local: The path of inside the docker container, you want to persist images from.
                 CAUTION: When setting this to "/", it will effectively erase your container.
                 Only supported for jobs with a docker image compute environment.
@@ -83,11 +84,11 @@ class MMTCLI:
         teamspace: Optional[str] = None,
         org: Optional[str] = None,
         user: Optional[str] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
     ) -> None:
@@ -103,8 +104,8 @@ class MMTCLI:
         machine_enum = Machine(machine.upper())
 
         teamspace = Teamspace(name=teamspace, org=org, user=user)
-        if cluster is None:
-            cluster = teamspace.default_cluster
+        if cloud_account is None:
+            cloud_account = teamspace.default_cloud_account
 
         if image is None:
             raise RuntimeError("Currently only docker images are specified")
@@ -118,11 +119,11 @@ class MMTCLI:
             teamspace=teamspace,
             org=org,
             user=user,
-            cluster=cluster,
+            cloud_account=cloud_account,
             env=env,
             interruptible=interruptible,
             image_credentials=image_credentials,
-            cluster_auth=cluster_auth,
+            cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
         )

lightning_sdk/cli/run.py

@@ -32,14 +32,15 @@ class _Run:
             teamspace: The teamspace the job should be associated with. Defaults to the current teamspace.
             org: The organization owning the teamspace (if any). Defaults to the current organization.
             user: The user owning the teamspace (if any). Defaults to the current user.
-            cluster: The cluster to run the job on. Defaults to the studio cluster if running with studio compute env.
-                If not provided will fall back to the teamspaces default cluster.
+            cloud_account: The cloud account to run the job on.
+                Defaults to the studio cloud account if running with studio compute env.
+                If not provided will fall back to the teamspaces default cloud account.
             env: Environment variables to set inside the job.
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             image_credentials: The credentials used to pull the image. Required if the image is private.
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
-            cluster_auth: Whether to authenticate with the cluster to pull the image.
-                Required if the registry is part of a cluster provider (e.g. ECR).
+            cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
+                Required if the registry is part of a cloud provider (e.g. ECR).
             artifacts_local: The path of inside the docker container, you want to persist images from.
                 CAUTION: When setting this to "/", it will effectively erase your container.
                 Only supported for jobs with a docker image compute environment.
@@ -68,11 +69,11 @@ class _Run:
         teamspace: Optional[str] = None,
         org: Optional[str] = None,
         user: Optional[str] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
     ) -> None:
@@ -86,11 +87,11 @@ class _Run:
             teamspace=teamspace,
             org=org,
             user=user,
-            cluster=cluster,
+            cloud_account=cloud_account,
             env=env,
             interruptible=interruptible,
             image_credentials=image_credentials,
-            cluster_auth=cluster_auth,
+            cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
         )

lightning_sdk/cli/serve.py

@@ -0,0 +1,130 @@
+import os
+import subprocess
+import warnings
+from pathlib import Path
+from typing import Union
+
+from rich.console import Console
+
+
+class _LitServe:
+    """Serve a LitServe model.
+
+    Example:
+        lightning serve api server.py
+    """
+
+    def api(
+        self,
+        script_path: Union[str, Path],
+        easy: bool = False,
+    ) -> None:
+        """Deploy a LitServe model script.
+
+        Args:
+            script_path: Path to the script to serve
+            easy: If True, generates a client for the model
+
+        Raises:
+            FileNotFoundError: If script_path doesn't exist
+            ImportError: If litserve is not installed
+            subprocess.CalledProcessError: If the script fails to run
+            IOError: If client.py generation fails
+        """
+        console = Console()
+        script_path = Path(script_path)
+        if not script_path.exists():
+            raise FileNotFoundError(f"Script not found: {script_path}")
+        if not script_path.is_file():
+            raise ValueError(f"Path is not a file: {script_path}")
+
+        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
+
+        if easy:
+            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
+
+        try:
+            subprocess.run(
+                ["python", str(script_path)],
+                check=True,
+                text=True,
+            )
+        except subprocess.CalledProcessError as e:
+            error_msg = f"Script execution failed with exit code {e.returncode}\nstdout: {e.stdout}\nstderr: {e.stderr}"
+            raise RuntimeError(error_msg) from None
+
+
+class _Docker:
+    """Generate a Dockerfile for a LitServe model."""
+
+    def api(self, server_filename: str, port: int = 8000, gpu: bool = False) -> None:
+        """Generate a Dockerfile for the given server code.
+
+        Example:
+            lightning litserve dockerize server.py --port 8000 --gpu
+
+        Args:
+            server_filename (str): The path to the server file. Example sever.py or app.py.
+            port (int, optional): The port to expose in the Docker container.
+            gpu (bool, optional): Whether to use a GPU-enabled Docker image.
+        """
+        import litserve as ls
+        from litserve import docker_builder
+
+        console = Console()
+        requirements = ""
+        if os.path.exists("requirements.txt"):
+            requirements = "-r requirements.txt"
+        else:
+            warnings.warn(
+                f"requirements.txt not found at {os.getcwd()}. "
+                f"Make sure to install the required packages in the Dockerfile.",
+                UserWarning,
+            )
+
+        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} litserve-model:latest"
+            docker_template = docker_builder.CUDA_DOCKER_TEMPLATE
+        else:
+            run_cmd = f"docker run -p {port}:{port} litserve-model: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)
+
+        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 litserve-model .[/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 litserve-model[/underline]
+"""
+        console.print(success_msg)

lightning_sdk/deployment/deployment.py

@@ -32,7 +32,7 @@ from lightning_sdk.organization import Organization
 from lightning_sdk.services.utilities import _get_cluster
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _resolve_org, _resolve_teamspace, _resolve_user
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_org, _resolve_teamspace, _resolve_user
 
 
 class Deployment:
@@ -81,7 +81,7 @@ class Deployment:
             raise ValueError("You need to pass a teamspace or an org for your deployment.")
 
         self._deployment_api = DeploymentApi()
-        self._cluster = _get_cluster(client=self._deployment_api._client, project_id=self._teamspace.id)
+        self._cloud_account = _get_cluster(client=self._deployment_api._client, project_id=self._teamspace.id)
         self._is_created = False
         deployment = self._deployment_api.get_deployment_by_name(name, self._teamspace.id)
         if deployment:
@@ -102,8 +102,9 @@ class Deployment:
         replicas: Optional[int] = None,
         health_check: Optional[Union[HttpHealthCheck, ExecHealthCheck]] = None,
         auth: Optional[Union[BasicAuth, TokenAuth]] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         custom_domain: Optional[str] = None,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> None:
         """The Lightning AI Deployment.
 
@@ -124,7 +125,7 @@ class Deployment:
             replicas: The number of replicas to deploy with.
             health_check: The health check config to know whether your service is ready to receive traffic.
             auth: The auth config to protect your services. Only Basic and Token supported.
-            cluster: The name of the cluster, the studio should be created on.
+            cloud_account: The name of the cloud account, the studio should be created on.
                 Doesn't matter when the studio already exists.
             custom_domain: Whether your service would be referenced under a custom doamin.
 
@@ -136,9 +137,11 @@ class Deployment:
         if self._is_created:
             raise RuntimeError("This deployment has already been started.")
 
-        if cluster is None and self._cluster is not None:
-            print(f"No cluster was provided, defaulting to {self._cluster.cluster_id}")
-            cluster = os.getenv("LIGHTNING_CLUSTER_ID") or self._cluster.cluster_id
+        cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
+
+        if cloud_account is None and self._cloud_account is not None:
+            print(f"No cloud account was provided, defaulting to {self._cloud_account.cluster_id}")
+            cloud_account = os.getenv("LIGHTNING_CLUSTER_ID") or self._cloud_account.cluster_id
 
         self._deployment = self._deployment_api.create_deployment(
             V1Deployment(
@@ -148,7 +151,7 @@ class Deployment:
                 project_id=self._teamspace.id,
                 replicas=replicas,
                 spec=to_spec(
-                    cluster_id=cluster,
+                    cloud_account=cloud_account,
                     command=command,
                     entrypoint=entrypoint,
                     env=env,
@@ -171,7 +174,7 @@ class Deployment:
         command: Optional[str] = None,
         env: Optional[List[Union[Env, Secret]]] = None,
         spot: Optional[bool] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         health_check: Optional[Union[HttpHealthCheck, ExecHealthCheck]] = None,
         # Changing those arguments don't create a new release
         min_replicas: Optional[int] = None,
@@ -182,7 +185,10 @@ class Deployment:
         replicas: Optional[int] = None,
         auth: Optional[Union[BasicAuth, TokenAuth]] = None,
         custom_domain: Optional[str] = None,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> None:
+        cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
+
         self._deployment = self._deployment_api.update_deployment(
             self._deployment,
             name=name or self._name,
@@ -190,7 +196,7 @@ class Deployment:
             replicas=replicas,
             min_replicas=min_replicas,
             max_replicas=max_replicas,
-            cluster_id=cluster,
+            cloud_account=cloud_account,
             machine=machine,
             environment=environment,
             entrypoint=entrypoint,
@@ -312,8 +318,8 @@ class Deployment:
         return None
 
     @property
-    def cluster(self) -> Optional[str]:
-        """The cluster of the replicas."""
+    def cloud_account(self) -> Optional[str]:
+        """The cloud_account of the replicas."""
         if self._deployment:
             self._deployment = self._deployment_api.get_deployment_by_name(self._name, self._teamspace.id)
             return self._deployment.spec.cluster_id