lightning-sdk 2025.8.6rc2__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

@@ -66,10 +66,18 @@ def _resolve_deprecated_provider(
         )
         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]) -> Optional[str]:
+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(
@@ -84,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
 
 
@@ -118,6 +140,11 @@ def _resolve_org(org: Optional[Union[str, "Organization"]]) -> Optional["Organiz
 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
 
 
@@ -137,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
 
 
@@ -165,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"]:
@@ -198,12 +249,36 @@ 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, Optional[str]]:

{lightning_sdk-2025.8.6rc2.dist-info → lightning_sdk-2025.11.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 2025.8.6rc2
+Version: 2025.11.5
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License
@@ -39,6 +39,7 @@ Requires-Dist: docker
 Requires-Dist: fastapi
 Requires-Dist: packaging
 Requires-Dist: pyjwt
+Requires-Dist: pyyaml
 Requires-Dist: requests
 Requires-Dist: rich
 Requires-Dist: simple-term-menu