lightning-sdk 2025.7.17__py3-none-any.whl → 2025.7.30rc0__py3-none-any.whl

lightning_sdk/pipeline/steps.py

@@ -25,11 +25,11 @@ from lightning_sdk.lightning_cloud.openapi.models import (
 from lightning_sdk.machine import Machine
 from lightning_sdk.mmt.v2 import MMTApiV2
 from lightning_sdk.pipeline.utils import DEFAULT, _get_studio, _to_wait_for, _validate_cloud_account
-from lightning_sdk.studio import Studio
+from lightning_sdk.studio import CloudAccountApi, Studio
 
 if TYPE_CHECKING:
     from lightning_sdk.organization import Organization
-    from lightning_sdk.teamspace import Teamspace
+    from lightning_sdk.teamspace import CloudProvider, Teamspace
     from lightning_sdk.user import User
 
 
@@ -57,6 +57,7 @@ class DeploymentStep:
         custom_domain: Optional[str] = None,
         quantity: Optional[int] = None,
         include_credentials: Optional[bool] = None,
+        max_runtime: Optional[int] = None,
         wait_for: Optional[Union[str, List[str]]] = DEFAULT,
     ) -> None:
         self.name = name
@@ -96,6 +97,7 @@ class DeploymentStep:
         self.custom_domain = custom_domain
         self.quantity = quantity
         self.include_credentials = include_credentials or True
+        self.max_runtime = max_runtime
         self.wait_for = wait_for
 
     def to_proto(
@@ -124,6 +126,7 @@ class DeploymentStep:
                     quantity=self.quantity,
                     cloudspace_id=self.studio._studio.id if self.studio else None,
                     include_credentials=self.include_credentials,
+                    max_runtime=self.max_runtime,
                 ),
                 strategy=to_strategy(self.release_strategy),
             ),
@@ -144,12 +147,14 @@ class JobStep:
         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,
         wait_for: Union[str, List[str], None] = DEFAULT,
     ) -> None:
         self.name = name
@@ -168,12 +173,14 @@ class JobStep:
         self.org = org
         self.user = user
         self.cloud_account = cloud_account or "" if self.studio is None else self.studio.cloud_account
+        self.cloud_provider = cloud_provider
         self.env = env
         self.interruptible = interruptible
         self.image_credentials = image_credentials
         self.cloud_account_auth = cloud_account_auth
         self.entrypoint = entrypoint
         self.path_mappings = path_mappings
+        self.max_runtime = max_runtime
         self.wait_for = wait_for
 
     def to_proto(
@@ -186,12 +193,16 @@ class JobStep:
             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)
+        resolved_cloud_account = CloudAccountApi().resolve_cloud_account(
+            teamspace.id, self.cloud_account, self.cloud_provider, teamspace.default_cloud_account
+        )
+
+        _validate_cloud_account(cloud_account, resolved_cloud_account, shared_filesystem)
 
         body = JobApiV2._create_job_body(
             name=self.name,
             command=self.command,
-            cloud_account=self.cloud_account or cloud_account,
+            cloud_account=resolved_cloud_account or cloud_account,
             studio_id=studio._studio.id if isinstance(studio, Studio) else None,
             image=self.image,
             machine=self.machine,
@@ -203,6 +214,7 @@ class JobStep:
             path_mappings=self.path_mappings,
             artifacts_local=None,
             artifacts_remote=None,
+            max_runtime=self.max_runtime,
         )
 
         return V1PipelineStep(
@@ -234,6 +246,7 @@ class MMTStep:
         cloud_account_auth: bool = False,
         entrypoint: str = "sh -c",
         path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
         wait_for: Optional[Union[str, List[str]]] = DEFAULT,
     ) -> None:
         self.machine = machine or Machine.CPU
@@ -258,6 +271,7 @@ class MMTStep:
         self.cloud_account_auth = cloud_account_auth
         self.entrypoint = entrypoint
         self.path_mappings = path_mappings
