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

lightning_sdk/machine.py

@@ -13,60 +13,101 @@ class CloudProvider(Enum):
     NEBIUS = "NEBIUS"
     LIGHTNING = "LIGHTNING"
 
+    def __str__(self) -> str:
+        """Converts the CloudProvider to a str."""
+        return self.value
+
 
 @dataclass(frozen=True)
 class Machine:
-    # Default Machines
-    CPU: ClassVar["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: 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"]
-    A10G: ClassVar["Machine"]
-    A10G_X_4: ClassVar["Machine"]
-    A10G_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"]
-    B200_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
-    instance_type: 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 Machine."""
-        return str(self.name) if self.name else str(self.instance_type)
+        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 self.instance_type == other.instance_type
+            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 == Machine.CPU
-            or self == Machine.CPU_SMALL
-            or self == Machine.DATA_PREP
-            or self == Machine.DATA_PREP_MAX
-            or self == Machine.DATA_PREP_ULTRA
-        )
+        return self.family in ("CPU", "DATA_PREP")
 
     @classmethod
     def from_str(cls, machine: str, *additional_machine_ids: Any) -> "Machine":
@@ -75,34 +116,89 @@ class Machine:
         )
         for m in possible_values:
             for machine_id in [machine, *additional_machine_ids]:
-                if machine_id in (getattr(m, "name", None), getattr(m, "instance_type", None)):
+                if 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.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")
-Machine.A10G_X_4 = Machine(name="A10G_X_4", instance_type="g5.12xlarge")
-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="p5en.48xlarge")
+        return cls(machine, machine, machine)
+
+
+# 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 = 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)

lightning_sdk/mmt/v2.py

@@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 from lightning_sdk.api.mmt_api import MMTApiV2
 from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.status import Status
+from lightning_sdk.utils.resolve import _get_org_id
 
 if TYPE_CHECKING:
     from lightning_sdk.job.job import Job
@@ -191,7 +192,11 @@ class _MMTV2(_BaseMMT):
     @property
     def machine(self) -> Union["Machine", str]:
         """Returns the machine type this job is running on."""
-        return self._job_api._get_job_machine_from_spec(self._guaranteed_job.spec)
+        return self._job_api._get_job_machine_from_spec(
+            self._guaranteed_job.spec,
+            self.teamspace.id,
+            _get_org_id(self.teamspace),
+        )
 
     def _update_internal_job(self) -> None:
         if getattr(self, "_job", None) is None:

lightning_sdk/models.py

@@ -127,7 +127,7 @@ def download_model(
         return api.download_model_files(
             name=model_name,
             version=version,
-            download_dir=Path(download_dir),
+            download_dir=Path(download_dir).expanduser().resolve(),
             teamspace_name=teamspace_name,
             teamspace_owner_name=teamspace_owner_name,
             progress_bar=progress_bar,

lightning_sdk/studio.py

@@ -7,7 +7,6 @@ from tqdm.auto import tqdm
 
 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.machine import CloudProvider, Machine
 from lightning_sdk.organization import Organization
@@ -16,6 +15,7 @@ from lightning_sdk.status import Status
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
 from lightning_sdk.utils.resolve import (
+    _get_org_id,
     _resolve_deprecated_cluster,
     _resolve_deprecated_provider,
     _resolve_teamspace,
@@ -176,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:
@@ -223,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")

lightning_sdk/teamspace.py

@@ -205,8 +205,11 @@ class Teamspace:
         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,

lightning_sdk/utils/resolve.py

@@ -240,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.22.dist-info → lightning_sdk-2025.7.31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 2025.7.22
+Version: 2025.7.31
 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")