lightning-sdk 2025.7.17__py3-none-any.whl → 2025.7.22__py3-none-any.whl

lightning_sdk/mmt/mmt.py

@@ -1,12 +1,13 @@
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.mmt.base import MMTMachine, _BaseMMT
 from lightning_sdk.mmt.v1 import _MMTV1
 from lightning_sdk.mmt.v2 import _MMTV2
 from lightning_sdk.utils.resolve import _setup_logger
 
 if TYPE_CHECKING:
-    from lightning_sdk.machine import Machine
+    from lightning_sdk.machine import CloudProvider, Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.status import Status
     from lightning_sdk.studio import Studio
@@ -75,6 +76,7 @@ class MMT(_BaseMMT):
             )
 
         self._internal_mmt = mmt
+        self._cloud_account_api = CloudAccountApi()
 
     @classmethod
     def run(
@@ -89,12 +91,14 @@ class MMT(_BaseMMT):
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
@@ -114,7 +118,11 @@ class MMT(_BaseMMT):
             user: The user owning the teamspace (if any). Defaults to the current user.
             cloud_account: The cloud account to run the job 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.
+                If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+                will fall back to the teamspaces default cloud account.
+            cloud_account_provider: The provider to select the cloud-account from.
+                If set, must be in agreement with the provider from the cloud_account (if specified).
+                If not specified, falls backto the teamspace default cloud account.
             env: Environment variables to set inside the job.
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             image_credentials: The credentials used to pull the image. Required if the image is private.
@@ -145,6 +153,7 @@ class MMT(_BaseMMT):
             org=org,
             user=user,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             env=env,
             interruptible=interruptible,
             image_credentials=image_credentials,
@@ -154,6 +163,7 @@ class MMT(_BaseMMT):
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
             cluster=cluster,  # deprecated in favor of cloud_account
+            max_runtime=max_runtime,
         )
         # required for typing with "MMT"
         assert isinstance(ret_val, cls)
@@ -173,10 +183,12 @@ class MMT(_BaseMMT):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
     ) -> "MMT":
@@ -193,7 +205,11 @@ class MMT(_BaseMMT):
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             cloud_account: The cloud account to run the job 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.
+                If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+                will fall back to the teamspaces default cloud account.
+            cloud_account_provider: The provider to select the cloud-account from.
+                If set, must be in agreement with the provider from the cloud_account (if specified).
+                If not specified, falls backto the teamspace default cloud account.
             image_credentials: The credentials used to pull the image. Required if the image is private.
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
             cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
@@ -211,11 +227,16 @@ class MMT(_BaseMMT):
                     }
                 If the path inside the connection is omitted it's assumed to be the root path of that connection.
                 Only applicable when submitting docker jobs.
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
         """
         self._job = self._internal_mmt._submit(
             num_machines=num_machines,
             machine=machine,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             command=command,
             studio=studio,
             image=image,
@@ -227,6 +248,7 @@ class MMT(_BaseMMT):
             path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            max_runtime=max_runtime,
         )
         return self
 

lightning_sdk/mmt/v1.py

@@ -7,7 +7,7 @@ from lightning_sdk.job.work import Work
 from lightning_sdk.status import Status
 
 if TYPE_CHECKING:
-    from lightning_sdk.machine import Machine
+    from lightning_sdk.machine import CloudProvider, Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.studio import Studio
     from lightning_sdk.teamspace import Teamspace
@@ -50,10 +50,12 @@ class _MMTV1(_BaseMMT):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
     ) -> "_MMTV1":
@@ -90,6 +92,10 @@ class _MMTV1(_BaseMMT):
                 To use the pre-defined entrypoint of the provided image, set this to an empty string.
                 Only applicable when submitting docker jobs.
             path_mappings: The mappings from data connection inside your container (not supported)
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
 
         """
         raise NotImplementedError("Cannot submit new mmts with MMTV1!")

lightning_sdk/mmt/v2.py

@@ -6,7 +6,7 @@ from lightning_sdk.status import Status
 
 if TYPE_CHECKING:
     from lightning_sdk.job.job import Job
-    from lightning_sdk.machine import Machine
+    from lightning_sdk.machine import CloudProvider, Machine
     from lightning_sdk.organization import Organization
     from lightning_sdk.studio import Studio
     from lightning_sdk.teamspace import Teamspace
@@ -49,10 +49,12 @@ class _MMTV2(_BaseMMT):
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
     ) -> "_MMTV2":
@@ -69,7 +71,11 @@ class _MMTV2(_BaseMMT):
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
             cloud_account: The cloud account to run the job 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.
+                If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+                will fall back to the teamspaces default cloud account.
+            cloud_account_provider: The provider to select the cloud-account from.
+                If set, must be in agreement with the provider from the cloud_account (if specified).
+                If not specified, falls backto the teamspace default cloud account.
             image_credentials: The credentials used to pull the image. Required if the image is private.
                 This should be the name of the respective credentials secret created on the Lightning AI platform.
             cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