+        self.max_runtime = max_runtime
         self.wait_for = wait_for
 
     def to_proto(
@@ -288,6 +302,7 @@ class MMTStep:
             path_mappings=self.path_mappings,
             artifacts_local=None,  # deprecated in favor of path_mappings
             artifacts_remote=None,  # deprecated in favor of path_mappings
+            max_runtime=self.max_runtime,
         )
 
         return V1PipelineStep(

lightning_sdk/sandbox.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Optional, Type, Union
 
-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
@@ -31,6 +31,7 @@ class _Sandbox:
         org: The organization to use for the sandbox.
         user: The user to use for the sandbox.
         cloud_account: The cloud account to use for the sandbox.
+        cloud_provider: Selects the cloud account based on the available cloud accounts and the specified provider.
         disable_secrets: If true, user secrets such as LIGHTNING_API_KEY are not stored in the sandbox.
 
     Example:
@@ -48,6 +49,7 @@ class _Sandbox:
         org: Optional[Union[str, Organization]] = None,
         user: Optional[Union[str, User]] = None,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union[CloudProvider, str]] = None,
         disable_secrets: bool = True,
     ) -> None:
         if name is None:
@@ -62,6 +64,7 @@ class _Sandbox:
             org=org,
             user=user,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             disable_secrets=disable_secrets,
         )
 

lightning_sdk/serve.py

@@ -101,6 +101,8 @@ Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additio
 
 [bold]To push the container to a registry:[/bold]
 > [underline]docker push {tag}[/underline]
+
+Check out [blue][link=https://lightning.ai/docs/litserve/features]the docs[/link][/blue] for more details.
 """
             console.print(success_msg)
         return os.path.abspath("Dockerfile")

lightning_sdk/studio.py

@@ -1,23 +1,26 @@
 import glob
 import os
 import warnings
-from enum import Enum
 from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
 
 from tqdm.auto import tqdm
 
-from lightning_sdk.api.cluster_api import ClusterApi
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.api.studio_api import StudioApi
-from lightning_sdk.api.utils import _machine_to_compute_name
 from lightning_sdk.constants import _LIGHTNING_DEBUG
-from lightning_sdk.lightning_cloud.openapi import V1CloudSpaceSourceType
-from lightning_sdk.machine import Machine
+from lightning_sdk.machine import CloudProvider, Machine
 from lightning_sdk.organization import Organization
 from lightning_sdk.owner import Owner
 from lightning_sdk.status import Status
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace, _setup_logger
+from lightning_sdk.utils.resolve import (
+    _get_org_id,
+    _resolve_deprecated_cluster,
+    _resolve_deprecated_provider,
+    _resolve_teamspace,
+    _setup_logger,
+)
 
 if TYPE_CHECKING:
     from lightning_sdk.job import Job
@@ -27,19 +30,6 @@ if TYPE_CHECKING:
 _logger = _setup_logger(__name__)
 
 
-class Provider(Enum):
-    # Machine providers based on v1CloudProvider
-    AWS = "AWS"
-    GCP = "GCP"
-    VULTR = "VULTR"
-    LAMBDA_LABS = "LAMBDA_LABS"
-    DGX = "DGX"
-    VOLTAGE_PARK = "VOLTAGE_PARK"
-    NEBIUS = "NEBIUS"
-    CLOUDFLARE = "CLOUDFLARE"
-    LIGHTNING = "LIGHTNING"
-
-
 class Studio:
     """A single Lightning AI Studio.
 
@@ -53,6 +43,9 @@ class Studio:
         user: the name of the user owning the :param`teamspace` in case it is owned directly by a user instead of an org
         cloud_account: the name of the cloud account, the studio should be created on.
             Doesn't matter when the studio already exists.
+        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.
         create_ok: whether the studio will be created if it does not yet exist. Defaults to True
         provider: the provider of the machine, the studio should be created on.
 
@@ -72,34 +65,37 @@ class Studio:
         org: Optional[Union[str, Organization]] = None,
         user: Optional[Union[str, User]] = None,
         cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union[CloudProvider, str]] = None,
         create_ok: bool = True,
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
-        provider: Optional[str] = None,
-        source: Optional[V1CloudSpaceSourceType] = None,
+        source: Optional[str] = None,
         disable_secrets: bool = False,
+        provider: Optional[Union[CloudProvider, str]] = None,  # deprecated in favor of cloud_provider
     ) -> None:
         self._studio_api = StudioApi()
