lightning-sdk 2025.8.6rc2__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/plugin.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Generator, Optional, Protocol, Union, run
 from lightning_sdk.job import Job
 from lightning_sdk.machine import Machine
 from lightning_sdk.studio import Studio
+from lightning_sdk.utils.logging import TrackCallsABCMeta
 from lightning_sdk.utils.resolve import (
     _LIGHTNING_SERVICE_EXECUTION_ID_KEY,
     _resolve_deprecated_cloud_compute,
@@ -22,7 +23,7 @@ if TYPE_CHECKING:
 _logger = _setup_logger(__name__)
 
 
-class _Plugin(ABC):
+class _Plugin(ABC, metaclass=TrackCallsABCMeta):
     """Abstract Plugin class defining the API.
 
     Args:
@@ -128,6 +129,7 @@ class JobsPlugin(_Plugin):
         machine: Machine = Machine.CPU,
         cloud_compute: Optional[Machine] = None,
         interruptible: bool = False,
+        reuse_snapshot: bool = True,
     ) -> Job:
         """Launches an asynchronous job.
 
@@ -137,6 +139,8 @@ class JobsPlugin(_Plugin):
             machine: The machine to run the job on.
             interruptible: Whether to run the job on an interruptible machine.
                 These are cheaper but can be preempted at any time.
+            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.
         """
         if not name:
             name = _run_name("job")
@@ -151,6 +155,7 @@ class JobsPlugin(_Plugin):
             teamspace=self._studio.teamspace,
             cloud_account=self._studio.cloud_account,
             interruptible=interruptible,
+            reuse_snapshot=reuse_snapshot,
         )
 
 

lightning_sdk/studio.py

@@ -1,5 +1,6 @@
 import glob
 import os
+import threading
 import warnings
 from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
 
@@ -7,13 +8,17 @@ 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.base_studio import BaseStudio
 from lightning_sdk.constants import _LIGHTNING_DEBUG
-from lightning_sdk.machine import CloudProvider, Machine
+from lightning_sdk.lightning_cloud.openapi import V1ClusterType
+from lightning_sdk.machine import DEFAULT_MACHINE, 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.logging import TrackCallsMeta
+from lightning_sdk.utils.names import random_unique_name
 from lightning_sdk.utils.resolve import (
     _get_org_id,
     _resolve_deprecated_cluster,
@@ -30,7 +35,7 @@ if TYPE_CHECKING:
 _logger = _setup_logger(__name__)
 
 
-class Studio:
+class Studio(metaclass=TrackCallsMeta):
     """A single Lightning AI Studio.
 
     Allows to fully control a studio, including retrieving the status, running commands
@@ -45,9 +50,11 @@ class Studio:
             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.
+            If not specified, falls back to 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.
+        studio_type: Type of studio to create. Only effective during initial creation;
+            ignored for existing studios.
 
     Note:
         Since a teamspace can either be owned by an org or by a user directly,
@@ -56,7 +63,11 @@ class Studio:
     """
 
     # skips init of studio, only set when using this as a shell for names, ids etc.
-    _skip_init = False
+    _skip_init = threading.local()
+    _skip_setup = threading.local()
+
+    # whether to show progress bars during operations
+    show_progress = False
 
     def __init__(
         self,
@@ -71,56 +82,117 @@ class Studio:
         source: Optional[str] = None,
         disable_secrets: bool = False,
         provider: Optional[Union[CloudProvider, str]] = None,  # deprecated in favor of cloud_provider
+        studio_type: Optional[str] = None,  # for base studio templates
     ) -> None:
         self._studio_api = StudioApi()
         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._prevent_refetch = False
+        self._teamspace = None
+
+        # don't resolve anything if we're skipping init
+        if not getattr(self._skip_init, "value", False):
+            _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 = _teamspace
-        self._cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
+            self._teamspace = _teamspace
 
-        self._setup_done = False
+        self._setup_done = getattr(self._skip_setup, "value", False)
         self._disable_secrets = disable_secrets
 
         self._plugins = {}
+        self._studio = None
+
+        # Check to see if we're inside a studio
+        current_studio = None
+        studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
+        if studio_id is not None and self._teamspace is not None:
+            # We're inside a studio, get it by ID
+            current_studio = self._studio_api.get_studio_by_id(studio_id=studio_id, teamspace_id=self._teamspace.id)
+
+        if cloud_account or not cloud_provider:
+            cloud_account = _resolve_deprecated_cluster(
+                cloud_account, cluster, current_studio.cluster_id if current_studio else None
+            )
+            cloud_provider = _resolve_deprecated_provider(cloud_provider, provider)
+        else:
+            cloud_provider = _resolve_deprecated_provider(cloud_provider, provider)
 
-        cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
-        cloud_provider = _resolve_deprecated_provider(cloud_provider, provider)
+        cls_name = self._cls_name
 
-        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 we're skipping init, we don't need to resolve the cloud account as then we're not creating a studio
+        if self._teamspace is not None:
+            _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)
-            if studio_id is None:
-                raise ValueError("Cannot autodetect Studio. Either use the SDK from within a Studio or pass a name!")
-            self._studio = self._studio_api.get_studio_by_id(studio_id=studio_id, teamspace_id=self._teamspace.id)
+        self._studio_type = None
+        if studio_type:
+            self._base_studio = BaseStudio(teamspace=self._teamspace)
+            self._available_base_studios = self._base_studio.list()
+            for bst in self._available_base_studios:
+                if (
+                    bst.id == studio_type
+                    or bst.name == studio_type
+                    or bst.name.lower().replace(" ", "-") == studio_type
+                ):
+                    self._studio_type = bst.id
+
+            if not self._studio_type:
+                raise ValueError(
+                    f"Could not find studio type with ID or name '{studio_type}'. "
+                    f"Available studio types: "
+                    f"{[bst.name.lower().replace(' ', '-') for bst in self._available_base_studios]}"
+                )
         else:
+            if current_studio:
+                self._studio_type = current_studio.environment_template_id
+
+        # Resolve studio name if not provided: explicit → env (LIGHTNING_CLOUD_SPACE_ID) → config defaults
+        if name is None and not getattr(self._skip_init, "value", False):
+            if current_studio:
+                name = current_studio.name
+            else:
+                # Try config defaults
+                from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+                config = Config()
+                name = config.get_value(DefaultConfigKeys.studio)
+                if name is None and not create_ok:
+                    raise ValueError(
+                        f"Cannot autodetect {cls_name}. Either use the SDK from within a {cls_name} or pass a name!"
+                    )
+
+        if self._studio is None and not getattr(self._skip_init, "value", False):
+            # If we have a name (explicit or from config), get studio by name
             try:
+                if name is None:
+                    # if we don't have a name, raise an error to get
+                    # to the exception path and optionally create a studio
+                    raise ValueError(
+                        f"Cannot autodetect {cls_name}. Either use the SDK from within a {cls_name} or pass a name!"
+                    )
                 self._studio = self._studio_api.get_studio(name, self._teamspace.id)
             except ValueError as e:
                 if create_ok:
+                    name = name or random_unique_name()
                     self._studio = self._studio_api.create_studio(
                         name,
                         self._teamspace.id,
-                        cloud_account=self._cloud_account,
+                        cloud_account=_cloud_account,
                         source=source,
                         disable_secrets=self._disable_secrets,
+                        cloud_space_environment_template_id=self._studio_type,
                     )
                 else:
-                    raise ValueError(f"Studio {name} does not exist.") from e
-
-        self._cloud_account = self._studio.cluster_id
+                    raise e
 
         if (
-            not self._skip_init
+            not getattr(self._skip_init, "value", False)
             and _internal_status_to_external_status(
                 self._studio_api._get_studio_instance_status_from_object(self._studio)
             )
@@ -183,6 +255,14 @@ class Studio:
             _get_org_id(self._teamspace),
         )
 
+    @property
+    def public_ip(self) -> Optional[str]:
+        """Returns the public IP address of the machine the Studio is running on."""
+        return self._studio_api.get_public_ip(
+            self._studio.id,
+            self._teamspace.id,
+        )
+
     @property
     def interruptible(self) -> bool:
         """Returns whether the Studio is running on a interruptible instance."""
@@ -194,7 +274,9 @@ class Studio:
     @property
     def cluster(self) -> str:
         """Returns the cluster the Studio is running on."""
-        warnings.warn("Studio.cluster is deprecated. Use Studio.cloud_account instead", DeprecationWarning)
+        warnings.warn(
+            f"{self._cls_name}.cluster is deprecated. Use {self._cls_name}.cloud_account instead", DeprecationWarning
+        )
         return self.cloud_account
 
     @property
@@ -203,7 +285,7 @@ class Studio:
 
     def start(
         self,
-        machine: Union[Machine, str] = Machine.CPU,
+        machine: Optional[Union[Machine, str]] = None,
         interruptible: Optional[bool] = None,
         max_runtime: Optional[int] = None,
     ) -> None:
@@ -218,6 +300,22 @@ class Studio:
                 Defaults to 3h
 
         """
+        # Check to see if we're inside a studio and if its running
+        current_studio_machine = None
+        studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
+        if studio_id is not None:
+            # We're inside a studio, get the machine if it is running
+            current_studio = self._studio_api.get_studio_by_id(studio_id=studio_id, teamspace_id=self._teamspace.id)
+            current_status = self._studio_api._get_studio_instance_status_from_object(current_studio)
+
+            if current_status and _internal_status_to_external_status(current_status) == Status.Running:
+                current_studio_machine = self._studio_api.get_machine(
+                    current_studio.id,
+                    self._teamspace.id,
+                    current_studio.cluster_id,
+                    _get_org_id(self._teamspace),
+                )
+
         status = self.status
 
         if interruptible is None:
@@ -227,23 +325,53 @@ class Studio:
             else:
                 interruptible = self.teamspace.start_studios_on_interruptible
 
-        if status == Status.Running:
+        new_machine = DEFAULT_MACHINE
+        if machine is not None:
             new_machine = machine
-            if not isinstance(machine, Machine):
-                new_machine = Machine.from_str(machine)
+        elif current_studio_machine is not None:
+            new_machine = current_studio_machine
+
+        if not isinstance(new_machine, Machine):
+            new_machine = Machine.from_str(new_machine)
+
+        if status == Status.Running:
             if new_machine != self.machine:
                 raise RuntimeError(
-                    f"Requested to start studio on {new_machine}, but studio is already running on {self.machine}."
+                    f"Requested to start {self._cls_name} on {new_machine}, "
+                    "but {self._cls_name} is already running on {self.machine}."
                     " Consider switching instead!"
                 )
-            _logger.info(f"Studio {self.name} is already running")
+            _logger.info(f"{self._cls_name} {self.name} is already running")
             return
 
         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, max_runtime=max_runtime
-        )
+            raise RuntimeError(
+                f"Cannot start a {self._cls_name} that is not stopped. {self._cls_name} {self.name} is {status}."
+            )
+
+        # Show progress bar during startup
+        if self.show_progress:
+            from lightning_sdk.utils.progress import StudioProgressTracker
+
+            with StudioProgressTracker("start", show_progress=True) as progress:
+                # Start the studio without blocking
+                self._studio_api.start_studio_async(
+                    self._studio.id,
+                    self._teamspace.id,
+                    new_machine,
+                    interruptible=interruptible,
+                    max_runtime=max_runtime,
+                )
+
+                # Track progress through completion
+                progress.track_startup_phases(
+                    lambda: self._studio_api.get_studio_status(self._studio.id, self._teamspace.id)
+                )
+        else:
+            # Use the blocking version if no progress is needed
+            self._studio_api.start_studio(
+                self._studio.id, self._teamspace.id, new_machine, interruptible=interruptible, max_runtime=max_runtime
+            )
 
         self._setup()
 
