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

lightning_sdk/lightning_cloud/openapi/models/v1_volume.py

@@ -57,6 +57,7 @@ class V1Volume(object):
         'server_id': 'str',
         'size_gb': 'str',
         'state': 'V1VolumeState',
+        'subnet_id': 'str',
         'throughput': 'str',
         'type': 'str',
         'updated_at': 'datetime',
@@ -80,13 +81,14 @@ class V1Volume(object):
         'server_id': 'serverId',
         'size_gb': 'sizeGb',
         'state': 'state',
+        'subnet_id': 'subnetId',
         'throughput': 'throughput',
         'type': 'type',
         'updated_at': 'updatedAt',
         'user_id': 'userId'
     }
 
-    def __init__(self, attached_at: 'datetime' =None, availability_zone: 'str' =None, detached_at: 'datetime' =None, encrypted: 'bool' =None, iops: 'str' =None, metadata: 'V1Metadata' =None, path: 'str' =None, provider: 'str' =None, provider_id: 'str' =None, region: 'str' =None, resource_id: 'str' =None, resource_type: 'str' =None, retention_period: 'str' =None, server_id: 'str' =None, size_gb: 'str' =None, state: 'V1VolumeState' =None, throughput: 'str' =None, type: 'str' =None, updated_at: 'datetime' =None, user_id: 'str' =None):  # noqa: E501
+    def __init__(self, attached_at: 'datetime' =None, availability_zone: 'str' =None, detached_at: 'datetime' =None, encrypted: 'bool' =None, iops: 'str' =None, metadata: 'V1Metadata' =None, path: 'str' =None, provider: 'str' =None, provider_id: 'str' =None, region: 'str' =None, resource_id: 'str' =None, resource_type: 'str' =None, retention_period: 'str' =None, server_id: 'str' =None, size_gb: 'str' =None, state: 'V1VolumeState' =None, subnet_id: 'str' =None, throughput: 'str' =None, type: 'str' =None, updated_at: 'datetime' =None, user_id: 'str' =None):  # noqa: E501
         """V1Volume - a model defined in Swagger"""  # noqa: E501
         self._attached_at = None
         self._availability_zone = None
@@ -104,6 +106,7 @@ class V1Volume(object):
         self._server_id = None
         self._size_gb = None
         self._state = None
+        self._subnet_id = None
         self._throughput = None
         self._type = None
         self._updated_at = None
@@ -141,6 +144,8 @@ class V1Volume(object):
             self.size_gb = size_gb
         if state is not None:
             self.state = state
+        if subnet_id is not None:
+            self.subnet_id = subnet_id
         if throughput is not None:
             self.throughput = throughput
         if type is not None:
@@ -486,6 +491,27 @@ class V1Volume(object):
 
         self._state = state
 
+    @property
+    def subnet_id(self) -> 'str':
+        """Gets the subnet_id of this V1Volume.  # noqa: E501
+
+
+        :return: The subnet_id of this V1Volume.  # noqa: E501
+        :rtype: str
+        """
+        return self._subnet_id
+
+    @subnet_id.setter
+    def subnet_id(self, subnet_id: 'str'):
+        """Sets the subnet_id of this V1Volume.
+
+
+        :param subnet_id: The subnet_id of this V1Volume.  # noqa: E501
+        :type: str
+        """
+
+        self._subnet_id = subnet_id
+
     @property
     def throughput(self) -> 'str':
         """Gets the throughput of this V1Volume.  # noqa: E501

lightning_sdk/llm/llm.py

@@ -9,6 +9,7 @@ PUBLIC_MODEL_PROVIDERS: Dict[str, str] = {
     "openai": "OpenAI",
     "anthropic": "Anthropic",
     "google": "Google",
+    "lightning-ai": "lightning-ai",
 }
 
 
@@ -54,8 +55,18 @@ class LLM:
         Raises:
             ValueError: If teamspace information cannot be resolved.
         """
-        # TODO support user input teamspace
-        self._get_auth_info()
+        teamspace_name = None
+        if teamspace:
+            try:
+                owner, teamspace_name = teamspace.split("/", maxsplit=1)
+            except ValueError as e:
+                raise ValueError(
+                    f"Invalid teamspace format: '{teamspace}'. "
+                    "Teamspace should be specified as '{teamspace_owner}/{teamspace_name}' "
+                    "(e.g., 'my-org/my-teamspace')."
+                ) from e
+
+        self._get_auth_info(teamspace_name)
 
         self._model_provider, self._model_name = self._parse_model_name(name)
         self._enable_async = enable_async
