lightning-sdk 2025.11.13.post0__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

@@ -0,0 +1,325 @@
+import logging
+import os
+import warnings
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Generator, List, Optional, Tuple, Union
+
+from lightning_sdk.api import TeamspaceApi, UserApi
+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
+
+
+_LIGHTNING_SERVICE_EXECUTION_ID_KEY = "LIGHTNING_SERVICE_EXECUTION_ID"
+
+
+def _setup_logger(name: str, level: int = logging.INFO) -> logging.Logger:
+    _logger = logging.getLogger(name)
+    _handler = logging.StreamHandler()
+    _handler.setLevel(level)
+    _logger.setLevel(level)
+    _formatter = logging.Formatter("%(levelname)s - %(message)s")
+    _handler.setFormatter(_formatter)
+    _logger.addHandler(_handler)
+    return _logger
+
+
+def _resolve_deprecated_cloud_compute(machine: Machine, cloud_compute: Optional[Machine]) -> Machine:
+    if cloud_compute is not None:
+        if machine == Machine.CPU:
+            # user explicitly set cloud_compute and not machine, so use cloud_compute
+            warnings.warn(
+                "The 'cloud_compute' argument will be deprecated in the future! "
+                "Please consider using the 'machine' argument instead!",
+                DeprecationWarning,
+            )
+            return cloud_compute
+
+        raise ValueError(
+            "Cannot use both 'cloud_compute' and 'machine' at the same time."
+            "Please don't set the 'cloud_compute' as it will be deprecated!"
+        )
+
+    return machine
+
+
+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(
+                "Cannot use both 'cluster' and 'cloud_account' at the same time."
+                "Please don't set the 'cluster' as it will be deprecated!"
+            )
+
+        warnings.warn(
+            "The 'cluster' argument will be deprecated in the future! "
+            "Please consider using the 'cloud_account' argument instead!",
+            DeprecationWarning,
+        )
+        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
+
+
+def _resolve_org(org: Optional[Union[str, "Organization"]]) -> Optional["Organization"]:
+    from lightning_sdk.organization import Organization
+
+    if isinstance(org, Organization):
+        return org
+
+    org = _resolve_org_name(org)
+
+    if org is None:
+        return None
+
+    from lightning_sdk.organization import Organization
+
+    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
+
+
+def _resolve_user(user: Optional[Union[str, "User"]]) -> Optional["User"]:
+    from lightning_sdk.user import User
+
+    if isinstance(user, User):
+        return user
+
+    user = _resolve_user_name(user)
+    if user is None:
+        return None
+
+    return User(name=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
+
+
+def _resolve_teamspace(
+    teamspace: Optional[Union[str, "Teamspace"]],
+    org: Optional[Union[str, "Organization"]],
+    user: Optional[Union[str, "User"]],
+) -> Optional["Teamspace"]:
+    from lightning_sdk.teamspace import Teamspace
+
+    if isinstance(teamspace, Teamspace):
+        return teamspace
+
+    teamspace = _resolve_teamspace_name(teamspace)
+    if teamspace is None:
+        return None
+
+    # if user was specified explicitly, use that, else resolve
+    if user is not None:
+        user = _resolve_user(user=user)
+        return Teamspace(name=teamspace, user=user)
+
+    org = _resolve_org(org)
+
+    if org is not None:
+        return Teamspace(name=teamspace, org=org)
+
+    user = _resolve_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"]:
+    """Returns Organizations the current Authed user is a member of."""
+    from lightning_sdk.organization import Organization
+
+    _orgs = UserApi()._get_organizations_for_authed_user()
+    return [Organization(_org.name) for _org in _orgs]
+
+
+def _get_teamspace_names_for_authed_user() -> List[str]:
+    """Returns Teamspace's names the current Authed user is a member of."""
+    teamspaces = UserApi()._get_all_teamspace_memberships("")
+    return sorted([ts.name for ts in teamspaces])
+
+
+def _get_authed_user() -> "User":
+    from lightning_sdk.user import User
+
+    user_id = TeamspaceApi()._get_authed_user_id()
+    _user = UserApi()._get_user_by_id(user_id)
+    return User(name=_user.username)
+
+
+@contextmanager
+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 = 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._prevent_refetch = prev_prevent_refetch_state
+
+
+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], 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 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-2025.11.13.post0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Grid.ai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.