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

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/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.4.dist-info → lightning_sdk-0.2.5.dist-info}/METADATA

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