@@ -76,13 +87,15 @@ class LLM:
     def provider(self) -> str:
         return self._model_provider
 
-    def _get_auth_info(self) -> None:
+    def _get_auth_info(self, teamspace_name: Optional[str] = None) -> None:
+        # TODO: Validate user input teamspace name
         if not LLM._auth_info_cached:
-            teamspace_name = os.environ.get("LIGHTNING_TEAMSPACE", None)
+            if teamspace_name is None:
+                teamspace_name = os.environ.get("LIGHTNING_TEAMSPACE", None)
             if teamspace_name is None:
                 raise ValueError(
                     "Teamspace name must be provided either through "
-                    "the environment variable LIGHTNING_TEAMSPACE or as an argument."
+                    "the environment variable LIGHTNING_TEAMSPACE or as an argument - LLM(..., teamspace=...)"
                 )
             LLM._cached_auth_info = {
                 "teamspace_name": teamspace_name,
@@ -138,6 +151,18 @@ class LLM:
                     "Please check the model name or provider."
                 ) from e
 
+        if self._model_provider == "lightning-ai":
+            # Try model provider model
+            try:
+                return self._llm_api.get_assistant(
+                    model_provider=self._model_provider,
+                    model_name=self._model_name,
+                    user_name="",
+                    org_name="",
+                )
+            except Exception:
+                pass
+
         # Try organization model
         try:
             return self._llm_api.get_assistant(
@@ -241,6 +266,7 @@ class LLM:
         metadata: Optional[Dict[str, str]] = None,
         stream: bool = False,
         full_response: bool = False,
+        tools: Optional[List[Dict[str, Any]]] = None,
         **kwargs: Any,
     ) -> Union[
         V1ConversationResponseChunk, Generator[V1ConversationResponseChunk, None, None], str, Generator[str, None, None]
@@ -281,6 +307,7 @@ class LLM:
             metadata=metadata,
             name=conversation,
             stream=stream,
+            tools=tools,
             **kwargs,
         )
         if not stream:

lightning_sdk/llm/public_assistants.json

@@ -4,5 +4,7 @@
     "openai/o3-mini": "ast_01jz3t13fhnjhh11t1k8b5gyp1",
     "anthropic/claude-3-5-sonnet-20240620": "ast_01jd3923a6p98rqwh3dpj686pq",
     "google/gemini-2.5-pro": "ast_01jz3tdb1fhey798k95pv61v57",
-    "google/gemini-2.5-flash": "ast_01jz3thxskg4fcdk4xhkjkym5a"
+    "google/gemini-2.5-flash": "ast_01jz3thxskg4fcdk4xhkjkym5a",
+    "google/gemini-2.5-flash-lite-preview-06-17": "ast_01jz3thxskg4fcdk4xhkjkym5b",
+    "lightning-ai/llama4-maverick": "ast_01k073vsqs66tenpns02cf5jnq"
 }

lightning_sdk/machine.py

@@ -1,17 +1,31 @@
 from dataclasses import dataclass
+from enum import Enum
 from typing import Any, ClassVar, Optional, Tuple
 
 
+class CloudProvider(Enum):
+    AWS = "AWS"
+    GCP = "GCP"
+    VULTR = "VULTR"
+    LAMBDA_LABS = "LAMBDA_LABS"
+    DGX = "DGX"
+    VOLTAGE_PARK = "VOLTAGE_PARK"
+    NEBIUS = "NEBIUS"
+    LIGHTNING = "LIGHTNING"
+
+
 @dataclass(frozen=True)
 class Machine:
     # Default Machines
     CPU: ClassVar["Machine"]
+    CPU_SMALL: ClassVar["Machine"]
     DATA_PREP: ClassVar["Machine"]
     DATA_PREP_MAX: ClassVar["Machine"]
     DATA_PREP_ULTRA: ClassVar["Machine"]
     T4: ClassVar["Machine"]
     T4_X_4: ClassVar["Machine"]
     L4: ClassVar["Machine"]
+    L4_X_2: ClassVar["Machine"]
     L4_X_4: ClassVar["Machine"]
     L4_X_8: ClassVar["Machine"]
     A10G: ClassVar["Machine"]
@@ -20,7 +34,10 @@ class Machine:
     L40S: ClassVar["Machine"]
     L40S_X_4: ClassVar["Machine"]
     L40S_X_8: ClassVar["Machine"]
