lightning-sdk 0.1.49__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/machine.py

@@ -1,32 +1,223 @@
-from lightning_sdk.utils.enum import DeprecationEnum
-
-
-class Machine(DeprecationEnum):
-    """Enum holding all supported machine types for studios."""
-
-    CPU_SMALL = "CPU_SMALL"
-    CPU = "CPU"
-    DATA_PREP = "DATA_PREP"
-    DATA_PREP_MAX = "DATA_PREP_MAX"
-    DATA_PREP_ULTRA = "DATA_PREP_ULTRA"
-    T4 = "T4"
-    T4_X_4 = "T4_X_4"
-    L4 = "L4"
-    L4_X_4 = "L4_X_4"
-    L4_X_8 = "L4_X_8"
-    A10G = "A10G"
-    A10G_X_4 = "A10G_X_4"
-    A10G_X_8 = "A10G_X_8"
-    L40S = "L40S"
-    L40 = "L40", "L40S"  # deprecated
-    L40S_X_4 = "L40S_X_4"
-    L40_X_4 = "L40_X_4", "L40S_X_4"  # deprecated
-    L40S_X_8 = "L40S_X_8"
-    L40_X_8 = "L40_X_8", "L40S_X_8"  # deprecated
-    A100_X_8 = "A100_X_8"
-    H100_X_8 = "H100_X_8"
-    H200_X_8 = "H200_X_8"
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, ClassVar, Optional, Tuple
+
+
+class CloudProvider(Enum):
+    AWS = "AWS"
+    GCP = "GCP"
+    LAMBDA_LABS = "LAMBDA_LABS"
+    DGX = "DGX"
+    VOLTAGE_PARK = "VOLTAGE_PARK"
+    NEBIUS = "NEBIUS"
+    LIGHTNING = "LIGHTNING"
+
+    def __str__(self) -> str:
+        """Converts the CloudProvider to a str."""
+        return self.value
+
+
+@dataclass(frozen=True)
+class Machine:
+    # supported CPU variations
+    CPU_X_2: ClassVar["Machine"]
+    CPU_X_4: ClassVar["Machine"]
+    CPU_X_8: ClassVar["Machine"]
+    CPU_X_16: ClassVar["Machine"]
+    # default CPU machines
+    CPU_SMALL: ClassVar["Machine"]
+    CPU: ClassVar["Machine"]
+    # supported data-prep variations (big disk)
+    DATA_PREP: ClassVar["Machine"]
+    DATA_PREP_MAX: ClassVar["Machine"]
+    DATA_PREP_ULTRA: ClassVar["Machine"]
+
+    # supported GPU types
+    # supported T4 variations
+    T4_SMALL: ClassVar["Machine"]
+    T4: ClassVar["Machine"]
+    T4_X_2: ClassVar["Machine"]
+    T4_X_4: ClassVar["Machine"]
+    T4_X_8: ClassVar["Machine"]
+    # supported L4 variations
+    L4: ClassVar["Machine"]
+    L4_X_2: ClassVar["Machine"]
+    L4_X_4: ClassVar["Machine"]
+    L4_X_8: ClassVar["Machine"]
+    # supported L40S variations
+    L40S: ClassVar["Machine"]
+    L40S_X_2: ClassVar["Machine"]
+    L40S_X_4: ClassVar["Machine"]
+    L40S_X_8: ClassVar["Machine"]
+    # supported A100 variations
+    # defaults, can be either A100 type depending on cloud provider availability
+    A100: ClassVar["Machine"]
+    A100_X_2: ClassVar["Machine"]
+    A100_X_4: ClassVar["Machine"]
+    A100_X_8: ClassVar["Machine"]
+    # A100 40GB versions
+    A100_40GB: ClassVar["Machine"]
+    A100_40GB_X_2: ClassVar["Machine"]
+    A100_40GB_X_4: ClassVar["Machine"]
+    A100_40GB_X_8: ClassVar["Machine"]
+    # A100 80GB versions
+    A100_80GB: ClassVar["Machine"]
+    A100_80GB_X_2: ClassVar["Machine"]
+    A100_80GB_X_4: ClassVar["Machine"]
+    A100_80GB_X_8: ClassVar["Machine"]
+
+    H100: ClassVar["Machine"]
+    H100_X_2: ClassVar["Machine"]
+    H100_X_4: ClassVar["Machine"]
+    H100_X_8: ClassVar["Machine"]
+
+    H200: ClassVar["Machine"]
+    H200_X_8: ClassVar["Machine"]
+    B200_X_8: ClassVar["Machine"]
+
+    # Specialized Machines
+
+    name: str
+    slug: str
+    instance_type: Optional[str] = None
+    family: Optional[str] = None
+    accelerator_count: Optional[int] = None
+    cost: Optional[float] = None
+    interruptible_cost: Optional[float] = None
+    wait_time: Optional[float] = None
+    interruptible_wait_time: Optional[float] = None
+    _include_in_cli: bool = True
 
     def __str__(self) -> str:
