lightning-sdk 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

lightning_sdk/__init__.py

@@ -31,6 +31,6 @@ __all__ = [
     "User",
 ]
 
-__version__ = "0.2.3"
+__version__ = "0.2.5"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/ai_hub.py

@@ -3,6 +3,7 @@ from urllib.parse import quote
 
 from lightning_sdk.api import AIHubApi
 from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.lightning_cloud.openapi.models import V1Deployment
 from lightning_sdk.user import User
 from lightning_sdk.utils.resolve import _resolve_teamspace
 
@@ -109,7 +110,8 @@ class AIHub:
         org: Optional[Union[str, "Organization"]] = None,
         user: Optional[Union[str, "User"]] = None,
         machine: Optional[Union[str, "Machine"]] = None,
-    ) -> Dict[str, Union[str, bool]]:
+        quantity: Optional[int] = None,
+    ) -> V1Deployment:
         """Deploy an API from the AI Hub.
 
         Example:
@@ -126,14 +128,14 @@ class AIHub:
             api_arguments: Additional API argument, such as model name, or batch size.
             name: Name for the deployed API. Defaults to None.
             cloud_account: The cloud account where you want to run the template, such as "lightning-public-prod".
-                Defaults to None.
+                Defaults to None, which will use the cloud account of the teamspace.
             teamspace: The team or group for deployment. Defaults to None.
             org: The organization for deployment. Don't pass user with this. Defaults to None.
             user: The user for deployment. Don't pass org with this. Defaults to None.
             machine: The machine to run the deployment on. Defaults to the first option set in the AI Hub template.
 
         Returns:
-            A dictionary containing the name of the deployed API,
+            A V1Deployment object containing the name of the deployed API,
             the URL to access it, and whether it is interruptible.
 
         Raises:
@@ -157,6 +159,7 @@ class AIHub:
             name=name,
             api_arguments=api_arguments,
             machine=machine,
+            quantity=quantity,
         )
 
         url = (
@@ -169,16 +172,9 @@ class AIHub:
 
         print("Deployment available at:", url)
 
-        return {
-            "id": deployment.id,
-            "name": deployment.name,
-            "deployment_url": url,
-            "api_endpoint": deployment.status.urls[0],
-            "interruptible": deployment.spec.spot,
-            "teamspace id": teamspace_id,
-        }
+        return deployment
 
-    def delete_deployment(self, deployment: Dict[str, Union[str, bool]]) -> None:
+    def delete_deployment(self, deployment: V1Deployment) -> None:
         """Delete a deployment from the AI Hub.
 
         Example:
@@ -188,9 +184,6 @@ class AIHub:
             hub.delete_deployment(deployment)
 
         Args:
-            deployment: The deployment dictionary returned by the run method.
+            deployment: The deployment object returned by the run method.
         """
-        if "teamspace id" not in deployment or "id" not in deployment:
-            raise ValueError("Deployment dictionary must contain 'teamspace id' and 'id' keys.")
-
-        self._api.delete_api(deployment["id"], deployment["teamspace id"])
+        self._api.delete_api(deployment.id, deployment.project_id)

lightning_sdk/api/ai_hub_api.py

@@ -1,4 +1,5 @@
 import traceback
+import warnings
 from typing import Dict, List, Optional, Tuple, Union
 
 import backoff
@@ -95,15 +96,17 @@ class AIHubApi:
 
                 if p.type == V1DeploymentTemplateParameterType.CHECKBOX and p.checkbox:
                     api_arguments[p.name] = (
-                        (p.checkbox.true_value or "True")
-                        if p.checkbox.is_checked
-                        else (p.checkbox.false_value or "False")
+                        (p.checkbox.true_value or "") if p.checkbox.is_checked else (p.checkbox.false_value or "")
                     )
 
         for p in parameters:
             name = p.name
             pattern = f"${{{name}}}"
             if name in api_arguments:
+                if p.type == V1DeploymentTemplateParameterType.CHECKBOX and p.checkbox:
+                    api_arguments[p.name] = (
+                        (p.checkbox.true_value or "") if api_arguments[name] is True else (p.checkbox.false_value or "")
+                    )
                 AIHubApi._update_parameters(job, p.placements, pattern, api_arguments[name])
             elif not p.required:
                 AIHubApi._update_parameters(job, p.placements, pattern, "")
@@ -120,6 +123,7 @@ class AIHubApi:
         name: Optional[str],
         api_arguments: Dict[str, str],
         machine: Optional[Union[str, Machine]],
+        quantity: Optional[int],
     ) -> V1Deployment:
         template = self._client.deployment_templates_service_get_deployment_template(template_id)
         name = name or template.name
@@ -139,6 +143,19 @@ class AIHubApi:
             apply_change(template.spec_v2.job, "instance_name", machine)
             apply_change(template.spec_v2.job, "instance_type", machine)
 
+        if quantity != template.spec_v2.job.quantity:
+            # If the quantity is different from the published template quantity, override it with warnging
+            warnings.warn(
+                "Overriding the quantity of the template with the provided quantity. "
+                "This may result in unexpected behavior. "
+                f"Please verify the template (https://lightning.ai/lightning-ai/ai-hub/{template_id}) "
+                "and asscoiated parameters before running."
+            )
+            apply_change(template.spec_v2.job, "quantity", quantity)
+
+        # Override the cluster_id with the cloud_account if it is provided
+        if len(cloud_account) > 0:
+            apply_change(template.spec_v2.job, "cluster_id", cloud_account)
         return self._client.jobs_service_create_deployment(
             project_id=project_id,
             body=CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs(

lightning_sdk/api/studio_api.py

@@ -140,11 +140,6 @@ class StudioApi:
         startup_status = self.get_studio_status(studio_id, teamspace_id).in_use.startup_status
         return startup_status and startup_status.top_up_restore_finished
 
-    @backoff.on_exception(backoff.expo, AttributeError, max_tries=10)
-    def _check_code_status_sync_in_progress(self, studio_id: str, teamspace_id: str) -> bool:
-        """Retries checking the sync_in_progress value of the code status when there's an AttributeError."""
-        return self.get_studio_status(studio_id, teamspace_id).in_use.sync_in_progress
-
     def start_studio(
         self, studio_id: str, teamspace_id: str, machine: Union[Machine, str], interruptible: False
     ) -> None:
