dh-cli 0.12.3__tar.gz → 0.12.4__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {dh_cli-0.12.3 → dh_cli-0.12.4}/PKG-INFO +1 -1
- {dh_cli-0.12.3 → dh_cli-0.12.4}/pyproject.toml +1 -1
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/engine_commands.py +94 -4
- dh_cli-0.12.4/src/dh_cli/engines_studios/retry.py +79 -0
- dh_cli-0.12.4/tests/engines_studios/test_engine_retry_commands.py +133 -0
- dh_cli-0.12.4/tests/engines_studios/test_retry.py +182 -0
- dh_cli-0.12.4/tests/github/__init__.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/.gitignore +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/LICENSE +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/README.md +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/__init__.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/_identity.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/__init__.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/aws_batch.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/__init__.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/boltz.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/cancel.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/clean.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/disco.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/embed_t5.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/equilibrator.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/esmfold2.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/finalize.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/list_jobs.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/local.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/logs.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/orca.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/protmpnn.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/protmpnn_to_boltz.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/retry.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/rfd3.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/status.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/submit.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/wait_for.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/fasta_utils.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/h5_utils.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/job_id.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/manifest.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/s3_transport.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/bedrock/__init__.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/bedrock/commands.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/bedrock/cost_report.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/bedrock/pricing.yaml +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/cloud_commands.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/codeartifact.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/__init__.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/api_client.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/auth.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/progress.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/ssh_config.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/studio_commands.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/github_commands.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/__init__.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/deploy.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/local.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/test.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/tf.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/users.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/main.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/sweep.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/utility_commands.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/warehouse.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/__init__.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_aws_batch_resources.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_image_override.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_no_train_command.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_submit_cpu_only.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_submit_image_validation.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_submit_merge.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_submit_mlflow_env.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/conftest.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/A_cache_write.json +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/B_cache_read.json +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/C_plain.json +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/D_cursor_user.json +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/E_service_role.json +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/F_legacy_shared.json +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/G_unknown_model.json +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_build_report.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_classify_arn.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_cli_exit_codes.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_cost_calc.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_cost_command.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_cur_reconciliation.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_key_command.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_render_formats.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_resolve_base_model.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_rotate.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_s3_walker.py +0 -0
- {dh_cli-0.12.3/tests/github → dh_cli-0.12.4/tests/engines_studios}/__init__.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/conftest.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_engine_role_cannot_read_github_pat.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_identity.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_login.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_login_error_paths.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_login_security.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_logout.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_rotate.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_status.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/hz/test_init.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/hz/test_suites.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/hz/test_users.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_cloud_gcp.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_disco_command.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_disco_jobdef.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_equilibrator_command.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_esmfold2_cli.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_boltz_tar.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_disco.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_disco_dispatch.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_equilibrator.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_esmfold2.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_protmpnn.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_rfd3.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_rfd3_command.py +0 -0
- {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_sweep.py +0 -0
|
@@ -8,8 +8,65 @@ import click
|
|
|
8
8
|
from .api_client import StudioManagerClient
|
|
9
9
|
from .auth import check_aws_auth, detect_aws_environment, get_aws_username
|
|
10
10
|
from .progress import format_idle_state, format_time_ago, wait_with_progress
|
|
11
|
+
from .retry import retry_on_capacity
|
|
11
12
|
from .ssh_config import update_ssh_config_silent
|
|
12
13
|
|
|
14
|
+
# Default capacity-retry budget: transient InsufficientInstanceCapacity in our
|
|
15
|
+
# G-family AZs usually clears within a few minutes (see topic-quota-and-caps.md).
|
|
16
|
+
_DEFAULT_RETRY_TIMEOUT = 600 # seconds (~10 min)
|
|
17
|
+
_DEFAULT_RETRY_INTERVAL = 30 # seconds
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def _retry_options(func):
|
|
21
|
+
"""Shared --retry / --no-retry / interval options for launch and start."""
|
|
22
|
+
func = click.option(
|
|
23
|
+
"--no-retry",
|
|
24
|
+
is_flag=True,
|
|
25
|
+
help="Fail immediately on capacity errors instead of retrying.",
|
|
26
|
+
)(func)
|
|
27
|
+
func = click.option(
|
|
28
|
+
"--retry-timeout",
|
|
29
|
+
type=int,
|
|
30
|
+
default=_DEFAULT_RETRY_TIMEOUT,
|
|
31
|
+
show_default=True,
|
|
32
|
+
help="Max seconds to keep retrying transient capacity errors (0 disables).",
|
|
33
|
+
)(func)
|
|
34
|
+
func = click.option(
|
|
35
|
+
"--retry-interval",
|
|
36
|
+
type=int,
|
|
37
|
+
default=_DEFAULT_RETRY_INTERVAL,
|
|
38
|
+
show_default=True,
|
|
39
|
+
help="Seconds to wait between capacity-retry attempts.",
|
|
40
|
+
)(func)
|
|
41
|
+
return func
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _make_capacity_retry(engine_type: str, no_retry: bool, timeout: int, interval: int):
|
|
45
|
+
"""Build a retry wrapper that echoes status between capacity retries.
|
|
46
|
+
|
|
47
|
+
Returns a callable taking a zero-arg `action` and running it with bounded
|
|
48
|
+
capacity-retry. With `no_retry` (or timeout 0) it runs the action once.
|
|
49
|
+
"""
|
|
50
|
+
effective_timeout = 0 if no_retry else timeout
|
|
51
|
+
|
|
52
|
+
def on_retry(attempt: int, error: BaseException, wait_seconds: float) -> None:
|
|
53
|
+
click.echo(
|
|
54
|
+
f"⏳ Attempt {attempt} failed: no {engine_type} capacity available "
|
|
55
|
+
f"right now (AWS-side). Retrying in {int(wait_seconds)}s… "
|
|
56
|
+
f"(Ctrl+C to cancel)",
|
|
57
|
+
err=True,
|
|
58
|
+
)
|
|
59
|
+
|
|
60
|
+
def run(action):
|
|
61
|
+
return retry_on_capacity(
|
|
62
|
+
action,
|
|
63
|
+
timeout=effective_timeout,
|
|
64
|
+
interval=interval,
|
|
65
|
+
on_retry=on_retry,
|
|
66
|
+
)
|
|
67
|
+
|
|
68
|
+
return run
|
|
69
|
+
|
|
13
70
|
|
|
14
71
|
@click.group()
|
|
15
72
|
def engine_cli():
|
|
@@ -44,6 +101,7 @@ def engine_cli():
|
|
|
44
101
|
default=None,
|
|
45
102
|
help="Environment (dev, sand, prod) - auto-detected if not specified",
|
|
46
103
|
)
|
|
104
|
+
@_retry_options
|
|
47
105
|
def launch_engine(
|
|
48
106
|
name: str,
|
|
49
107
|
engine_type: str,
|
|
@@ -53,6 +111,9 @@ def launch_engine(
|
|
|
53
111
|
no_wait: bool,
|
|
54
112
|
skip_ssh_config: bool,
|
|
55
113
|
env: Optional[str],
|
|
114
|
+
retry_interval: int,
|
|
115
|
+
retry_timeout: int,
|
|
116
|
+
no_retry: bool,
|
|
56
117
|
):
|
|
57
118
|
"""Launch a new engine for the current user (or specified user with --user flag)."""
|
|
58
119
|
|
|
@@ -83,8 +144,13 @@ def launch_engine(
|
|
|
83
144
|
click.echo(f"🚀 Launching {engine_type} engine '{name}' for {user}...")
|
|
84
145
|
|
|
85
146
|
try:
|
|
86
|
-
# Launch the engine
|
|
87
|
-
|
|
147
|
+
# Launch the engine, retrying transient AWS capacity errors by default.
|
|
148
|
+
retry = _make_capacity_retry(engine_type, no_retry, retry_timeout, retry_interval)
|
|
149
|
+
engine = retry(
|
|
150
|
+
lambda: client.launch_engine(
|
|
151
|
+
name=name, user=user, engine_type=engine_type, boot_disk_size=boot_disk_size
|
|
152
|
+
)
|
|
153
|
+
)
|
|
88
154
|
|
|
89
155
|
engine_id = engine["instance_id"]
|
|
90
156
|
click.echo(f"✓ EC2 instance launched: {engine_id}")
|
|
@@ -128,6 +194,10 @@ def launch_engine(
|
|
|
128
194
|
if update_ssh_config_silent(client, env):
|
|
129
195
|
click.echo("✓ SSH config updated")
|
|
130
196
|
|
|
197
|
+
except KeyboardInterrupt:
|
|
198
|
+
click.echo("\n✗ Cancelled while waiting for capacity.", err=True)
|
|
199
|
+
raise click.Abort()
|
|
200
|
+
|
|
131
201
|
except Exception as e:
|
|
132
202
|
error_msg = str(e)
|
|
133
203
|
|
|
@@ -197,7 +267,17 @@ def launch_engine(
|
|
|
197
267
|
default=None,
|
|
198
268
|
help="Environment (dev, sand, prod) - auto-detected if not specified",
|
|
199
269
|
)
|
|
200
|
-
|
|
270
|
+
@_retry_options
|
|
271
|
+
def start_engine(
|
|
272
|
+
name_or_id: str,
|
|
273
|
+
no_wait: bool,
|
|
274
|
+
skip_ssh_config: bool,
|
|
275
|
+
yes: bool,
|
|
276
|
+
env: Optional[str],
|
|
277
|
+
retry_interval: int,
|
|
278
|
+
retry_timeout: int,
|
|
279
|
+
no_retry: bool,
|
|
280
|
+
):
|
|
201
281
|
"""Start a stopped engine."""
|
|
202
282
|
|
|
203
283
|
# Check AWS auth first to provide clear error messages
|
|
@@ -225,7 +305,13 @@ def start_engine(name_or_id: str, no_wait: bool, skip_ssh_config: bool, yes: boo
|
|
|
225
305
|
engine_id = engine["instance_id"]
|
|
226
306
|
engine_name = engine.get("name", engine_id)
|
|
227
307
|
|
|
228
|
-
|
|
308
|
+
# Start the instance, retrying transient AWS capacity errors by default.
|
|
309
|
+
# Note: a stopped engine restarts in its ORIGINAL AZ, so retry here is a
|
|
310
|
+
# time-wait for that one AZ's capacity to return — it cannot fall to
|
|
311
|
+
# another AZ (unlike launch).
|
|
312
|
+
engine_type_label = engine.get("engine_type", "engine")
|
|
313
|
+
retry = _make_capacity_retry(engine_type_label, no_retry, retry_timeout, retry_interval)
|
|
314
|
+
result = retry(lambda: client.start_engine(engine_id))
|
|
229
315
|
|
|
230
316
|
if "error" in result:
|
|
231
317
|
click.echo(f"✗ Error: {result['error']}", err=True)
|
|
@@ -305,6 +391,10 @@ def start_engine(name_or_id: str, no_wait: bool, skip_ssh_config: bool, yes: boo
|
|
|
305
391
|
if update_ssh_config_silent(client, env):
|
|
306
392
|
click.echo("✓ SSH config updated")
|
|
307
393
|
|
|
394
|
+
except KeyboardInterrupt:
|
|
395
|
+
click.echo("\n✗ Cancelled while waiting for capacity.", err=True)
|
|
396
|
+
raise click.Abort()
|
|
397
|
+
|
|
308
398
|
except Exception as e:
|
|
309
399
|
click.echo(f"✗ Error: {e}", err=True)
|
|
310
400
|
raise click.Abort()
|
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
"""Bounded capacity-retry for engine launch/start.
|
|
2
|
+
|
|
3
|
+
AWS periodically has no on-demand G-family capacity in our AZs, surfacing as
|
|
4
|
+
`InsufficientInstanceCapacity` (and, when a subnet is IP-starved,
|
|
5
|
+
`InsufficientFreeAddressesInSubnet`). These are transient: capacity returns
|
|
6
|
+
within minutes. This helper retries such failures within a wall-clock budget
|
|
7
|
+
so users (and agents) don't have to re-run the command by hand.
|
|
8
|
+
|
|
9
|
+
Quota errors (`VcpuLimitExceeded`), instance-count limits
|
|
10
|
+
(`InstanceLimitExceeded`), and validation errors are NOT capacity problems and
|
|
11
|
+
are never retried — they will not resolve by waiting.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
from __future__ import annotations
|
|
15
|
+
|
|
16
|
+
import time
|
|
17
|
+
from typing import Callable, Optional, TypeVar
|
|
18
|
+
|
|
19
|
+
T = TypeVar("T")
|
|
20
|
+
|
|
21
|
+
# Substrings that mark a transient, retryable AWS capacity/IP shortage.
|
|
22
|
+
_CAPACITY_MARKERS = (
|
|
23
|
+
"InsufficientInstanceCapacity",
|
|
24
|
+
"InsufficientFreeAddressesInSubnet",
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def is_capacity_error(error: BaseException) -> bool:
|
|
29
|
+
"""True if `error` is a transient AWS capacity/IP shortage worth retrying."""
|
|
30
|
+
message = str(error)
|
|
31
|
+
return any(marker in message for marker in _CAPACITY_MARKERS)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def retry_on_capacity(
|
|
35
|
+
action: Callable[[], T],
|
|
36
|
+
*,
|
|
37
|
+
timeout: float,
|
|
38
|
+
interval: float,
|
|
39
|
+
on_retry: Optional[Callable[[int, BaseException, float], None]] = None,
|
|
40
|
+
monotonic: Callable[[], float] = time.monotonic,
|
|
41
|
+
sleep: Callable[[float], None] = time.sleep,
|
|
42
|
+
) -> T:
|
|
43
|
+
"""Call `action`, retrying transient capacity errors within a budget.
|
|
44
|
+
|
|
45
|
+
Args:
|
|
46
|
+
action: Zero-arg callable performing the launch/start attempt.
|
|
47
|
+
timeout: Total wall-clock budget in seconds. A new attempt is only
|
|
48
|
+
started while elapsed time is strictly below this. ``0`` yields a
|
|
49
|
+
single attempt (no retries).
|
|
50
|
+
interval: Seconds to wait between attempts.
|
|
51
|
+
on_retry: Optional callback invoked before each sleep with
|
|
52
|
+
``(attempt_number, error, wait_seconds)`` — used to print status.
|
|
53
|
+
monotonic: Monotonic clock source (injected for testing).
|
|
54
|
+
sleep: Sleep function (injected for testing).
|
|
55
|
+
|
|
56
|
+
Returns:
|
|
57
|
+
Whatever `action` returns on its first success.
|
|
58
|
+
|
|
59
|
+
Raises:
|
|
60
|
+
The last capacity error if the budget is exhausted, or immediately
|
|
61
|
+
re-raises any non-capacity error from `action`.
|
|
62
|
+
"""
|
|
63
|
+
start = monotonic()
|
|
64
|
+
attempt = 0
|
|
65
|
+
last_error: Optional[BaseException] = None
|
|
66
|
+
while True:
|
|
67
|
+
attempt += 1
|
|
68
|
+
try:
|
|
69
|
+
return action()
|
|
70
|
+
except Exception as error: # noqa: BLE001 — re-raised below unless retryable
|
|
71
|
+
if not is_capacity_error(error):
|
|
72
|
+
raise
|
|
73
|
+
last_error = error
|
|
74
|
+
# Stop if the next attempt would not begin within the budget.
|
|
75
|
+
if (monotonic() - start) + interval >= timeout:
|
|
76
|
+
raise last_error
|
|
77
|
+
if on_retry is not None:
|
|
78
|
+
on_retry(attempt, error, interval)
|
|
79
|
+
sleep(interval)
|
|
@@ -0,0 +1,133 @@
|
|
|
1
|
+
"""Command-level tests for capacity-retry wiring in `dh engine launch/start`.
|
|
2
|
+
|
|
3
|
+
These invoke the Click commands with a mocked `StudioManagerClient` and a
|
|
4
|
+
patched sleep, asserting that a transient `InsufficientInstanceCapacity`
|
|
5
|
+
is retried by default, that `--no-retry` disables it, and that non-capacity
|
|
6
|
+
errors still fail fast.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
from unittest.mock import MagicMock, patch
|
|
12
|
+
|
|
13
|
+
import pytest
|
|
14
|
+
from click.testing import CliRunner
|
|
15
|
+
|
|
16
|
+
from dh_cli.engines_studios import engine_commands as ec
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@pytest.fixture
|
|
20
|
+
def runner():
|
|
21
|
+
return CliRunner()
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@pytest.fixture(autouse=True)
|
|
25
|
+
def _fast_no_side_effects():
|
|
26
|
+
"""Neutralize auth/env detection, SSH-config writes, and real sleeping."""
|
|
27
|
+
with (
|
|
28
|
+
patch.object(ec, "check_aws_auth", lambda: None),
|
|
29
|
+
patch.object(ec, "detect_aws_environment", lambda: "dev"),
|
|
30
|
+
patch.object(ec, "get_aws_username", lambda: "tester"),
|
|
31
|
+
patch.object(ec, "update_ssh_config_silent", lambda *a, **k: False),
|
|
32
|
+
# Retry helper sleeps via time.sleep in retry module; make it instant.
|
|
33
|
+
patch("dh_cli.engines_studios.retry.time.sleep", lambda s: None),
|
|
34
|
+
):
|
|
35
|
+
yield
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _capacity_error():
|
|
39
|
+
return RuntimeError(
|
|
40
|
+
"Failed to launch instance: An error occurred "
|
|
41
|
+
"(InsufficientInstanceCapacity) when calling the RunInstances operation"
|
|
42
|
+
)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def test_launch_retries_capacity_then_succeeds(runner):
|
|
46
|
+
client = MagicMock()
|
|
47
|
+
# Fail twice with capacity, then succeed.
|
|
48
|
+
client.launch_engine.side_effect = [
|
|
49
|
+
_capacity_error(),
|
|
50
|
+
_capacity_error(),
|
|
51
|
+
{"instance_id": "i-abc", "instance_type": "g5.2xlarge", "state": "pending"},
|
|
52
|
+
]
|
|
53
|
+
|
|
54
|
+
with patch.object(ec, "StudioManagerClient", return_value=client):
|
|
55
|
+
result = runner.invoke(
|
|
56
|
+
ec.launch_engine,
|
|
57
|
+
["myeng", "--type", "a10g", "--no-wait", "--retry-interval", "1"],
|
|
58
|
+
)
|
|
59
|
+
|
|
60
|
+
assert result.exit_code == 0, result.output
|
|
61
|
+
assert client.launch_engine.call_count == 3
|
|
62
|
+
assert "i-abc" in result.output
|
|
63
|
+
assert "Retrying" in result.output
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def test_launch_no_retry_fails_fast_on_capacity(runner):
|
|
67
|
+
client = MagicMock()
|
|
68
|
+
client.launch_engine.side_effect = _capacity_error()
|
|
69
|
+
|
|
70
|
+
with patch.object(ec, "StudioManagerClient", return_value=client):
|
|
71
|
+
result = runner.invoke(
|
|
72
|
+
ec.launch_engine,
|
|
73
|
+
["myeng", "--type", "a10g", "--no-wait", "--no-retry"],
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
assert result.exit_code != 0
|
|
77
|
+
assert client.launch_engine.call_count == 1
|
|
78
|
+
assert "insufficient EC2 capacity" in result.output
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def test_launch_does_not_retry_quota_error(runner):
|
|
82
|
+
client = MagicMock()
|
|
83
|
+
client.launch_engine.side_effect = RuntimeError("VcpuLimitExceeded: quota")
|
|
84
|
+
|
|
85
|
+
with patch.object(ec, "StudioManagerClient", return_value=client):
|
|
86
|
+
result = runner.invoke(
|
|
87
|
+
ec.launch_engine,
|
|
88
|
+
["myeng", "--type", "a10g", "--no-wait", "--retry-interval", "1"],
|
|
89
|
+
)
|
|
90
|
+
|
|
91
|
+
assert result.exit_code != 0
|
|
92
|
+
assert client.launch_engine.call_count == 1
|
|
93
|
+
assert "quota exceeded" in result.output
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def test_start_retries_capacity_then_succeeds(runner):
|
|
97
|
+
client = MagicMock()
|
|
98
|
+
client.get_engine_by_name.return_value = {
|
|
99
|
+
"instance_id": "i-xyz",
|
|
100
|
+
"name": "myeng",
|
|
101
|
+
}
|
|
102
|
+
client.start_engine.side_effect = [
|
|
103
|
+
_capacity_error(),
|
|
104
|
+
{"instance_id": "i-xyz", "status": "pending"},
|
|
105
|
+
]
|
|
106
|
+
|
|
107
|
+
with patch.object(ec, "StudioManagerClient", return_value=client):
|
|
108
|
+
result = runner.invoke(
|
|
109
|
+
ec.start_engine,
|
|
110
|
+
["myeng", "--no-wait", "--retry-interval", "1"],
|
|
111
|
+
)
|
|
112
|
+
|
|
113
|
+
assert result.exit_code == 0, result.output
|
|
114
|
+
assert client.start_engine.call_count == 2
|
|
115
|
+
assert "Retrying" in result.output
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
def test_start_no_retry_fails_fast_on_capacity(runner):
|
|
119
|
+
client = MagicMock()
|
|
120
|
+
client.get_engine_by_name.return_value = {
|
|
121
|
+
"instance_id": "i-xyz",
|
|
122
|
+
"name": "myeng",
|
|
123
|
+
}
|
|
124
|
+
client.start_engine.side_effect = _capacity_error()
|
|
125
|
+
|
|
126
|
+
with patch.object(ec, "StudioManagerClient", return_value=client):
|
|
127
|
+
result = runner.invoke(
|
|
128
|
+
ec.start_engine,
|
|
129
|
+
["myeng", "--no-wait", "--no-retry"],
|
|
130
|
+
)
|
|
131
|
+
|
|
132
|
+
assert result.exit_code != 0
|
|
133
|
+
assert client.start_engine.call_count == 1
|
|
@@ -0,0 +1,182 @@
|
|
|
1
|
+
"""Tests for the capacity-retry helper used by `dh engine launch` / `start`.
|
|
2
|
+
|
|
3
|
+
The helper's whole job is: keep calling an action while it fails with an
|
|
4
|
+
AWS *capacity* error, within a bounded wall-clock budget, and re-raise the
|
|
5
|
+
last error once the budget is spent. Non-capacity errors must propagate
|
|
6
|
+
immediately. Time is injected so the tests are instant and deterministic.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import pytest
|
|
12
|
+
|
|
13
|
+
from dh_cli.engines_studios.retry import is_capacity_error, retry_on_capacity
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class FakeClock:
|
|
17
|
+
"""Deterministic monotonic clock; `sleep` advances it, records intervals."""
|
|
18
|
+
|
|
19
|
+
def __init__(self) -> None:
|
|
20
|
+
self.now = 0.0
|
|
21
|
+
self.slept: list[float] = []
|
|
22
|
+
|
|
23
|
+
def monotonic(self) -> float:
|
|
24
|
+
return self.now
|
|
25
|
+
|
|
26
|
+
def sleep(self, seconds: float) -> None:
|
|
27
|
+
self.slept.append(seconds)
|
|
28
|
+
self.now += seconds
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def test_returns_immediately_on_success():
|
|
32
|
+
clock = FakeClock()
|
|
33
|
+
calls = []
|
|
34
|
+
|
|
35
|
+
def action():
|
|
36
|
+
calls.append(1)
|
|
37
|
+
return "ok"
|
|
38
|
+
|
|
39
|
+
result = retry_on_capacity(
|
|
40
|
+
action, timeout=600, interval=30, monotonic=clock.monotonic, sleep=clock.sleep
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
assert result == "ok"
|
|
44
|
+
assert len(calls) == 1
|
|
45
|
+
assert clock.slept == []
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_retries_capacity_error_then_succeeds():
|
|
49
|
+
clock = FakeClock()
|
|
50
|
+
attempts = {"n": 0}
|
|
51
|
+
|
|
52
|
+
def action():
|
|
53
|
+
attempts["n"] += 1
|
|
54
|
+
if attempts["n"] < 3:
|
|
55
|
+
raise RuntimeError(
|
|
56
|
+
"Failed to launch instance: An error occurred "
|
|
57
|
+
"(InsufficientInstanceCapacity) when calling RunInstances"
|
|
58
|
+
)
|
|
59
|
+
return "launched"
|
|
60
|
+
|
|
61
|
+
result = retry_on_capacity(
|
|
62
|
+
action, timeout=600, interval=30, monotonic=clock.monotonic, sleep=clock.sleep
|
|
63
|
+
)
|
|
64
|
+
|
|
65
|
+
assert result == "launched"
|
|
66
|
+
assert attempts["n"] == 3
|
|
67
|
+
# Slept between the two failures (not after success).
|
|
68
|
+
assert clock.slept == [30, 30]
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def test_non_capacity_error_propagates_without_retry():
|
|
72
|
+
clock = FakeClock()
|
|
73
|
+
attempts = {"n": 0}
|
|
74
|
+
|
|
75
|
+
def action():
|
|
76
|
+
attempts["n"] += 1
|
|
77
|
+
raise RuntimeError("VcpuLimitExceeded: vCPU quota exceeded")
|
|
78
|
+
|
|
79
|
+
with pytest.raises(RuntimeError, match="VcpuLimitExceeded"):
|
|
80
|
+
retry_on_capacity(
|
|
81
|
+
action, timeout=600, interval=30, monotonic=clock.monotonic, sleep=clock.sleep
|
|
82
|
+
)
|
|
83
|
+
|
|
84
|
+
assert attempts["n"] == 1
|
|
85
|
+
assert clock.slept == []
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
def test_reraises_last_capacity_error_after_budget_exhausted():
|
|
89
|
+
clock = FakeClock()
|
|
90
|
+
attempts = {"n": 0}
|
|
91
|
+
|
|
92
|
+
def action():
|
|
93
|
+
attempts["n"] += 1
|
|
94
|
+
raise RuntimeError("InsufficientInstanceCapacity: no capacity")
|
|
95
|
+
|
|
96
|
+
with pytest.raises(RuntimeError, match="InsufficientInstanceCapacity"):
|
|
97
|
+
retry_on_capacity(
|
|
98
|
+
action, timeout=60, interval=30, monotonic=clock.monotonic, sleep=clock.sleep
|
|
99
|
+
)
|
|
100
|
+
|
|
101
|
+
# attempt@t=0 fails; next attempt (t=30) fits in the 60s budget -> sleep 30.
|
|
102
|
+
# attempt@t=30 fails; next attempt (t=60) would NOT be < 60 -> stop, no sleep.
|
|
103
|
+
assert attempts["n"] == 2
|
|
104
|
+
assert clock.slept == [30]
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def test_timeout_zero_means_single_attempt():
|
|
108
|
+
clock = FakeClock()
|
|
109
|
+
attempts = {"n": 0}
|
|
110
|
+
|
|
111
|
+
def action():
|
|
112
|
+
attempts["n"] += 1
|
|
113
|
+
raise RuntimeError("InsufficientInstanceCapacity")
|
|
114
|
+
|
|
115
|
+
with pytest.raises(RuntimeError, match="InsufficientInstanceCapacity"):
|
|
116
|
+
retry_on_capacity(
|
|
117
|
+
action, timeout=0, interval=30, monotonic=clock.monotonic, sleep=clock.sleep
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
assert attempts["n"] == 1
|
|
121
|
+
assert clock.slept == []
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def test_on_retry_callback_receives_attempt_and_error():
|
|
125
|
+
clock = FakeClock()
|
|
126
|
+
seen = []
|
|
127
|
+
attempts = {"n": 0}
|
|
128
|
+
|
|
129
|
+
def action():
|
|
130
|
+
attempts["n"] += 1
|
|
131
|
+
if attempts["n"] < 2:
|
|
132
|
+
raise RuntimeError("InsufficientInstanceCapacity")
|
|
133
|
+
return "ok"
|
|
134
|
+
|
|
135
|
+
def on_retry(attempt, error, wait_seconds):
|
|
136
|
+
seen.append((attempt, str(error), wait_seconds))
|
|
137
|
+
|
|
138
|
+
retry_on_capacity(
|
|
139
|
+
action,
|
|
140
|
+
timeout=600,
|
|
141
|
+
interval=30,
|
|
142
|
+
on_retry=on_retry,
|
|
143
|
+
monotonic=clock.monotonic,
|
|
144
|
+
sleep=clock.sleep,
|
|
145
|
+
)
|
|
146
|
+
|
|
147
|
+
assert len(seen) == 1
|
|
148
|
+
attempt, msg, wait = seen[0]
|
|
149
|
+
assert attempt == 1
|
|
150
|
+
assert "InsufficientInstanceCapacity" in msg
|
|
151
|
+
assert wait == 30
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
def test_keyboard_interrupt_during_sleep_propagates():
|
|
155
|
+
class InterruptingClock(FakeClock):
|
|
156
|
+
def sleep(self, seconds):
|
|
157
|
+
raise KeyboardInterrupt
|
|
158
|
+
|
|
159
|
+
clock = InterruptingClock()
|
|
160
|
+
|
|
161
|
+
def action():
|
|
162
|
+
raise RuntimeError("InsufficientInstanceCapacity")
|
|
163
|
+
|
|
164
|
+
with pytest.raises(KeyboardInterrupt):
|
|
165
|
+
retry_on_capacity(
|
|
166
|
+
action, timeout=600, interval=30, monotonic=clock.monotonic, sleep=clock.sleep
|
|
167
|
+
)
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
@pytest.mark.parametrize(
|
|
171
|
+
"message,expected",
|
|
172
|
+
[
|
|
173
|
+
("InsufficientInstanceCapacity", True),
|
|
174
|
+
("Failed to launch instance: (InsufficientInstanceCapacity) ...", True),
|
|
175
|
+
("InsufficientFreeAddressesInSubnet", True),
|
|
176
|
+
("VcpuLimitExceeded", False),
|
|
177
|
+
("InstanceLimitExceeded", False),
|
|
178
|
+
("some validation error", False),
|
|
179
|
+
],
|
|
180
|
+
)
|
|
181
|
+
def test_is_capacity_error_classification(message, expected):
|
|
182
|
+
assert is_capacity_error(RuntimeError(message)) is expected
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|