+    A100_X_2: ClassVar["Machine"]
+    A100_X_4: ClassVar["Machine"]
     A100_X_8: ClassVar["Machine"]
+    B200_X_8: ClassVar["Machine"]
     H100_X_8: ClassVar["Machine"]
     H200_X_8: ClassVar["Machine"]
 
@@ -45,6 +62,7 @@ class Machine:
         """Whether the machine is a CPU."""
         return (
             self == Machine.CPU
+            or self == Machine.CPU_SMALL
             or self == Machine.DATA_PREP
             or self == Machine.DATA_PREP_MAX
             or self == Machine.DATA_PREP_ULTRA
@@ -66,12 +84,14 @@ class Machine:
 
 
 Machine.CPU = Machine(name="CPU", instance_type="cpu-4")
+Machine.CPU_SMALL = Machine(name="CPU_SMALL", instance_type="n2d-standard-2")  # GCP
 Machine.DATA_PREP = Machine(name="DATA_PREP", instance_type="data-large")
 Machine.DATA_PREP_MAX = Machine(name="DATA_PREP_MAX", instance_type="data-max")
 Machine.DATA_PREP_ULTRA = Machine(name="DATA_PREP_ULTRA", instance_type="data-ultra")
 Machine.T4 = Machine(name="T4", instance_type="g4dn.2xlarge")
 Machine.T4_X_4 = Machine(name="T4_X_4", instance_type="g4dn.12xlarge")
 Machine.L4 = Machine(name="L4", instance_type="g6.4xlarge")
+Machine.L4_X_2 = Machine(name="L4_X_2", instance_type="g2-standard-24")  # GCP
 Machine.L4_X_4 = Machine(name="L4_X_4", instance_type="g6.12xlarge")
 Machine.L4_X_8 = Machine(name="L4_X_8", instance_type="g6.48xlarge")
 Machine.A10G = Machine(name="A10G", instance_type="g5.8xlarge")
@@ -80,6 +100,9 @@ Machine.A10G_X_8 = Machine(name="A10G_X_8", instance_type="g5.48xlarge")
 Machine.L40S = Machine(name="L40S", instance_type="g6e.4xlarge")
 Machine.L40S_X_4 = Machine(name="L40S_X_4", instance_type="g6e.12xlarge")
 Machine.L40S_X_8 = Machine(name="L40S_X_8", instance_type="g6e.48xlarge")
+Machine.A100_X_2 = Machine(name="A100_X_2", instance_type="a2-ultragpu-2g")  # GCP
+Machine.A100_X_4 = Machine(name="A100_X_4", instance_type="a2-ultragpu-4g")  # GCP
 Machine.A100_X_8 = Machine(name="A100_X_8", instance_type="p4d.24xlarge")
+Machine.B200_X_8 = Machine(name="B200_X_8", instance_type="a4-highgpu-8g")  # GCP
 Machine.H100_X_8 = Machine(name="H100_X_8", instance_type="p5.48xlarge")
-Machine.H200_X_8 = Machine(name="H200_X_8", instance_type="p5e.48xlarge")
+Machine.H200_X_8 = Machine(name="H200_X_8", instance_type="p5en.48xlarge")

lightning_sdk/mmt/base.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Protocol, Tuple, Union
 
 if TYPE_CHECKING:
     from lightning_sdk.job.base import MachineDict
-    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
@@ -64,12 +64,14 @@ class _BaseMMT(_BaseJob):
         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,  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
@@ -89,7 +91,11 @@ class _BaseMMT(_BaseJob):
             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.
@@ -109,6 +115,10 @@ class _BaseMMT(_BaseJob):
                     }
                 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
         """
         from lightning_sdk.lightning_cloud.openapi.rest import ApiException
         from lightning_sdk.studio import Studio
@@ -191,6 +201,7 @@ class _BaseMMT(_BaseJob):
             num_machines=num_machines,
             machine=machine,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             command=command,
             studio=studio,
             image=image,
@@ -202,6 +213,7 @@ class _BaseMMT(_BaseJob):
             path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            max_runtime=max_runtime,
         )
         return inst
 
@@ -216,12 +228,14 @@ class _BaseMMT(_BaseJob):
         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,
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
+        max_runtime: Optional[int] = None,
     ) -> None:
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -253,6 +267,10 @@ class _BaseMMT(_BaseJob):
                     }
                 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
         """
 
     @property

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."""