lightning-sdk 2025.12.9__py3-none-any.whl → 2025.12.17__py3-none-any.whl

lightning_sdk/organization.py

@@ -1,9 +1,12 @@
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from lightning_sdk.api import OrgApi
 from lightning_sdk.owner import Owner
 from lightning_sdk.utils.resolve import _resolve_org_name
 
+if TYPE_CHECKING:
+    from lightning_sdk.teamspace import Teamspace
+
 
 class Organization(Owner):
     """Represents an organization owner of teamspaces and studios.
@@ -44,6 +47,12 @@ class Organization(Owner):
     def default_cloud_account(self) -> Optional[str]:
         return self._org.preferred_cluster or None
 
+    def create_teamspace(self, name: str) -> "Teamspace":
+        from lightning_sdk.teamspace import Teamspace
+
+        self._org_api.create_teamspace(name, self.id)
+        return Teamspace(name=name, org=self)
+
     def __repr__(self) -> str:
         """Returns reader friendly representation."""
         return f"Organization(name={self.name})"

lightning_sdk/owner.py

@@ -24,6 +24,10 @@ class Owner(ABC, metaclass=TrackCallsABCMeta):
     def id(self) -> str:
         """The owner's ID."""
 
+    @abstractmethod
+    def create_teamspace(self, name: str) -> "Teamspace":
+        """Creates a new teamspace."""
+
     @property
     def teamspaces(self) -> List["Teamspace"]:
         """All teamspaces by this owner."""

lightning_sdk/pipeline/steps.py

@@ -176,6 +176,7 @@ class JobStep:
         max_runtime: Optional[int] = None,
         wait_for: Union[str, List[str], None] = DEFAULT,
         reuse_snapshot: bool = True,
+        scratch_disks: Optional[Dict[str, int]] = None,
     ) -> None:
         self.name = name
         self.machine = machine or Machine.CPU
@@ -203,6 +204,7 @@ class JobStep:
         self.max_runtime = max_runtime
         self.wait_for = wait_for
         self.reuse_snapshot = reuse_snapshot
+        self.scratch_disks = scratch_disks
 
     def to_proto(
         self, teamspace: "Teamspace", cloud_account: str, shared_filesystem: Union[bool, V1SharedFilesystem]
@@ -242,6 +244,7 @@ class JobStep:
             max_runtime=self.max_runtime,
             machine_image_version=machine_image_version,
             reuse_snapshot=self.reuse_snapshot,
+            scratch_disks=self.scratch_disks,
         )
 
         return V1PipelineStep(

lightning_sdk/plugin.py

@@ -411,12 +411,12 @@ class CustomPortPlugin(_Plugin):
         if name is None:
             name = _run_name("port")
 
-        return self._studio._studio_api.start_new_port(
+        return self._studio._studio_api.add_port(
             teamspace_id=self._studio._teamspace.id,
             studio_id=self._studio._studio.id,
             name=name,
             port=port,
-        )
+        ).urls[0]
 
 
 @runtime_checkable

lightning_sdk/studio.py

@@ -2,7 +2,7 @@ import glob
 import os
 import threading
 import warnings
-from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union
 
 from tqdm.auto import tqdm
 
@@ -12,7 +12,7 @@ from lightning_sdk.api.utils import AccessibleResource, raise_access_error_if_no
 from lightning_sdk.base_studio import BaseStudio
 from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.exceptions import OutOfCapacityError
-from lightning_sdk.lightning_cloud.openapi import V1ClusterType
+from lightning_sdk.lightning_cloud.openapi import V1ClusterType, V1Endpoint
 from lightning_sdk.machine import DEFAULT_MACHINE, CloudProvider, Machine
 from lightning_sdk.organization import Organization
 from lightning_sdk.owner import Owner
@@ -707,6 +707,37 @@ class Studio(metaclass=TrackCallsMeta):
             interruptible=interruptible,
         )
 