@@ -258,13 +386,20 @@ class Studio:
         """Deletes the current Studio."""
         self._studio_api.delete_studio(self._studio.id, self._teamspace.id)
 
-    def duplicate(self, target_teamspace: Optional[Union["Teamspace", str]] = None) -> "Studio":
+    def duplicate(
+        self,
+        target_teamspace: Optional[Union["Teamspace", str]] = None,
+        machine: Machine = Machine.CPU,
+        name: Optional[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.
+            machine: the machine to start the duplicated studio on.
+                Defaults to CPU
         """
         if target_teamspace is None:
             target_teamspace_id = self._teamspace.id
@@ -284,16 +419,23 @@ class Studio:
             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
+            studio_id=self._studio.id,
+            teamspace_id=self._teamspace.id,
+            target_teamspace_id=target_teamspace_id,
+            machine=machine,
+            new_name=name,
         )
         return Studio(**kwargs)
 
-    def switch_machine(self, machine: Union[Machine, str], interruptible: bool = False) -> None:
+    def switch_machine(
+        self, machine: Union[Machine, str], interruptible: bool = False, cloud_provider: Optional[CloudProvider] = None
+    ) -> None:
         """Switches machine to the provided machine type/.
 
         Args:
             machine: the new machine type to switch to
             interruptible: determines whether to switch to an interruptible instance
+            cloud_provider: the cloud provider to switch to, has no effect if the Studio is not on Lightning Cloud
 
         Note:
             this call is blocking until the new machine is provisioned
@@ -302,12 +444,49 @@ class Studio:
         status = self.status
         if status != Status.Running:
             raise RuntimeError(
-                f"Cannot switch machine on a studio that is not running. Studio {self.name} is {status}."
+                f"Cannot switch machine on a {self._cls_name} that is not running. "
+                "{self._cls_name} {self.name} is {status}."
             )
-        self._studio_api.switch_studio_machine(
-            self._studio.id, self._teamspace.id, machine, interruptible=interruptible
+
+        current_cloud = self._cloud_account_api.get_cloud_account_non_org(
+            self._teamspace.id,
+            self._studio.cluster_id,
         )
 
+        cloud_account = ""
+        if cloud_provider is not None and current_cloud.spec.cluster_type == V1ClusterType.GLOBAL:
+            cloud_account = self._cloud_account_api.resolve_cloud_account(
+                self._teamspace.id,
+                cloud_account=None,
+                cloud_provider=cloud_provider,
+                default_cloud_account=None,
+            )
+
+        if self.show_progress:
+            from lightning_sdk.utils.progress import StudioProgressTracker
+
+            with StudioProgressTracker("switch", show_progress=True) as progress:
+                # Update progress before starting the switch
+                progress.update_progress(5, "Initiating machine switch...")
+
+                # Start the switch operation with progress tracking
+                self._studio_api.switch_studio_machine_with_progress(
+                    self._studio.id,
+                    self._teamspace.id,
+                    machine,
+                    interruptible=interruptible,
+                    progress=progress,
+                    cloud_account=cloud_account,
+                )
+        else:
+            self._studio_api.switch_studio_machine(
+                self._studio.id, self._teamspace.id, machine, interruptible=interruptible, cloud_account=cloud_account
+            )
+
+        if self._studio and cloud_account:
+            # TODO: get this from the API
+            self._studio.cluster_id = cloud_account
+
     def run_and_detach(self, *commands: str, timeout: float = 10, check_interval: float = 1) -> str:
         """Runs given commands on the Studio and returns immediately.
 
@@ -324,7 +503,10 @@ class Studio:
             print(f"Running {commands=}")
         status = self.status
         if status != Status.Running:
-            raise RuntimeError(f"Cannot run a command in a studio that is not running. Studio {self.name} is {status}.")
+            raise RuntimeError(
+                f"Cannot run a command in a {self._cls_name} that is not running. "
+                "{self._cls_name} {self.name} is {status}."
+            )
 
         iter_output = self._studio_api.run_studio_commands_and_yield(
             self._studio.id, self._teamspace.id, *commands, timeout=timeout, check_interval=check_interval
@@ -350,7 +532,10 @@ class Studio:
 
         status = self.status
         if status != Status.Running:
-            raise RuntimeError(f"Cannot run a command in a studio that is not running. Studio {self.name} is {status}.")
+            raise RuntimeError(
+                f"Cannot run a command in a {self._cls_name} that is not running. "
+                "{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()
 
@@ -446,6 +631,7 @@ class Studio:
         command: str,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
+        reuse_snapshot: bool = True,
     ) -> "Job":
         """Run async workloads using the compute environment from your studio.
 
@@ -455,6 +641,8 @@ class Studio:
             command: The command to run inside your job.
             env: Environment variables to set inside the job.
             interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+            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.job import Job
 
@@ -468,6 +656,7 @@ class Studio:
             cloud_account=self.cloud_account,
             env=env,
             interruptible=interruptible,
+            reuse_snapshot=reuse_snapshot,
         )
 
     def run_mmt(
@@ -512,6 +701,14 @@ class Studio:
         self._assistant_id = assistant.id
         _logger.info(assistant_info)
 
+    def rename(self, new_name: str) -> None:
+        """Renames the current Studio to the provided new name."""
+        if new_name == self._studio.name:
+            return
+
+        self._studio_api._update_cloudspace(self._studio, self._teamspace.id, "display_name", new_name)
+        self._update_studio_reference()
+
     @property
     def auto_sleep(self) -> bool:
         """Returns if a Studio has auto-sleep enabled."""
@@ -520,8 +717,8 @@ class Studio:
     @auto_sleep.setter
     def auto_sleep(self, value: bool) -> None:
         if not value and self.machine == Machine.CPU:
-            warnings.warn("Disabling auto-sleep will convert the Studio from free to paid!")
-        self._studio_api.update_autoshutdown(self._studio.id, self._teamspace.id, enabled=value, studio=self._studio)
+            warnings.warn(f"Disabling auto-sleep will convert the {self._cls_name} from free to paid!")
+        self._studio_api.update_autoshutdown(self._studio.id, self._teamspace.id, enabled=value)
         self._update_studio_reference()
 
     @property
@@ -531,10 +728,8 @@ class Studio:
 
     @auto_sleep_time.setter
     def auto_sleep_time(self, value: int) -> None:
-        warnings.warn("Setting auto-sleep time will convert the Studio from free to paid!")
-        self._studio_api.update_autoshutdown(
-            self._studio.id, self._teamspace.id, idle_shutdown_seconds=value, studio=self._studio
-        )
+        warnings.warn(f"Setting auto-sleep time will convert the {self._cls_name} from free to paid!")
+        self._studio_api.update_autoshutdown(self._studio.id, self._teamspace.id, idle_shutdown_seconds=value)
         self._update_studio_reference()
 
     @property
@@ -557,6 +752,22 @@ class Studio:
         warnings.warn("auto_shutdown_time is deprecated. Use auto_sleep_time instead", DeprecationWarning)
         self.auto_sleep_time = value
 
+    @property
+    def env(self) -> Dict[str, str]:
+        self._update_studio_reference()
+        return self._studio_api.get_env(self._studio)
+
+    def set_env(self, new_env: Dict[str, str], partial: bool = True) -> None:
+        """Set the environment variables for the Studio.
+
+        Args:
+            new_env: The new environment variables to set.
+            partial: Whether to only set the environment variables that are provided.
+                If False, existing environment variables that are not in new_env will be removed.
+                If True, existing environment variables that are not in new_env will be kept.
+        """
+        self._studio_api.set_env(self._studio, self._teamspace.id, new_env, partial=partial)
+
     @property
     def available_plugins(self) -> Mapping[str, str]:
         """All available plugins to install in the current Studio."""
@@ -648,6 +859,36 @@ class Studio:
     def _update_studio_reference(self) -> None:
         self._studio = self._studio_api.get_studio_by_id(studio_id=self._studio.id, teamspace_id=self._teamspace.id)
 
+    @property
+    def _cls_name(self) -> str:
+        return self.__class__.__qualname__
+
+
+class VM(Studio):
+    """A single Lightning AI VM.
+
+    Allows to fully control a vm, including retrieving the status, running commands
+    and switching machine types.
+
+    Args:
+        name: the name of the vm
+        teamspace: the name of the teamspace the vm is contained by
+        org: the name of the organization owning the :param`teamspace` in case it is owned by an org
+        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 vm should be created on.
+            Doesn't matter when the vm 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 vm will be created if it does not yet exist. Defaults to True
+        provider: the provider of the machine, the vm should be created on.
+
+    Note:
+        Since a teamspace can either be owned by an org or by a user directly,
+        only one of the arguments can be provided.
+
+    """
+
 
 def _internal_status_to_external_status(internal_status: str) -> Status:
     """Converts internal status strings from HTTP requests to external enums."""