lightning-sdk 2025.8.19.post0__py3-none-any.whl → 2025.8.26__py3-none-any.whl

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
 
@@ -14,6 +15,7 @@ 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.names import random_unique_name
 from lightning_sdk.utils.resolve import (
     _get_org_id,
     _resolve_deprecated_cluster,
@@ -56,8 +58,11 @@ class Studio:
     """
 
     # skips init of studio, only set when using this as a shell for names, ids etc.
-    _skip_init = False
-    _skip_setup = False
+    _skip_init = threading.local()
+    _skip_setup = threading.local()
+
+    # whether to show progress bars during operations
+    show_progress = False
 
     def __init__(
         self,
@@ -76,64 +81,77 @@ class Studio:
         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._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._teamspace = _teamspace
 
-        self._setup_done = self._skip_setup
+        self._setup_done = getattr(self._skip_setup, "value", False)
         self._disable_secrets = disable_secrets
 
         self._plugins = {}
+        self._studio = None
 
         cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
         cloud_provider = _resolve_deprecated_provider(cloud_provider, provider)
 
-        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,
+            )
 
         # Resolve studio name if not provided: explicit → env (LIGHTNING_CLOUD_SPACE_ID) → config defaults
-        if name is None:
+        if name is None and not getattr(self._skip_init, "value", False):
             studio_id = os.environ.get("LIGHTNING_CLOUD_SPACE_ID", None)
             if studio_id is not None:
                 # We're inside a studio, get it by ID
                 self._studio = self._studio_api.get_studio_by_id(studio_id=studio_id, teamspace_id=self._teamspace.id)
+                name = self._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:
+                if name is None and not create_ok:
                     raise ValueError(
                         "Cannot autodetect Studio. Either use the SDK from within a Studio or pass a name!"
                     )
 
-        # If we have a name (explicit or from config), get studio by name
-        if name is not None:
+        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(
+                        "Cannot autodetect Studio. Either use the SDK from within a Studio 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,
                     )
                 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)
             )
@@ -254,9 +272,26 @@ class Studio:
 
         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
-        )
+
+        # 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, 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, machine, interruptible=interruptible, max_runtime=max_runtime
+            )
 
         self._setup()
 
@@ -324,9 +359,22 @@ class Studio:
             raise RuntimeError(
                 f"Cannot switch machine on a studio that is not running. Studio {self.name} is {status}."
             )
-        self._studio_api.switch_studio_machine(
-            self._studio.id, self._teamspace.id, machine, interruptible=interruptible
-        )
+
+        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
+                )
+        else:
+            self._studio_api.switch_studio_machine(
+                self._studio.id, self._teamspace.id, machine, interruptible=interruptible
+            )
 
     def run_and_detach(self, *commands: str, timeout: float = 10, check_interval: float = 1) -> str:
         """Runs given commands on the Studio and returns immediately.
@@ -575,6 +623,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."""

lightning_sdk/teamspace.py

@@ -2,7 +2,7 @@ import glob
 import os
 import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
 from tqdm.auto import tqdm
 
@@ -21,6 +21,7 @@ from lightning_sdk.utils.resolve import (
     _resolve_org,
     _resolve_teamspace_name,
     _resolve_user,
+    skip_studio_init,
 )
 
 if TYPE_CHECKING:
@@ -125,7 +126,11 @@ class Teamspace:
         for cl in cloud_accounts:
             _studios = self._teamspace_api.list_studios(teamspace_id=self.id, cloud_account=cl.cluster_id)
             for s in _studios:
-                studios.append(Studio(name=s.name, teamspace=self, cluster=cl.cluster_name, create_ok=False))
+                with skip_studio_init():
+                    studio = Studio(name=s.name, teamspace=self, cluster=cl.cluster_name, create_ok=False)
+                    studio._studio = s
+                    studio._teamspace = self
+                    studios.append(studio)
 
         return studios
 
@@ -209,6 +214,24 @@ class Teamspace:
 
         return tuple(mmts)
 
+    @property
+    def secrets(self) -> Dict[str, str]:
+        """All (encrypted) secrets for the teamspace.
+
+        Note:
+            Once created, the secret values are encrypted and cannot be viewed here anymore.
+        """
+        return self._teamspace_api.get_secrets(self.id)
+
+    def set_secret(self, key: str, value: str) -> None:
+        """Set the (encrypted) secrets for the teamspace."""
+        if not self._teamspace_api.verify_secret_name(key):
+            raise ValueError(
+                "Secret keys must only contain alphanumeric characters and underscores and not begin with a number."
+            )
+
+        self._teamspace_api.set_secret(self.id, key, value)
+
     def list_machines(self, cloud_account: Optional[str] = None) -> List[Machine]:
         if cloud_account is None:
             cloud_account = os.getenv("LIGHTNING_CLUSTER_ID") or self.default_cloud_account

lightning_sdk/user.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Dict, Optional
 
 from lightning_sdk.api import UserApi
 from lightning_sdk.owner import Owner
@@ -37,6 +37,24 @@ class User(Owner):
         """The user's ID."""
         return self._user.id
 
+    @property
+    def secrets(self) -> Dict[str, str]:
+        """All (encrypted) secrets for the user.
+
+        Note:
+            Once created, the secret values are encrypted and cannot be viewed here anymore.
+        """
+        return self._user_api.get_secrets()
+
+    def set_secret(self, key: str, value: str) -> None:
+        """Set a (encrypted) secret for the user."""
+        if not self._user_api.verify_secret_name(key):
+            raise ValueError(
+                "Secret keys must only contain alphanumeric characters and underscores and not begin with a number."
+            )
+
+        self._user_api.set_secret(key, value)
+
     def __repr__(self) -> str:
         """Returns reader friendly representation."""
         return f"User(name={self.name})"

lightning_sdk/utils/config.py

@@ -144,6 +144,12 @@ class Config:
             sort_keys=True,
         )
 
+    def get(self, key: str) -> Optional[str]:
+        return self.get_value(key)
+
+    def set(self, key: str, value: str) -> None:
+        self._set_nested([key], value)
+
 
 def _unflatten_dict(flat_dict: Dict[str, Any]) -> Dict[str, Any]:
     unflattened_dict = {}