lightning-sdk 2025.12.16__py3-none-any.whl → 2026.1.22__py3-none-any.whl

lightning_sdk/machine.py

@@ -162,12 +162,12 @@ Machine.CPU_X_4 = Machine(name="CPU_X_4", slug="cpu-4", family="CPU", accelerato
 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 = 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
+    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
+    name="DATA_PREP_ULTRA", slug="data-prep-ultra-extra-large", family="DATA-PREP", accelerator_count=96
 )
 
 # GPU machines

lightning_sdk/mmt/base.py

@@ -69,7 +69,7 @@ class _BaseMMT(_BaseJob):
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
-        entrypoint: str = "sh -c",
+        entrypoint: Optional[str] = None,
         path_mappings: Optional[Dict[str, str]] = None,
         max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
@@ -104,8 +104,10 @@ class _BaseMMT(_BaseJob):
             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).
             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.
+                just runs the provided command in a standard shell if a command is provided.
+                If no command is provided, it will run the pre-defined entrypoint of the provided image.
+                To use the pre-defined entrypoint of the provided image with a specified command,
+                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.
@@ -173,7 +175,7 @@ class _BaseMMT(_BaseJob):
                     "Other jobs will automatically persist artifacts to the teamspace distributed filesystem."
                 )
 
-            if entrypoint != "sh -c":
+            if entrypoint is not None:
                 raise ValueError("Specifying the entrypoint has no effect for jobs with Studio envs.")
 
         else:
@@ -199,6 +201,20 @@ class _BaseMMT(_BaseJob):
                     f"<CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>, got {artifacts_local}"
                 )
 
+            # command specified, so use the default entrypoint of sh -c
+            if command is not None and entrypoint is None:
+                entrypoint = "sh -c"
+
+            # entrypoint specifically set to empty string, so set to None here to fall back to the image entrypoint
+            elif entrypoint == "":  # noqa: SIM114
+                entrypoint = None
+
+            # entrypoint not specified, but also no command specified, so use the image entrypoint
+            elif entrypoint is None:
+                entrypoint = None
+
+            # all other cases, the entrypoint has been specifically set, so use it as is
+
         inst = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False)
         inst._submit(
             num_machines=num_machines,
@@ -235,7 +251,7 @@ class _BaseMMT(_BaseJob):
         cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
-        entrypoint: str = "sh -c",
+        entrypoint: Optional[str] = None,
         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
@@ -260,8 +276,11 @@ 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).
-            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.
+            entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
+                just runs the provided command in a standard shell if a command is provided.
+                If no command is provided, it will run the pre-defined entrypoint of the provided image.
+                To use the pre-defined entrypoint of the provided image with a specified command,
+                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.

lightning_sdk/mmt/mmt.py

@@ -100,7 +100,7 @@ class MMT(_BaseMMT):
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
-        entrypoint: str = "sh -c",
+        entrypoint: Optional[str] = None,
         path_mappings: Optional[Dict[str, str]] = None,
         max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,
@@ -134,8 +134,11 @@ class MMT(_BaseMMT):
                 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).
-            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.
+            entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
+                just runs the provided command in a standard shell if a command is provided.
+                If no command is provided, it will run the pre-defined entrypoint of the provided image.
+                To use the pre-defined entrypoint of the provided image with a specified command,
+                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.
@@ -194,7 +197,7 @@ class MMT(_BaseMMT):
         cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
-        entrypoint: str = "sh -c",
+        entrypoint: Optional[str] = None,
         path_mappings: Optional[Dict[str, str]] = None,
         max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
@@ -224,8 +227,10 @@ class MMT(_BaseMMT):
             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).
             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.
+                just runs the provided command in a standard shell if a command is provided.
+                If no command is provided, it will run the pre-defined entrypoint of the provided image.
+                To use the pre-defined entrypoint of the provided image with a specified command,
+                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.

lightning_sdk/mmt/v1.py

@@ -53,7 +53,7 @@ class _MMTV1(_BaseMMT):
         cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
-        entrypoint: str = "sh -c",
+        entrypoint: Optional[str] = None,
         path_mappings: Optional[Dict[str, str]] = None,
         max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,
@@ -89,7 +89,10 @@ class _MMTV1(_BaseMMT):
                 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.
+                just runs the provided command in a standard shell if a command is provided.
+                If no command is provided, it will run the pre-defined entrypoint of the provided image.
+                To use the pre-defined entrypoint of the provided image with a specified command,
+                set this to an empty string.
                 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)

lightning_sdk/mmt/v2.py

@@ -53,7 +53,7 @@ class _MMTV2(_BaseMMT):
         cloud_provider: Optional[Union["CloudProvider", str]] = None,
         image_credentials: Optional[str] = None,
         cloud_account_auth: bool = False,
