lightning-sdk 0.2.10__py3-none-any.whl → 0.2.11__py3-none-any.whl

lightning_sdk/__init__.py

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

lightning_sdk/api/deployment_api.py

@@ -224,6 +224,7 @@ class DeploymentApi:
     def create_deployment(
         self,
         deployment: V1Deployment,
+        from_onboarding: Optional[bool] = None,
     ) -> V1Deployment:
         return self._client.jobs_service_create_deployment(
             project_id=deployment.project_id,
@@ -236,6 +237,7 @@ class DeploymentApi:
                 replicas=deployment.replicas,
                 spec=deployment.spec,
                 strategy=deployment.strategy,
+                from_onboarding=from_onboarding,
             ),
         )
 

lightning_sdk/api/teamspace_api.py

@@ -17,6 +17,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1ClusterAccelerator,
     V1Endpoint,
     V1Job,
+    V1Model,
     V1ModelVersionArchive,
     V1MultiMachineJob,
     V1Project,
@@ -199,9 +200,7 @@ class TeamspaceApi:
 
     def delete_model(self, name: str, version: Optional[str], teamspace_id: str) -> None:
         """Delete a model or a version from the model store."""
-        models = self.models_api.models_store_list_models(project_id=teamspace_id, name=name).models
-        assert len(models) == 1, "Multiple models with the same name found"
-        model = models[0]
+        model = self.get_model(teamspace_id=teamspace_id, model_name=name)
         # decide if delete only version of whole model
         if version:
             if version == "default":
@@ -257,7 +256,7 @@ class TeamspaceApi:
             if main_pbar:
                 main_pbar.update(1)
 
-    def complete_model_upload(self, model_id: str, version: str, teamspace_id: str) -> None:
+    def _complete_model_upload(self, model_id: str, version: str, teamspace_id: str) -> None:
         self.models_api.models_store_complete_model_upload(
             body=_DummyBody(),
             project_id=teamspace_id,
@@ -306,3 +305,29 @@ class TeamspaceApi:
             project_id=teamspace_id, id=cloud_account
         )
         return response.accelerator
+
+    def get_model(self, teamspace_id: str, model_id: Optional[str] = None, model_name: Optional[str] = None) -> V1Model:
+        if model_id:
+            return self.models_api.models_store_get_model(project_id=teamspace_id, model_id=model_id)
+        if not model_name:
+            raise ValueError("Either `model_id` or `model_name` must be provided.")
+        # list models with specific name
+        models = self.models_api.models_store_list_models(project_id=teamspace_id, name=model_name).models
+        if len(models) == 0:
+            raise ValueError(f"Model '{model_name}' does not exist.")
+        if len(models) > 1:
+            raise RuntimeError(f"Model name '{model_name}' is not a unique with this teamspace.")
+        # if there is only one model with the name, return it
+        return models[0]
+
+    def list_models(self, teamspace_id: str) -> List[V1Model]:
+        response = self.models_api.models_store_list_models(project_id=teamspace_id)
+        return response.models
+
+    def list_model_versions(
+        self, teamspace_id: str, model_id: Optional[str] = None, model_name: Optional[str] = None
+    ) -> List[V1ModelVersionArchive]:
+        if model_name and not model_id:
+            model_id = self.get_model(teamspace_id=teamspace_id, model_name=model_name).id
+        response = self.models_api.models_store_list_model_versions(project_id=teamspace_id, model_id=model_id)
+        return response.versions

lightning_sdk/cli/serve.py

@@ -333,12 +333,13 @@ class _Onboarding:
             return _OnboardingStatus.ONBOARDED
         return _OnboardingStatus.ONBOARDING
 
-    def _wait(self, timeout: int = _POLL_TIMEOUT) -> None:
+    def _wait_user_onboarding(self, timeout: int = _POLL_TIMEOUT) -> None:
         """Wait for user onboarding if they can join the teamspace otherwise move to select a teamspace."""
         status = self.status
         if status == _OnboardingStatus.ONBOARDED:
             return
 
+        self.console.print("Waiting for account setup. Visit lightning.ai")
         start_time = datetime.now()
         while self.status != _OnboardingStatus.ONBOARDED:
             time.sleep(5)
@@ -351,9 +352,13 @@ class _Onboarding:
 
     def get_cloudspace_id(self, teamspace: Teamspace) -> Optional[str]:
         cloudspaces: List[V1CloudSpace] = self.client.cloud_space_service_list_cloud_spaces(teamspace.id).cloudspaces