-        """String representation of the enum."""
-        return str(self.value)
+        """String representation of the Machine."""
+        return str(self.name) if self.name else (self.slug if self.slug else str(self.instance_type))
+
+    def __eq__(self, other: object) -> bool:
+        """Machines are equal if the instance type is equal."""
+        if isinstance(other, Machine):
+            return (
+                # equality based on raw instance type (provider specific)
+                (self.instance_type and self.instance_type == other.instance_type)
+                # equality based on slug (provider agnostic)
+                or self.slug == other.slug
+                # equality based on machine specs (e.g. A100_80GB_X_8 == A100_X_8)
+                or (self.family == other.family and self.accelerator_count == other.accelerator_count)
+            )
+        return False
+
+    def is_cpu(self) -> bool:
+        """Whether the machine is a CPU."""
+        return self.family in ("CPU", "DATA_PREP")
+
+    @classmethod
+    def from_str(cls, machine: str, *additional_machine_ids: Any) -> "Machine":
+        possible_values: Tuple["Machine", ...] = tuple(
+            [machine for machine in cls.__dict__.values() if isinstance(machine, cls)]
+        )
+        for m in possible_values:
+            for machine_id in [machine, *additional_machine_ids]:
+                if machine_id and machine_id in (
+                    getattr(m, "name", None),
+                    getattr(m, "instance_type", None),
+                    getattr(m, "slug", None),
+                ):
+                    return m
+
+        if additional_machine_ids:
+            return cls(machine, *additional_machine_ids)
+        return cls(machine, machine, machine)
+
+    @classmethod
+    def _from_accelerator(cls, accelerator: Any) -> "Machine":
+        if accelerator.accelerator_type == "GPU":
+            accelerator_resources_count = accelerator.resources.gpu
+        else:
+            accelerator_resources_count = accelerator.resources.cpu
+
+        return Machine.from_str(
+            accelerator.slug_multi_cloud,
+            accelerator.slug,
+            accelerator.instance_id,
+            accelerator.secondary_instance_id,
+            f"lit-{accelerator.family.lower()}-{accelerator_resources_count}",
+        )
+
+
+# CPU machines
+# default CPU machines
+Machine.CPU_SMALL = Machine(name="CPU_SMALL", slug="cpu-2", family="CPU", accelerator_count=2)
+Machine.CPU = Machine(name="CPU", slug="cpu-4", family="CPU", accelerator_count=4)
+# available CPU variations
+Machine.CPU_X_2 = Machine(name="CPU_X_2", slug="cpu-2", family="CPU", accelerator_count=2)
+Machine.CPU_X_4 = Machine(name="CPU_X_4", slug="cpu-4", family="CPU", accelerator_count=4)
+Machine.CPU_X_8 = Machine(name="CPU_X_8", slug="cpu-8", family="CPU", accelerator_count=8)
+Machine.CPU_X_16 = Machine(name="CPU_X_16", slug="cpu-16", family="CPU", accelerator_count=16)
+# available data-prep (big disk) machines
+Machine.DATA_PREP = Machine(name="DATA_PREP", slug="data-prep-mid", family="DATA_PREP", accelerator_count=32)
+Machine.DATA_PREP_MAX = Machine(
+    name="DATA_PREP_MAX", slug="data-prep-max-large", family="DATA_PREP", accelerator_count=64
+)
+Machine.DATA_PREP_ULTRA = Machine(
+    name="DATA_PREP_ULTRA", slug="data-prep-ultra-extra-large", family="DATA_PREP", accelerator_count=96
+)
+
+# GPU machines
+# available T4 machines
+Machine.T4_SMALL = Machine(name="T4_SMALL", slug="lit-t4-1-small", family="T4", accelerator_count=1)
+Machine.T4 = Machine(name="T4", slug="lit-t4-1", family="T4", accelerator_count=1)
+Machine.T4_X_2 = Machine(name="T4_X_2", slug="lit-t4-2", family="T4", accelerator_count=2)
+Machine.T4_X_4 = Machine(name="T4_X_4", slug="lit-t4-4", family="T4", accelerator_count=4)
+Machine.T4_X_8 = Machine(name="T4_X_8", slug="lit-t4-8", family="T4", accelerator_count=8)
+# available L4 machines
+Machine.L4 = Machine(name="L4", slug="lit-l4-1", family="L4", accelerator_count=1)
+Machine.L4_X_2 = Machine(name="L4_X_2", slug="lit-l4-2", family="L4", accelerator_count=2)
+Machine.L4_X_4 = Machine(name="L4_X_4", slug="lit-l4-4", family="L4", accelerator_count=4)
+Machine.L4_X_8 = Machine(name="L4_X_8", slug="lit-l4-8", family="L4", accelerator_count=8)
+# available L40S machines
+Machine.L40S = Machine(name="L40S", slug="lit-l40s-1", family="L40S", accelerator_count=1)
+Machine.L40S_X_2 = Machine(name="L40S_X_2", slug="lit-l40s-2", family="L40S", accelerator_count=2)
+Machine.L40S_X_4 = Machine(name="L40S_X_4", slug="lit-l40s-4", family="L40S", accelerator_count=4)
+Machine.L40S_X_8 = Machine(name="L40S_X_8", slug="lit-l40s-8", family="L40S", accelerator_count=8)
+# available A100 Machines
+Machine.A100 = Machine(name="A100", slug="lit-a100-1", family="A100", accelerator_count=1)
+Machine.A100_X_2 = Machine(name="A100_X_2", slug="lit-a100-2", family="A100", accelerator_count=2)
+Machine.A100_X_4 = Machine(name="A100_X_4", slug="lit-a100-4", family="A100", accelerator_count=4)
+Machine.A100_X_8 = Machine(name="A100_X_8", slug="lit-a100-8", family="A100", accelerator_count=8)
+# don't include variants in cli, only default types that can match for all variants
+Machine.A100_40GB = Machine(
+    name="A100_40GB", slug="lit-a100-40gb-1", family="A100", accelerator_count=1, _include_in_cli=False
+)
+Machine.A100_40GB_X_2 = Machine(
+    name="A100_40GB_X_2", slug="lit-a100-40gb-2", family="A100", accelerator_count=2, _include_in_cli=False
+)
+Machine.A100_40GB_X_4 = Machine(
+    name="A100_40GB_X_4", slug="lit-a100-40gb-4", family="A100", accelerator_count=4, _include_in_cli=False
+)
+Machine.A100_40GB_X_8 = Machine(
+    name="A100_40GB_X_8", slug="lit-a100-40gb-8", family="A100", accelerator_count=8, _include_in_cli=False
+)
+Machine.A100_80GB = Machine(
+    name="A100_80GB", slug="lit-a100-80gb-1", family="A100", accelerator_count=1, _include_in_cli=False
+)
+Machine.A100_80GB_X_2 = Machine(
+    name="A100_80GB_X_2", slug="lit-a100-80gb-2", family="A100", accelerator_count=2, _include_in_cli=False
+)
+Machine.A100_80GB_X_4 = Machine(
+    name="A100_80GB_X_4", slug="lit-a100-80gb-4", family="A100", accelerator_count=4, _include_in_cli=False
+)
+Machine.A100_80GB_X_8 = Machine(
+    name="A100_80GB_X_8", slug="lit-a100-80gb-8", family="A100", accelerator_count=8, _include_in_cli=False
+)
+# available H100 machines
+Machine.H100 = Machine(name="H100", slug="lit-h100-1", family="H100", accelerator_count=1)
+Machine.H100_X_2 = Machine(name="H100_X_2", slug="lit-h100-2", family="H100", accelerator_count=2)
+Machine.H100_X_4 = Machine(name="H100_X_4", slug="lit-h100-4", family="H100", accelerator_count=4)
+Machine.H100_X_8 = Machine(name="H100_X_8", slug="lit-h100-8", family="H100", accelerator_count=8)
+# available H200 machines
+Machine.H200 = Machine(name="H200", slug="lit-h200x-1", family="H200", accelerator_count=1)
+Machine.H200_X_8 = Machine(name="H200_X_8", slug="lit-h200x-8", family="H200", accelerator_count=8)
+# available B200 machines
+Machine.B200_X_8 = Machine(name="B200_X_8", slug="lit-b200x-8", family="B200", accelerator_count=8)
+
+
+DEFAULT_MACHINE = Machine.CPU.name

