lightning-sdk 2025.8.18.post0__py3-none-any.whl → 2025.8.21__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py

@@ -41,7 +41,6 @@ class V1UserFeatures(object):
                             and the value is json key in definition.
     """
     swagger_types = {
-        'add_data_v2': 'bool',
         'affiliate_links': 'bool',
         'agents_v2': 'bool',
         'ai_hub_monetization': 'bool',
@@ -75,13 +74,11 @@ class V1UserFeatures(object):
         'f241': 'bool',
         'f242': 'bool',
         'f243': 'bool',
-        'f244': 'bool',
         'f245': 'bool',
         'fair_share': 'bool',
         'featured_studios_admin': 'bool',
         'gcp_overprovisioning': 'bool',
         'gcs_connections_optimized': 'bool',
-        'gcs_folders': 'bool',
         'instant_capacity_reservation': 'bool',
         'job_artifacts_v2': 'bool',
         'kubernetes_cluster_ui': 'bool',
@@ -109,7 +106,6 @@ class V1UserFeatures(object):
         'product_license': 'bool',
         'project_selector': 'bool',
         'publish_pipelines': 'bool',
-        'r2_data_connections': 'bool',
         'reserved_machines_tab': 'bool',
         'restartable_jobs': 'bool',
         'runnable_public_studio_page': 'bool',
@@ -123,7 +119,6 @@ class V1UserFeatures(object):
         'studio_sharing_v2': 'bool',
         'studio_version_visibility': 'bool',
         'trainium2': 'bool',
-        'use_internal_data_connection_mounts': 'bool',
         'use_rclone_mounts_only': 'bool',
         'vultr': 'bool',
         'weka': 'bool',
@@ -131,7 +126,6 @@ class V1UserFeatures(object):
     }
 
     attribute_map = {
-        'add_data_v2': 'addDataV2',
         'affiliate_links': 'affiliateLinks',
         'agents_v2': 'agentsV2',
         'ai_hub_monetization': 'aiHubMonetization',
@@ -165,13 +159,11 @@ class V1UserFeatures(object):
         'f241': 'f241',
         'f242': 'f242',
         'f243': 'f243',
-        'f244': 'f244',
         'f245': 'f245',
         'fair_share': 'fairShare',
         'featured_studios_admin': 'featuredStudiosAdmin',
         'gcp_overprovisioning': 'gcpOverprovisioning',
         'gcs_connections_optimized': 'gcsConnectionsOptimized',
-        'gcs_folders': 'gcsFolders',
         'instant_capacity_reservation': 'instantCapacityReservation',
         'job_artifacts_v2': 'jobArtifactsV2',
         'kubernetes_cluster_ui': 'kubernetesClusterUi',
@@ -199,7 +191,6 @@ class V1UserFeatures(object):
         'product_license': 'productLicense',
         'project_selector': 'projectSelector',
         'publish_pipelines': 'publishPipelines',
-        'r2_data_connections': 'r2DataConnections',
         'reserved_machines_tab': 'reservedMachinesTab',
         'restartable_jobs': 'restartableJobs',
         'runnable_public_studio_page': 'runnablePublicStudioPage',
@@ -213,16 +204,14 @@ class V1UserFeatures(object):
         'studio_sharing_v2': 'studioSharingV2',
         'studio_version_visibility': 'studioVersionVisibility',
         'trainium2': 'trainium2',
-        'use_internal_data_connection_mounts': 'useInternalDataConnectionMounts',
         'use_rclone_mounts_only': 'useRcloneMountsOnly',
         'vultr': 'vultr',
         'weka': 'weka',
         'writable_s3_connections': 'writableS3Connections'
     }
 
-    def __init__(self, add_data_v2: 'bool' =None, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, byo_machine_type: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, chat_models: 'bool' =None, cloudspace_schedules: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, control_center_monitoring: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_persistent_disk: 'bool' =None, drive_v2: 'bool' =None, enterprise_compute_admin: 'bool' =None, f227: 'bool' =None, f234: 'bool' =None, f236: 'bool' =None, f237: 'bool' =None, f238: 'bool' =None, f239: 'bool' =None, f240: 'bool' =None, f241: 'bool' =None, f242: 'bool' =None, f243: 'bool' =None, f244: 'bool' =None, f245: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, gcp_overprovisioning: 'bool' =None, gcs_connections_optimized: 'bool' =None, gcs_folders: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, kubernetes_cluster_ui: 'bool' =None, kubernetes_clusters: 'bool' =None, landing_studios: 'bool' =None, lit_logger: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, model_api_dashboard: 'bool' =None, multiple_studio_versions: 'bool' =None, nerf_fs_nonpaying: 'bool' =None, onboarding_v2: 'bool' =None, org_level_member_permissions: 'bool' =None, org_usage_limits: 'bool' =None, persistent_disk: 'bool' =None, plugin_distributed: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_sweeps: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, product_license: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, r2_data_connections: 'bool' =None, reserved_machines_tab: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, single_wallet: 'bool' =None, slurm: 'bool' =None, specialised_studios: 'bool' =None, storage_overuse_deletion: 'bool' =None, studio_config: 'bool' =None, studio_sharing_v2: 'bool' =None, studio_version_visibility: 'bool' =None, trainium2: 'bool' =None, use_internal_data_connection_mounts: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: 'bool' =None, weka: 'bool' =None, writable_s3_connections: 'bool' =None):  # noqa: E501
+    def __init__(self, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, byo_machine_type: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, chat_models: 'bool' =None, cloudspace_schedules: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, control_center_monitoring: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_persistent_disk: 'bool' =None, drive_v2: 'bool' =None, enterprise_compute_admin: 'bool' =None, f227: 'bool' =None, f234: 'bool' =None, f236: 'bool' =None, f237: 'bool' =None, f238: 'bool' =None, f239: 'bool' =None, f240: 'bool' =None, f241: 'bool' =None, f242: 'bool' =None, f243: 'bool' =None, f245: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, gcp_overprovisioning: 'bool' =None, gcs_connections_optimized: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, kubernetes_cluster_ui: 'bool' =None, kubernetes_clusters: 'bool' =None, landing_studios: 'bool' =None, lit_logger: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, model_api_dashboard: 'bool' =None, multiple_studio_versions: 'bool' =None, nerf_fs_nonpaying: 'bool' =None, onboarding_v2: 'bool' =None, org_level_member_permissions: 'bool' =None, org_usage_limits: 'bool' =None, persistent_disk: 'bool' =None, plugin_distributed: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_sweeps: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, product_license: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, reserved_machines_tab: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, single_wallet: 'bool' =None, slurm: 'bool' =None, specialised_studios: 'bool' =None, storage_overuse_deletion: 'bool' =None, studio_config: 'bool' =None, studio_sharing_v2: 'bool' =None, studio_version_visibility: 'bool' =None, trainium2: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: 'bool' =None, weka: 'bool' =None, writable_s3_connections: 'bool' =None):  # noqa: E501
         """V1UserFeatures - a model defined in Swagger"""  # noqa: E501
-        self._add_data_v2 = None
         self._affiliate_links = None
         self._agents_v2 = None
         self._ai_hub_monetization = None
@@ -256,13 +245,11 @@ class V1UserFeatures(object):
         self._f241 = None
         self._f242 = None
         self._f243 = None
-        self._f244 = None
         self._f245 = None
         self._fair_share = None
         self._featured_studios_admin = None
         self._gcp_overprovisioning = None
         self._gcs_connections_optimized = None
-        self._gcs_folders = None
         self._instant_capacity_reservation = None
         self._job_artifacts_v2 = None
         self._kubernetes_cluster_ui = None
@@ -290,7 +277,6 @@ class V1UserFeatures(object):
         self._product_license = None
         self._project_selector = None
         self._publish_pipelines = None
-        self._r2_data_connections = None
         self._reserved_machines_tab = None
         self._restartable_jobs = None
         self._runnable_public_studio_page = None
@@ -304,14 +290,11 @@ class V1UserFeatures(object):
         self._studio_sharing_v2 = None
         self._studio_version_visibility = None
         self._trainium2 = None
-        self._use_internal_data_connection_mounts = None
         self._use_rclone_mounts_only = None
         self._vultr = None
         self._weka = None
         self._writable_s3_connections = None
         self.discriminator = None
-        if add_data_v2 is not None:
-            self.add_data_v2 = add_data_v2
         if affiliate_links is not None:
             self.affiliate_links = affiliate_links
         if agents_v2 is not None:
@@ -378,8 +361,6 @@ class V1UserFeatures(object):
             self.f242 = f242
         if f243 is not None:
             self.f243 = f243
-        if f244 is not None:
-            self.f244 = f244
         if f245 is not None:
             self.f245 = f245
         if fair_share is not None:
@@ -390,8 +371,6 @@ class V1UserFeatures(object):
             self.gcp_overprovisioning = gcp_overprovisioning
         if gcs_connections_optimized is not None:
             self.gcs_connections_optimized = gcs_connections_optimized
-        if gcs_folders is not None:
-            self.gcs_folders = gcs_folders
         if instant_capacity_reservation is not None:
             self.instant_capacity_reservation = instant_capacity_reservation
         if job_artifacts_v2 is not None:
@@ -446,8 +425,6 @@ class V1UserFeatures(object):
             self.project_selector = project_selector
         if publish_pipelines is not None:
             self.publish_pipelines = publish_pipelines
-        if r2_data_connections is not None:
-            self.r2_data_connections = r2_data_connections
         if reserved_machines_tab is not None:
             self.reserved_machines_tab = reserved_machines_tab
         if restartable_jobs is not None:
@@ -474,8 +451,6 @@ class V1UserFeatures(object):
             self.studio_version_visibility = studio_version_visibility
         if trainium2 is not None:
             self.trainium2 = trainium2
-        if use_internal_data_connection_mounts is not None:
-            self.use_internal_data_connection_mounts = use_internal_data_connection_mounts
         if use_rclone_mounts_only is not None:
             self.use_rclone_mounts_only = use_rclone_mounts_only
         if vultr is not None:
@@ -485,27 +460,6 @@ class V1UserFeatures(object):
         if writable_s3_connections is not None:
             self.writable_s3_connections = writable_s3_connections
 
-    @property
-    def add_data_v2(self) -> 'bool':
-        """Gets the add_data_v2 of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The add_data_v2 of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._add_data_v2
-
-    @add_data_v2.setter
-    def add_data_v2(self, add_data_v2: 'bool'):
-        """Sets the add_data_v2 of this V1UserFeatures.
-
-
-        :param add_data_v2: The add_data_v2 of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._add_data_v2 = add_data_v2
-
     @property
     def affiliate_links(self) -> 'bool':
         """Gets the affiliate_links of this V1UserFeatures.  # noqa: E501
