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

lightning_sdk/lit_container.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 from typing import Dict, List, Optional
 
 from lightning_sdk.api.lit_container_api import LitContainerApi
@@ -29,11 +30,17 @@ class LitContainer:
         repositories = self._api.list_containers(project_id)
         table = []
         for repo in repositories:
+            created_date = repo.creation_time
+            if isinstance(repo.creation_time, str):
+                created_date = datetime.fromisoformat(created_date)
+
+            created = created_date.strftime("%Y-%m-%d %H:%M:%S")
+
             table.append(
                 {
                     "REPOSITORY": repo.name,
                     "IMAGE ID": repo.id,
-                    "CREATED": repo.creation_time.strftime("%Y-%m-%d %H:%M:%S"),
+                    "CREATED": created,
                 }
             )
         return table

lightning_sdk/mmt/mmt.py

@@ -1,7 +1,5 @@
-from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
-from lightning_sdk.api.user_api import UserApi
 from lightning_sdk.mmt.base import MMTMachine, _BaseMMT
 from lightning_sdk.mmt.v1 import _MMTV1
 from lightning_sdk.mmt.v2 import _MMTV2
@@ -18,21 +16,6 @@ if TYPE_CHECKING:
 _logger = _setup_logger(__name__)
 
 
-@lru_cache(maxsize=None)
-def _has_mmt_v2() -> bool:
-    api = UserApi()
-    try:
-        feature_flags = api._get_feature_flags()
-    except Exception:
-        return False
-
-    try:
-        return feature_flags.mmt_v2
-    except AttributeError:
-        # Feature flag doesn't exist anymore, so return True
-        return True
-
-
 class MMT(_BaseMMT):
     """Class to submit and manage multi-machine jobs on the Lightning AI Platform."""
 
@@ -60,7 +43,7 @@ class MMT(_BaseMMT):
         """
         from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 
-        if _has_mmt_v2() and not self._force_v1:
+        if not self._force_v1:
             # try with v2 and fall back to v1
             try:
                 mmt = _MMTV2(

lightning_sdk/mmt/v1.py

@@ -92,34 +92,7 @@ class _MMTV1(_BaseMMT):
             path_mappings: The mappings from data connection inside your container (not supported)
 
         """
-        if studio is None:
-            raise ValueError("Studio is required for submitting jobs")
-        if image is not None or image_credentials is not None or cloud_account_auth or entrypoint != "sh -c":
-            raise ValueError("Image is not supported for submitting jobs")
-
-        if artifacts_local is not None or artifacts_remote is not None:
-            raise ValueError("Specifying how to persist artifacts is not yet supported with jobs")
-
-        if env is not None:
-            raise ValueError("Environment variables are not supported for submitting jobs")
-        if command is None:
-            raise ValueError("Command is required for submitting multi-machine jobs")
-
-        _submitted = self._job_api.submit_job(
-            name=self._name,
-            num_machines=num_machines,
-            command=command,
-            studio_id=studio._studio.id,
-            teamspace_id=self._teamspace.id,
-            cloud_account=cloud_account or "",
-            machine=machine,
-            interruptible=interruptible,
-            strategy="parallel",
-        )
-
-        self._name = _submitted.name
-        self._job = _submitted
-        return self
+        raise NotImplementedError("Cannot submit new mmts with MMTV1!")
 
     def _update_internal_job(self) -> None:
         try:

lightning_sdk/models.py

@@ -72,6 +72,19 @@ def _get_teamspace(name: str, organization: str) -> "Teamspace":
     raise RuntimeError(f"Teamspace `{requested_teamspace}` not found. Available teamspaces: {os.linesep}\t{options}")
 
 
+def _extend_model_name_with_teamspace(name: str) -> str:
+    """Extend the model name with the teamspace if it can be determined from env. variables."""
+    if "/" in name:
+        return name
+    # do some magic if you run studio
+    teamspace = _resolve_teamspace(None, None, None)
+    if not teamspace:
+        raise ValueError(
+            f"Model name must be in the format `organization/teamspace/model_name` but you provided '{name}'."
+        )
+    return f"{teamspace.owner.name}/{teamspace.name}/{name}"
+
+
 def _parse_model_name_and_version(name: str) -> Tuple[str, str, str, str]:
     """Parse the name argument into its components."""
     try:
@@ -105,9 +118,7 @@ def download_model(
         download_dir: The directory where the Model should be downloaded.
         progress_bar: Whether to show a progress bar when downloading.
     """
-    if "/" not in name:  # do some magic if you run studio
-        teamspace = _resolve_teamspace(None, None, None)
-        name = f"{teamspace.owner.name}/{teamspace.name}/{name}"
+    name = _extend_model_name_with_teamspace(name)
     teamspace_owner_name, teamspace_name, model_name, version = _parse_model_name_and_version(name)
 
     download_dir = Path(download_dir)
@@ -147,9 +158,7 @@ def upload_model(
             If not provided, the default cloud account for the Teamspace will be used.
         progress_bar: Whether to show a progress bar for the upload.
     """
-    if "/" not in name:  # do some magic if you run studio
-        teamspace = _resolve_teamspace(None, None, None)
-        name = f"{teamspace.owner.name}/{teamspace.name}/{name}"
+    name = _extend_model_name_with_teamspace(name)
     org_name, teamspace_name, model_name, _ = _parse_model_name_and_version(name)
     teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
     return teamspace.upload_model(

lightning_sdk/pipeline/pipeline.py

@@ -49,8 +49,8 @@ class Pipeline:
 
         self._teamspace = _resolve_teamspace(
             teamspace=teamspace,
-            org=self._org,
-            user=self._user,
+            org=org,
+            user=user,
         )
 
         self._pipeline_api = PipelineApi()

lightning_sdk/pipeline/types.py

@@ -152,12 +152,20 @@ class Job:
         self.wait_for = wait_for
 
     def to_proto(self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: bool) -> V1PipelineStep:
+        studio = _get_studio(self.studio)
+        if isinstance(studio, Studio):
+            if self.cloud_account is None:
+                self.cloud_account = studio.cloud_account
+            elif studio.cloud_account != self.cloud_account:
+                raise ValueError("The provided cloud account doesn't match the studio")
+
         _validate_cloud_account(cloud_account, self.cloud_account, shared_filesystem)
+
         body = JobApiV2._create_job_body(
             name=self.name,
             command=self.command,
             cloud_account=self.cloud_account or cloud_account,
-            studio_id=None,
+            studio_id=studio._studio.id if isinstance(studio, Studio) else None,
             image=self.image,
             machine=self.machine,
             interruptible=self.interruptible,
@@ -220,13 +228,21 @@ class MMT:
         self.wait_for = wait_for
 
     def to_proto(self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: bool) -> V1PipelineStep:
+        studio = _get_studio(self.studio)
+        if isinstance(studio, Studio):
+            if self.cloud_account is None:
+                self.cloud_account = studio.cloud_account
+            elif studio.cloud_account != self.cloud_account:
+                raise ValueError("The provided cloud account doesn't match the studio")
+
         _validate_cloud_account(cloud_account, self.cloud_account, shared_filesystem)
+
         body = MMTApiV2._create_mmt_body(
             name=self.name,
             num_machines=self.num_machines,
             command=self.command,
             cloud_account=self.cloud_account or cloud_account,
-            studio_id=self.studio.studio_id if isinstance(self.studio, Studio) else None,
+            studio_id=studio._studio.id if isinstance(studio, Studio) else None,
             image=self.image,
             machine=self.machine,
             interruptible=self.interruptible,
@@ -266,3 +282,13 @@ def _validate_cloud_account(pipeline_cloud_account: str, step_cloud_account: str
             "With shared filesystem enabled, all the pipeline steps wait_for to be on the same cluster."
             f" Found {pipeline_cloud_account} and {step_cloud_account}"
         )
+
+
+def _get_studio(studio: Union["Studio", str, None]) -> Union[Studio, None]:
+    if studio is None:
+        return None
+
+    if isinstance(studio, Studio):
+        return studio
+
+    return Studio(studio)

lightning_sdk/pipeline/utils.py

@@ -43,7 +43,7 @@ def prepare_steps(steps: List["V1PipelineStep"]) -> List["V1PipelineStep"]:
         else:
             for name in current_step.wait_for:
                 if current_step.name == name:
-                    raise ValueError("You can only reference prior steps")
+                    raise ValueError(f"You can only reference prior steps. Found {current_step.name}")
 
                 if name not in name_to_step:
                     raise ValueError(f"The step {current_step_idx} doesn't have a valid wait_for. Found {name}")

lightning_sdk/plugin.py

@@ -303,7 +303,6 @@ class SlurmJobsPlugin(_Plugin):
         name: Optional[str] = None,
         cluster_id: Optional[str] = None,
         work_dir: str = "/home/lightning_manager",
-        num_gpus: int = 1,
         sync_env: bool = True,
         cache_id: Optional[str] = None,
     ) -> "Externalv1LightningappInstance":
@@ -315,7 +314,6 @@ class SlurmJobsPlugin(_Plugin):
             cluster_id: The name of the SLURM Cluster to submit the job on.
                 If the cluster_id isn't provided, the oldest running SLURM cluster will be selected.
             work_dir: The position where the the files will be created on the SLURM cluster.
-            num_gpus: The number of GPUs requested.
             sync_env: Whether to force an environement sync.
             cache_id: A string to avoid re-downloading the Studio files to the SLURM cluster.
                 If you update your files and don't change the cache_id, they won't be used.
@@ -326,9 +324,6 @@ class SlurmJobsPlugin(_Plugin):
         if work_dir == "":
             raise ValueError("The argument `work_dir` needs to be a proper path on the SLURM Cluster.")
 
-        if num_gpus <= 0:
-            raise ValueError("The argument `num_gpus` needs to be strictly positive.")
-
         if name is None:
             name = _run_name("slurm")
 
@@ -392,7 +387,6 @@ class SlurmJobsPlugin(_Plugin):
                 sync_env=sync_env,
                 work_dir=work_dir,
                 service_id=service_id,
-                num_gpus=num_gpus,
                 cache_id=cache_id,
             ),
         )

lightning_sdk/serve.py

@@ -1,9 +1,9 @@
 import os
 import shlex
 import subprocess
-import warnings
 from pathlib import Path
-from typing import Generator, List, Optional
+from typing import Generator, Optional
+from urllib.parse import urlencode
 
 import docker
 from rich.console import Console
@@ -13,6 +13,26 @@ from lightning_sdk import Deployment, Machine, Teamspace
 from lightning_sdk.api.deployment_api import AutoScaleConfig
 from lightning_sdk.api.lit_container_api import LitContainerApi
 from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.lightning_cloud import env
+from lightning_sdk.lightning_cloud.login import Auth, AuthServer
+
+_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 _AuthServer(AuthServer):
+    def get_auth_url(self, port: int) -> str:
+        redirect_uri = f"http://localhost:{port}/login-complete"
+        params = urlencode({"redirectTo": redirect_uri, "inviteCode": "litserve"})
+        return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
+
+
+class _Auth(Auth):
+    def _run_server(self) -> None:
+        _AuthServer().login_with_browser(self)
 
 
 class _LitServeDeployer:
@@ -20,38 +40,49 @@ class _LitServeDeployer:
         self._console = Console()
         self._client = None
 
+    @staticmethod
+    def authenticate() -> None:
+        auth = _Auth()
+        auth.authenticate()
+
     @property
     def client(self) -> docker.DockerClient:
-        os.environ["DOCKER_BUILDKIT"] = "1"
-
         if self._client is None:
             try:
+                os.environ["DOCKER_BUILDKIT"] = "1"
                 self._client = docker.from_env()
                 self._client.ping()
-            except docker.errors.DockerException as e:
-                raise RuntimeError(f"Failed to connect to Docker daemon: {e!s}. Is Docker running?") from None
+            except docker.errors.DockerException:
+                raise RuntimeError(_DOCKER_NOT_RUNNING_MSG) from None
         return self._client
 
     def dockerize_api(
-        self, server_filename: str, port: int = 8000, gpu: bool = False, tag: str = "litserve-model"
+        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
-        if os.path.exists("Dockerfile"):
-            console.print("Dockerfile already exists. Skipping generation.")
-            return os.path.abspath("Dockerfile")
-
         requirements = ""
         if os.path.exists("requirements.txt"):
             requirements = "-r requirements.txt"
         else:
-            warnings.warn(
+            console.print(
                 f"requirements.txt not found at {os.getcwd()}. "
                 f"Make sure to install the required packages in the Dockerfile.",
-                UserWarning,
+                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()}")
@@ -72,7 +103,8 @@ class _LitServeDeployer:
         with open("Dockerfile", "w") as f:
             f.write(dockerfile_content)
 
-        success_msg = f"""[bold]Dockerfile created successfully[/bold]
+        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]
@@ -84,7 +116,7 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
 [bold]To push the container to a registry:[/bold]
 > [underline]docker push {tag}[/underline]
 """
-        console.print(success_msg)
+            console.print(success_msg)
         return os.path.abspath("Dockerfile")
 
     @staticmethod
@@ -187,19 +219,19 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
         deployment_name: str,
         teamspace: Teamspace,
         image: str,
-        ports: List[int],
-        gpu: bool = False,
         metric: Optional[str] = None,
         machine: Optional[Machine] = None,
-        min_replica: Optional[int] = 1,
+        min_replica: Optional[int] = 0,
         max_replica: Optional[int] = 1,
         spot: Optional[bool] = None,
-        replicas: Optional[int] = None,
+        replicas: Optional[int] = 1,
         cloud_account: Optional[str] = None,
+        port: Optional[int] = 8000,
+        include_credentials: Optional[bool] = True,
     ) -> dict:
-        machine = machine or Machine.CPU
-        metric = metric or "GPU" if gpu else "CPU"
         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:
             raise RuntimeError(
@@ -211,11 +243,12 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
         deployment.start(
             machine=machine,
             image=image,
-            ports=ports,
             autoscale=autoscale,
             spot=spot,
             replicas=replicas,
             cloud_account=cloud_account,
+            ports=[port],
+            include_credentials=include_credentials,
         )
 
         return {"deployment": deployment, "url": url}

lightning_sdk/utils/resolve.py

@@ -77,6 +77,7 @@ def _resolve_org(org: Optional[Union[str, "Organization"]]) -> Optional["Organiz
 
     if isinstance(org, Organization):
         return org
+
     org = _resolve_org_name(org)
 
     if org is None:

{lightning_sdk-0.2.3.dist-info → lightning_sdk-0.2.5.dist-info}/METADATA

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