-        for cloudspace in cloudspaces:
-            if "scratch-studio" in cloudspace.name or "scratch-studio" in cloudspace.display_name:
-                return cloudspace.id
+        cloudspaces = sorted(cloudspaces, key=lambda cloudspace: cloudspace.created_at, reverse=True)
+        if len(cloudspaces) == 0:
+            raise RuntimeError("Error creating deployment! Finish account setup at lightning.ai first.")
+        # get the first cloudspace
+        cloudspace = cloudspaces[0]
+        if "scratch-studio" in cloudspace.name or "scratch-studio" in cloudspace.display_name:
+            return cloudspace.id
         return None
 
     def select_teamspace(self, teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace:
@@ -368,15 +373,8 @@ class _Onboarding:
 
         # Run only when user hasn't completed onboarding yet.
         menu = _TeamspacesMenu()
-        possible_teamspaces = menu._get_possible_teamspaces(self.user)
-        can_join_org = self.can_join_org
-
-        if len(possible_teamspaces) == 1 and can_join_org:
-            # wait for onboarding to complete so that user can join an org
-            # create deployment in the org default teamspace
-            self.console.print("Waiting for account setup. Visit lightning.ai")
-            self._wait()
-
+        self._wait_user_onboarding()
+        # Onboarding has been completed - user already selected organization if they could
         possible_teamspaces = menu._get_possible_teamspaces(self.user)
         if len(possible_teamspaces) == 1:
             # User didn't select any org
@@ -465,13 +463,16 @@ def _handle_cloud(
     authenticate(shall_confirm=not non_interactive)
     user_status = poll_verified_status()
     cloudspace_id: Optional[str] = None
+    from_onboarding = False
     if not user_status["verified"]:
         console.print("❌ Verify phone number to continue. Visit lightning.ai.", style="red")
         return
     if not user_status["onboarded"]:
+        console.print("onboarding user")
         onboarding = _Onboarding(console)
         resolved_teamspace = onboarding.select_teamspace(teamspace, org, user)
         cloudspace_id = onboarding.get_cloudspace_id(resolved_teamspace)
+        from_onboarding = True
     else:
         resolved_teamspace = select_teamspace(teamspace, org, user)
 
@@ -517,6 +518,7 @@ def _handle_cloud(
         replicas=replicas,
         include_credentials=include_credentials,
         cloudspace_id=cloudspace_id,
+        from_onboarding=from_onboarding,
     )
     console.print(f"🚀 Deployment started, access at [i]{deployment_status.get('url')}[/i]")
     if user_status["onboarded"]:

lightning_sdk/deployment/deployment.py

@@ -120,6 +120,7 @@ class Deployment:
         cloudspace_id: Optional[str] = None,
         quantity: Optional[int] = None,
         include_credentials: Optional[bool] = None,
+        from_onboarding: Optional[bool] = None,
     ) -> None:
         """The Lightning AI Deployment.
 
@@ -146,6 +147,7 @@ class Deployment:
             cloudspace_id: Connect deployment to a Studio.
             quantity: The number of machines per replica to deploy.
             include_credentials: Whether to include the environment variables for the SDK to authenticate
+            from_onboarding: Whether the deployment is from onboarding.
 
         Note:
             Since a teamspace can either be owned by an org or by a user directly,
@@ -182,7 +184,8 @@ class Deployment:
                     include_credentials=include_credentials if include_credentials is not None else True,
                 ),
                 strategy=to_strategy(release_strategy),
-            )
+            ),
+            from_onboarding=from_onboarding,
         )
 
         # Overrides the name

lightning_sdk/models.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 from lightning_sdk.api import OrgApi, TeamspaceApi, UserApi
+from lightning_sdk.lightning_cloud.openapi import V1ModelVersionArchive
 from lightning_sdk.lightning_cloud.openapi.models import V1Membership, V1OwnerType
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.user import User
@@ -184,4 +185,19 @@ def delete_model(
     name = _extend_model_name_with_teamspace(name)
     org_name, teamspace_name, model_name, version = _parse_org_teamspace_model_version(name)
     teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
-    teamspace.delete_model(name=model_name)
+    teamspace.delete_model(name=f"{model_name}:{version}" if version else model_name)
+
+
+def list_model_versions(
+    name: str,
+) -> List[V1ModelVersionArchive]:
+    """List all versions of a model.
+
+    Args:
+        name: The name of the model you want to list versions for.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+    """
+    name = _extend_model_name_with_teamspace(name)
+    org_name, teamspace_name, model_name, _ = _parse_org_teamspace_model_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    return teamspace.list_model_versions(name=model_name)

lightning_sdk/serve.py

@@ -249,6 +249,7 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
         port: Optional[int] = 8000,
         include_credentials: Optional[bool] = True,
         cloudspace_id: Optional[str] = None,
+        from_onboarding: bool = False,
     ) -> dict:
         """Run a deployment on the cloud. If the deployment already exists, it will be updated.
 
@@ -266,6 +267,7 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
             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.
@@ -300,6 +302,7 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
             ports=[port],
             include_credentials=include_credentials,
             cloudspace_id=cloudspace_id,
+            from_onboarding=from_onboarding,
         )
 
         return {"deployment": deployment, "url": url}

