dh-cli 0.7.1__tar.gz → 0.8.0__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 (91) hide show
  1. {dh_cli-0.7.1 → dh_cli-0.8.0}/PKG-INFO +1 -1
  2. {dh_cli-0.7.1 → dh_cli-0.8.0}/pyproject.toml +1 -1
  3. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/aws_batch.py +1 -1
  4. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/submit.py +3 -2
  5. dh_cli-0.8.0/src/dh_cli/github_commands.py +275 -0
  6. dh_cli-0.8.0/tests/batch/__init__.py +0 -0
  7. dh_cli-0.8.0/tests/batch/test_aws_batch_resources.py +92 -0
  8. dh_cli-0.8.0/tests/batch/test_submit_cpu_only.py +162 -0
  9. dh_cli-0.7.1/src/dh_cli/_identity.py +0 -82
  10. dh_cli-0.7.1/src/dh_cli/github_commands.py +0 -794
  11. dh_cli-0.7.1/tests/github/__init__.py +0 -1
  12. dh_cli-0.7.1/tests/github/conftest.py +0 -197
  13. dh_cli-0.7.1/tests/github/test_engine_role_cannot_read_github_pat.py +0 -46
  14. dh_cli-0.7.1/tests/github/test_identity.py +0 -74
  15. dh_cli-0.7.1/tests/github/test_login.py +0 -272
  16. dh_cli-0.7.1/tests/github/test_login_error_paths.py +0 -152
  17. dh_cli-0.7.1/tests/github/test_login_security.py +0 -152
  18. dh_cli-0.7.1/tests/github/test_logout.py +0 -39
  19. dh_cli-0.7.1/tests/github/test_rotate.py +0 -226
  20. dh_cli-0.7.1/tests/github/test_status.py +0 -93
  21. {dh_cli-0.7.1 → dh_cli-0.8.0}/.gitignore +0 -0
  22. {dh_cli-0.7.1 → dh_cli-0.8.0}/LICENSE +0 -0
  23. {dh_cli-0.7.1 → dh_cli-0.8.0}/README.md +0 -0
  24. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/__init__.py +0 -0
  25. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/__init__.py +0 -0
  26. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/__init__.py +0 -0
  27. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/boltz.py +0 -0
  28. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/cancel.py +0 -0
  29. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/clean.py +0 -0
  30. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/embed_t5.py +0 -0
  31. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/finalize.py +0 -0
  32. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/list_jobs.py +0 -0
  33. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/local.py +0 -0
  34. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/logs.py +0 -0
  35. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/orca.py +0 -0
  36. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/protmpnn.py +0 -0
  37. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/protmpnn_to_boltz.py +0 -0
  38. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/retry.py +0 -0
  39. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/status.py +0 -0
  40. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/train.py +0 -0
  41. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/commands/wait_for.py +0 -0
  42. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/fasta_utils.py +0 -0
  43. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/h5_utils.py +0 -0
  44. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/job_id.py +0 -0
  45. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/manifest.py +0 -0
  46. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/batch/s3_transport.py +0 -0
  47. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/bedrock/__init__.py +0 -0
  48. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/bedrock/commands.py +0 -0
  49. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/bedrock/cost_report.py +0 -0
  50. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/bedrock/pricing.yaml +0 -0
  51. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/cloud_commands.py +0 -0
  52. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/codeartifact.py +0 -0
  53. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/engines_studios/__init__.py +0 -0
  54. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/engines_studios/api_client.py +0 -0
  55. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/engines_studios/auth.py +0 -0
  56. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/engines_studios/engine_commands.py +0 -0
  57. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/engines_studios/progress.py +0 -0
  58. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/engines_studios/ssh_config.py +0 -0
  59. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/engines_studios/studio_commands.py +0 -0
  60. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/hz/__init__.py +0 -0
  61. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/hz/deploy.py +0 -0
  62. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/hz/local.py +0 -0
  63. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/hz/test.py +0 -0
  64. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/hz/tf.py +0 -0
  65. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/hz/users.py +0 -0
  66. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/main.py +0 -0
  67. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/utility_commands.py +0 -0
  68. {dh_cli-0.7.1 → dh_cli-0.8.0}/src/dh_cli/warehouse.py +0 -0
  69. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/conftest.py +0 -0
  70. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/fixtures/A_cache_write.json +0 -0
  71. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/fixtures/B_cache_read.json +0 -0
  72. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/fixtures/C_plain.json +0 -0
  73. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/fixtures/D_cursor_user.json +0 -0
  74. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/fixtures/E_service_role.json +0 -0
  75. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/fixtures/F_legacy_shared.json +0 -0
  76. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/fixtures/G_unknown_model.json +0 -0
  77. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/test_build_report.py +0 -0
  78. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/test_classify_arn.py +0 -0
  79. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/test_cli_exit_codes.py +0 -0
  80. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/test_cost_calc.py +0 -0
  81. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/test_cost_command.py +0 -0
  82. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/test_cur_reconciliation.py +0 -0
  83. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/test_key_command.py +0 -0
  84. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/test_render_formats.py +0 -0
  85. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/test_resolve_base_model.py +0 -0
  86. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/bedrock/test_s3_walker.py +0 -0
  87. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/hz/test_init.py +0 -0
  88. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/hz/test_suites.py +0 -0
  89. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/hz/test_users.py +0 -0
  90. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/test_cloud_gcp.py +0 -0
  91. {dh_cli-0.7.1 → dh_cli-0.8.0}/tests/test_finalize_protmpnn.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dh-cli
