lightning-sdk 0.1.53__py3-none-any.whl → 0.1.54__py3-none-any.whl

lightning_sdk/__init__.py

@@ -29,5 +29,5 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.53"
+__version__ = "0.1.54"
 _check_version_and_prompt_upgrade(__version__)

lightning_sdk/api/deployment_api.py

@@ -213,6 +213,14 @@ class DeploymentApi:
                 return None
             raise ex
 
+    def get_deployment_by_id(self, deployment_id: str, teamspace_id: str) -> Optional[V1Deployment]:
+        try:
+            return self._client.jobs_service_get_deployment(project_id=teamspace_id, id=deployment_id)
+        except ApiException as ex:
+            if "Reason: Not Found" in str(ex):
+                return None
+            raise ex
+
     def create_deployment(
         self,
         deployment: V1Deployment,
@@ -269,7 +277,7 @@ class DeploymentApi:
         requires_release |= apply_change(deployment.spec, "entrypoint", entrypoint)
         requires_release |= apply_change(deployment.spec, "command", command)
         requires_release |= apply_change(deployment.spec, "env", to_env(env))
-        requires_release |= apply_change(deployment.spec, "env", to_health_check(health_check))
+        requires_release |= apply_change(deployment.spec, "readiness_probe", to_health_check(health_check))
         requires_release |= apply_change(deployment.spec, "cluster_id", cloud_account)
         requires_release |= apply_change(deployment.spec, "spot", spot)
 

lightning_sdk/api/job_api.py

@@ -275,7 +275,7 @@ class JobApiV2:
             return
 
         if current_state != Status.Stopping:
-            update_body = JobsIdBody1(cloudspace_id=current_job.spec.cloudspace_id, state=self.v2_job_state_stop)
+            update_body = JobsIdBody1(state=self.v2_job_state_stop)
             self._client.jobs_service_update_job(body=update_body, project_id=teamspace_id, id=job_id)
 
         while True:

lightning_sdk/api/lit_container_api.py

@@ -48,7 +48,8 @@ class LitContainerApi:
             raise ValueError(f"Container {container} does not exist") from None
 
         registry_url = _get_registry_url()
-        repository = f"{registry_url}/lit-container/{teamspace.owner.name}/{teamspace.name}/{container}"
+        container_basename = container.split("/")[-1]
+        repository = f"{registry_url}/lit-container/{teamspace.owner.name}/{teamspace.name}/{container_basename}"
         tagged = self._docker_client.api.tag(container, repository, tag)
         if not tagged:
             raise ValueError(f"Could not tag container {container} with {repository}:{tag}")

lightning_sdk/deployment/deployment.py

@@ -42,7 +42,7 @@ class Deployment:
     and switching machine types, etc..
 
     Args:
-        name: The name of the deployment.
+        name: The name or the id of the deployment.
         teamspace: The teamspace in which you want to deploy.
         org: The name of the organization owning the :param`teamspace` in case it is owned by an org
         user: The name of the user owning the :param`teamspace` in case it is owned directly by a user instead of an org
@@ -55,7 +55,7 @@ class Deployment:
 
     def __init__(
         self,
-        name: str,  # Only the name is required in case a deployment already exist.
+        name: str,
         teamspace: Optional[Union[str, Teamspace]] = None,
         org: Optional[Union[str, Organization]] = None,
         user: Optional[Union[str, User]] = None,
@@ -83,8 +83,14 @@ class Deployment:
         self._deployment_api = DeploymentApi()
         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 name.startswith("dep_"):
+            deployment = self._deployment_api.get_deployment_by_id(name, self._teamspace.id)
+        else:
+            deployment = self._deployment_api.get_deployment_by_name(name, self._teamspace.id)
+
         if deployment:
+            self._name = deployment.name
             self._is_created = True
             self._deployment = deployment
 
@@ -163,6 +169,9 @@ class Deployment:
                 strategy=to_strategy(release_strategy),
             )
         )