lightning_sdk/mmt/base.py

@@ -1,9 +1,10 @@
+import warnings
 from abc import abstractmethod
 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
@@ -11,7 +12,7 @@ if TYPE_CHECKING:
     from lightning_sdk.user import User
 
 from lightning_sdk.job.base import _BaseJob
-from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, in_studio
 
 
 class MMTMachine(Protocol):
@@ -23,7 +24,7 @@ class MMTMachine(Protocol):
         ...
 
     @property
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """The actual machine type this node is running on."""
         ...
 
@@ -54,7 +55,7 @@ class _BaseMMT(_BaseJob):
     def run(
         cls,
         name: str,
-        machine: "Machine",
+        machine: Union["Machine", str],
         num_machines: int,
         command: Optional[str] = None,
         studio: Union["Studio", str, None] = None,
@@ -63,14 +64,18 @@ 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,
-        artifacts_local: Optional[str] = None,
-        artifacts_remote: Optional[str] = None,
         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
+        reuse_snapshot: bool = True,
     ) -> "_BaseMMT":
         """Run async workloads using a docker image across multiple machines.
 
@@ -87,28 +92,38 @@ 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.
                 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.
                 Required if the registry is part of a cloud provider (e.g. ECR).
-            artifacts_local: The path of inside the docker container, you want to persist images from.
-                CAUTION: When setting this to "/", it will effectively erase your container.
-                Only supported for jobs with a docker image compute environment.
-            artifacts_remote: The remote storage to persist your artifacts to.
-                Should be of format <CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>.
-                PATH_WITHIN_CONNECTION hereby is a path relative to the connection's root.
-                E.g. efs:data:some-path would result in an EFS connection named `data` and to the path `some-path`
-                within it.
-                Note that the connection needs to be added to the teamspace already in order for it to be found.
-                Only supported for jobs with a docker image compute environment.
             entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
                 just runs the provided command in a standard shell.
                 To use the pre-defined entrypoint of the provided image, set this to an empty string.
                 Only applicable when submitting docker jobs.
+            path_mappings: Dictionary of path mappings. The keys are the path inside the container whereas the value
+                represents the data-connection name and the path inside that connection.
+                Should be of form
+                    {
+                        "<CONTAINER_PATH_1>": "<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1>",
+                        "<CONTAINER_PATH_2>": "<CONNECTION_NAME_2>"
+                    }
+                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
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
         """
