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.
Files changed (116) hide show
  1. {dh_cli-0.12.3 → dh_cli-0.12.4}/PKG-INFO +1 -1
  2. {dh_cli-0.12.3 → dh_cli-0.12.4}/pyproject.toml +1 -1
  3. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/engine_commands.py +94 -4
  4. dh_cli-0.12.4/src/dh_cli/engines_studios/retry.py +79 -0
  5. dh_cli-0.12.4/tests/engines_studios/test_engine_retry_commands.py +133 -0
  6. dh_cli-0.12.4/tests/engines_studios/test_retry.py +182 -0
  7. dh_cli-0.12.4/tests/github/__init__.py +0 -0
  8. {dh_cli-0.12.3 → dh_cli-0.12.4}/.gitignore +0 -0
  9. {dh_cli-0.12.3 → dh_cli-0.12.4}/LICENSE +0 -0
  10. {dh_cli-0.12.3 → dh_cli-0.12.4}/README.md +0 -0
  11. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/__init__.py +0 -0
  12. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/_identity.py +0 -0
  13. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/__init__.py +0 -0
  14. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/aws_batch.py +0 -0
  15. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/__init__.py +0 -0
  16. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/boltz.py +0 -0
  17. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/cancel.py +0 -0
  18. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/clean.py +0 -0
  19. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/disco.py +0 -0
  20. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/embed_t5.py +0 -0
  21. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/equilibrator.py +0 -0
  22. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/esmfold2.py +0 -0
  23. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/finalize.py +0 -0
  24. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/list_jobs.py +0 -0
  25. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/local.py +0 -0
  26. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/logs.py +0 -0
  27. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/orca.py +0 -0
  28. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/protmpnn.py +0 -0
  29. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/protmpnn_to_boltz.py +0 -0
  30. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/retry.py +0 -0
  31. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/rfd3.py +0 -0
  32. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/status.py +0 -0
  33. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/submit.py +0 -0
  34. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/commands/wait_for.py +0 -0
  35. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/fasta_utils.py +0 -0
  36. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/h5_utils.py +0 -0
  37. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/job_id.py +0 -0
  38. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/manifest.py +0 -0
  39. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/batch/s3_transport.py +0 -0
  40. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/bedrock/__init__.py +0 -0
  41. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/bedrock/commands.py +0 -0
  42. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/bedrock/cost_report.py +0 -0
  43. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/bedrock/pricing.yaml +0 -0
  44. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/cloud_commands.py +0 -0
  45. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/codeartifact.py +0 -0
  46. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/__init__.py +0 -0
  47. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/api_client.py +0 -0
  48. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/auth.py +0 -0
  49. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/progress.py +0 -0
  50. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/ssh_config.py +0 -0
  51. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/engines_studios/studio_commands.py +0 -0
  52. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/github_commands.py +0 -0
  53. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/__init__.py +0 -0
  54. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/deploy.py +0 -0
  55. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/local.py +0 -0
  56. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/test.py +0 -0
  57. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/tf.py +0 -0
  58. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/hz/users.py +0 -0
  59. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/main.py +0 -0
  60. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/sweep.py +0 -0
  61. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/utility_commands.py +0 -0
  62. {dh_cli-0.12.3 → dh_cli-0.12.4}/src/dh_cli/warehouse.py +0 -0
  63. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/__init__.py +0 -0
  64. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_aws_batch_resources.py +0 -0
  65. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_image_override.py +0 -0
  66. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_no_train_command.py +0 -0
  67. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_submit_cpu_only.py +0 -0
  68. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_submit_image_validation.py +0 -0
  69. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_submit_merge.py +0 -0
  70. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/batch/test_submit_mlflow_env.py +0 -0
  71. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/conftest.py +0 -0
  72. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/A_cache_write.json +0 -0
  73. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/B_cache_read.json +0 -0
  74. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/C_plain.json +0 -0
  75. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/D_cursor_user.json +0 -0
  76. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/E_service_role.json +0 -0
  77. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/F_legacy_shared.json +0 -0
  78. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/fixtures/G_unknown_model.json +0 -0
  79. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_build_report.py +0 -0
  80. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_classify_arn.py +0 -0
  81. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_cli_exit_codes.py +0 -0
  82. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_cost_calc.py +0 -0
  83. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_cost_command.py +0 -0
  84. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_cur_reconciliation.py +0 -0
  85. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_key_command.py +0 -0
  86. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_render_formats.py +0 -0
  87. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_resolve_base_model.py +0 -0
  88. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_rotate.py +0 -0
  89. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/bedrock/test_s3_walker.py +0 -0
  90. {dh_cli-0.12.3/tests/github → dh_cli-0.12.4/tests/engines_studios}/__init__.py +0 -0
  91. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/conftest.py +0 -0
  92. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_engine_role_cannot_read_github_pat.py +0 -0
  93. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_identity.py +0 -0
  94. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_login.py +0 -0
  95. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_login_error_paths.py +0 -0
  96. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_login_security.py +0 -0
  97. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_logout.py +0 -0
  98. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_rotate.py +0 -0
  99. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/github/test_status.py +0 -0
  100. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/hz/test_init.py +0 -0
  101. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/hz/test_suites.py +0 -0
  102. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/hz/test_users.py +0 -0
  103. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_cloud_gcp.py +0 -0
  104. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_disco_command.py +0 -0
  105. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_disco_jobdef.py +0 -0
  106. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_equilibrator_command.py +0 -0
  107. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_esmfold2_cli.py +0 -0
  108. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_boltz_tar.py +0 -0
  109. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_disco.py +0 -0
  110. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_disco_dispatch.py +0 -0
  111. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_equilibrator.py +0 -0
  112. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_esmfold2.py +0 -0
  113. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_protmpnn.py +0 -0
  114. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_finalize_rfd3.py +0 -0
  115. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_rfd3_command.py +0 -0
  116. {dh_cli-0.12.3 → dh_cli-0.12.4}/tests/test_sweep.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dh-cli
3
- Version: 0.12.3
3
+ Version: 0.12.4
4
4
  Summary: Dayhoff Labs developer CLI
5
5
  Author-email: Dayhoff Labs <dev@dayhofflabs.com>
6
6
  License: # PolyForm Noncommercial License 1.0.0
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "dh-cli"
7
- version = "0.12.3"
7
+ version = "0.12.4"
8
8
  description = "Dayhoff Labs developer CLI"
9
9
  requires-python = ">=3.11"
10
10
  readme = "README.md"
@@ -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
- engine = client.launch_engine(name=name, user=user, engine_type=engine_type, boot_disk_size=boot_disk_size)
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
- def start_engine(name_or_id: str, no_wait: bool, skip_ssh_config: bool, yes: bool, env: Optional[str]):
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
- result = client.start_engine(engine_id)
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