@@ -1199,27 +1153,6 @@ class V1UserFeatures(object):
 
         self._f243 = f243
 
-    @property
-    def f244(self) -> 'bool':
-        """Gets the f244 of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The f244 of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._f244
-
-    @f244.setter
-    def f244(self, f244: 'bool'):
-        """Sets the f244 of this V1UserFeatures.
-
-
-        :param f244: The f244 of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._f244 = f244
-
     @property
     def f245(self) -> 'bool':
         """Gets the f245 of this V1UserFeatures.  # noqa: E501
@@ -1325,27 +1258,6 @@ class V1UserFeatures(object):
 
         self._gcs_connections_optimized = gcs_connections_optimized
 
-    @property
-    def gcs_folders(self) -> 'bool':
-        """Gets the gcs_folders of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The gcs_folders of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._gcs_folders
-
-    @gcs_folders.setter
-    def gcs_folders(self, gcs_folders: 'bool'):
-        """Sets the gcs_folders of this V1UserFeatures.
-
-
-        :param gcs_folders: The gcs_folders of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._gcs_folders = gcs_folders
-
     @property
     def instant_capacity_reservation(self) -> 'bool':
         """Gets the instant_capacity_reservation of this V1UserFeatures.  # noqa: E501
@@ -1913,27 +1825,6 @@ class V1UserFeatures(object):
 
         self._publish_pipelines = publish_pipelines
 