+    def add_ports(self, ports: Union[int, List[int], Dict[str, int]]) -> List[V1Endpoint]:
+        """Add one or more ports to the studio and return their endpoints.
+
+        Args:
+            ports: Port to add. Can be:
+                - int: Single port (e.g., 8080)
+                - List[int]: Multiple ports ()
+                - dict[str, int]: Named ports (e.g., {"web": 8080})
+
+        Returns:
+            List of V1Endpoint objects. Access endpoint properties like:
+                - endpoint.name: Port name (None for unnamed ports)
+                - endpoint.ports: List of port numbers
+                - endpoint.urls: List of accessible URLs
+        """
+        if isinstance(ports, dict):
+            port_items = ports.items()
+        elif isinstance(ports, list):
+            port_items = ((None, port) for port in ports)
+        else:
+            port_items = [(None, ports)]
+
+        return [
+            self._studio_api.add_port(self._teamspace.id, self._studio.id, name=name, port=port)
+            for name, port in port_items
+        ]
+
+    def list_ports(self) -> List[V1Endpoint]:
+        """List ports that are exposed in the Studio."""
+        return self._studio_api.list_ports(self._teamspace.id, self._studio.id)
+
     def create_assistant(self, name: str, port: int) -> None:
         assistant = self._studio_api.create_assistant(
             studio_id=self._studio.id, teamspace_id=self._teamspace.id, port=port, assistant_name=name

lightning_sdk/teamspace.py

@@ -11,7 +11,12 @@ import lightning_sdk
 from lightning_sdk.agents import Agent
 from lightning_sdk.api import CloudAccountApi, TeamspaceApi
 from lightning_sdk.api.utils import AccessibleResource, raise_access_error_if_not_allowed
-from lightning_sdk.lightning_cloud.openapi import V1ClusterType, V1Model, V1ModelVersionArchive, V1ProjectClusterBinding
+from lightning_sdk.lightning_cloud.openapi import (
+    V1ClusterType,
+    V1Model,
+    V1ModelVersionArchive,
+    V1ProjectClusterBinding,
+)
 from lightning_sdk.machine import CloudProvider, Machine
 from lightning_sdk.models import UploadedModelInfo
 from lightning_sdk.organization import Organization
@@ -275,14 +280,48 @@ class Teamspace(metaclass=TrackCallsMeta):
 
         self._teamspace_api.set_secret(self.id, key, value)
 
-    def list_machines(self, cloud_account: Optional[str] = None) -> List[Machine]:
+    def list_machines(self, cloud_account: Optional[str] = None, machine: Optional[str] = None) -> List[Machine]:
+        """List available machines across cloud accounts.
+
+        Args:
+            cloud_account: The cloud account from which to list available machines. If None, uses LIGHTNING_CLUSTER_ID
+                environment variable. If that's also None, queries all global cloud accounts.
+            machine: Specific machine name to filter by. If provided, only returns that
+                machine type. Must be a valid Machine enum value.
+
+        Returns:
+            List of available machines, excluding out-of-capacity machines.
+        """
         if cloud_account is None:
-            cloud_account = os.getenv("LIGHTNING_CLUSTER_ID") or self.default_cloud_account
+            cloud_account = os.getenv("LIGHTNING_CLUSTER_ID", None)
 
+        # if cloud_account is not given as a paramter and as a env var, use global cloud_accounts
         if cloud_account is None:
+            global_cloud_accounts = self._cloud_account_api.list_global_cloud_accounts(teamspace_id=self.id)
+            cloud_accounts = [cm.id for cm in global_cloud_accounts]
+        else:
+            cloud_accounts = [cloud_account]
+
+        if cloud_accounts is None:
             raise RuntimeError("Could not resolve cloud account")
 
-        cloud_machines = self._teamspace_api.list_machines(self.id, cloud_account=cloud_account)
+        if machine:
+            _machine_values = tuple(
+                [
+                    machine.name
+                    for machine in Machine.__dict__.values()
+                    if isinstance(machine, Machine) and machine._include_in_cli
+                ]
+            )
+            if machine not in _machine_values:
+                raise ValueError(f"Machine '{machine}' is not valid. Valid machines are: {_machine_values}")
+            machine = getattr(Machine, machine.upper(), Machine(machine, machine))
+
+        cloud_machines = self._teamspace_api.list_machines(
+            self.id, cloud_accounts=cloud_accounts, machine=machine, org_id=self._org.id
+        )
+        # filter out of capacity machines
+        cloud_machines = [cm for cm in cloud_machines if not cm.out_of_capacity]
         return [
             Machine(
                 name=cluster_machine.instance_id,
@@ -292,6 +331,7 @@ class Teamspace(metaclass=TrackCallsMeta):
                 accelerator_count=cluster_machine.resources.gpu or cluster_machine.resources.cpu,
                 cost=cluster_machine.cost,
                 interruptible_cost=cluster_machine.spot_price,
+                provider=cluster_machine.provider,
                 wait_time=float(cluster_machine.available_in_seconds) if cluster_machine.available_in_seconds else None,
                 interruptible_wait_time=float(cluster_machine.available_in_seconds_spot)
                 if cluster_machine.available_in_seconds_spot

lightning_sdk/user.py

@@ -1,8 +1,12 @@
-from typing import Dict, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional
 
 from lightning_sdk.api import UserApi
 from lightning_sdk.owner import Owner
-from lightning_sdk.utils.resolve import _resolve_user_name
+from lightning_sdk.utils.resolve import _get_authed_user, _get_organizations_for_authed_user, _resolve_user_name
+
+if TYPE_CHECKING:
+    from lightning_sdk.organization import Organization
+    from lightning_sdk.teamspace import Teamspace
 
 
 class User(Owner):
@@ -55,6 +59,22 @@ class User(Owner):
 
         self._user_api.set_secret(key, value)
 
+    def create_teamspace(self, name: str) -> "Teamspace":
+        from lightning_sdk.teamspace import Teamspace
+
+        if not _get_authed_user().id == self.id:
+            raise ValueError("Can only create teamspaces for currently authenticated user")
+
+        self._user_api.create_teamspace(name)
+        return Teamspace(name=name, user=self)
+
+    @property
+    def organizations(self) -> List["Organization"]:
+        if not _get_authed_user().id == self.id:
+            raise ValueError("Can only list organizations for currently authenticated user")
+
+        return _get_organizations_for_authed_user(user_api=self._user_api)
+
     def __repr__(self) -> str:
         """Returns reader friendly representation."""
         return f"User(name={self.name})"

lightning_sdk/utils/resolve.py

@@ -2,6 +2,7 @@ import logging
 import os
 import warnings
 from contextlib import contextmanager
+from functools import lru_cache
 from typing import TYPE_CHECKING, Generator, List, Optional, Tuple, Union
 
 from lightning_sdk.api import TeamspaceApi, UserApi
@@ -222,25 +223,26 @@ def _resolve_teamspace(
     raise RuntimeError("Neither user nor org provided, but one of them needs to be provided")
 
 
-def _get_organizations_for_authed_user() -> List["Organization"]:
+def _get_organizations_for_authed_user(user_api: Optional[UserApi] = None) -> List["Organization"]:
     """Returns Organizations the current Authed user is a member of."""
     from lightning_sdk.organization import Organization
 
-    _orgs = UserApi()._get_organizations_for_authed_user()
+    _orgs = (user_api or UserApi())._get_organizations_for_authed_user()
     return [Organization(_org.name) for _org in _orgs]
 
 
-def _get_teamspace_names_for_authed_user() -> List[str]:
+def _get_teamspace_names_for_authed_user(user_api: Optional[UserApi] = None) -> List[str]:
     """Returns Teamspace's names the current Authed user is a member of."""
-    teamspaces = UserApi()._get_all_teamspace_memberships("")
+    teamspaces = (user_api or UserApi())._get_all_teamspace_memberships("")
     return sorted([ts.name for ts in teamspaces])
 
 
-def _get_authed_user() -> "User":
+@lru_cache(maxsize=1)
+def _get_authed_user(user_api: Optional[UserApi] = None, teamspace_api: Optional[TeamspaceApi] = None) -> "User":
     from lightning_sdk.user import User
 
-    user_id = TeamspaceApi()._get_authed_user_id()
-    _user = UserApi()._get_user_by_id(user_id)
+    user_id = (teamspace_api or TeamspaceApi())._get_authed_user_id()
+    _user = (user_api or UserApi())._get_user_by_id(user_id)
     return User(name=_user.username)
 
 

{lightning_sdk-2025.12.9.dist-info → lightning_sdk-2025.12.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 2025.12.9
+Version: 2025.12.17
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License