Skip to content
This repository was archived by the owner on Jan 23, 2026. It is now read-only.
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
150 changes: 119 additions & 31 deletions packages/jumpstarter/jumpstarter/client/lease.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import logging
import os
import sys
from collections.abc import AsyncGenerator, Generator
from contextlib import (
ExitStack,
Expand All @@ -20,6 +22,7 @@
from anyio.from_thread import BlockingPortal
from grpc.aio import AioRpcError, Channel
from jumpstarter_protocol import jumpstarter_pb2, jumpstarter_pb2_grpc
from rich.console import Console

from .exceptions import LeaseError
from jumpstarter.client import client_from_path
Expand Down Expand Up @@ -112,44 +115,68 @@ async def request_async(self):

return await self._acquire()

def _update_spinner_status(self, spinner, result):
"""Update spinner with appropriate status message based on lease conditions."""
if condition_true(result.conditions, "Pending"):
pending_message = condition_message(result.conditions, "Pending")
if pending_message:
spinner.update_status(f"Waiting for lease: {pending_message}")
else:
spinner.update_status("Waiting for lease to be ready...")
else:
spinner.update_status("Waiting for server to provide status updates...")

async def _acquire(self):
"""Acquire a lease.

Makes sure the lease is ready, and returns the lease object.
"""
try:
with fail_after(self.acquisition_timeout):
while True:
logger.debug("Polling Lease %s", self.name)
result = await self.get()
# lease ready
if condition_true(result.conditions, "Ready"):
logger.debug("Lease %s acquired", self.name)
self.exporter_name = result.exporter
return self
# lease unsatisfiable
if condition_true(result.conditions, "Unsatisfiable"):
message = condition_message(result.conditions, "Unsatisfiable")
logger.debug("Lease %s cannot be satisfied: %s", self.name, message)
raise LeaseError(f"the lease cannot be satisfied: {message}")

# lease invalid
if condition_true(result.conditions, "Invalid"):
message = condition_message(result.conditions, "Invalid")
logger.debug("Lease %s is invalid: %s", self.name, message)
raise LeaseError(f"the lease is invalid: {message}")

# lease not pending
if condition_false(result.conditions, "Pending"):
raise LeaseError(
f"Lease {self.name} is not in pending, but it isn't in Ready or Unsatisfiable state either"
)

# lease released
if condition_present_and_equal(result.conditions, "Ready", "False", "Released"):
raise LeaseError(f"lease {self.name} released")

await sleep(5)
with LeaseAcquisitionSpinner(self.name) as spinner:
while True:
logger.debug("Polling Lease %s", self.name)
result = await self.get()

# lease ready
if condition_true(result.conditions, "Ready"):
logger.debug("Lease %s acquired", self.name)
spinner.update_status(f"Lease {self.name} acquired successfully!")
self.exporter_name = result.exporter
break

# lease unsatisfiable
if condition_true(result.conditions, "Unsatisfiable"):
message = condition_message(result.conditions, "Unsatisfiable")
logger.debug("Lease %s cannot be satisfied: %s", self.name, message)
raise LeaseError(f"the lease cannot be satisfied: {message}")

# lease invalid
if condition_true(result.conditions, "Invalid"):
message = condition_message(result.conditions, "Invalid")
logger.debug("Lease %s is invalid: %s", self.name, message)
raise LeaseError(f"the lease is invalid: {message}")

# lease not pending
if condition_false(result.conditions, "Pending"):
raise LeaseError(
f"Lease {self.name} is not in pending, but it isn't in Ready or "
f"Unsatisfiable state either"
)

# lease released
if condition_present_and_equal(result.conditions, "Ready", "False", "Released"):
raise LeaseError(f"lease {self.name} released")

# Update spinner with appropriate status message
self._update_spinner_status(spinner, result)

# Wait in 1-second increments with tick updates for better UX
for _ in range(5):
await sleep(1)
spinner.tick()
return self

except TimeoutError:
logger.debug(f"Lease {self.name} acquisition timed out after {self.acquisition_timeout} seconds")
raise LeaseError(
Expand Down Expand Up @@ -269,3 +296,64 @@ def serve_unix(self):
def monitor(self, threshold: timedelta = timedelta(minutes=5)):
with self.portal.wrap_async_context_manager(self.monitor_async(threshold)):
yield


class LeaseAcquisitionSpinner:
"""Context manager for displaying a spinner during lease acquisition."""

def __init__(self, lease_name: str | None = None):
self.lease_name = lease_name
self.console = Console()
self.spinner = None
self.start_time = None
self._should_show_spinner = self._is_terminal_available() and not self._is_non_interactive()
self._current_message = None

def _is_non_interactive(self) -> bool:
"""Check if the user desires a NONINTERACTIVE environment."""
return os.environ.get("NONINTERACTIVE", "false").lower() in ["true", "1"]

def _is_terminal_available(self) -> bool:
"""Check if we're running in a terminal/TTY."""
return (
hasattr(sys.stdout, 'isatty') and
sys.stdout.isatty() and
hasattr(sys.stderr, 'isatty') and
sys.stderr.isatty()
)

def __enter__(self):
self.start_time = datetime.now()
if self._should_show_spinner:
self.spinner = self.console.status(
f"Acquiring lease {self.lease_name or '...'}...",
spinner="dots",
spinner_style="blue"
)
self.spinner.start()
return self

def __exit__(self, exc_type, exc_val, exc_tb):
if self.spinner:
self.spinner.stop()

def update_status(self, message: str):
"""Update the spinner status message."""
if self.spinner and self._should_show_spinner:
self._current_message = f"[blue]{message}[/blue]"
elapsed = datetime.now() - self.start_time
elapsed_str = str(elapsed).split('.')[0] # Remove microseconds
self.spinner.update(f"{self._current_message} [dim]({elapsed_str})[/dim]")
else:
# Log info message when no console is available
elapsed = datetime.now() - self.start_time
elapsed_str = str(elapsed).split('.')[0] # Remove microseconds
logger.info(f"{message} ({elapsed_str})")

def tick(self):
"""Update the spinner with current elapsed time without changing the message."""
if self.spinner and self._should_show_spinner and self._current_message:
elapsed = datetime.now() - self.start_time
elapsed_str = str(elapsed).split('.')[0] # Remove microseconds
# Use the stored current message and update with new elapsed time
self.spinner.update(f"{self._current_message} [dim]({elapsed_str})[/dim]")
Loading
Loading