@@ -87,6 +93,10 @@ class _MMTV2(_BaseMMT):
                     }
                 If the path inside the connection is omitted it's assumed to be the root path of that connection.
                 Only applicable when submitting docker jobs.
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
         """
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -104,6 +114,14 @@ class _MMTV2(_BaseMMT):
             studio_id = None
             if image is None:
                 raise ValueError("either image or studio must be provided")
+
+        cloud_account = self._cloud_account_api.resolve_cloud_account(
+            self._teamspace.id,
+            cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
+            default_cloud_account=self._teamspace.default_cloud_account,
+        )
+
         submitted = self._job_api.submit_job(
             name=self.name,
             num_machines=num_machines,
@@ -121,6 +139,7 @@ class _MMTV2(_BaseMMT):
             path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            max_runtime=max_runtime,
         )
         self._job = submitted
         self._name = submitted.name

lightning_sdk/organization.py

@@ -40,6 +40,10 @@ class Organization(Owner):
         """The organization's ID."""
         return self._org.id
 
+    @property
+    def default_cloud_account(self) -> Optional[str]:
+        return self._org.preferred_cluster or None
+
     def __repr__(self) -> str:
         """Returns reader friendly representation."""
         return f"Organization(name={self.name})"

lightning_sdk/pipeline/pipeline.py

@@ -1,9 +1,12 @@
 import os
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import List, Optional, Union
 
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.api.pipeline_api import PipelineApi
+from lightning_sdk.machine import CloudProvider
 from lightning_sdk.organization import Organization
 from lightning_sdk.pipeline.printer import PipelinePrinter
+from lightning_sdk.pipeline.schedule import _TIMEZONES, Schedule
 from lightning_sdk.pipeline.steps import DeploymentStep, JobStep, MMTStep, _get_studio
 from lightning_sdk.pipeline.utils import prepare_steps
 from lightning_sdk.services.utilities import _get_cluster
@@ -12,9 +15,6 @@ from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
 from lightning_sdk.utils.resolve import _resolve_teamspace
 
-if TYPE_CHECKING:
-    from lightning_sdk.pipeline.schedule import Schedule
-
 
 class Pipeline:
     def __init__(
@@ -24,6 +24,7 @@ class Pipeline:
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union[CloudProvider, str]] = None,
         shared_filesystem: Optional[bool] = None,
         studio: Optional[Union[Studio, str]] = None,
     ) -> None:
@@ -45,9 +46,14 @@ class Pipeline:
             org=org,
             user=user,
         )
+        if self._teamspace is None:
+            raise RuntimeError("Could not resolve teamspace")
 
         self._pipeline_api = PipelineApi()
-        self._cloud_account = cloud_account
+        self._cloud_account_api = CloudAccountApi()
+        self._cloud_account = self._cloud_account_api.resolve_cloud_account(
+            self._teamspace.id, cloud_account, cloud_provider, self._teamspace.default_cloud_account
+        )
         self._default_cluster = _get_cluster(
             client=self._pipeline_api._client, project_id=self._teamspace.id, cluster_id=cloud_account
         )
@@ -108,6 +114,11 @@ class Pipeline:
             if schedule.name is None:
                 schedule.name = f"schedule-{schedule_idx}"
 
+            if schedule.timezone is not None and schedule.timezone not in _TIMEZONES:
+                raise ValueError(
+                    f"The schedule {schedule.name} timezone isn't supported. Available list is {_TIMEZONES}. Found {schedule.timezone}."  # noqa: E501
+                )
+
         parent_pipeline_id = None if self._pipeline is None else self._pipeline.id
 
         self._pipeline = self._pipeline_api.create_pipeline(

lightning_sdk/pipeline/printer.py

@@ -1,7 +1,7 @@
 import os
 from typing import Any, ClassVar, Dict, List
 
-from lightning_sdk.lightning_cloud.openapi.models import V1Pipeline, V1PipelineStepType
+from lightning_sdk.lightning_cloud.openapi.models import V1Pipeline, V1PipelineStepType, V1Schedule
 from lightning_sdk.pipeline.utils import _get_spec
 
 
@@ -21,7 +21,7 @@ class PipelinePrinter:
         pipeline: V1Pipeline,
         teamspace: Any,
         proto_steps: List[Any],
-        schedules: List[Any],
+        schedules: List[V1Schedule],
     ) -> None:
         self._name = name
         self._initial = initial
@@ -78,7 +78,9 @@ class PipelinePrinter:
             return
 
         for schedule in self._schedules:
-            self._print(f"  - '{schedule.name}' runs on cron schedule: `{schedule.cron_expression}`")
+            self._print(
+                f"  - '{schedule.name}' runs on cron schedule: `{schedule.cron_expression} in timezone {schedule.timezone or 'UTC'}`"  # noqa: E501
+            )
 
     def _print_footer(self) -> None:
         """Prints the final link and closing message."""