-    @property
-    def r2_data_connections(self) -> 'bool':
-        """Gets the r2_data_connections of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The r2_data_connections of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._r2_data_connections
-
-    @r2_data_connections.setter
-    def r2_data_connections(self, r2_data_connections: 'bool'):
-        """Sets the r2_data_connections of this V1UserFeatures.
-
-
-        :param r2_data_connections: The r2_data_connections of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._r2_data_connections = r2_data_connections
-
     @property
     def reserved_machines_tab(self) -> 'bool':
         """Gets the reserved_machines_tab of this V1UserFeatures.  # noqa: E501
@@ -2207,27 +2098,6 @@ class V1UserFeatures(object):
 
         self._trainium2 = trainium2
 
-    @property
-    def use_internal_data_connection_mounts(self) -> 'bool':
-        """Gets the use_internal_data_connection_mounts of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The use_internal_data_connection_mounts of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._use_internal_data_connection_mounts
-
-    @use_internal_data_connection_mounts.setter
-    def use_internal_data_connection_mounts(self, use_internal_data_connection_mounts: 'bool'):
-        """Sets the use_internal_data_connection_mounts of this V1UserFeatures.
-
-
-        :param use_internal_data_connection_mounts: The use_internal_data_connection_mounts of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._use_internal_data_connection_mounts = use_internal_data_connection_mounts
-
     @property
     def use_rclone_mounts_only(self) -> 'bool':
         """Gets the use_rclone_mounts_only of this V1UserFeatures.  # noqa: E501