+        from lightning_sdk.lightning_cloud.openapi.rest import ApiException
         from lightning_sdk.studio import Studio
 
         cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
@@ -167,6 +182,13 @@ class _BaseMMT(_BaseJob):
                     "image and studio are mutually exclusive as both define the environment to run the job in"
                 )
 
+            if cloud_account is None and in_studio():
+                try:
+                    resolve_studio = Studio(teamspace=teamspace, user=user, org=org)
+                    cloud_account = resolve_studio.cloud_account
+                except (ValueError, ApiException):
+                    warnings.warn("Could not infer cloud account from studio. Using teamspace default.")
+
             # they either need to specified both or none of them
             if bool(artifacts_local) != bool(artifacts_remote):
                 raise ValueError("Artifact persistence requires both artifacts_local and artifacts_remote to be set")
@@ -182,6 +204,7 @@ class _BaseMMT(_BaseJob):
             num_machines=num_machines,
             machine=machine,
             cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
             command=command,
             studio=studio,
             image=image,
@@ -189,9 +212,12 @@ class _BaseMMT(_BaseJob):
             interruptible=interruptible,
             image_credentials=image_credentials,
             cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
-            entrypoint=entrypoint,
+            max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
         )
         return inst
 
@@ -199,18 +225,22 @@ class _BaseMMT(_BaseJob):
     def _submit(
         self,
         num_machines: int,
-        machine: "Machine",
+        machine: Union["Machine", str],
         command: Optional[str] = None,
         studio: Optional["Studio"] = None,
         image: Optional[str] = None,
         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,
-        artifacts_local: Optional[str] = None,
-        artifacts_remote: Optional[str] = None,
         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,
+        reuse_snapshot: bool = True,
     ) -> None:
         """Submit a new multi-machine job to the Lightning AI platform.
 
@@ -230,19 +260,24 @@ class _BaseMMT(_BaseJob):
                 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.
                 Required if the registry is part of a cloud provider (e.g. ECR).
-            artifacts_local: The path of inside the docker container, you want to persist images from.
-                CAUTION: When setting this to "/", it will effectively erase your container.
-                Only supported for jobs with a docker image compute environment.
-            artifacts_remote: The remote storage to persist your artifacts to.
-                Should be of format <CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>.
-                PATH_WITHIN_CONNECTION hereby is a path relative to the connection's root.
-                E.g. efs:data:some-path would result in an EFS connection named `data` and to the path `some-path`
-                within it.
-                Note that the connection needs to be added to the teamspace already in order for it to be found.
-                Only supported for jobs with a docker image compute environment.
             entrypoint: The entrypoint of your docker container. Defaults to sh -c.
                 To use the pre-defined entrypoint of the provided image, set this to an empty string.
                 Only applicable when submitting docker jobs.
+            path_mappings: Dictionary of path mappings. The keys are the path inside the container whereas the value
+                represents the data-connection name and the path inside that connection.
+                Should be of form
+                    {
+                        "<CONTAINER_PATH_1>": "<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1>",
+                        "<CONTAINER_PATH_2>": "<CONNECTION_NAME_2>"
+                    }
+                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
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
         """
 
     @property
@@ -257,7 +292,7 @@ class _BaseMMT(_BaseJob):
 
     @property
     @abstractmethod
-    def machine(self) -> "Machine":
+    def machine(self) -> Union["Machine", str]:
         """Returns the machine type this job is running on."""
 
     @abstractmethod