diff --git a/packages/jumpstarter-cli/jumpstarter_cli/common.py b/packages/jumpstarter-cli/jumpstarter_cli/common.py index cd0968d3c..5e9bbe719 100644 --- a/packages/jumpstarter-cli/jumpstarter_cli/common.py +++ b/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", diff --git a/packages/jumpstarter-cli/jumpstarter_cli/shell.py b/packages/jumpstarter-cli/jumpstarter_cli/shell.py index 4e460dd71..e28a052d4 100644 --- a/packages/jumpstarter-cli/jumpstarter_cli/shell.py +++ b/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) diff --git a/packages/jumpstarter-cli/pyproject.toml b/packages/jumpstarter-cli/pyproject.toml index c395bb589..8d8afa4a8 100644 --- a/packages/jumpstarter-cli/pyproject.toml +++ b/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] diff --git a/packages/jumpstarter/jumpstarter/config/client.py b/packages/jumpstarter/jumpstarter/config/client.py index 8cb24c504..1df373e09 100644 --- a/packages/jumpstarter/jumpstarter/config/client.py +++ b/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 diff --git a/uv.lock b/uv.lock index b8ad24798..e6febc9e6 100644 --- a/uv.lock +++ b/uv.lock @@ -1201,6 +1201,7 @@ source = { editable = "packages/jumpstarter-cli" } dependencies = [ { name = "jumpstarter-cli-admin" }, { name = "jumpstarter-cli-driver" }, + { name = "pytimeparse2" }, ] [package.dev-dependencies] @@ -1215,6 +1216,7 @@ dev = [ requires-dist = [ { name = "jumpstarter-cli-admin", editable = "packages/jumpstarter-cli-admin" }, { name = "jumpstarter-cli-driver", editable = "packages/jumpstarter-cli-driver" }, + { name = "pytimeparse2", specifier = ">=1.7.1" }, ] [package.metadata.requires-dev] @@ -3435,6 +3437,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5f/ed/539768cf28c661b5b068d66d96a2f155c4971a5d55684a514c1a0e0dec2f/python_dotenv-1.1.1-py3-none-any.whl", hash = "sha256:31f23644fe2602f88ff55e1f5c79ba497e01224ee7737937930c448e4d0e24dc", size = 20556, upload-time = "2025-06-24T04:21:06.073Z" }, ] +[[package]] +name = "pytimeparse2" +version = "1.7.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/19/10/cc63fecd69905eb4d300fe71bd580e4a631483e9f53fdcb8c0ad345ce832/pytimeparse2-1.7.1.tar.gz", hash = "sha256:98668cdcba4890e1789e432e8ea0059ccf72402f13f5d52be15bdfaeb3a8b253", size = 10431, upload-time = "2023-05-11T21:40:55.774Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/17/9e/85abf91ef5df452f56498927affdb7128194d15644084f6c6722477c305b/pytimeparse2-1.7.1-py3-none-any.whl", hash = "sha256:a162ea6a7707fd0bb82dd99556efb783935f51885c8bdced0fce3fffe85ab002", size = 6136, upload-time = "2023-05-11T21:40:46.051Z" }, +] + [[package]] name = "pyudev" version = "0.24.3"