-        self._cluster_api = ClusterApi()
+        self._cloud_account_api = CloudAccountApi()
+
+        _teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+        if _teamspace is None:
+            raise ValueError("Couldn't resolve teamspace from the provided name, org, or user")
 
-        self._teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+        self._teamspace = _teamspace
         self._cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
+
         self._setup_done = False
         self._disable_secrets = disable_secrets
 
         self._plugins = {}
 
-        if self._teamspace is None:
-            raise ValueError("Couldn't resolve teamspace from the provided name, org, or user")
+        cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
+        cloud_provider = _resolve_deprecated_provider(cloud_provider, provider)
 
-        if provider is not None:
-            if isinstance(provider, str) and provider in Provider.__members__:
-                provider = Provider(provider)
-            else:
-                raise ValueError(f"Invalid provider: {provider}. Must be one of {Provider.__members__.keys()}.")
-            self._cloud_account = self._cluster_api.get_cluster_provider_mapping(
-                self._teamspace.id,
-                self._teamspace.owner.id,
-            )[provider.value]
+        self._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,
+        )
 
         if name is None:
             studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
@@ -121,6 +117,8 @@ class Studio:
                 else:
                     raise ValueError(f"Studio {name} does not exist.") from e
 
+        self._cloud_account = self._studio.cluster_id
+
         if (
             not self._skip_init
             and _internal_status_to_external_status(
@@ -178,7 +176,12 @@ class Studio:
         """Returns the current machine type the Studio is running on."""
         if self.status != Status.Running:
             return None
-        return self._studio_api.get_machine(self._studio.id, self._teamspace.id)
+        return self._studio_api.get_machine(
+            self._studio.id,
+            self._teamspace.id,
+            self.cloud_account,
+            _get_org_id(self._teamspace),
+        )
 
     @property
     def interruptible(self) -> bool:
@@ -198,8 +201,23 @@ class Studio:
     def cloud_account(self) -> str:
         return self._studio.cluster_id
 
-    def start(self, machine: Union[Machine, str] = Machine.CPU, interruptible: Optional[bool] = None) -> None:
-        """Starts a Studio on the specified machine type (default: CPU-4)."""
+    def start(
+        self,
+        machine: Union[Machine, str] = Machine.CPU,
+        interruptible: Optional[bool] = None,
+        max_runtime: Optional[int] = None,
+    ) -> None:
+        """Starts a Studio on the specified machine type (default: CPU-4).
+
+        Args:
+            machine: the machine type to start the studio on. Defaults to CPU-4
+            interruptible: whether to use interruptible machines
+            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
+
+        """
         status = self.status
 
         if interruptible is None:
@@ -210,10 +228,12 @@ class Studio:
                 interruptible = self.teamspace.start_studios_on_interruptible
 
         if status == Status.Running:
-            curr_machine = _machine_to_compute_name(self.machine) if self.machine is not None else None
-            if curr_machine != _machine_to_compute_name(machine):
+            new_machine = machine
+            if not isinstance(machine, Machine):
+                new_machine = Machine.from_str(machine)
+            if new_machine != self.machine:
                 raise RuntimeError(
-                    f"Requested to start studio on {machine}, but studio is already running on {self.machine}."
+                    f"Requested to start studio on {new_machine}, but studio is already running on {self.machine}."
                     " Consider switching instead!"
                 )
             _logger.info(f"Studio {self.name} is already running")
@@ -221,7 +241,9 @@ class Studio:
 
         if status != Status.Stopped:
             raise RuntimeError(f"Cannot start a studio that is not stopped. Studio {self.name} is {status}.")
-        self._studio_api.start_studio(self._studio.id, self._teamspace.id, machine, interruptible=interruptible)
+        self._studio_api.start_studio(
+            self._studio.id, self._teamspace.id, machine, interruptible=interruptible, max_runtime=max_runtime
+        )
 
         self._setup()
 
@@ -236,9 +258,34 @@ class Studio:
         """Deletes the current Studio."""
         self._studio_api.delete_studio(self._studio.id, self._teamspace.id)
 
-    def duplicate(self) -> "Studio":
-        """Duplicates the existing Studio to the same teamspace."""
-        kwargs = self._studio_api.duplicate_studio(self._studio.id, self._teamspace.id, self._teamspace.id)
+    def duplicate(self, target_teamspace: Optional[Union["Teamspace", str]] = None) -> "Studio":
+        """Duplicates the existing Studio.
+
+        Args:
+            target_teamspace: the teamspace to duplicate the studio to.
+                Must have the same owner as the source teamspace.
+                If not provided, defaults to current teamspace.
+        """
+        if target_teamspace is None:
+            target_teamspace_id = self._teamspace.id
+        else:
+            target_teamspace = _resolve_teamspace(
+                target_teamspace,
+                org=self._teamspace.owner if isinstance(self._teamspace.owner, Organization) else None,
+                user=self._teamspace.owner if isinstance(self._teamspace.owner, User) else None,
+            )
+
+            if target_teamspace is None:
+                raise ValueError(
+                    f"Could not resolve target teamspace {target_teamspace} "
+                    f"with owner {self.teamspace.owner} for duplication!"
+                )
+
+            target_teamspace_id = target_teamspace.id
+
+        kwargs = self._studio_api.duplicate_studio(
+            studio_id=self._studio.id, teamspace_id=self._teamspace.id, target_teamspace_id=target_teamspace_id
+        )
         return Studio(**kwargs)
 
     def switch_machine(self, machine: Union[Machine, str], interruptible: bool = False) -> None:

lightning_sdk/teamspace.py

@@ -107,8 +107,8 @@ class Teamspace:
         from lightning_sdk.studio import Studio
 
         studios = []
-        clusters = self._teamspace_api.list_cloud_accounts(teamspace_id=self.id)
-        for cl in clusters:
+        cloud_accounts = self._teamspace_api.list_cloud_accounts(teamspace_id=self.id)
+        for cl in cloud_accounts:
             _studios = self._teamspace_api.list_studios(teamspace_id=self.id, cloud_account=cl.cluster_id)
             for s in _studios:
                 studios.append(Studio(name=s.name, teamspace=self, cluster=cl.cluster_name, create_ok=False))
@@ -116,8 +116,11 @@ class Teamspace:
         return studios
 
     @property
-    def default_cloud_account(self) -> str:
-        return self._teamspace.project_settings.preferred_cluster
+    def default_cloud_account(self) -> Optional[str]:
+        owner_preferred_cluster = (
+            getattr(self.owner, "default_cloud_account", None) if isinstance(self.owner, Organization) else None
+        )
+        return self._teamspace.project_settings.preferred_cluster or owner_preferred_cluster
 
     @property
     def start_studios_on_interruptible(self) -> bool:
@@ -126,8 +129,8 @@ class Teamspace:
     @property
     def cloud_accounts(self) -> List[str]:
         """All cloud accounts associated with that teamspace."""
-        clusters = self._teamspace_api.list_cloud_accounts(teamspace_id=self.id)
-        return [cl.cluster_name for cl in clusters]
+        cloud_accounts = self._teamspace_api.list_cloud_accounts(teamspace_id=self.id)
+        return [cl.cluster_name for cl in cloud_accounts]
 
     @property
     def cloud_account_objs(self) -> List[V1ProjectClusterBinding]:
@@ -196,11 +199,17 @@ class Teamspace:
         if cloud_account is None:
             cloud_account = os.getenv("LIGHTNING_CLUSTER_ID") or self.default_cloud_account
 
-        cluster_machines = self._teamspace_api.list_machines(self.id, cloud_account=cloud_account)
+        if cloud_account is None:
+            raise RuntimeError("Could not resolve cloud account")
+
+        cloud_machines = self._teamspace_api.list_machines(self.id, cloud_account=cloud_account)
         return [
             Machine(
-                cluster_machine.instance_id,
-                cluster_machine.instance_id,
+                name=cluster_machine.instance_id,
+                slug=cluster_machine.slug_multi_cloud,
+                instance_type=cluster_machine.instance_id,
+                family=cluster_machine.family,
+                accelerator_count=cluster_machine.resources.gpu or cluster_machine.resources.cpu,
                 cost=cluster_machine.cost,
                 interruptible_cost=cluster_machine.spot_price,
                 wait_time=float(cluster_machine.available_in_seconds) if cluster_machine.available_in_seconds else None,
@@ -208,7 +217,7 @@ class Teamspace:
                 if cluster_machine.available_in_seconds_spot
                 else None,
             )
-            for cluster_machine in cluster_machines
+            for cluster_machine in cloud_machines
         ]
 
     def __eq__(self, other: "Teamspace") -> bool:

lightning_sdk/utils/resolve.py

@@ -6,7 +6,8 @@ from typing import TYPE_CHECKING, Generator, List, Optional, Tuple, Union
 
 from lightning_sdk.api import TeamspaceApi, UserApi
 from lightning_sdk.api.utils import _get_cloud_url
-from lightning_sdk.machine import Machine
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import CloudProvider, Machine
 
 if TYPE_CHECKING:
     from lightning_sdk.organization import Organization
@@ -48,6 +49,26 @@ def _resolve_deprecated_cloud_compute(machine: Machine, cloud_compute: Optional[
     return machine
 
 
+def _resolve_deprecated_provider(
+    cloud_provider: Optional[Union[CloudProvider, str]], provider: Optional[Union[CloudProvider, str]]
+) -> Optional[Union[CloudProvider, str]]:
+    if provider is not None:
+        if cloud_provider is not None:
+            raise ValueError(
+                "Cannot use both 'provider' and 'cloud_provider' at the same time."
+                "Please don't set the 'provider' as it will be deprecated!"
+            )
+
+        warnings.warn(
+            "The 'provider' argument will be deprecated in the future! "
+            "Please consider using the 'cloud_provider' argument instead!",
+            DeprecationWarning,
+        )
+        return provider
+
+    return cloud_provider
+
+
 def _resolve_deprecated_cluster(cloud_account: Optional[str], cluster: Optional[str]) -> Optional[str]:
     if cluster is not None:
         if cloud_account is not None:
@@ -85,7 +106,13 @@ def _resolve_org(org: Optional[Union[str, "Organization"]]) -> Optional["Organiz
 
     from lightning_sdk.organization import Organization
 
-    return Organization(name=org)
+    try:
+        return Organization(name=org)
+    # Handle case where user name is mistakenly used as organization name
+    except ApiException as ae:
+        if ae.status == 404:
+            raise ValueError(f"Organization '{org}' does not exist or you are not a member of it.") from ae
+        raise RuntimeError(f"Failed to resolve organization '{org}': {ae}") from ae
 
 
 def _resolve_user_name(name: Optional[str]) -> Optional[str]:
@@ -213,3 +240,11 @@ def _get_studio_url(studio: "Studio", turn_on: bool = False) -> str:
     if turn_on:
         return f"{base_url}?turnOn=true"
     return base_url
+
+
+def _get_org_id(teamspace: "Teamspace") -> str:
+    from lightning_sdk.organization import Organization
+
+    if isinstance(teamspace.owner, Organization):
+        return teamspace.owner.id
+    return ""

{lightning_sdk-2025.7.17.dist-info → lightning_sdk-2025.7.30rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 2025.7.17
+Version: 2025.7.30rc0
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License
@@ -25,6 +25,8 @@ License: MIT License
         OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
         SOFTWARE.
         
+Project-URL: Homepage, https://lightning.ai
+Project-URL: Documentation, https://lightning.ai/docs/overview/sdk-reference
 Keywords: deep learning,machine learning,pytorch,AI
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -79,11 +81,11 @@ s.start()
 # prints Machine.CPU-4
 print(s.machine)
 
-# or start directly on this machine with s.start(Machine.A10G)
+# or start directly on this machine with s.start(Machine.L4)
 print("switching Studio machine...")
-s.switch_machine(Machine.A10G)
+s.switch_machine(Machine.L4)
 
-# prints Machine.A10G
+# prints Machine.L4
 print(s.machine)
 
 # prints Status.Running
@@ -96,7 +98,7 @@ s.install_plugin("jobs")
 s.install_plugin("multi-machine-training")
 
 # run the resulting plugins to start 1 job and 1 multi-machine training
-s.installed_plugins["jobs"].run("python my_dummy_file", name="my_first_job", machine=Machine.A10G)
+s.installed_plugins["jobs"].run("python my_dummy_file", name="my_first_job", machine=Machine.L4)
 s.installed_plugins["multi-machine-training"].run("python my_dummy_file", name="my_first_mmt", machine=Machine.T4, num_instances=42)
 
 print("Stopping Studio")