Merge pull request #740 from jumpstarter-dev/acquisition-timeout

acquisition timeout support and humanized durations

Author: mangelajo

packages/jumpstarter-cli/jumpstarter_cli/common.py

@@ -3,6 +3,7 @@
 
 import click
 from pydantic import TypeAdapter, ValidationError
+from pytimeparse2 import parse as parse_duration
 
 opt_selector = click.option(
     "-l",
@@ -15,14 +16,59 @@
 class DurationParamType(click.ParamType):
     name = "duration"
 
+    def __init__(self, minimum: timedelta | None = None):
+        super().__init__()
+        self.minimum = minimum
+
     def convert(self, value, param, ctx):
         if isinstance(value, timedelta):
-            return value
+            td = value
+        elif isinstance(value, int):
+            # Integer as seconds (backward compatibility)
+            td = timedelta(seconds=value)
+        elif isinstance(value, str):
+            # Try parsing as plain integer first (backward compatibility)
+            try:
+                int_value = int(value)
+                td = timedelta(seconds=int_value)
+            except ValueError:
+                # Parse with pytimeparse2 first (supports human-readable formats)
+                td = None
+                try:
+                    seconds = parse_duration(value)
+                    if seconds is not None:
+                        td = timedelta(seconds=seconds)
+                except (ValueError, TypeError):
+                    pass
+
+                # Fall back to pydantic/speedate for ISO 8601 and other formats
+                if td is None:
+                    try:
+                        td = TypeAdapter(timedelta).validate_python(value)
+                    except (ValueError, ValidationError):
+                        self.fail(
+                            (
+                                f"{value!r} is not a valid duration "
+                                "(e.g., '30m', '3h30m', '1d', '1d3h40m', 'PT1H30M', '01:30:00')"
+                            ),
+                            param,
+                            ctx,
+                        )
+        else:
+            self.fail(
+                f"{value!r} is not a valid duration (e.g., '30m', '3h30m', '1d', '1d3h40m')",
+                param,
+                ctx,
+            )
 
-        try:
-            return TypeAdapter(timedelta).validate_python(value)
-        except (ValueError, ValidationError):
-            self.fail(f"{value!r} is not a valid duration", param, ctx)
+        # Validate minimum if specified
+        if self.minimum is not None and td < self.minimum:
+            min_seconds = int(self.minimum.total_seconds())
+            self.fail(
+                f"{value!r} must be at least {min_seconds} seconds", param, ctx
+            )
+
+        return td
 
 
 DURATION = DurationParamType()
@@ -36,12 +82,11 @@ def convert(self, value, param, ctx):
 Accepted duration formats:
 
 \b
-PnYnMnDTnHnMnS - ISO 8601 duration format
-HH:MM:SS - time in hours, minutes, seconds
-D days, HH:MM:SS - time prefixed by X days
-D d, HH:MM:SS - time prefixed by X d
+Human-readable: 30m, 3h30m, 1d, 1d3h40m, etc.
+ISO 8601: PT1H30M, P1DT2H30M, etc.
+Time format: 01:30:00, 2 days, 01:30:00, etc.
 
-See https://docs.rs/speedate/latest/speedate/ for details
+See https://github.com/wroberts/pytimeparse2 for details
 """,
 )
 
@@ -67,6 +112,21 @@ def convert(self, value, param, ctx):
 
 DATETIME = DateTimeParamType()
 
+
+ACQUISITION_TIMEOUT = DurationParamType(minimum=timedelta(seconds=5))
+
+opt_acquisition_timeout = partial(
+    click.option,
+    "--acquisition-timeout",
+    "acquisition_timeout",
+    type=ACQUISITION_TIMEOUT,
+    default=None,
+    help=(
+        "Override acquisition timeout (e.g., '30m', '3h30m', '1d', '1d3h40m', "
+        "or seconds as integer). Must be >= 5 seconds."
+    ),
+)
+
 opt_begin_time = click.option(
     "--begin-time",
     "begin_time",

packages/jumpstarter-cli/jumpstarter_cli/shell.py

@@ -8,7 +8,7 @@
 from jumpstarter_cli_common.exceptions import handle_exceptions_with_reauthentication
 from jumpstarter_cli_common.signal import signal_handler
 
-from .common import opt_duration_partial, opt_selector
+from .common import opt_acquisition_timeout, opt_duration_partial, opt_selector
 from .login import relogin_client
 from jumpstarter.common.utils import launch_shell
 from jumpstarter.config.client import ClientConfigV1Alpha1
@@ -33,7 +33,9 @@ def launch_remote_shell(path: str) -> int:
                 return launch_remote_shell(path)
 
 
-async def _shell_with_signal_handling(config, selector, lease_name, duration, exporter_logs, command):
+async def _shell_with_signal_handling(
+    config, selector, lease_name, duration, exporter_logs, command, acquisition_timeout
+):
     """Handle lease acquisition and shell execution with signal handling."""
     exit_code = 0
     cancelled_exc_class = get_cancelled_exc_class()
@@ -43,7 +45,7 @@ async def _shell_with_signal_handling(config, selector, lease_name, duration, ex
         try:
             try:
                 async with anyio.from_thread.BlockingPortal() as portal:
-                    async with config.lease_async(selector, lease_name, duration, portal) as lease:
+                    async with config.lease_async(selector, lease_name, duration, portal, acquisition_timeout) as lease:
                         exit_code = await anyio.to_thread.run_sync(
                             _run_shell_with_lease, lease, exporter_logs, config, command
                         )
@@ -70,9 +72,10 @@ async def _shell_with_signal_handling(config, selector, lease_name, duration, ex
 @opt_selector
 @opt_duration_partial(default=timedelta(minutes=30), show_default="00:30:00")
 @click.option("--exporter-logs", is_flag=True, help="Enable exporter log streaming")
+@opt_acquisition_timeout()
 # end client specific
 @handle_exceptions_with_reauthentication(relogin_client)
-def shell(config, command: tuple[str, ...], lease_name, selector, duration, exporter_logs):
+def shell(config, command: tuple[str, ...], lease_name, selector, duration, exporter_logs, acquisition_timeout):
     """
     Spawns a shell (or custom command) connecting to a local or remote exporter
 
@@ -88,7 +91,14 @@ def shell(config, command: tuple[str, ...], lease_name, selector, duration, expo
     match config:
         case ClientConfigV1Alpha1():
             exit_code = anyio.run(
-                _shell_with_signal_handling, config, selector, lease_name, duration, exporter_logs, command
+                _shell_with_signal_handling,
+                config,
+                selector,
+                lease_name,
+                duration,
+                exporter_logs,
+                command,
+                acquisition_timeout,
             )
             sys.exit(exit_code)
 

packages/jumpstarter-cli/pyproject.toml

@@ -13,6 +13,7 @@ requires-python = ">=3.11"
 dependencies = [
   "jumpstarter-cli-admin",
   "jumpstarter-cli-driver",
+  "pytimeparse2>=1.7.1",
 ]
 
 [dependency-groups]

packages/jumpstarter/jumpstarter/config/client.py

@@ -255,6 +255,7 @@ async def lease_async(
         lease_name: str | None,
         duration: timedelta,
         portal: BlockingPortal,
+        acquisition_timeout: timedelta | None = None,
     ):
         from jumpstarter.client import Lease
 
@@ -263,6 +264,12 @@ async def lease_async(
         # when no lease name is provided, release the lease on exit
         release_lease = lease_name == ""
         try:
+            # Convert timedelta to seconds for acquisition_timeout
+            acquisition_timeout_seconds = (
+                int(acquisition_timeout.total_seconds())
+                if acquisition_timeout is not None
+                else self.leases.acquisition_timeout
+            )
             async with Lease(
                 channel=await self.channel(),
                 namespace=self.metadata.namespace,
@@ -275,7 +282,7 @@ async def lease_async(
                 release=release_lease,
                 tls_config=self.tls,
                 grpc_options=self.grpcOptions,
-                acquisition_timeout=self.leases.acquisition_timeout,
+                acquisition_timeout=acquisition_timeout_seconds,
             ) as lease:
                 yield lease