lightning_sdk/llm/llm.py

@@ -285,7 +285,7 @@ class LLM:
         self,
         prompt: str,
         system_prompt: Optional[str] = None,
-        max_completion_tokens: Optional[int] = 500,
+        max_completion_tokens: Optional[int] = None,
         images: Optional[Union[List[str], str]] = None,
         conversation: Optional[str] = None,
         metadata: Optional[Dict[str, str]] = None,
@@ -319,7 +319,7 @@ class LLM:
         self,
         prompt: str,
         system_prompt: Optional[str] = None,
-        max_completion_tokens: Optional[int] = 500,
+        max_completion_tokens: Optional[int] = None,
         images: Optional[Union[List[str], str]] = None,
         conversation: Optional[str] = None,
         metadata: Optional[Dict[str, str]] = None,

lightning_sdk/studio.py

@@ -59,6 +59,9 @@ class Studio:
     _skip_init = False
     _skip_setup = False
 
+    # whether to show progress bars during operations
+    show_progress = False
+
     def __init__(
         self,
         name: Optional[str] = None,
@@ -254,9 +257,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 +344,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.

lightning_sdk/utils/progress.py

@@ -0,0 +1,284 @@
+"""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."""
+
+    ALLOCATING_MACHINE = "Allocating machine from cloud provider..."
+    STUDIO_STARTING = "Studio is starting up..."
+    SETTING_UP_ENVIRONMENT = "Setting up Studio environment..."
+    RESTORING_STATE = "Restoring Studio state..."
+    FINALIZING_SETUP = "Finalizing Studio setup..."
+    COMPLETED = "Studio started successfully"
+
+
+def get_switching_progress_message(percentage: int, is_base_studio: bool, is_new_cloud_space: bool) -> str:
+    """Get progress message for switching studios."""
+    percentage = max(0, min(100, round(percentage)))
+
+    if percentage > 98:
+        message = "Done"
+    elif percentage > 80:
+        if is_new_cloud_space:
+            message = "Setting up Base Studio..." if is_base_studio else "Preparing Studio..."
+        else:
+            message = "Restoring Studio..."
+    elif percentage > 60:
+        message = "Setting up machine from the cloud provider"
+    else:
+        message = "Allocating machine from the cloud provider"
+
+    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, is_new_cloud_space: 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, is_new_cloud_space)
+        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.ALLOCATING_MACHINE.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.ALLOCATING_MACHINE.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.COMPLETED.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, 85)
+                            current_message = StartupPhase.FINALIZING_SETUP.value
+                        elif hasattr(startup_status, "container_ready") and startup_status.container_ready:
+                            current_progress = max(current_progress, 70)
+                            current_message = StartupPhase.RESTORING_STATE.value
+                        elif hasattr(startup_status, "machine_ready") and startup_status.machine_ready:
+                            current_progress = max(current_progress, 60)
+                            current_message = StartupPhase.SETTING_UP_ENVIRONMENT.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.STUDIO_STARTING.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.ALLOCATING_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.ALLOCATING_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.ALLOCATING_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.ALLOCATING_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-2025.8.18.post0.dist-info → lightning_sdk-2025.8.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 2025.8.18.post0
+Version: 2025.8.21
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License