3
- Version: 0.7.1
3
+ Version: 0.8.0
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.7.1"
7
+ version = "0.8.0"
8
8
  description = "Dayhoff Labs developer CLI"
9
9
  requires-python = ">=3.11"
10
10
  readme = "README.md"
@@ -124,7 +124,7 @@ class BatchClient:
124
124
  resource_requirements.append({"type": "VCPU", "value": str(vcpus)})
125
125
  if memory_mb is not None:
126
126
  resource_requirements.append({"type": "MEMORY", "value": str(memory_mb)})
127
- if gpus is not None:
127
+ if gpus not in (None, 0):
128
128
  resource_requirements.append({"type": "GPU", "value": str(gpus)})
129
129
 
130
130
  container_overrides: dict[str, Any] = {}
@@ -71,7 +71,7 @@ def submit(
71
71
  queue: t4-1x-spot
72
72
  memory: 30G
73
73
  vcpus: 8
74
- gpus: 1
74
+ gpus: 1 # omit, set to 0, or null for CPU-only
75
75
  array: 10
76
76
  retry: 3
77
77
  timeout: 6h
@@ -113,11 +113,12 @@ def submit(
113
113
  timeout_seconds = _parse_timeout(job_timeout)
114
114
 
115
115
  # Show plan
116
+ gpus_display = "no GPUs" if job_gpus in (None, 0) else f"{job_gpus} GPUs"
116
117
  click.echo()
117
118
  click.echo(f"Job ID: {job_id}")
118
119
  click.echo(f"Command: {job_command}")
119
120
  click.echo(f"Queue: {job_queue}")
120
- click.echo(f"Resources: {job_vcpus} vCPUs, {job_memory} memory, {job_gpus} GPUs")
121
+ click.echo(f"Resources: {job_vcpus} vCPUs, {job_memory} memory, {gpus_display}")
121
122
  click.echo(f"Array Size: {job_array}")
122
123
  click.echo(f"Retry: {job_retry}")
123
124
  click.echo(f"Timeout: {job_timeout} ({timeout_seconds}s)")
@@ -0,0 +1,275 @@
1
+ """CLI commands for GitHub authentication.
2
+
3
+ This module provides commands for authenticating with GitHub from within
4
+ development containers using the GitHub CLI (gh).
5
+
6
+ The implementation:
7
+ 1. Wraps `gh auth login` with sensible defaults for devcontainer environments
8
+ 2. Automatically configures git to use GitHub CLI for credential management
9
+ 3. Uses HTTPS protocol and device flow authentication (works in headless envs)
10
+ """
11
+
12
+ import shutil
13
+ import subprocess
14
+ import sys
15
+ from typing import List, Tuple
16
+
17
+ import typer
18
+
19
+ # --- Configuration ---
20
+ # OAuth scopes to request during login:
21
+ # - repo: Full access to private/public repositories
22
+ # - read:org: Read-only access to organization membership
23
+ # - workflow: Ability to update GitHub Actions workflow files
24
+ GITHUB_SCOPES = "repo,read:org,workflow"
25
+ GITHUB_PROTOCOL = "https"
26
+
27
+ # --- Color constants for formatted output ---
28
+ RED = "\033[0;31m"
29
+ GREEN = "\033[0;32m"
30
+ YELLOW = "\033[0;33m"
31
+ BLUE = "\033[0;36m"
32
+ NC = "\033[0m" # No Color
33
+
34
+
35
+ # --- Helper Functions ---
36
+ def _find_executable(name: str) -> str:
37
+ """Find the full path to an executable in PATH."""
38
+ path = shutil.which(name)
39
+ if not path:
40
+ raise FileNotFoundError(f"{name} command not found. Please ensure it's installed.")
41
+ return path
42
+
43
+
44
+ def _run_command(
45
+ cmd_list: List[str],
46
+ capture: bool = False,
47
+ check: bool = True,
48
+ suppress_output: bool = False,
49
+ ) -> Tuple[int, str, str]:
50
+ """Run a command and return its result.
51
+
52
+ Args:
53
+ cmd_list: List of command arguments
54
+ capture: Whether to capture output
55
+ check: Whether to raise on non-zero exit code
56
+ suppress_output: Whether to hide output even if not captured
57
+
58
+ Returns:
59
+ Tuple of (return_code, stdout_str, stderr_str)
60
+ """
61
+ stdout_opt = subprocess.PIPE if capture else subprocess.DEVNULL if suppress_output else None
62
+ stderr_opt = subprocess.PIPE if capture else subprocess.DEVNULL if suppress_output else None
63
+
64
+ try:
65
+ result = subprocess.run(cmd_list, stdout=stdout_opt, stderr=stderr_opt, check=check, text=True)
66
+ return (
67
+ result.returncode,
68
+ result.stdout if capture else "",
69
+ result.stderr if capture else "",
70
+ )
71
+ except subprocess.CalledProcessError as e:
72
+ if capture:
73
+ return (e.returncode, e.stdout or "", e.stderr or "")
74
+ return (e.returncode, "", "")
75
+
76
+
77
+ def _is_gh_authenticated() -> bool:
78
+ """Check if GitHub CLI is authenticated.
79
+
80
+ Returns:
81
+ True if `gh auth status` succeeds, False otherwise.
82
+ """
83
+ try:
84
+ gh_path = _find_executable("gh")
85
+ returncode, _, _ = _run_command(
86
+ [gh_path, "auth", "status"],
87
+ capture=True,
88
+ check=False,
89
+ suppress_output=True,
90
+ )
91
+ return returncode == 0
92
+ except FileNotFoundError:
93
+ return False
94
+
95
+
96
+ def _get_gh_user() -> str:
97
+ """Get the currently authenticated GitHub username.
98
+
99
+ Returns:
100
+ The username or 'Not authenticated' if not logged in.
101
+ """
102
+ try:
103
+ gh_path = _find_executable("gh")
104
+ returncode, stdout, _ = _run_command(
105
+ [gh_path, "auth", "status", "--show-token"],
106
+ capture=True,
107
+ check=False,
108
+ suppress_output=True,
109
+ )
110
+ if returncode != 0:
111
+ return "Not authenticated"
112
+
113
+ # Parse the output to find the logged in account
114
+ # Format: "Logged in to github.com account username (keyring)"
115
+ for line in stdout.split("\n"):
116
+ if "Logged in to" in line and "account" in line:
117
+ parts = line.split("account")
118
+ if len(parts) > 1:
119
+ # Extract username from "account username (keyring)" or similar
120
+ user_part = parts[1].strip().split()[0]
121
+ return user_part
122
+ return "Unknown"
123
+ except FileNotFoundError:
124
+ return "gh CLI not found"
125
+
126
+
127
+ def _get_gh_scopes() -> str:
128
+ """Get the OAuth scopes for the current GitHub authentication.
129
+
130
+ Returns:
131
+ Comma-separated list of scopes or 'Unknown' if not available.
132
+ """
133
+ try:
134
+ gh_path = _find_executable("gh")
135
+ returncode, stdout, _ = _run_command(
136
+ [gh_path, "auth", "status"],
137
+ capture=True,
138
+ check=False,
139
+ suppress_output=True,
140
+ )
141
+ if returncode != 0:
142
+ return "N/A"
143
+
144
+ # Parse output for scopes - format varies by gh version
145
+ # Look for lines containing "Token scopes:" or similar
146
+ for line in stdout.split("\n"):
147
+ if "scopes" in line.lower():
148
+ # Extract everything after the colon
149
+ if ":" in line:
150
+ scopes = line.split(":", 1)[1].strip()
151
+ return scopes if scopes else "none"
152
+ return "Unknown"
153
+ except FileNotFoundError:
154
+ return "N/A"
155
+
156
+
157
+ # --- Typer Application ---
158
+ gh_app = typer.Typer(
159
+ help="Manage GitHub authentication using the gh CLI.",
160
+ context_settings={"help_option_names": ["-h", "--help"]},
161
+ )
162
+
163
+
164
+ @gh_app.command("status")
165
+ def gh_status():
166
+ """Show current GitHub authentication status."""
167
+ print(f"{BLUE}--- GitHub Authentication Status ---{NC}")
168
+
169
+ try:
170
+ gh_path = _find_executable("gh")
171
+ except FileNotFoundError:
172
+ print(f"{RED}Error: GitHub CLI (gh) is not installed.{NC}")
173
+ print(f"Install it with: {YELLOW}brew install gh{NC} (Mac) or see https://cli.github.com")
174
+ sys.exit(1)
175
+
176
+ if _is_gh_authenticated():
177
+ user = _get_gh_user()
178
+ print(f" Status: {GREEN}Authenticated{NC}")
179
+ print(f" User: {GREEN}{user}{NC}")
180
+
181
+ # Show detailed status
182
+ print(f"\n{BLUE}Detailed status:{NC}")
183
+ _run_command([gh_path, "auth", "status"], check=False)
184
+ else:
185
+ print(f" Status: {RED}Not authenticated{NC}")
186
+ print("\nTo authenticate, run:")
187
+ print(f" {YELLOW}dh gh login{NC}")
188
+
189
+
190
+ @gh_app.command("login")
191
+ def gh_login():
192
+ """Authenticate with GitHub and configure git credential helper.
193
+
194
+ This command:
195
+ 1. Authenticates with GitHub using device flow (works in headless environments)
196
+ 2. Requests scopes: repo, read:org, workflow
197
+ 3. Configures git to use the GitHub CLI for credential management
198
+ """
199
+ try:
200
+ gh_path = _find_executable("gh")
201
+ except FileNotFoundError:
202
+ print(f"{RED}Error: GitHub CLI (gh) is not installed.{NC}")
203
+ print(f"Install it with: {YELLOW}brew install gh{NC} (Mac) or see https://cli.github.com")
204
+ sys.exit(1)
205
+
206
+ # Step 1: Authenticate with GitHub
207
+ print(f"{BLUE}Authenticating with GitHub...{NC}")
208
+ print(f"{YELLOW}Requesting scopes: {GITHUB_SCOPES}{NC}")
209
+ print(f"{YELLOW}Using protocol: {GITHUB_PROTOCOL}{NC}")
210
+ print()
211
+
212
+ login_cmd = [
213
+ gh_path,
214
+ "auth",
215
+ "login",
216
+ "--web", # Use device flow (works in devcontainers)
217
+ "--git-protocol",
218
+ GITHUB_PROTOCOL,
219
+ "--scopes",
220
+ GITHUB_SCOPES,
221
+ ]
222
+
223
+ returncode, _, _ = _run_command(login_cmd, capture=False, check=False)
224
+
225
+ if returncode != 0:
226
+ print(f"\n{RED}Authentication failed. Please check the output above.{NC}")
227
+ sys.exit(1)
228
+
229
+ # Step 2: Configure git credential helper
230
+ print(f"\n{BLUE}Configuring git to use GitHub CLI for credentials...{NC}")
231
+ setup_cmd = [gh_path, "auth", "setup-git"]
232
+ setup_rc, _, setup_err = _run_command(setup_cmd, capture=True, check=False)
233
+
234
+ if setup_rc != 0:
235
+ print(f"{YELLOW}Warning: Failed to configure git credential helper: {setup_err}{NC}")
236
+ print(f"You may need to run manually: {YELLOW}gh auth setup-git{NC}")
237
+ else:
238
+ print(f"{GREEN}Git credential helper configured.{NC}")
239
+
240
+ # Step 3: Show final status
241
+ print(f"\n{GREEN}GitHub authentication complete!{NC}")
242
+ print(f"\n{BLUE}--- Current Status ---{NC}")
243
+ gh_status()
244
+
245
+
246
+ @gh_app.command("logout")
247
+ def gh_logout():
248
+ """Log out from GitHub and clear credentials.
249
+
250
+ This removes the GitHub authentication token and unconfigures
251
+ the git credential helper.
252
+ """
253
+ try:
254
+ gh_path = _find_executable("gh")
255
+ except FileNotFoundError:
256
+ print(f"{RED}Error: GitHub CLI (gh) is not installed.{NC}")
257
+ sys.exit(1)
258
+
259
+ if not _is_gh_authenticated():
260
+ print(f"{YELLOW}Not currently authenticated with GitHub.{NC}")
261
+ return
262
+
263
+ print(f"{BLUE}Logging out from GitHub...{NC}")
264
+
265
+ # Log out from github.com
266
+ logout_cmd = [gh_path, "auth", "logout", "--hostname", "github.com"]
267
+ returncode, _, _ = _run_command(logout_cmd, capture=False, check=False)
268
+
269
+ if returncode != 0:
270
+ print(f"{RED}Logout may have failed. Check the output above.{NC}")
271
+ sys.exit(1)
272
+
273
+ print(f"\n{GREEN}Successfully logged out from GitHub.{NC}")
274
+ print(f"\n{BLUE}To log back in:{NC}")
275
+ print(f" {YELLOW}dh gh login{NC}")
File without changes
@@ -0,0 +1,92 @@
1
+ """Tests for BatchClient resource-requirement construction.
2
+
3
+ These tests mock only boto3.client("batch") and assert on the
4
+ submit_args passed to submit_job. They pin the GPU-handling contract:
5
+ gpus=None and gpus=0 both mean "no GPU requirement"; positive values
6
+ are forwarded verbatim.
7
+ """
8
+
9
+ from unittest.mock import MagicMock, patch
10
+
11
+
12
+ def _make_client():
13
+ """Construct a BatchClient with mocked boto3 batch/logs clients."""
14
+ with patch("dh_cli.batch.aws_batch.boto3") as mock_boto3:
15
+ mock_batch = MagicMock()
16
+ mock_logs = MagicMock()
17
+ mock_boto3.client.side_effect = [mock_batch, mock_logs]
18
+ mock_batch.submit_job.return_value = {"jobId": "aws-uuid"}
19
+
20
+ from dh_cli.batch.aws_batch import BatchClient
21
+
22
+ client = BatchClient()
23
+ return client, mock_batch
24
+
25
+
26
+ def _get_gpu_resource(submit_args):
27
+ """Extract the GPU resourceRequirement entry from submit_args, or None."""
28
+ overrides = submit_args.get("containerOverrides", {})
29
+ for req in overrides.get("resourceRequirements", []):
30
+ if req.get("type") == "GPU":
31
+ return req
32
+ return None
33
+
34
+
35
+ class TestGpuResourceRequirement:
36
+ """gpus=None and gpus=0 must both omit the GPU resource requirement."""
37
+
38
+ def test_gpus_none_omits_gpu_resource(self):
39
+ client, mock_batch = _make_client()
40
+ client.submit_job(
41
+ job_name="j",
42
+ job_definition="dayhoff-generic",
43
+ job_queue="cpu-spot",
44
+ vcpus=4,
45
+ memory_mb=8192,
46
+ gpus=None,
47
+ command=["sh", "-c", "echo hi"],
48
+ )
49
+ submit_args = mock_batch.submit_job.call_args[1]
50
+ assert _get_gpu_resource(submit_args) is None
51
+
52
+ def test_gpus_zero_omits_gpu_resource(self):
53
+ client, mock_batch = _make_client()
54
+ client.submit_job(
55
+ job_name="j",
56
+ job_definition="dayhoff-generic",
57
+ job_queue="cpu-spot",
58
+ vcpus=4,
59
+ memory_mb=8192,
60
+ gpus=0,
61
+ command=["sh", "-c", "echo hi"],
62
+ )
63
+ submit_args = mock_batch.submit_job.call_args[1]
64
+ assert _get_gpu_resource(submit_args) is None
65
+
66
+ def test_gpus_one_includes_gpu_resource(self):
67
+ client, mock_batch = _make_client()
68
+ client.submit_job(
69
+ job_name="j",
70
+ job_definition="dayhoff-generic",
71
+ job_queue="t4-1x-spot",
72
+ vcpus=4,
73
+ memory_mb=8192,
74
+ gpus=1,
75
+ command=["sh", "-c", "echo hi"],
76
+ )
77
+ submit_args = mock_batch.submit_job.call_args[1]
78
+ assert _get_gpu_resource(submit_args) == {"type": "GPU", "value": "1"}
79
+
80
+ def test_gpus_two_includes_gpu_resource(self):
81
+ client, mock_batch = _make_client()
82
+ client.submit_job(
83
+ job_name="j",
84
+ job_definition="dayhoff-generic",
85
+ job_queue="t4-1x-spot",
86
+ vcpus=4,
87
+ memory_mb=8192,
88
+ gpus=2,
89
+ command=["sh", "-c", "echo hi"],
90
+ )
91
+ submit_args = mock_batch.submit_job.call_args[1]
92
+ assert _get_gpu_resource(submit_args) == {"type": "GPU", "value": "2"}
@@ -0,0 +1,162 @@
1
+ """Tests for CPU-only submission paths through the CLI.
2
+
3
+ Uses Click's CliRunner with BatchClient.submit_job mocked. Pins the
4
+ contract: users can express "no GPU" via --gpus 0, YAML gpus: 0, or
5
+ YAML gpus: null; the implicit --gpus 1 default still applies when
6
+ nothing is specified; and the dry-run echo reads "no GPUs" rather
7
+ than "None GPUs" or "0 GPUs".
8
+ """
9
+
10
+ from unittest.mock import MagicMock, patch
11
+
12
+ import pytest
13
+ import yaml
14
+ from click.testing import CliRunner
15
+
16
+
17
+ @pytest.fixture
18
+ def cli_runner():
19
+ return CliRunner()
20
+
21
+
22
+ def _invoke(cli_runner, args, tmp_path):
23
+ """Invoke submit with BatchClient mocked and base-path tucked under tmp_path."""
24
+ base = tmp_path / "jobs"
25
+ with (
26
+ patch("dh_cli.batch.commands.submit.get_aws_username", return_value="jason"),
27
+ patch("dh_cli.batch.commands.submit.BatchClient") as mock_batch_cls,
28
+ patch(
29
+ "dh_cli.batch.commands.submit.generate_job_id",
30
+ return_value="jason-batch-20260511-cpu00001",
31
+ ),
32
+ ):
33
+ mock_client = MagicMock()
34
+ mock_client.submit_job.return_value = "aws-uuid-cpu"
35
+ mock_batch_cls.return_value = mock_client
36
+
37
+ from dh_cli.batch.commands.submit import submit
38
+
39
+ result = cli_runner.invoke(submit, args + ["--base-path", str(base)])
40
+ return result, mock_client
41
+
42
+
43
+ class TestCpuOnlySubmissionPaths:
44
+ """--gpus 0, YAML gpus: 0, and YAML gpus: null must all submit with no GPU."""
45
+
46
+ def test_cli_gpus_zero_sends_no_gpu_requirement(self, cli_runner, tmp_path):
47
+ result, mock_client = _invoke(
48
+ cli_runner,
49
+ ["--command", "echo hi", "--queue", "cpu-spot", "--gpus", "0"],
50
+ tmp_path,
51
+ )
52
+ assert result.exit_code == 0, result.output
53
+ call_kwargs = mock_client.submit_job.call_args[1]
54
+ assert call_kwargs["gpus"] == 0
55
+
56
+ def test_yaml_gpus_zero_sends_no_gpu_requirement(self, cli_runner, tmp_path):
57
+ config_path = tmp_path / "job.yaml"
58
+ config_path.write_text(
59
+ yaml.dump(
60
+ {
61
+ "command": "echo hi",
62
+ "queue": "cpu-spot",
63
+ "gpus": 0,
64
+ }
65
+ )
66
+ )
67
+ result, mock_client = _invoke(cli_runner, ["-f", str(config_path)], tmp_path)
68
+ assert result.exit_code == 0, result.output
69
+ call_kwargs = mock_client.submit_job.call_args[1]
70
+ assert call_kwargs["gpus"] == 0
71
+
72
+ def test_yaml_gpus_null_sends_no_gpu_requirement(self, cli_runner, tmp_path):
73
+ config_path = tmp_path / "job.yaml"
74
+ config_path.write_text(
75
+ yaml.dump(
76
+ {
77
+ "command": "echo hi",
78
+ "queue": "cpu-spot",
79
+ "gpus": None,
80
+ }
81
+ )
82
+ )
83
+ result, mock_client = _invoke(cli_runner, ["-f", str(config_path)], tmp_path)
84
+ assert result.exit_code == 0, result.output
85
+ call_kwargs = mock_client.submit_job.call_args[1]
86
+ assert call_kwargs["gpus"] is None
87
+
88
+
89
+ class TestGpuDefaultPreserved:
90
+ """Omitting --gpus and YAML gpus still gives the implicit GPU=1 default."""
91
+
92
+ def test_cli_gpus_omitted_gets_implicit_one(self, cli_runner, tmp_path):
93
+ result, mock_client = _invoke(cli_runner, ["--command", "echo hi"], tmp_path)
94
+ assert result.exit_code == 0, result.output
95
+ call_kwargs = mock_client.submit_job.call_args[1]
96
+ assert call_kwargs["gpus"] == 1
97
+
98
+ def test_cli_gpus_one_still_works(self, cli_runner, tmp_path):
99
+ result, mock_client = _invoke(cli_runner, ["--command", "echo hi", "--gpus", "1"], tmp_path)
100
+ assert result.exit_code == 0, result.output
101
+ call_kwargs = mock_client.submit_job.call_args[1]
102
+ assert call_kwargs["gpus"] == 1
103
+
104
+ def test_yaml_gpus_one_still_works(self, cli_runner, tmp_path):
105
+ config_path = tmp_path / "job.yaml"
106
+ config_path.write_text(yaml.dump({"command": "echo hi", "gpus": 1}))
107
+ result, mock_client = _invoke(cli_runner, ["-f", str(config_path)], tmp_path)
108
+ assert result.exit_code == 0, result.output
109
+ call_kwargs = mock_client.submit_job.call_args[1]
110
+ assert call_kwargs["gpus"] == 1
111
+
112
+
113
+ class TestCliBeatsYamlOnZero:
114
+ """--gpus 0 on the CLI must beat YAML gpus: 1."""
115
+
116
+ def test_cli_gpus_zero_overrides_yaml_one(self, cli_runner, tmp_path):
117
+ config_path = tmp_path / "job.yaml"
118
+ config_path.write_text(yaml.dump({"command": "echo hi", "gpus": 1}))
119
+ result, mock_client = _invoke(cli_runner, ["-f", str(config_path), "--gpus", "0"], tmp_path)
120
+ assert result.exit_code == 0, result.output
121
+ call_kwargs = mock_client.submit_job.call_args[1]
122
+ assert call_kwargs["gpus"] == 0
123
+
124
+
125
+ class TestDryRunEcho:
126
+ """The dry-run plan line must read 'no GPUs' or 'N GPUs', never 'None'/'0 GPUs'."""
127
+
128
+ def test_dry_run_echo_reads_no_gpus_when_gpus_zero(self, cli_runner, tmp_path):
129
+ result, _ = _invoke(
130
+ cli_runner,
131
+ [
132
+ "--command",
133
+ "echo hi",
134
+ "--queue",
135
+ "cpu-spot",
136
+ "--gpus",
137
+ "0",
138
+ "--dry-run",
139
+ ],
140
+ tmp_path,
141
+ )
142
+ assert result.exit_code == 0, result.output
143
+ assert "no GPUs" in result.output
144
+ assert "None GPUs" not in result.output
145
+ assert "0 GPUs" not in result.output
146
+
147
+ def test_dry_run_echo_reads_no_gpus_when_yaml_gpus_null(self, cli_runner, tmp_path):
148
+ config_path = tmp_path / "job.yaml"
149
+ config_path.write_text(yaml.dump({"command": "echo hi", "queue": "cpu-spot", "gpus": None}))
150
+ result, _ = _invoke(cli_runner, ["-f", str(config_path), "--dry-run"], tmp_path)
151
+ assert result.exit_code == 0, result.output
152
+ assert "no GPUs" in result.output
153
+ assert "None GPUs" not in result.output
154
+
155
+ def test_dry_run_echo_reads_n_gpus_when_gpus_positive(self, cli_runner, tmp_path):
156
+ result, _ = _invoke(
157
+ cli_runner,
158
+ ["--command", "echo hi", "--gpus", "2", "--dry-run"],
159
+ tmp_path,
160
+ )
161
+ assert result.exit_code == 0, result.output
162
+ assert "2 GPUs" in result.output
@@ -1,82 +0,0 @@
1
- """Identity resolution for `dh` commands that read per-developer secrets.
2
-
3
- The `github_commands` and (future) `bedrock` commands both key per-user
4
- Secrets Manager entries on the caller's Dayhoff handle. This module
5
- resolves that handle from the current SSO session in a way that matches
6
- the server-side resource policy (which keys on
7
- `aws:PrincipalTag/Email`).
8
-
9
- Design note — why not session tags directly?
10
-
11
- AWS Secrets Manager resource policies evaluate
12
- `aws:PrincipalTag/Email` automatically because IAM Identity Center
13
- attaches an `Email` principal tag to DeveloperAccess sessions. But
14
- there is no SDK API that lets a session read its own principal
15
- tags back — session tags are visible *only* to IAM condition
16
- evaluation, not to callers. The closest observable we have is the
17
- assumed-role ARN's RoleSessionName, which Identity Center sets to
18
- the user's email by default.
19
-
20
- So the resolution below extracts the RoleSessionName and strips the
21
- `@<domain>` suffix if present. This matches what the policy's
22
- `aws:PrincipalTag/Email` condition will also evaluate to at
23
- GetSecretValue time — i.e. "which secret can I read?" and "what
24
- handle am I?" are answered by the same identity fact by
25
- construction.
26
-
27
- CANARY (resolve pre-code-freeze, plan §"Pre-implementation canaries"
28
- Canary 1): confirm the RoleSessionName on DeveloperAccess is
29
- `<handle>@dayhoff.com` and not `<handle>@dayhofflabs.com`. The
30
- DEFAULT_DOMAIN constant below is the single place to flip that.
31
- """
32
- from __future__ import annotations
33
-
34
- import re
35
-
36
- DEFAULT_DOMAIN = "dayhoff.com"
37
-
38
- _SSO_ASSUMED_ROLE_RE = re.compile(
39
- r"^arn:aws:sts::\d+:assumed-role/AWSReservedSSO_[^/]+/(?P<session>.+)$"
40
- )
41
-
42
-
43
- class HandleResolutionError(RuntimeError):
44
- """Raised when the current session's handle can't be determined.
45
-
46
- The caller is expected to turn this into a user-facing error
47
- pointing at `awslogin dev-devaccess`.
48
- """
49
-
50
-
51
- def resolve_handle_from_session(session, *, domain: str = DEFAULT_DOMAIN) -> str:
52
- """Return the dev handle matching the SSO session's Email principal tag.
53
-
54
- Args:
55
- session: a boto3 Session configured with the caller's SSO
56
- credentials. The function calls `sts:GetCallerIdentity` on
57
- it; if that call would fall through to the engine instance
58
- role instead of the dev's SSO creds, the caller must have
59
- already detected and errored on that before getting here
60
- (see `github_commands._sso_session`).
61
- domain: email domain to strip from the session name. Defaults to
62
- `dayhoff.com` to match `cursor_bedrock_role.tf`'s tag.
63
-
64
- Returns:
65
- The handle (e.g. `"dma"`).
66
-
67
- Raises:
68
- HandleResolutionError: the caller's ARN doesn't look like an
69
- Identity Center DeveloperAccess session.
70
- """
71
- arn = session.client("sts").get_caller_identity()["Arn"]
72
- match = _SSO_ASSUMED_ROLE_RE.match(arn)
73
- if not match:
74
- raise HandleResolutionError(
75
- f"Caller ARN does not look like an AWS SSO session: {arn}. "
76
- f"Run `awslogin dev-devaccess` (or pass --handle explicitly)."
77
- )
78
- session_name = match.group("session")
79
- suffix = f"@{domain}"
80
- if session_name.endswith(suffix):
81
- return session_name[: -len(suffix)]
82
- return session_name