lightning_sdk/teamspace.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 import lightning_sdk
 from lightning_sdk.agents import Agent
 from lightning_sdk.api import TeamspaceApi
-from lightning_sdk.lightning_cloud.openapi import V1ProjectClusterBinding
+from lightning_sdk.lightning_cloud.openapi import V1Model, V1ModelVersionArchive, V1ProjectClusterBinding
 from lightning_sdk.machine import Machine
 from lightning_sdk.models import UploadedModelInfo
 from lightning_sdk.organization import Organization
@@ -317,7 +317,7 @@ class Teamspace:
             teamspace_id=self.id,
             progress_bar=progress_bar,
         )
-        self._teamspace_api.complete_model_upload(
+        self._teamspace_api._complete_model_upload(
             model_id=model.model_id,
             version=model.version,
             teamspace_id=self.id,
@@ -389,6 +389,18 @@ class Teamspace:
         name, version = _parse_model_and_version(name)
         self._teamspace_api.delete_model(name=name, version=version, teamspace_id=self.id)
 
+    def list_models(self) -> List[V1Model]:
+        """List all models in the model store."""
+        return self._teamspace_api.list_models(teamspace_id=self.id)
+
+    def list_model_versions(self, name: str) -> List[V1ModelVersionArchive]:
+        """List all versions of a model in the model store."""
+        if ":" in name:
+            raise ValueError(
+                "Model name should not contain a version tag. Please provide the model name without a version."
+            )
+        return self._teamspace_api.list_model_versions(teamspace_id=self.id, model_name=name)
+
 
 def _list_files(path: Union[str, Path]) -> Tuple[List[Path], List[str]]:
     """List all folders in a directory and return them as a list and relative path."""

{lightning_sdk-0.2.10.dist-info → lightning_sdk-0.2.11.dist-info}/METADATA

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

{lightning_sdk-0.2.10.dist-info → lightning_sdk-0.2.11.dist-info}/RECORD

@@ -1,32 +1,32 @@
 docs/source/conf.py,sha256=r8yX20eC-4mHhMTd0SbQb5TlSWHhO6wnJ0VJ_FBFpag,13249
-lightning_sdk/__init__.py,sha256=-t2JkyMlxbkPfbSjgk84YQIBy91o4GStvVhwQxiDHAc,1105
+lightning_sdk/__init__.py,sha256=CndcMqG7IuvuyazHxh3jvRLoKhhy72_ST2Y_mfjGIeM,1105
 lightning_sdk/agents.py,sha256=ly6Ma1j0ZgGPFyvPvMN28JWiB9dATIstFa5XM8pMi6I,1577
 lightning_sdk/ai_hub.py,sha256=iI1vNhgcz_Ff1c3rN1ogN7dK-r-HXRj6NMtS2cA14UA,6925
 lightning_sdk/constants.py,sha256=ztl1PTUBULnqTf3DyKUSJaV_O20hNtUYT6XvAYIrmIk,749
 lightning_sdk/helpers.py,sha256=KWMWnORHItIIA3PGn71YPs-7RjzGi8IXa2kQ5Qo4U8M,2459
 lightning_sdk/lit_container.py,sha256=8ys49TXI9MO89jLTA7MwDrKrssTRARAIF9OVmolDLq0,5273
 lightning_sdk/machine.py,sha256=EMr-ulyYEYhIkKFnGBOofnFf4asndTSeoOQxlyd3xW4,3632
-lightning_sdk/models.py,sha256=1v0eNZ-KQvXBAG3IOpyYbCuErqmT5gWK_mtwKOarkBI,7163
+lightning_sdk/models.py,sha256=tPFiiDQiDsFhfStcVPH5_PcujbBEfKS2IB5FS9NfwZk,7838
 lightning_sdk/organization.py,sha256=WCfzdgjtvY1_A07DnxOpp74V2JR2gQwtXbIEcFDnoVU,1232
 lightning_sdk/owner.py,sha256=t5svD2it4C9pbSpVuG9WJL46CYi37JXNziwnXxhiU5U,1361
 lightning_sdk/plugin.py,sha256=_K0m-PSyXE86B-SFQVQm0JEvfdIVXeARmnlATCAiMdo,14632
-lightning_sdk/serve.py,sha256=tv7x86T682EE3j6a88bpnR5J__COtPT6rMOctyqBIIg,11663
+lightning_sdk/serve.py,sha256=vtnxjJRLcpipTgM_64WNZls_zW1uOdSiPX4-OHFYQhk,11814
 lightning_sdk/status.py,sha256=lLGAuSvXBoXQFEEsEYwdCi0RcSNatUn5OPjJVjDtoM0,386
 lightning_sdk/studio.py,sha256=hyvAiVhkETAtbu0RRF1Aw6F8Y__E1SSAmsB8PHfmqHo,19935
-lightning_sdk/teamspace.py,sha256=mCUwNJdQsl-hcu8q1qs7jj-0dheHXrNHuybaDVh6dL0,15921
+lightning_sdk/teamspace.py,sha256=GL0fiZ7oz3Uz1qftDmSFZScpuUm3H4kBcm4abv06xLM,16533
 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=azqDZ-PzasVAcoQHno7k7OO_xFOHQ4NDozxF8jEh83Y,7864
 lightning_sdk/api/cluster_api.py,sha256=OxAhEdNt-9bYQ-u3MCADdVGsilXf2e7Wcxvb73Cmn8k,916
-lightning_sdk/api/deployment_api.py,sha256=-zoIuMhZBsgona2__msh_Vm6e2m1Z323J-43by37hj8,22639
+lightning_sdk/api/deployment_api.py,sha256=a3QIlQDQBAUv-acU_pLdBGbGQMsfPvfvex2ShlxfotY,22736
 lightning_sdk/api/job_api.py,sha256=_mMAI_BG_48i-BLwCP_U72zgmM5zYa2KUZ7u66HWkIc,13568
 lightning_sdk/api/lit_container_api.py,sha256=WQRNTOu74wSTfg0TKLxGPDRvnu2o0hCOw9fOTr-56Tw,11206
 lightning_sdk/api/mmt_api.py,sha256=-v7ATab-ThAM-HRClS92Ehxuu9MlBfdKWWFCGvVUHiM,8962
 lightning_sdk/api/org_api.py,sha256=Ze3z_ATVrukobujV5YdC42DKj45Vuwl7X52q_Vr-o3U,803
 lightning_sdk/api/pipeline_api.py,sha256=P5P9C6qOpyBGU0t5N68h1LuFAsAKmPPgkac6uObrYKw,1676
 lightning_sdk/api/studio_api.py,sha256=BqmGmotUJrM-RoMfwYSdymAir491tBkPsIlvdJqOoUw,25858
-lightning_sdk/api/teamspace_api.py,sha256=4sCKenZOVIueMLwtCPZseMtEfRwI4PIznC2o0fnMdeY,11823
+lightning_sdk/api/teamspace_api.py,sha256=syN_FK8nHhLl3FG9Oxe8L6aK5YbOMn_GQlkijhTRJA8,13147
 lightning_sdk/api/user_api.py,sha256=sL7RIjjtmZmvCZWx7BBZslhj1BeNh4Idn-RVcdmf7M0,2598
 lightning_sdk/api/utils.py,sha256=3KvVmucfrUIPV0HC0y5DHhWB0b3D5tvLFXSrDU6XjcI,23833
 lightning_sdk/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -49,7 +49,7 @@ lightning_sdk/cli/list.py,sha256=bH_rX9J_Ubb9dkF2fpoZf0zMXH9YoEjcm63gk5v0zs4,109
 lightning_sdk/cli/mmts_menu.py,sha256=HUXo3ZoZ3fWOCNWTQWoJgUlFXYq5uVm_6uFjAq7BDe8,2219
 lightning_sdk/cli/open.py,sha256=sgMLWBnkXdIq8H9XK_rph2bye3b07AKTJBQIk9fCGVc,1937
 lightning_sdk/cli/run.py,sha256=8JZiDrKwDhlaTOJd6qq2mCWJRqKm6shCWLzpbmFYIkE,13929
-lightning_sdk/cli/serve.py,sha256=L6sHbC0sPt47ns0VeGluhrbkq5bxJuegg66eSJebrew,18756
+lightning_sdk/cli/serve.py,sha256=lNSa2tg66vLCrtYikNmFAiG2fMX-luhpnizyyB1Xlqk,18968
 lightning_sdk/cli/start.py,sha256=jUk52lkEFC_fqiEPkwM8GwE68WMNEtzBuzjkvr3POd0,1840
 lightning_sdk/cli/stop.py,sha256=5nCrUe1BONpX1nKNhbSFqLaXXKaRhSO7PvM1BVYLgn4,2864
 lightning_sdk/cli/studios_menu.py,sha256=TA9rO6_fFHGMz0Nt4rJ6iV80X5pZE4xShrSiyXoU-oQ,4129
@@ -57,7 +57,7 @@ lightning_sdk/cli/switch.py,sha256=qLvDoQRldCNgO1XvIGg-i-GyqnkHVb_U1-X4svqeNNU,1
 lightning_sdk/cli/teamspace_menu.py,sha256=C3g3spTKgtMwoK7pnooy0MBPz4AKhFjcObkvZyZ4v04,3797
 lightning_sdk/cli/upload.py,sha256=jQbWVlwxsDp3YnFtZTXFeK_PqDssNpyFzL8wV3NkWk8,13849
 lightning_sdk/deployment/__init__.py,sha256=dXsa4psDzFYFklsq3JC-2V_L4FQjGZnQAf-ZiVlqG9c,545
-lightning_sdk/deployment/deployment.py,sha256=dPLXxPNSi86oZiCLNJFxaiMoPe1PJyLlZfG0L_sBMfQ,18356
+lightning_sdk/deployment/deployment.py,sha256=KF7YtYQm9GCiUrPtSNnJLIUIi6BqQ2dvLRSJ-xpbBzQ,18522
 lightning_sdk/job/__init__.py,sha256=1MxjQ6rHkyUHCypSW9RuXuVMVH11WiqhIXcU2LCFMwE,64
 lightning_sdk/job/base.py,sha256=TS5KavtfBAFbLbNqqumEizY9edjO1joSmtUdcO5CThQ,17748
 lightning_sdk/job/job.py,sha256=1Xf0ne4wwXpkb_GJgWD8mfueJZoKR8G4lC5T6OXijlw,13075
@@ -1004,9 +1004,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=6qlBUkOmcFTjhQx_CAGfnvWBbMYp6XrCV_sX_IqplLE,6748
-lightning_sdk-0.2.10.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
-lightning_sdk-0.2.10.dist-info/METADATA,sha256=cF9VhuDfRfSM1wx3EMN0XgcSuVL94cAUo-Z5pNrDwUA,3992
-lightning_sdk-0.2.10.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-lightning_sdk-0.2.10.dist-info/entry_points.txt,sha256=msB9PJWIJ784dX-OP8by51d4IbKYH3Fj1vCuA9oXjHY,68
-lightning_sdk-0.2.10.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
-lightning_sdk-0.2.10.dist-info/RECORD,,
+lightning_sdk-0.2.11.dist-info/LICENSE,sha256=uFIuZwj5z-4TeF2UuacPZ1o17HkvKObT8fY50qN84sg,1064
+lightning_sdk-0.2.11.dist-info/METADATA,sha256=c2dtRhsFy2j9rD0ioEa85TOrK3eSSlHCPmFGtBfiuzY,3992
+lightning_sdk-0.2.11.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+lightning_sdk-0.2.11.dist-info/entry_points.txt,sha256=msB9PJWIJ784dX-OP8by51d4IbKYH3Fj1vCuA9oXjHY,68
+lightning_sdk-0.2.11.dist-info/top_level.txt,sha256=ps8doKILFXmN7F1mHncShmnQoTxKBRPIcchC8TpoBw4,19
+lightning_sdk-0.2.11.dist-info/RECORD,,