@@ -165,9 +160,6 @@ class StudioApi:
             studio_id,
         )
 
-        while self._check_code_status_sync_in_progress(studio_id, teamspace_id):
-            time.sleep(1)
-
         while True:
             if self._check_code_status_top_up_restore_finished(studio_id, teamspace_id):
                 break

lightning_sdk/cli/serve.py

@@ -4,17 +4,19 @@ from pathlib import Path
 from typing import Optional, Union
 
 import click
-import docker
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from rich.prompt import Confirm
 
+from lightning_sdk import Machine, Teamspace
 from lightning_sdk.api.deployment_api import DeploymentApi
 from lightning_sdk.api.lit_container_api import LitContainerApi
 from lightning_sdk.cli.exceptions import StudioCliError
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
 from lightning_sdk.serve import _LitServeDeployer
 
+_MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
+
 
 @click.group("serve")
 def serve() -> None:
@@ -24,11 +26,10 @@ def serve() -> None:
         lightning serve api server.py  # serve locally
 
     Example:
-        lightning serve api server.py --cloud  # deploy to the cloud
+        lightning serve api server.py --cloud --name litserve-api  # deploy to the cloud
 
     You can deploy the API to the cloud by running `lightning serve api server.py --cloud`.
-    This will generate a Dockerfile, build the image, and push it to the image registry.
-    Deploying to the cloud requires pre-login to the docker registry.
+    This will build a docker container for the server.py script and deploy it to the Lightning AI platform.
     """
 
 
@@ -48,7 +49,6 @@ def serve() -> None:
     flag_value=True,
     help="Deploy the model to the Lightning AI platform",
 )
-@click.option("--gpu", is_flag=True, default=False, flag_value=True, help="Use GPU for serving")
 @click.option("--name", default=None, help="Name of the deployed API (e.g., 'classification-api', 'Llama-api')")
 @click.option(
     "--non-interactive",
@@ -58,17 +58,88 @@ def serve() -> None:
     flag_value=True,
     help="Do not prompt for confirmation",
 )
+@click.option(
+    "--machine",
+    default="CPU",
+    show_default=True,
+    type=click.Choice(_MACHINE_VALUES),
+    help="The machine type to deploy the API on.",
+)
+@click.option(
+    "--interruptible",
+    is_flag=True,
+    default=False,
+    flag_value=True,
+    help="Whether the machine should be interruptible (spot) or not.",
+)
+@click.option(
+    "--teamspace",
+    default=None,
+    help="The teamspace the deployment should be associated with. Defaults to the current teamspace.",
+)
+@click.option(
+    "--org",
+    default=None,
+    help="The organization owning the teamspace (if any). Defaults to the current organization.",
+)
+@click.option("--user", default=None, help="The user owning the teamspace (if any). Defaults to the current user.")
+@click.option(
+    "--cloud-account",
+    "--cloud_account",
+    default=None,
+    help=(
+        "The cloud account to run the deployment 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."
+    ),
+)
+@click.option("--port", default=8000, help="The port to expose the API on.")
+@click.option("--min_replica", "--min-replica", default=0, help="Number of replicas to start with.")
+@click.option("--max_replica", "--max-replica", default=1, help="Number of replicas to scale up to.")
+@click.option("--replicas", "--replicas", default=1, help="Deployment will start with this many replicas.")
+@click.option(
+    "--no_credentials",
+    "--no-credentials",
+    is_flag=True,
+    default=False,
+    flag_value=True,
+    help="Whether to include credentials in the deployment.",
+)
 def api(
     script_path: str,
     easy: bool,
     cloud: bool,
-    gpu: bool,
     name: str,
     non_interactive: bool,
+    machine: str,
+    interruptible: bool,
+    teamspace: Optional[str],
+    org: Optional[str],
+    user: Optional[str],
+    cloud_account: Optional[str],
+    port: Optional[int],
+    min_replica: Optional[int],
+    max_replica: Optional[int],
+    replicas: Optional[int],
+    no_credentials: Optional[bool],
 ) -> None:
     """Deploy a LitServe model script."""
     return api_impl(
-        script_path=script_path, easy=easy, cloud=cloud, gpu=gpu, repository=name, non_interactive=non_interactive
+        script_path=script_path,
+        easy=easy,
+        cloud=cloud,
+        repository=name,
+        non_interactive=non_interactive,
+        machine=machine,
+        interruptible=interruptible,
+        teamspace=teamspace,
+        org=org,
+        user=user,
+        cloud_account=cloud_account,
+        port=port,
+        min_replica=min_replica,
+        max_replica=max_replica,
+        include_credentials=not no_credentials,
     )
 
 
@@ -76,10 +147,20 @@ def api_impl(
     script_path: Union[str, Path],
     easy: bool = False,
     cloud: bool = False,
-    gpu: bool = False,
     repository: [str] = None,
     tag: Optional[str] = None,
     non_interactive: bool = False,
+    machine: str = "CPU",
+    interruptible: bool = False,
+    teamspace: Optional[str] = None,
+    org: Optional[str] = None,
+    user: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    port: Optional[int] = 8000,
+    min_replica: Optional[int] = 0,
+    max_replica: Optional[int] = 1,
+    replicas: Optional[int] = 1,
+    include_credentials: Optional[bool] = True,
 ) -> None:
     """Deploy a LitServe model script."""
     console = Console()
@@ -93,8 +174,24 @@ def api_impl(
 
     if cloud:
         repository = repository or "litserve-model"
+        machine = Machine.from_str(machine)
         return _handle_cloud(
-            script_path, console, gpu=gpu, repository=repository, tag=tag, non_interactive=non_interactive
+            script_path,
+            console,
+            repository=repository,
+            tag=tag,
+            non_interactive=non_interactive,
+            machine=machine,
+            interruptible=interruptible,
+            teamspace=teamspace,
+            org=org,
+            user=user,
+            cloud_account=cloud_account,
+            port=port,
+            min_replica=min_replica,
+            max_replica=max_replica,
+            replicas=replicas,
+            include_credentials=include_credentials,
         )
 
     try:
@@ -111,22 +208,30 @@ def api_impl(
 def _handle_cloud(
     script_path: Union[str, Path],
     console: Console,
-    gpu: bool,
     repository: str = "litserve-model",
     tag: Optional[str] = None,
-    teamspace: Optional[str] = None,
     non_interactive: bool = False,
+    machine: Machine = "CPU",
+    interruptible: bool = False,
+    teamspace: Optional[str] = None,
+    org: Optional[str] = None,
+    user: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    port: Optional[int] = 8000,
+    min_replica: Optional[int] = 0,
+    max_replica: Optional[int] = 1,
+    replicas: Optional[int] = 1,
+    include_credentials: Optional[bool] = True,
 ) -> None:
-    menu = _TeamspacesMenu()
-    teamspace = menu._resolve_teamspace(teamspace)
-    try:
-        client = docker.from_env()
-        client.ping()
-    except docker.errors.DockerException as e:
-        raise StudioCliError(f"Failed to connect to Docker daemon: {e!s}. Is Docker running?") from None
+    if teamspace is None:
+        menu = _TeamspacesMenu()
+        resolved_teamspace = menu._resolve_teamspace(teamspace)
+    else:
+        resolved_teamspace = Teamspace(name=teamspace, org=org, user=user)
 
+    port = port or 8000
     ls_deployer = _LitServeDeployer()
-    path = ls_deployer.dockerize_api(script_path, port=8000, gpu=gpu, tag=tag)
+    path = ls_deployer.dockerize_api(script_path, port=port, gpu=not machine.is_cpu(), tag=tag, print_success=False)
     console.clear()
     if non_interactive:
         console.print("[italic]non-interactive[/italic] mode enabled, skipping confirmation prompts", style="blue")
@@ -142,7 +247,8 @@ def _handle_cloud(
     lit_cr = LitContainerApi()
     deployment_name = os.path.basename(repository)
 
-    if DeploymentApi().get_deployment_by_name(deployment_name, teamspace.id):
+    ls_deployer.authenticate()
+    if DeploymentApi().get_deployment_by_name(deployment_name, resolved_teamspace.id):
         raise StudioCliError(f"Deployment {deployment_name} already exists. Please choose a different name.") from None
 
     with Progress(
@@ -154,7 +260,7 @@ def _handle_cloud(
     ) as progress:
         try:
             ls_deployer.build_container(path, repository, tag, console, progress)
-            push_status = ls_deployer.push_container(repository, tag, teamspace, lit_cr, progress)
+            push_status = ls_deployer.push_container(repository, tag, resolved_teamspace, lit_cr, progress)
         except Exception as e:
             console.print(f"❌ Deployment failed: {e}", style="red")
             return
@@ -163,6 +269,17 @@ def _handle_cloud(
     repository = push_status.get("repository")
 
     deployment_status = ls_deployer._run_on_cloud(
-        deployment_name=deployment_name, image=repository, teamspace=teamspace, ports=[8000], gpu=gpu, metric=None
+        deployment_name=deployment_name,
+        image=repository,
+        teamspace=resolved_teamspace,
+        metric=None,
+        machine=machine,
+        spot=interruptible,
+        cloud_account=cloud_account,
+        port=port,
+        min_replica=min_replica,
+        max_replica=max_replica,
+        replicas=replicas,
+        include_credentials=include_credentials,
     )
     console.print(f"🚀 Deployment started, access at [i]{deployment_status.get('url')}[/i]")

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_deprecated_cluster, _resolve_teamspace, _resolve_user
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_org, _resolve_teamspace, _resolve_user
 
 
 class Deployment:
@@ -74,6 +74,7 @@ class Deployment:
 
         self._name = name
         self._user = _resolve_user(self._user or user)
+        self._org = _resolve_org(org)
 
         self._teamspace = _resolve_teamspace(
             teamspace=teamspace,
@@ -96,6 +97,8 @@ class Deployment:
             self._name = deployment.name
             self._is_created = True
             self._deployment = deployment
+        else:
+            self._deployment = None
 
     def start(
         self,
@@ -235,6 +238,13 @@ class Deployment:
         """All the deployment replicas will be stopped and all their traffic blocked."""
         self._deployment = self._deployment_api.stop(self._deployment)
 
+    @property
+    def name(self) -> Optional[str]:
+        if self._deployment:
+            self._deployment = self._deployment_api.get_deployment_by_name(self._name, self._teamspace.id)
+            self._name = self._deployment.name
+        return self._name
+
     @property
     def replicas(self) -> Optional[int]:
         """The default number of replicas the release starts with."""
@@ -248,7 +258,7 @@ class Deployment:
         """The minimum number of replicas."""
         if self._deployment:
             self._deployment = self._deployment_api.get_deployment_by_name(self._name, self._teamspace.id)
-            return self._deployment.autoscaling.min_replicas
+            return int(self._deployment.autoscaling.min_replicas)
         return None
 
     @property
@@ -256,7 +266,7 @@ class Deployment:
         """The maximum number of replicas."""
         if self._deployment:
             self._deployment = self._deployment_api.get_deployment_by_name(self._name, self._teamspace.id)
-            return self._deployment.autoscaling.max_replicas
+            return int(self._deployment.autoscaling.max_replicas)
         return None
 
     @property
@@ -371,9 +381,13 @@ class Deployment:
             return self._deployment.spec.include_credentials
         return None
 
+    @property
+    def org(self) -> Optional[Organization]:
+        return self._org
+
     @property
     def user(self) -> Optional[User]:
-        """The teamspace of the deployment."""
+        """The user of the deployment."""
         return self._user
 
     @property
@@ -385,6 +399,20 @@ class Deployment:
     def is_started(self) -> bool:
         return self._is_created
 
+    @property
+    def is_stopped(self) -> Optional[bool]:
+        if self._deployment:
+            self._deployment = self._deployment_api.get_deployment_by_name(self._name, self._teamspace.id)
+            return self._deployment.autoscaling.max_replicas == 0
+        return None
+
+    @property
+    def image(self) -> Optional[str]:
+        if self._deployment:
+            self._deployment = self._deployment_api.get_deployment_by_name(self._name, self._teamspace.id)
+            return self._deployment.spec.image
+        return None
+
     @property
     def _session(self) -> Any:
         if self._request_session is None: