lightning-sdk 0.1.49__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/utils/progress.py

@@ -0,0 +1,283 @@
+"""Studio startup/switch progress bar utilities."""
+
+import time
+import types
+from enum import Enum
+from typing import Any, Callable, List, Optional, Type, Union
+
+from rich.console import Console
+from rich.progress import BarColumn, Progress, SpinnerColumn, TaskID, TextColumn, TimeElapsedColumn
+
+from lightning_sdk.lightning_cloud.openapi.models.v1_cluster_accelerator import V1ClusterAccelerator
+from lightning_sdk.lightning_cloud.openapi.models.v1_get_cloud_space_instance_status_response import (
+    V1GetCloudSpaceInstanceStatusResponse,
+)
+
+
+class StartupPhase(Enum):
+    """Studio startup phase messages."""
+
+    STARTING_STUDIO = "Starting Studio..."
+    GETTING_MACHINE = "Getting a machine..."
+
+    SWITCHING_STUDIO = "Switching Studio..."
+
+    SETTING_UP_MACHINE = "Setting up machine..."
+    RESTORING_STUDIO = "Restoring Studio..."
+    PREPARING_STUDIO = "Preparing Studio..."
+    RESTORING_BASE_STUDIO = "Restoring Base Studio..."
+    SETTING_UP_BASE_STUDIO = "Setting up Base Studio..."
+    DONE = "Done"
+
+
+def get_switching_progress_message(percentage: int, is_base_studio: bool) -> str:
+    """Get progress message for switching studios."""
+    percentage = max(0, min(100, round(percentage)))
+
+    if percentage > 98:
+        message = StartupPhase.DONE.value
+    elif percentage > 80:
+        message = StartupPhase.RESTORING_BASE_STUDIO.value if is_base_studio else StartupPhase.RESTORING_STUDIO.value
+    elif percentage > 60:
+        message = StartupPhase.SETTING_UP_MACHINE.value
+    else:
+        message = StartupPhase.SWITCHING_STUDIO.value
+    return f"({percentage}%) {message}"
+
+
+def estimated_studio_ready_in_seconds(
+    cloud_space: Any, cloud_space_instance_status: Any, accelerators: Optional[List[V1ClusterAccelerator]] = None
+) -> int:
+    """Calculate estimated seconds until studio is ready."""
+    # Default estimate
+    return 120
+
+
+def progress_bar_growth(default_timeout: int, counter: float) -> int:
+    """Calculate progress bar growth based on timeout and counter."""
+    if default_timeout <= 0:
+        return 100
+
+    value = ((default_timeout - counter) / default_timeout) * 100
+    if value > 100:
+        value = 100
+    if value < 0:
+        value = 0
+    return int(value)
+
+
+class StudioProgressTracker:
+    """Tracks and displays progress for studio startup/switching operations."""
+
+    def __init__(self, operation_type: str = "start", show_progress: bool = True, check_interval: float = 1.0) -> None:
+        """Initialize progress tracker.
+
+        Args:
+            operation_type: Type of operation ('start' or 'switch')
+            show_progress: Whether to display progress bar
+            check_interval: Seconds between status checks
+        """
+        self.operation_type = operation_type
+        self.show_progress = show_progress
+        self.check_interval = check_interval
+        self.progress: Optional[Progress] = None
+        self.task_id: Optional[TaskID] = None
+        self.console = Console()
+        self._last_message = ""
+
+    def __enter__(self) -> "StudioProgressTracker":
+        """Enter context manager."""
+        if self.show_progress:
+            self.progress = Progress(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                BarColumn(),
+                TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+                TimeElapsedColumn(),
+                console=self.console,
+                transient=False,
+            )
+            self.progress.start()
+            self.task_id = self.progress.add_task(f"{self.operation_type.capitalize()}ing Studio...", total=100)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Union[Type[BaseException], None],
+        exc_val: Union[BaseException, None],
+        exc_tb: Union[types.TracebackType, None],
+    ) -> None:
+        """Exit context manager."""
+        if self.progress:
+            self.progress.stop()
+
+    def update_progress(self, percentage: int, message: str = "", is_base_studio: bool = False) -> None:
+        """Update progress bar with current percentage and message."""
+        if not self.progress or self.task_id is None:
+            return
+
+        if self.operation_type == "switch":
+            display_message = get_switching_progress_message(percentage, is_base_studio)
+        else:
+            display_message = message or f"{self.operation_type.capitalize()}ing Studio..."
+
+        # Update description if message changed
+        if display_message != self._last_message:
+            self.progress.update(self.task_id, description=display_message)
+            self._last_message = display_message
+
+        # Never show 100% until truly complete
+        completed = min(percentage, 98) if percentage < 100 else 100
+        self.progress.update(self.task_id, completed=completed)
+
+        # Force console refresh to ensure progress is visible
+        self.progress.refresh()
+
+    def complete(self, success_message: str = "") -> None:
+        """Mark operation as complete."""
+        if self.progress and self.task_id is not None:
+            self.progress.update(self.task_id, completed=100, description=success_message or "Done")
+
+    def track_startup_phases(
+        self,
+        status_getter: Callable[[], V1GetCloudSpaceInstanceStatusResponse],
+        accelerators: Optional[List[V1ClusterAccelerator]] = None,
+        timeout: int = 600,
+    ) -> None:
+        """Track startup phases and update progress accordingly."""
+        start_time = time.time()
+        last_progress = 5
+        phase_start_times = {}
+        last_message_time = 0
+        message_stability_delay = 3.0  # Seconds to wait before changing message
+
+        # Show initial progress immediately
+        self.update_progress(5, StartupPhase.STARTING_STUDIO.value)
+
+        while True:
+            try:
+                status = status_getter()
+                elapsed = time.time() - start_time
+
+                # Default fallback progress based on time
+                time_based_progress = min(95, int((elapsed / timeout) * 100))
+                current_progress = max(last_progress, time_based_progress)
+                current_message = StartupPhase.STARTING_STUDIO.value
+
+                # Check if we have detailed status information
+                if hasattr(status, "in_use") and status.in_use:
+                    in_use = status.in_use
+
+                    # Check startup status for detailed phases
+                    if hasattr(in_use, "startup_status") and in_use.startup_status:
+                        startup_status = in_use.startup_status
+
+                        # Check completion first
+                        if (
+                            hasattr(startup_status, "top_up_restore_finished")
+                            and startup_status.top_up_restore_finished
+                        ):
+                            self.complete(StartupPhase.DONE.value)
+                            break
+
+                        # Check other phases in descending priority
+                        if (
+                            hasattr(startup_status, "initial_restore_finished")
+                            and startup_status.initial_restore_finished
+                        ):
+                            current_progress = max(current_progress, 80)
+                            current_message = StartupPhase.PREPARING_STUDIO.value
+                        elif hasattr(startup_status, "container_ready") and startup_status.container_ready:
+                            current_progress = max(current_progress, 60)
+                            current_message = StartupPhase.SETTING_UP_MACHINE.value
+                        elif hasattr(startup_status, "machine_ready") and startup_status.machine_ready:
+                            current_progress = max(current_progress, 30)
+                            current_message = StartupPhase.GETTING_MACHINE.value
+
+                    # Check general phase information
+                    if hasattr(in_use, "phase") and in_use.phase:
+                        phase = in_use.phase
+
+                        if phase == "CLOUD_SPACE_INSTANCE_STATE_RUNNING":
+                            current_progress = max(current_progress, 80)
+                            current_message = StartupPhase.SETTING_UP_MACHINE.value
+                        elif phase == "CLOUD_SPACE_INSTANCE_STATE_PENDING":
+                            # Track time in pending phase for smoother progress
+                            if "pending" not in phase_start_times:
+                                phase_start_times["pending"] = time.time()
+
+                            pending_elapsed = time.time() - phase_start_times["pending"]
+                            # Progress more smoothly through pending phase (10-60%)
+                            pending_progress = 10 + min(50, int((pending_elapsed / 60) * 50))
+                            current_progress = max(current_progress, pending_progress)
+                            current_message = StartupPhase.GETTING_MACHINE.value
+
+                # Check for requested machine status (pre-allocation)
+                elif hasattr(status, "requested") and status.requested:
+                    if "allocation" not in phase_start_times:
+                        phase_start_times["allocation"] = time.time()
+
+                    allocation_elapsed = time.time() - phase_start_times["allocation"]
+                    # Progress through allocation phase (5-30%)
+                    allocation_progress = 5 + min(25, int((allocation_elapsed / 30) * 25))
+                    current_progress = max(current_progress, allocation_progress)
+                    current_message = StartupPhase.GETTING_MACHINE.value
+
+                # Ensure progress never decreases and moves smoothly
+                if current_progress > last_progress:
+                    # Smooth progress increases - don't jump too much at once
+                    max_increment = 3 if current_progress - last_progress > 10 else current_progress - last_progress
+                    current_progress = last_progress + max_increment
+
+                current_progress = min(98, current_progress)  # Never show 100% until truly complete
+
+                # Only update message if enough time has passed since last message change
+                # or if this is the first message
+                current_time = time.time()
+                should_update_message = current_message != self._last_message and (
+                    current_time - last_message_time >= message_stability_delay or last_message_time == 0
+                )
+
+                if should_update_message:
+                    self.update_progress(current_progress, current_message)
+                    last_message_time = current_time
+                else:
+                    # Update progress but keep existing message
+                    if self.progress and self.task_id is not None:
+                        self.progress.update(self.task_id, completed=current_progress)
+                        self.progress.refresh()
+
+                last_progress = current_progress
+
+                # Break if we timeout
+                if elapsed > timeout:
+                    self.complete("Studio start may still be in progress...")
+                    break
+
+            except Exception:
+                # Continue on API errors but still update progress
+                elapsed = time.time() - start_time
+                fallback_progress = min(95, max(last_progress, int((elapsed / timeout) * 100)))
+
+                # Only update message if enough time has passed
+                current_time = time.time()
+                should_update_message = StartupPhase.GETTING_MACHINE.value != self._last_message and (
+                    current_time - last_message_time >= message_stability_delay or last_message_time == 0
+                )
+
+                if should_update_message:
+                    self.update_progress(fallback_progress, StartupPhase.GETTING_MACHINE.value)
+                    last_message_time = current_time
+                else:
+                    # Update progress but keep existing message
+                    if self.progress and self.task_id is not None:
+                        self.progress.update(self.task_id, completed=fallback_progress)
+                        self.progress.refresh()
+
+                last_progress = fallback_progress
+
+                if elapsed > timeout:
+                    self.complete("Studio start may still be in progress...")
+                    break
+
+            time.sleep(self.check_interval)

lightning_sdk/utils/resolve.py

@@ -5,10 +5,13 @@ from contextlib import contextmanager
 from typing import TYPE_CHECKING, Generator, List, Optional, Tuple, Union
 
 from lightning_sdk.api import TeamspaceApi, UserApi
-from lightning_sdk.machine import Machine
+from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import CloudProvider, Machine
 
 if TYPE_CHECKING:
     from lightning_sdk.organization import Organization
+    from lightning_sdk.studio import Studio
     from lightning_sdk.teamspace import Teamspace
     from lightning_sdk.user import User
 
@@ -46,7 +49,35 @@ def _resolve_deprecated_cloud_compute(machine: Machine, cloud_compute: Optional[
     return machine
 
 
-def _resolve_deprecated_cluster(cloud_account: Optional[str], cluster: Optional[str]) -> Optional[str]:
+def _resolve_deprecated_provider(
+    cloud_provider: Optional[Union[CloudProvider, str]], provider: Optional[Union[CloudProvider, str]]
+) -> Optional[Union[CloudProvider, str]]:
+    if provider is not None:
+        if cloud_provider is not None:
+            raise ValueError(
+                "Cannot use both 'provider' and 'cloud_provider' at the same time."
+                "Please don't set the 'provider' as it will be deprecated!"
+            )
+
+        warnings.warn(
+            "The 'provider' argument will be deprecated in the future! "
+            "Please consider using the 'cloud_provider' argument instead!",
+            DeprecationWarning,
+        )
+        return provider
+
+    if cloud_provider is None:
+        from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+        config = Config()
+        cloud_provider = config.get_value(DefaultConfigKeys.cloud_provider)
+
+    return cloud_provider
+
+
+def _resolve_deprecated_cluster(
+    cloud_account: Optional[str], cluster: Optional[str], current_cloud_account: Optional[str] = None
+) -> Optional[str]:
     if cluster is not None:
         if cloud_account is not None:
             raise ValueError(
@@ -61,12 +92,26 @@ def _resolve_deprecated_cluster(cloud_account: Optional[str], cluster: Optional[
         )
         return cluster
 
+    if cloud_account is None:
+        from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+        config = Config()
+        cloud_account = config.get_value(DefaultConfigKeys.cloud_account)
+
+        if cloud_account is None:
+            cloud_account = current_cloud_account
+
     return cloud_account
 
 
 def _resolve_org_name(name: Optional[str]) -> Optional[str]:
     if name is None:
         name = os.environ.get("LIGHTNING_ORG", "") or None
+    if name is None:
+        from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+        config = Config()
+        name = config.get_value(DefaultConfigKeys.organization)
     return name
 
 
@@ -75,6 +120,7 @@ def _resolve_org(org: Optional[Union[str, "Organization"]]) -> Optional["Organiz
 
     if isinstance(org, Organization):
         return org
+
     org = _resolve_org_name(org)
 
     if org is None:
@@ -82,12 +128,23 @@ def _resolve_org(org: Optional[Union[str, "Organization"]]) -> Optional["Organiz
 
     from lightning_sdk.organization import Organization
 
-    return Organization(name=org)
+    try:
+        return Organization(name=org)
+    # Handle case where user name is mistakenly used as organization name
+    except ApiException as ae:
+        if ae.status == 404:
+            raise ValueError(f"Organization '{org}' does not exist or you are not a member of it.") from ae
+        raise RuntimeError(f"Failed to resolve organization '{org}': {ae}") from ae
 
 
 def _resolve_user_name(name: Optional[str]) -> Optional[str]:
     if name is None:
         name = os.environ.get("LIGHTNING_USERNAME", "") or None
+    if name is None:
+        from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+        config = Config()
+        name = config.get_value(DefaultConfigKeys.user)
     return name
 
 
@@ -107,6 +164,11 @@ def _resolve_user(user: Optional[Union[str, "User"]]) -> Optional["User"]:
 def _resolve_teamspace_name(name: Optional[str]) -> Optional[str]:
     if name is None:
         name = os.environ.get("LIGHTNING_TEAMSPACE", "") or None
+    if name is None:
+        from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+        config = Config()
+        name = config.get_value(DefaultConfigKeys.teamspace_name)
     return name
 
 
@@ -135,10 +197,29 @@ def _resolve_teamspace(
         return Teamspace(name=teamspace, org=org)
 
     user = _resolve_user(user)
-    if user is None:
-        raise RuntimeError("Neither user nor org provided, but one of them needs to be provided")
 
-    return Teamspace(name=teamspace, user=user)
+    # If still no user or org resolved, try config defaults
+    if user is None and org is None:
+        from lightning_sdk.utils.config import Config, DefaultConfigKeys
+
+        config = Config()
+        owner_type = config.get_value(DefaultConfigKeys.teamspace_owner_type)
+        owner_name = config.get_value(DefaultConfigKeys.teamspace_owner)
+
+        if owner_type and owner_name:
+            if owner_type.lower() == "organization":
+                org = _resolve_org(owner_name)
+            elif owner_type.lower() == "user":
+                user = _resolve_user(owner_name)
+
+    # Final resolution check
+    if org is not None:
+        return Teamspace(name=teamspace, org=org)
+
+    if user is not None:
+        return Teamspace(name=teamspace, user=user)
+
+    raise RuntimeError("Neither user nor org provided, but one of them needs to be provided")
 
 
 def _get_organizations_for_authed_user() -> List["Organization"]:
@@ -168,22 +249,77 @@ def skip_studio_init() -> Generator[None, None, None]:
     """Skip studio init based on current runtime."""
     from lightning_sdk.studio import Studio
 
-    prev_studio_init_state = Studio._skip_init
-    Studio._skip_init = True
+    prev_studio_init_state = getattr(Studio._skip_init, "value", False)
+    Studio._skip_init.value = True
+
+    yield
+
+    Studio._skip_init.value = prev_studio_init_state
+
+
+@contextmanager
+def skip_studio_setup() -> Generator[None, None, None]:
+    """Skip studio setup based on current runtime."""
+    from lightning_sdk.studio import Studio
+
+    prev_studio_setup_state = getattr(Studio._skip_setup, "value", False)
+    Studio._skip_setup.value = True
+
+    yield
+
+    Studio._skip_setup.value = prev_studio_setup_state
+
+
+@contextmanager
+def prevent_refetch_studio(studio: "Studio") -> Generator[None, None, None]:
+    """Prevent refetching the studio based on current runtime."""
+    prev_prevent_refetch_state = getattr(studio, "_prevent_refetch", False)
+    studio._prevent_refetch = True
 
     yield
 
-    Studio._skip_init = prev_studio_init_state
+    studio._prevent_refetch = prev_prevent_refetch_state
 
 
-def _parse_model_and_version(name: str) -> Tuple[str, str]:
+def _parse_model_and_version(name: str) -> Tuple[str, Optional[str]]:
+    """Parse the model name and version from the given string.
+
+    >>> _parse_model_and_version("org/teamspace/modelname")
+    ('org/teamspace/modelname', None)
+    >>> _parse_model_and_version("org/teamspace/modelname:version")
+    ('org/teamspace/modelname', 'version')
+    """
     parts = name.split(":")
     if len(parts) == 1:
-        return parts[0], "latest"
+        return parts[0], None
     if len(parts) == 2:
         return parts[0], parts[1]
     # The rest of the validation for name and version happens in the backend
     raise ValueError(
-        "Model version is expected to be in the format `entity/modelname:version` separated by a"
-        f" single colon, but got: {name}"
+        "Model version is expected to be in the format `entity/modelname:version` separated by a single colon,"
+        f" but got: {name}"
     )
+
+
+def in_studio() -> bool:
+    """Returns true if inside a studio, else false."""
+    has_cloudspace_id = bool(os.getenv("LIGHTNING_CLOUD_SPACE_ID", None))
+    is_interactive = os.getenv("LIGHTNING_INTERACTIVE", "false") == "true"
+    return has_cloudspace_id and is_interactive
+
+
+def _get_studio_url(studio: "Studio", turn_on: bool = False) -> str:
+    cloud_url = _get_cloud_url().replace(":443", "")
+    base_url = f"{cloud_url}/{studio.owner.name}/{studio.teamspace.name}/studios/{studio.name}/code"
+
+    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-0.1.49.dist-info → lightning_sdk-2025.11.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 0.1.49
+Version: 2025.11.5
 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
@@ -32,22 +34,23 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: backoff
+Requires-Dist: click
+Requires-Dist: docker
 Requires-Dist: fastapi
+Requires-Dist: packaging
 Requires-Dist: pyjwt
-Requires-Dist: python-multipart
+Requires-Dist: pyyaml
 Requires-Dist: requests
 Requires-Dist: rich
+Requires-Dist: simple-term-menu
 Requires-Dist: six
+Requires-Dist: tqdm
 Requires-Dist: urllib3
 Requires-Dist: uvicorn
 Requires-Dist: websocket-client
-Requires-Dist: tqdm
-Requires-Dist: fire
-Requires-Dist: simple-term-menu
-Requires-Dist: lightning-utilities
+Requires-Dist: wget
 Provides-Extra: serve
 Requires-Dist: litserve>=0.2.5; extra == "serve"
-Requires-Dist: docker; extra == "serve"
 
 # Lightning SDK
 
@@ -79,11 +82,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 +99,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")