-        entrypoint: str = "sh -c",
+        entrypoint: Optional[str] = None,
         path_mappings: Optional[Dict[str, str]] = None,
         max_runtime: Optional[int] = None,
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
@@ -83,7 +83,10 @@ class _MMTV2(_BaseMMT):
             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).
             entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
-                just runs the provided command in a standard shell.
+                just runs the provided command in a standard shell if a command is provided.
+                If no command is provided, it will run the pre-defined entrypoint of the provided image.
+                To use the pre-defined entrypoint of the provided image with a specified command,
+                set this to an empty string.
                 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

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/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
 
@@ -11,8 +11,8 @@ from lightning_sdk.api.studio_api import StudioApi
 from lightning_sdk.api.utils import AccessibleResource, raise_access_error_if_not_allowed
 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.exceptions import NotSupportedError, OutOfCapacityError
+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
@@ -341,12 +341,20 @@ class Studio(metaclass=TrackCallsMeta):
             if new_machine != self.machine:
                 raise RuntimeError(
                     f"Requested to start {self._cls_name} on {new_machine}, "
-                    "but {self._cls_name} is already running on {self.machine}."
+                    f"but {self._cls_name} is already running on {self.machine}."
                     " Consider switching instead!"
                 )
             _logger.info(f"{self._cls_name} {self.name} is already running")
             return
 
+        if not self._studio_api.machine_is_supported(
+            new_machine, self._teamspace.id, self.cloud_account, _get_org_id(self._teamspace)
+        ):
+            raise NotSupportedError(
+                "Requested machine is not supported in the selected cloud account. "
+                "Try a different machine or cloud account by setting the `machine` or `cloud_account` argument."
+            )
+
         if not self._studio_api.machine_has_capacity(
             new_machine,
             self._teamspace.id,
@@ -548,7 +556,7 @@ class Studio(metaclass=TrackCallsMeta):
         if status != Status.Running:
             raise RuntimeError(
                 f"Cannot run a command in a {self._cls_name} that is not running. "
-                "{self._cls_name} {self.name} is {status}."
+                f"{self._cls_name} {self.name} is {status}."
             )
         output, exit_code = self._studio_api.run_studio_commands(self._studio.id, self._teamspace.id, *commands)
         output = output.strip()
@@ -605,11 +613,13 @@ class Studio(metaclass=TrackCallsMeta):
             progress_bar = tqdm(total=len(all_files), desc="Uploading files", unit="file")
         for local_file, remote_path in sorted(all_files, key=lambda p: p[1]):
             if progress_bar:
-                progress_bar.set_description(f"Uploading {local_file}")
+                progress_bar.set_description("Uploading files")
             self.upload_file(local_file, remote_path=remote_path, progress_bar=False)
             if progress_bar:
                 progress_bar.update(1)
         if progress_bar:
+            progress_bar.set_description("Upload complete")
+            progress_bar.refresh()
             progress_bar.close()
 
     def download_file(self, remote_path: str, file_path: Optional[str] = None) -> None:
@@ -707,6 +717,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/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/logging.py

@@ -22,7 +22,7 @@ def track_calls() -> Callable[..., any]:
             bound_args.apply_defaults()
 
             args_str = ", ".join(f"{k}: {v}" for k, v in bound_args.arguments.items() if k != "self")
-            message = f"VERSION: {__version__} | ARGS: {args_str} "
+            message = f"ARGS: {args_str} "
 
             body = V1CreateSDKCommandHistoryRequest(
                 command=func.__qualname__,
@@ -30,6 +30,7 @@ def track_calls() -> Callable[..., any]:
                 project_id=None,
                 severity=V1SDKCommandHistorySeverity.INFO,
                 type=V1SDKCommandHistoryType.SDK,
+                version=__version__,
             )
 
             try:

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.16.dist-info → lightning_sdk-2026.1.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 2025.12.16
+Version: 2026.1.22
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License
@@ -38,8 +38,6 @@ Requires-Dist: click
 Requires-Dist: docker
 Requires-Dist: fastapi
 Requires-Dist: packaging
-Requires-Dist: pandas
-Requires-Dist: pyjwt
 Requires-Dist: pyyaml
 Requires-Dist: requests
 Requires-Dist: rich
@@ -48,8 +46,6 @@ Requires-Dist: six
 Requires-Dist: tqdm
 Requires-Dist: urllib3<=2.5.0
 Requires-Dist: uvicorn
-Requires-Dist: websocket-client
-Requires-Dist: wget
 Provides-Extra: serve
 Requires-Dist: litserve>=0.2.5; extra == "serve"