+
+        # Overrides the name
+        self._name = self._deployment._name
         self._is_created = True
 
     def update(

lightning_sdk/job/base.py

@@ -275,20 +275,25 @@ class _BaseJob(ABC):
         Caution: This also deletes all artifacts and snapshots associated with the job.
         """
 
-    def wait(self, interval: float = 5.0) -> None:
+    def wait(self, interval: float = 5.0, timeout: Optional[float] = None) -> None:
         """Waits for the job to be either completed, manually stopped or failed.
 
         Args:
-            interval: the number of seconds to spend in-between status checks.
+            interval: The number of seconds to spend in-between status checks.
+            timeout: The maximum number of seconds to wait before raising an error. If None, waits forever.
         """
         import time
 
         from lightning_sdk.status import Status
 
+        start = time.time()
         while True:
             if self.status in (Status.Completed, Status.Stopped, Status.Failed):
                 break
 
+            if timeout is not None and time.time() - start > timeout:
+                raise TimeoutError("Job didn't finish within the provided timeout.")
+
             time.sleep(interval)
 
     @property

lightning_sdk/job/job.py

@@ -1,7 +1,5 @@
-from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
-from lightning_sdk.api.user_api import UserApi
 from lightning_sdk.job.base import _BaseJob
 from lightning_sdk.job.v1 import _JobV1
 from lightning_sdk.job.v2 import _JobV2
@@ -18,15 +16,6 @@ if TYPE_CHECKING:
     from lightning_sdk.user import User
 
 
-@lru_cache(maxsize=None)
-def _has_jobs_v2() -> bool:
-    api = UserApi()
-    try:
-        return api._get_feature_flags().jobs_v2
-    except Exception:
-        return False
-
-
 class Job(_BaseJob):
     """Class to submit and manage single-machine jobs on the Lightning AI Platform."""
 
@@ -52,7 +41,7 @@ class Job(_BaseJob):
         """
         from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 
-        if _has_jobs_v2() and not self._force_v1:
+        if not self._force_v1:
             # try with v2 and fall back to v1
             try:
                 job = _JobV2(

lightning_sdk/job/v1.py

@@ -126,38 +126,7 @@ class _JobV1(_BaseJob):
             The submitted job.
 
         """
-        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:
-            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 jobs")
-
-        if entrypoint != "sh -c":
-            raise ValueError("Specifying the entrypoint is not yet supported with jobs")
-
-        if path_mappings is not None:
-            raise ValueError("Specifying path mappings is not yet supported with jobs")
-
-        # TODO: add support for empty names (will give an empty string)
-        _submitted = self._job_api.submit_job(
-            name=self._name,
-            command=command,
-            studio_id=studio._studio.id,
-            teamspace_id=self._teamspace.id,
-            cloud_account=cloud_account or "",
-            machine=machine,
-            interruptible=interruptible,
-        )
-        self._name = _submitted.name
-        self._job = _submitted
-        return self
+        raise NotImplementedError("Cannot submit new jobs with JobsV1!")
 
     def _update_internal_job(self) -> None:
         try:

lightning_sdk/job/v2.py

@@ -140,7 +140,12 @@ class _JobV2(_BaseJob):
     @property
     def status(self) -> "Status":
         """The current status of the job."""
-        return self._job_api._job_state_to_external(self._latest_job.state)
+        try:
+            return self._job_api._job_state_to_external(self._latest_job.state)
+        except Exception:
+            raise RuntimeError(
+                f"Job {self._name} does not exist in Teamspace {self.teamspace.name}. Did you delete it?"
+            ) from None
 
     @property
     def machine(self) -> Union["Machine", str]:

lightning_sdk/mmt/mmt.py

@@ -2,7 +2,6 @@ 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.job.job import _has_jobs_v2
 from lightning_sdk.mmt.base import MMTMachine, _BaseMMT
 from lightning_sdk.mmt.v1 import _MMTV1
 from lightning_sdk.mmt.v2 import _MMTV2
@@ -21,16 +20,18 @@ _logger = _setup_logger(__name__)
 
 @lru_cache(maxsize=None)
 def _has_mmt_v2() -> bool:
-    # users need both mmtv2 and jobsv2 flags in order for mmtv2 to work correctly
-    if not _has_jobs_v2():
-        return False
-
     api = UserApi()
     try:
-        return api._get_feature_flags().mmt_v2
+        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."""

lightning_sdk/plugin.py

@@ -432,6 +432,8 @@ def _success_message(resp: Union["Externalv1LightningappInstance", Job], plugin_
 def forced_v1(cls: Any) -> Generator[Any, None, None]:
     """Forces to use the v1 version of a class when using a class with multiple backends."""
     orig_val = getattr(cls, "_force_v1", False)
-    cls._force_v1 = True
-    yield cls
-    cls._force_v1 = orig_val
+    try:
+        cls._force_v1 = True
+        yield cls
+    finally:
+        cls._force_v1 = orig_val

{lightning_sdk-0.1.53.dist-info → lightning_sdk-0.1.54.dist-info}/METADATA

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

{lightning_sdk-0.1.53.dist-info → lightning_sdk-0.1.54.dist-info}/RECORD

@@ -1,5 +1,5 @@
 docs/source/conf.py,sha256=r8yX20eC-4mHhMTd0SbQb5TlSWHhO6wnJ0VJ_FBFpag,13249
-lightning_sdk/__init__.py,sha256=OspSXvFP3ft-tLCeHCDtd7YyPmS4jKx7hWDf0rVmHIs,970
+lightning_sdk/__init__.py,sha256=u3LjkfJQaHe9jKJXPyYO90NDPaQM7UezZxVhedZtvdA,970
 lightning_sdk/agents.py,sha256=ly6Ma1j0ZgGPFyvPvMN28JWiB9dATIstFa5XM8pMi6I,1577
 lightning_sdk/ai_hub.py,sha256=CpbpizTr3yzOCaEth0F1_mMifnR7Vf84quvthQnjzg4,6239
 lightning_sdk/constants.py,sha256=ztl1PTUBULnqTf3DyKUSJaV_O20hNtUYT6XvAYIrmIk,749
@@ -9,7 +9,7 @@ lightning_sdk/machine.py,sha256=qutfVwQJHC02c2ygB_O5wrFI-Eq9zbfA0E1OPEYcCO0,861
 lightning_sdk/models.py,sha256=Ssb88nspTKs_bPnZSixsTEi30FaKkmUEEVlf7ElQQ6s,5610
 lightning_sdk/organization.py,sha256=WCfzdgjtvY1_A07DnxOpp74V2JR2gQwtXbIEcFDnoVU,1232
 lightning_sdk/owner.py,sha256=t5svD2it4C9pbSpVuG9WJL46CYi37JXNziwnXxhiU5U,1361
-lightning_sdk/plugin.py,sha256=nWFL1l6Q9DE8Lr2krD5qRhZljwmYm00R2eR1tYRb20s,14902
+lightning_sdk/plugin.py,sha256=qv0oUuclcimQUcw4wsgNrbwKT-kjsDq65Vn78i3X_4c,14936
 lightning_sdk/status.py,sha256=lLGAuSvXBoXQFEEsEYwdCi0RcSNatUn5OPjJVjDtoM0,386
 lightning_sdk/studio.py,sha256=cDdVM70JXUj4tAjGBg6Purrd-_g92HNe_YgnOCxcMdU,19576
 lightning_sdk/teamspace.py,sha256=j6EghDJvxrYoWT1xMOA4qGAw3ezIpEytqGr-QlFnRXE,13247
@@ -17,9 +17,9 @@ lightning_sdk/user.py,sha256=vdn8pZqkAZO0-LoRsBdg0TckRKtd_H3QF4gpiZcl4iY,1130
 lightning_sdk/api/__init__.py,sha256=Qn2VVRvir_gO7w4yxGLkZY-R3T7kdiTPKgQ57BhIA9k,413
 lightning_sdk/api/agents_api.py,sha256=G47TbFo9kYqnBMqdw2RW-lfS1VAUBSXDmzs6fpIEMUs,4059
 lightning_sdk/api/ai_hub_api.py,sha256=NIFZcwibvicCK-oj9Vfm9ylGyXCFRn8gHIo-nmncKgY,5904
-lightning_sdk/api/deployment_api.py,sha256=T480Nej7LqmtkAx8SBkPGQ5JxeyQ-GVIDqUCc7Z1yfk,21448
-lightning_sdk/api/job_api.py,sha256=3ohMeDjeOaQAWU2qwbg4v5cDWF3U915hECMODxTdN6M,12541
-lightning_sdk/api/lit_container_api.py,sha256=zgmLcX6883nCWPLNMEANYTLB4zXAbMwXvGacHTzN5wI,3371
+lightning_sdk/api/deployment_api.py,sha256=AQh84-N2afeLIXvTgqkZUpOslt_7Tioi6_GUG_J1AWI,21809
+lightning_sdk/api/job_api.py,sha256=KhSZfjNlKtRFpT0yhVFyQ1tOOB16_C2UYWk_AqWrXro,12495
+lightning_sdk/api/lit_container_api.py,sha256=Cc5Whytka06xTzbDvAgnDj3a6ZBiY9vv0WV9I8RbvZY,3434
 lightning_sdk/api/mmt_api.py,sha256=km27mnZ7CT_TYa3ofjVrPlidjRMJk32pnznHxjVf1OM,7633
 lightning_sdk/api/org_api.py,sha256=Ze3z_ATVrukobujV5YdC42DKj45Vuwl7X52q_Vr-o3U,803
 lightning_sdk/api/studio_api.py,sha256=Cfsq8HFc4uUsj8hncnhnD_TLhw0cg-ryclGowj8S6Y0,26374
@@ -45,12 +45,12 @@ lightning_sdk/cli/studios_menu.py,sha256=0kQGqGel8gAbpdJtjOM1a6NEat_TnIqRNprNn8Q
 lightning_sdk/cli/teamspace_menu.py,sha256=GVTrMAqxxL3INX83alo5wybnrl6rBBKENo01ZNkWPik,3753
 lightning_sdk/cli/upload.py,sha256=cvirl8lfPiGV4THa5akeRy1KDoPx7Klj7tTrvWePwLc,11501
 lightning_sdk/deployment/__init__.py,sha256=BLu7_cVLp97TYxe6qe-J1zKUSZXAVcvCjgcA7plV2k4,497
-lightning_sdk/deployment/deployment.py,sha256=Dp15pn8rFAfMfaDhKn0v3bphFuvLgkPFs3KSNxW6eyc,15472
+lightning_sdk/deployment/deployment.py,sha256=Fn8iJVimaWB0Ngz53oJ03kBUzVe5_-3oHAbvEqFTf0s,15681
 lightning_sdk/job/__init__.py,sha256=1MxjQ6rHkyUHCypSW9RuXuVMVH11WiqhIXcU2LCFMwE,64
-lightning_sdk/job/base.py,sha256=VFMWFMZXUHxMD18DV3fljJr6gJd8RvT_IvW9OxkKSPo,17423
-lightning_sdk/job/job.py,sha256=TyueqMyci4rxmbDhqB3WUlkYsOHRFwl4lQrNv87I5As,13338
-lightning_sdk/job/v1.py,sha256=F9XqJTVx2JomxM6jYXqTpE394zTGZVGWv1JtAetj3uY,10378
-lightning_sdk/job/v2.py,sha256=7tamPEbgK6JuqsC3BFMi8Z2-wWR7qvnhjvdo0hTbNFw,10078
+lightning_sdk/job/base.py,sha256=TS5KavtfBAFbLbNqqumEizY9edjO1joSmtUdcO5CThQ,17748
+lightning_sdk/job/job.py,sha256=X4MLFZCZOtEhNlHIWZESR5mPrS4z6bhw3rSNXxuIHSA,13065
+lightning_sdk/job/v1.py,sha256=zgdZFnlyhzOBlbwdwx6Ob9YKBrzMnJflzLXy4qyb4So,9070
+lightning_sdk/job/v2.py,sha256=9O31CwIQjwm6qZSEZtFBWbZyBuCSmyFMsHCRbQvJdBU,10283
 lightning_sdk/job/work.py,sha256=Xm5byWTcOxYI8zZ3Sz6P8HxBTBAJ3S2KucndXA5KePQ,2332
 lightning_sdk/lightning_cloud/__init__.py,sha256=o91SMAlwr4Ke5ESe8fHjqXcj31_h7rT-MlFoXA-n2EI,173
 lightning_sdk/lightning_cloud/__version__.py,sha256=lOfmWHtjmiuSG28TbKQqd2B3nwmSGOlKVFwhaj_cRJk,23
@@ -869,7 +869,7 @@ lightning_sdk/lightning_cloud/utils/name_generator.py,sha256=MkciuA10332V0mcE2Px
 lightning_sdk/lightning_cloud/utils/network.py,sha256=axPgl8rhyPcPjxiztDxyksfxax3VNg2OXL5F5Uc81b4,406
 lightning_sdk/mmt/__init__.py,sha256=ExMu90-96bGBnyp5h0CErQszUGB1-PcjC4-R8_NYbeY,117
 lightning_sdk/mmt/base.py,sha256=B3HC-c82bPHprEZh1mhLCPCrCE8BOKqwIhY7xCF9CXg,15152
-lightning_sdk/mmt/mmt.py,sha256=Ha0Sp-m4_qBFuxDvASBt5ir_mVonPqcjyDS1QbIl0wQ,13894
+lightning_sdk/mmt/mmt.py,sha256=UbbK6UFEd-S2OPjezCG7cy-7kfMbyLOg3LAzs_Zq2Ds,13872
 lightning_sdk/mmt/v1.py,sha256=dROOXRWLzG2kUwGtvozB9Vywt_3QK8j4XFHVZFohWUM,9527
 lightning_sdk/mmt/v2.py,sha256=Em1XBoqViqUMKm-sshzdMcSH5UTtZZwbJcsgqY6-mw0,9625
 lightning_sdk/services/__init__.py,sha256=gSWUjccEhMI9CIWL_nbrFHUK2S6TM2725mEzrLMfK1Y,225
@@ -880,9 +880,9 @@ lightning_sdk/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 lightning_sdk/utils/dynamic.py,sha256=glUTO1JC9APtQ6Gr9SO02a3zr56-sPAXM5C3NrTpgyQ,1959
 lightning_sdk/utils/enum.py,sha256=h2JRzqoBcSlUdanFHmkj_j5DleBHAu1esQYUsdNI-hU,4106
 lightning_sdk/utils/resolve.py,sha256=idhRk4zg0jBOlySwKwjhsGgRorpsNVHOKOADOxMUYY0,6062
-lightning_sdk-0.1.53.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
-lightning_sdk-0.1.53.dist-info/METADATA,sha256=pjXt-rc8W-jTPmfq_ric7-NBi0t_82uks-e-z2yl3Ss,4013
-lightning_sdk-0.1.53.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-lightning_sdk-0.1.53.dist-info/entry_points.txt,sha256=msB9PJWIJ784dX-OP8by51d4IbKYH3Fj1vCuA9oXjHY,68
-lightning_sdk-0.1.53.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
-lightning_sdk-0.1.53.dist-info/RECORD,,
+lightning_sdk-0.1.54.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
+lightning_sdk-0.1.54.dist-info/METADATA,sha256=H9hzgevLDkHwM5VHEHnGB-7RHuIeNSOUtjp3ie_OhBw,4013
+lightning_sdk-0.1.54.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+lightning_sdk-0.1.54.dist-info/entry_points.txt,sha256=msB9PJWIJ784dX-OP8by51d4IbKYH3Fj1vCuA9oXjHY,68
+lightning_sdk-0.1.54.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
+lightning_sdk-0.1.54.dist-info/RECORD,,