dh-cli 0.4.5__tar.gz → 0.5.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.
- {dh_cli-0.4.5 → dh_cli-0.5.0}/PKG-INFO +1 -1
- {dh_cli-0.4.5 → dh_cli-0.5.0}/pyproject.toml +1 -1
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/__init__.py +2 -0
- dh_cli-0.5.0/src/dh_cli/batch/commands/orca.py +352 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/cloud_commands.py +23 -398
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/main.py +1 -1
- dh_cli-0.5.0/tests/test_cloud_gcp.py +67 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/.gitignore +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/LICENSE +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/README.md +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/__init__.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/aws_batch.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/__init__.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/boltz.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/cancel.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/clean.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/embed_t5.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/finalize.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/list_jobs.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/local.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/logs.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/protmpnn.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/protmpnn_to_boltz.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/retry.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/status.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/submit.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/train.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/commands/wait_for.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/fasta_utils.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/h5_utils.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/job_id.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/manifest.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/batch/s3_transport.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/codeartifact.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/engines_studios/__init__.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/engines_studios/api_client.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/engines_studios/auth.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/engines_studios/engine_commands.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/engines_studios/progress.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/engines_studios/ssh_config.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/engines_studios/studio_commands.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/github_commands.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/hz/__init__.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/hz/deploy.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/hz/local.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/hz/test.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/hz/tf.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/hz/users.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/utility_commands.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/src/dh_cli/warehouse.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/tests/hz/test_init.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/tests/hz/test_suites.py +0 -0
- {dh_cli-0.4.5 → dh_cli-0.5.0}/tests/hz/test_users.py +0 -0
|
@@ -12,6 +12,7 @@ import click
|
|
|
12
12
|
|
|
13
13
|
from .commands.boltz import boltz
|
|
14
14
|
from .commands.cancel import cancel
|
|
15
|
+
from .commands.orca import orca
|
|
15
16
|
from .commands.clean import clean
|
|
16
17
|
from .commands.embed_t5 import embed_t5
|
|
17
18
|
from .commands.finalize import finalize
|
|
@@ -99,6 +100,7 @@ batch_cli.add_command(wait_for, name="wait-for")
|
|
|
99
100
|
batch_cli.add_command(embed_t5, name="embed-t5")
|
|
100
101
|
batch_cli.add_command(boltz)
|
|
101
102
|
batch_cli.add_command(protmpnn)
|
|
103
|
+
batch_cli.add_command(orca)
|
|
102
104
|
batch_cli.add_command(protmpnn_to_boltz, name="protmpnn-to-boltz")
|
|
103
105
|
batch_cli.add_command(train)
|
|
104
106
|
|
|
@@ -0,0 +1,352 @@
|
|
|
1
|
+
"""ORCA quantum chemistry batch pipeline command.
|
|
2
|
+
|
|
3
|
+
Unlike Boltz/ProtMPNN, ORCA jobs are typically one-per-submission (each
|
|
4
|
+
calculation runs for hours). Multiple .inp files are submitted as separate
|
|
5
|
+
Batch jobs, not as an array job.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import os
|
|
9
|
+
import shutil
|
|
10
|
+
from pathlib import Path
|
|
11
|
+
|
|
12
|
+
import click
|
|
13
|
+
|
|
14
|
+
from ..aws_batch import BatchClient, BatchError, resolve_dependency
|
|
15
|
+
from ..job_id import generate_job_id, get_aws_username
|
|
16
|
+
from ..manifest import (
|
|
17
|
+
BATCH_JOBS_BASE,
|
|
18
|
+
BatchConfig,
|
|
19
|
+
InputConfig,
|
|
20
|
+
JobManifest,
|
|
21
|
+
JobStatus,
|
|
22
|
+
OutputConfig,
|
|
23
|
+
create_job_directory,
|
|
24
|
+
save_manifest,
|
|
25
|
+
save_manifest_s3,
|
|
26
|
+
save_local_stub,
|
|
27
|
+
)
|
|
28
|
+
from ..s3_transport import s3_job_prefix, upload_directory
|
|
29
|
+
|
|
30
|
+
DEFAULT_QUEUE = "cpu-ondemand"
|
|
31
|
+
DEFAULT_JOB_DEFINITION = "dayhoff-orca"
|
|
32
|
+
DEFAULT_IMAGE_URI = (
|
|
33
|
+
"074735440724.dkr.ecr.us-east-1.amazonaws.com/dayhoff:orca-latest"
|
|
34
|
+
)
|
|
35
|
+
DEFAULT_VCPUS = 16
|
|
36
|
+
DEFAULT_MEMORY_MB = 64000
|
|
37
|
+
DEFAULT_TIMEOUT_H = 24
|
|
38
|
+
|
|
39
|
+
S3_MAX_FILES = 100
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def _is_primordial_path(path: Path) -> bool:
|
|
43
|
+
return str(path).startswith("/primordial/")
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def _find_inp_files(input_path: Path) -> list[Path]:
|
|
47
|
+
"""Find .inp files — either a single file or all files in a directory."""
|
|
48
|
+
if input_path.is_file() and input_path.suffix == ".inp":
|
|
49
|
+
return [input_path]
|
|
50
|
+
if input_path.is_dir():
|
|
51
|
+
return sorted(input_path.glob("*.inp"))
|
|
52
|
+
return []
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
@click.command()
|
|
56
|
+
@click.argument("input_path", type=click.Path(exists=True))
|
|
57
|
+
@click.option(
|
|
58
|
+
"--queue", default=DEFAULT_QUEUE,
|
|
59
|
+
help=f"Batch queue [default: {DEFAULT_QUEUE}]",
|
|
60
|
+
)
|
|
61
|
+
@click.option(
|
|
62
|
+
"--vcpus", default=DEFAULT_VCPUS, type=int,
|
|
63
|
+
help=f"vCPUs per job [default: {DEFAULT_VCPUS}]",
|
|
64
|
+
)
|
|
65
|
+
@click.option(
|
|
66
|
+
"--memory", default=DEFAULT_MEMORY_MB, type=int,
|
|
67
|
+
help=f"Memory in MB per job [default: {DEFAULT_MEMORY_MB}]",
|
|
68
|
+
)
|
|
69
|
+
@click.option(
|
|
70
|
+
"--timeout", "timeout_h", default=DEFAULT_TIMEOUT_H, type=float,
|
|
71
|
+
help=f"Timeout in hours per job [default: {DEFAULT_TIMEOUT_H}]",
|
|
72
|
+
)
|
|
73
|
+
@click.option(
|
|
74
|
+
"--local", "run_local", is_flag=True,
|
|
75
|
+
help="Run ORCA locally in a Docker container instead of Batch",
|
|
76
|
+
)
|
|
77
|
+
@click.option(
|
|
78
|
+
"--shell", "run_shell", is_flag=True,
|
|
79
|
+
help="Drop into container shell for debugging",
|
|
80
|
+
)
|
|
81
|
+
@click.option("--dry-run", is_flag=True, help="Show plan without submitting")
|
|
82
|
+
@click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
|
|
83
|
+
@click.option(
|
|
84
|
+
"--after", multiple=True,
|
|
85
|
+
help="Job ID(s) to wait for before starting",
|
|
86
|
+
)
|
|
87
|
+
def orca(input_path, queue, vcpus, memory, timeout_h, run_local, run_shell,
|
|
88
|
+
dry_run, base_path, after):
|
|
89
|
+
"""Run ORCA DFT calculations on AWS Batch.
|
|
90
|
+
|
|
91
|
+
INPUT_PATH can be a single .inp file or a directory of .inp files.
|
|
92
|
+
Each .inp file is submitted as a separate Batch job (ORCA calculations
|
|
93
|
+
are long-running, so array jobs are not used).
|
|
94
|
+
|
|
95
|
+
\b
|
|
96
|
+
Examples:
|
|
97
|
+
# Submit a single calculation
|
|
98
|
+
dh batch orca /primordial/calculations/ts_opt.inp
|
|
99
|
+
|
|
100
|
+
# Submit all .inp files in a directory
|
|
101
|
+
dh batch orca /primordial/calculations/
|
|
102
|
+
|
|
103
|
+
# Use more CPUs for a heavy job
|
|
104
|
+
dh batch orca ts_opt.inp --vcpus 32 --memory 128000
|
|
105
|
+
|
|
106
|
+
# Test locally in container
|
|
107
|
+
dh batch orca ts_opt.inp --local
|
|
108
|
+
|
|
109
|
+
\b
|
|
110
|
+
After job completes:
|
|
111
|
+
dh batch status <job-id>
|
|
112
|
+
dh batch logs <job-id>
|
|
113
|
+
dh batch finalize <job-id> --output /primordial/results/
|
|
114
|
+
"""
|
|
115
|
+
path = Path(input_path).resolve()
|
|
116
|
+
|
|
117
|
+
if run_shell:
|
|
118
|
+
_run_shell_mode(path)
|
|
119
|
+
return
|
|
120
|
+
|
|
121
|
+
if run_local:
|
|
122
|
+
_run_local_mode(path, vcpus)
|
|
123
|
+
return
|
|
124
|
+
|
|
125
|
+
_submit_batch_jobs(path, queue, vcpus, memory, timeout_h, dry_run,
|
|
126
|
+
base_path, after)
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def _submit_batch_jobs(
|
|
130
|
+
input_path: Path,
|
|
131
|
+
queue: str,
|
|
132
|
+
vcpus: int,
|
|
133
|
+
memory: int,
|
|
134
|
+
timeout_h: float,
|
|
135
|
+
dry_run: bool,
|
|
136
|
+
base_path: str,
|
|
137
|
+
after: tuple[str, ...] = (),
|
|
138
|
+
):
|
|
139
|
+
"""Submit ORCA job(s) to AWS Batch."""
|
|
140
|
+
inp_files = _find_inp_files(input_path)
|
|
141
|
+
|
|
142
|
+
if not inp_files:
|
|
143
|
+
click.echo(
|
|
144
|
+
click.style("Error: No .inp files found", fg="red"), err=True,
|
|
145
|
+
)
|
|
146
|
+
raise SystemExit(1)
|
|
147
|
+
|
|
148
|
+
click.echo(f"Found {len(inp_files)} ORCA input file(s)")
|
|
149
|
+
|
|
150
|
+
use_s3 = not _is_primordial_path(input_path)
|
|
151
|
+
if use_s3 and len(inp_files) > S3_MAX_FILES:
|
|
152
|
+
click.echo(f"Error: {len(inp_files)} files exceeds S3 limit ({S3_MAX_FILES}).")
|
|
153
|
+
click.echo("Copy inputs to /primordial/ for large jobs.")
|
|
154
|
+
raise SystemExit(1)
|
|
155
|
+
|
|
156
|
+
timeout_s = int(timeout_h * 3600)
|
|
157
|
+
|
|
158
|
+
for inp_file in inp_files:
|
|
159
|
+
job_id = generate_job_id("orca")
|
|
160
|
+
|
|
161
|
+
click.echo()
|
|
162
|
+
click.echo(f"Job ID: {job_id}")
|
|
163
|
+
click.echo(f"Input: {inp_file.name}")
|
|
164
|
+
click.echo(f"Queue: {queue}")
|
|
165
|
+
click.echo(f"vCPUs: {vcpus}")
|
|
166
|
+
click.echo(f"Memory: {memory} MB")
|
|
167
|
+
click.echo(f"Timeout: {timeout_h}h")
|
|
168
|
+
click.echo(f"Job definition: {DEFAULT_JOB_DEFINITION}")
|
|
169
|
+
click.echo(f"Storage: {'S3' if use_s3 else 'Primordial'}")
|
|
170
|
+
|
|
171
|
+
if dry_run:
|
|
172
|
+
click.echo(click.style(" Dry run — not submitted", fg="yellow"))
|
|
173
|
+
continue
|
|
174
|
+
|
|
175
|
+
if len(inp_files) == 1 and not click.confirm("\nSubmit job?", default=True):
|
|
176
|
+
click.echo("Cancelled.")
|
|
177
|
+
raise SystemExit(0)
|
|
178
|
+
|
|
179
|
+
s3_prefix = s3_job_prefix(job_id) if use_s3 else None
|
|
180
|
+
job_dir_path = Path(base_path) / job_id
|
|
181
|
+
|
|
182
|
+
if use_s3:
|
|
183
|
+
click.echo("Uploading input to S3...")
|
|
184
|
+
upload_directory(inp_file.parent, f"{s3_prefix}input/", glob=inp_file.name)
|
|
185
|
+
else:
|
|
186
|
+
job_dir_path = create_job_directory(job_id, base_path)
|
|
187
|
+
inp_dest = job_dir_path / "input"
|
|
188
|
+
inp_dest.mkdir(parents=True, exist_ok=True)
|
|
189
|
+
shutil.copy2(inp_file, inp_dest / inp_file.name)
|
|
190
|
+
click.echo(f"Copied {inp_file.name} to {inp_dest}")
|
|
191
|
+
|
|
192
|
+
manifest = JobManifest(
|
|
193
|
+
job_id=job_id,
|
|
194
|
+
user=job_id.split("-")[0],
|
|
195
|
+
pipeline="orca",
|
|
196
|
+
status=JobStatus.PENDING,
|
|
197
|
+
image_uri=DEFAULT_IMAGE_URI,
|
|
198
|
+
storage_mode="s3" if use_s3 else "primordial",
|
|
199
|
+
s3_prefix=s3_prefix,
|
|
200
|
+
input=InputConfig(
|
|
201
|
+
source=str(inp_file),
|
|
202
|
+
num_sequences=1,
|
|
203
|
+
num_chunks=1,
|
|
204
|
+
),
|
|
205
|
+
batch=BatchConfig(
|
|
206
|
+
queue=queue,
|
|
207
|
+
job_definition=DEFAULT_JOB_DEFINITION,
|
|
208
|
+
array_size=None, # Single job, not array
|
|
209
|
+
),
|
|
210
|
+
output=OutputConfig(destination=None, finalized=False),
|
|
211
|
+
depends_on=list(after) if after else None,
|
|
212
|
+
)
|
|
213
|
+
|
|
214
|
+
if use_s3:
|
|
215
|
+
save_manifest_s3(manifest)
|
|
216
|
+
save_local_stub(job_id, s3_prefix)
|
|
217
|
+
else:
|
|
218
|
+
save_manifest(manifest, base_path)
|
|
219
|
+
|
|
220
|
+
try:
|
|
221
|
+
resolved = [resolve_dependency(jid, base_path) for jid in after]
|
|
222
|
+
depends_on = (
|
|
223
|
+
[{"jobId": aws_id} for aws_id in resolved if aws_id is not None]
|
|
224
|
+
or None
|
|
225
|
+
)
|
|
226
|
+
|
|
227
|
+
client = BatchClient()
|
|
228
|
+
environment = {
|
|
229
|
+
"JOB_DIR": str(job_dir_path),
|
|
230
|
+
"JOB_ID": job_id,
|
|
231
|
+
"BATCH_ARRAY_SIZE": "1",
|
|
232
|
+
"BATCH_NUM_FILES": "1",
|
|
233
|
+
}
|
|
234
|
+
if use_s3:
|
|
235
|
+
environment["STORAGE_MODE"] = "s3"
|
|
236
|
+
environment["S3_JOB_PREFIX"] = s3_prefix
|
|
237
|
+
|
|
238
|
+
batch_job_id = client.submit_job(
|
|
239
|
+
job_name=job_id,
|
|
240
|
+
job_definition=DEFAULT_JOB_DEFINITION,
|
|
241
|
+
job_queue=queue,
|
|
242
|
+
array_size=None, # Single job
|
|
243
|
+
environment=environment,
|
|
244
|
+
timeout_seconds=timeout_s,
|
|
245
|
+
retry_attempts=2,
|
|
246
|
+
depends_on=depends_on,
|
|
247
|
+
share_identifier=get_aws_username(),
|
|
248
|
+
vcpus=vcpus,
|
|
249
|
+
memory_mb=memory,
|
|
250
|
+
)
|
|
251
|
+
|
|
252
|
+
manifest.status = JobStatus.SUBMITTED
|
|
253
|
+
manifest.batch.job_id = batch_job_id
|
|
254
|
+
if use_s3:
|
|
255
|
+
save_manifest_s3(manifest)
|
|
256
|
+
else:
|
|
257
|
+
save_manifest(manifest, base_path)
|
|
258
|
+
|
|
259
|
+
click.echo()
|
|
260
|
+
click.echo(click.style("Job submitted!", fg="green"))
|
|
261
|
+
click.echo(f" AWS Job ID: {batch_job_id}")
|
|
262
|
+
click.echo(f" Status: dh batch status {job_id}")
|
|
263
|
+
click.echo(f" Logs: dh batch logs {job_id}")
|
|
264
|
+
|
|
265
|
+
except BatchError as e:
|
|
266
|
+
manifest.status = JobStatus.FAILED
|
|
267
|
+
manifest.error_message = str(e)
|
|
268
|
+
if use_s3:
|
|
269
|
+
save_manifest_s3(manifest)
|
|
270
|
+
else:
|
|
271
|
+
save_manifest(manifest, base_path)
|
|
272
|
+
click.echo(
|
|
273
|
+
click.style(f"Failed to submit: {e}", fg="red"), err=True,
|
|
274
|
+
)
|
|
275
|
+
raise SystemExit(1)
|
|
276
|
+
|
|
277
|
+
if dry_run:
|
|
278
|
+
click.echo()
|
|
279
|
+
click.echo(click.style("Dry run complete — no jobs submitted", fg="yellow"))
|
|
280
|
+
|
|
281
|
+
|
|
282
|
+
def _run_local_mode(input_path: Path, vcpus: int = DEFAULT_VCPUS):
|
|
283
|
+
"""Run ORCA locally in a Docker container."""
|
|
284
|
+
import subprocess
|
|
285
|
+
|
|
286
|
+
inp_files = _find_inp_files(input_path)
|
|
287
|
+
if not inp_files:
|
|
288
|
+
click.echo(click.style("Error: No .inp files found", fg="red"), err=True)
|
|
289
|
+
raise SystemExit(1)
|
|
290
|
+
|
|
291
|
+
inp_file = inp_files[0]
|
|
292
|
+
click.echo(f"Running ORCA locally: {inp_file.name}")
|
|
293
|
+
|
|
294
|
+
work_dir = inp_file.parent
|
|
295
|
+
job_dir = work_dir / ".local_orca_job"
|
|
296
|
+
inp_dir = job_dir / "input"
|
|
297
|
+
if job_dir.exists():
|
|
298
|
+
shutil.rmtree(job_dir)
|
|
299
|
+
inp_dir.mkdir(parents=True)
|
|
300
|
+
(job_dir / "output").mkdir()
|
|
301
|
+
shutil.copy2(inp_file, inp_dir / inp_file.name)
|
|
302
|
+
|
|
303
|
+
cmd = [
|
|
304
|
+
"docker", "run", "--rm",
|
|
305
|
+
"-v", "/primordial:/primordial:ro",
|
|
306
|
+
"-v", f"{job_dir}:{job_dir}",
|
|
307
|
+
"-e", f"JOB_DIR={job_dir}",
|
|
308
|
+
"-e", "AWS_BATCH_JOB_ARRAY_INDEX=0",
|
|
309
|
+
"-e", "BATCH_ARRAY_SIZE=1",
|
|
310
|
+
"-e", "BATCH_NUM_FILES=1",
|
|
311
|
+
DEFAULT_IMAGE_URI,
|
|
312
|
+
]
|
|
313
|
+
|
|
314
|
+
click.echo(f"Output: {job_dir / 'output'}")
|
|
315
|
+
click.echo()
|
|
316
|
+
|
|
317
|
+
result = subprocess.run(cmd)
|
|
318
|
+
if result.returncode != 0:
|
|
319
|
+
click.echo(
|
|
320
|
+
click.style(f"Container exited with code {result.returncode}", fg="red"),
|
|
321
|
+
err=True,
|
|
322
|
+
)
|
|
323
|
+
raise SystemExit(result.returncode)
|
|
324
|
+
|
|
325
|
+
output_dir = job_dir / "output"
|
|
326
|
+
outputs = list(output_dir.rglob("*.out"))
|
|
327
|
+
if outputs:
|
|
328
|
+
click.echo()
|
|
329
|
+
click.echo(click.style("ORCA completed!", fg="green"))
|
|
330
|
+
for f in outputs:
|
|
331
|
+
click.echo(f" {f}")
|
|
332
|
+
else:
|
|
333
|
+
click.echo(click.style("Warning: No .out files found", fg="yellow"))
|
|
334
|
+
|
|
335
|
+
|
|
336
|
+
def _run_shell_mode(input_path: Path):
|
|
337
|
+
"""Drop into container shell for debugging."""
|
|
338
|
+
import subprocess
|
|
339
|
+
|
|
340
|
+
click.echo("Dropping into ORCA container shell...")
|
|
341
|
+
|
|
342
|
+
mount_path = input_path if input_path.is_dir() else input_path.parent
|
|
343
|
+
cmd = [
|
|
344
|
+
"docker", "run", "--rm", "-it",
|
|
345
|
+
"-v", "/primordial:/primordial:ro",
|
|
346
|
+
"-v", f"{mount_path}:/work",
|
|
347
|
+
"--entrypoint", "/bin/bash",
|
|
348
|
+
DEFAULT_IMAGE_URI,
|
|
349
|
+
]
|
|
350
|
+
|
|
351
|
+
click.echo(f"Input available at: /work/")
|
|
352
|
+
subprocess.run(cmd)
|
|
@@ -1,9 +1,8 @@
|
|
|
1
1
|
"""CLI commands for cloud provider authentication and management.
|
|
2
2
|
|
|
3
3
|
This module provides commands for authenticating with GCP and AWS from within
|
|
4
|
-
development containers. It handles
|
|
5
|
-
|
|
6
|
-
shell RC files (AWS only) or gcloud config settings (GCP).
|
|
4
|
+
development containers. It handles persistent configuration via shell RC files
|
|
5
|
+
(AWS only) or gcloud config settings (GCP).
|
|
7
6
|
|
|
8
7
|
The implementation focuses on:
|
|
9
8
|
1. Unifying cloud authentication with the `dh` CLI tool
|
|
@@ -27,7 +26,6 @@ import questionary
|
|
|
27
26
|
import typer
|
|
28
27
|
|
|
29
28
|
# --- Configuration ---
|
|
30
|
-
GCP_DEVCON_SA = "devcon@enzyme-discovery.iam.gserviceaccount.com"
|
|
31
29
|
GCP_PROJECT_ID = "enzyme-discovery"
|
|
32
30
|
AWS_DEFAULT_PROFILE = "dev-devaccess"
|
|
33
31
|
AWS_CONFIG_FILE = Path.home() / ".aws" / "config"
|
|
@@ -134,26 +132,17 @@ def _get_env_var(variable: str) -> Optional[str]:
|
|
|
134
132
|
|
|
135
133
|
# --- GCP Functions ---
|
|
136
134
|
|
|
137
|
-
# New approach: Use gcloud config settings instead of environment variables
|
|
138
|
-
# for impersonation and project settings. This avoids modifying RC files
|
|
139
|
-
# and the need for `eval "$(dh gcp use-... --export)"`.
|
|
140
|
-
# ADC is updated during the initial `dh gcp login` for the user.
|
|
141
|
-
# Subsequent ADC updates for impersonation or user mode must be done manually
|
|
142
|
-
# if required by libraries, as the underlying gcloud commands can force interaction.
|
|
143
|
-
|
|
144
135
|
|
|
145
136
|
def _get_short_name(account: str) -> str:
|
|
146
|
-
"""Extracts a short name ('dma'
|
|
137
|
+
"""Extracts a short name ('dma') from a GCP account email.
|
|
147
138
|
|
|
148
139
|
Args:
|
|
149
140
|
account: The full account string (e.g., 'dma@dayhofflabs.com',
|
|
150
|
-
'
|
|
141
|
+
'None', 'Not authenticated').
|
|
151
142
|
|
|
152
143
|
Returns:
|
|
153
144
|
The short name or the original string if not a recognized email pattern.
|
|
154
145
|
"""
|
|
155
|
-
if account == GCP_DEVCON_SA:
|
|
156
|
-
return "devcon"
|
|
157
146
|
if "@" in account:
|
|
158
147
|
# Attempt to get the part before @, common for user accounts
|
|
159
148
|
user_part = account.split("@")[0]
|
|
@@ -168,7 +157,7 @@ def _gcloud_set_config(key: str, value: str) -> Tuple[int, str, str]:
|
|
|
168
157
|
"""Set a gcloud configuration value using `gcloud config set`.
|
|
169
158
|
|
|
170
159
|
Args:
|
|
171
|
-
key: The configuration key (e.g., 'project'
|
|
160
|
+
key: The configuration key (e.g., 'project').
|
|
172
161
|
value: The value to set for the key.
|
|
173
162
|
|
|
174
163
|
Returns:
|
|
@@ -179,20 +168,6 @@ def _gcloud_set_config(key: str, value: str) -> Tuple[int, str, str]:
|
|
|
179
168
|
return _run_command(cmd, capture=True, check=False, suppress_output=True)
|
|
180
169
|
|
|
181
170
|
|
|
182
|
-
def _gcloud_unset_config(key: str) -> Tuple[int, str, str]:
|
|
183
|
-
"""Unset a gcloud configuration value using `gcloud config unset`.
|
|
184
|
-
|
|
185
|
-
Args:
|
|
186
|
-
key: The configuration key to unset.
|
|
187
|
-
|
|
188
|
-
Returns:
|
|
189
|
-
Tuple of (return_code, stdout_str, stderr_str) from _run_command.
|
|
190
|
-
"""
|
|
191
|
-
gcloud_path = _find_executable("gcloud")
|
|
192
|
-
cmd = [gcloud_path, "config", "unset", key, "--quiet"]
|
|
193
|
-
return _run_command(cmd, capture=True, check=False, suppress_output=True)
|
|
194
|
-
|
|
195
|
-
|
|
196
171
|
def _get_adc_status() -> str:
|
|
197
172
|
"""Check the status and type of Application Default Credentials (ADC).
|
|
198
173
|
|
|
@@ -200,7 +175,7 @@ def _get_adc_status() -> str:
|
|
|
200
175
|
default GCE metadata server fallback if no explicit config is found.
|
|
201
176
|
|
|
202
177
|
Returns:
|
|
203
|
-
A short string describing the ADC principal ('dma',
|
|
178
|
+
A short string describing the ADC principal ('dma',
|
|
204
179
|
'default VM service account', 'Other/External', 'Not configured', etc.).
|
|
205
180
|
"""
|
|
206
181
|
adc_file = (
|
|
@@ -221,15 +196,6 @@ def _get_adc_status() -> str:
|
|
|
221
196
|
|
|
222
197
|
if cred_type == "authorized_user":
|
|
223
198
|
return "dma" # Assuming 'dma' is the likely user
|
|
224
|
-
elif cred_type == "impersonated_service_account":
|
|
225
|
-
sa_url = adc_data.get("service_account_impersonation_url", "")
|
|
226
|
-
sa_match = re.search(r"serviceAccounts/([^:]+)", sa_url)
|
|
227
|
-
if sa_match and sa_match.group(1) == GCP_DEVCON_SA:
|
|
228
|
-
return "devcon"
|
|
229
|
-
elif sa_match:
|
|
230
|
-
return f"Other SA ({_get_short_name(sa_match.group(1))})"
|
|
231
|
-
else:
|
|
232
|
-
return "devcon (?)" # Likely devcon but failed parse
|
|
233
199
|
elif cred_type == "external_account":
|
|
234
200
|
return "Other/External"
|
|
235
201
|
elif cred_type == "service_account":
|
|
@@ -325,21 +291,6 @@ def _get_current_gcp_user() -> str:
|
|
|
325
291
|
return "Not authenticated"
|
|
326
292
|
|
|
327
293
|
|
|
328
|
-
def _get_current_gcp_impersonation() -> str:
|
|
329
|
-
"""Get the current impersonated service account from gcloud config."""
|
|
330
|
-
gcloud_path = _find_executable("gcloud")
|
|
331
|
-
cmd = [
|
|
332
|
-
gcloud_path,
|
|
333
|
-
"config",
|
|
334
|
-
"get-value",
|
|
335
|
-
"auth/impersonate_service_account",
|
|
336
|
-
"--quiet",
|
|
337
|
-
]
|
|
338
|
-
returncode, stdout, _ = _run_command(cmd, capture=True, check=False)
|
|
339
|
-
sa = stdout.strip() if returncode == 0 else ""
|
|
340
|
-
return sa if sa else "None"
|
|
341
|
-
|
|
342
|
-
|
|
343
294
|
def _run_gcloud_login() -> None:
|
|
344
295
|
"""Run the gcloud auth login command, updating ADC using device flow.
|
|
345
296
|
|
|
@@ -372,11 +323,10 @@ def _run_gcloud_login() -> None:
|
|
|
372
323
|
print(f"{GREEN}User authentication complete. ADC updated for user account.{NC}")
|
|
373
324
|
|
|
374
325
|
|
|
375
|
-
def _test_gcp_credentials(user: str
|
|
326
|
+
def _test_gcp_credentials(user: str) -> None:
|
|
376
327
|
"""Test GCP credentials. Prints output on failure (to stderr) and success (to stdout)."""
|
|
377
328
|
gcloud_path = _find_executable("gcloud")
|
|
378
329
|
user_short = _get_short_name(user)
|
|
379
|
-
impersonation_short = _get_short_name(impersonation_sa)
|
|
380
330
|
|
|
381
331
|
if user != "Not authenticated" and "Not authenticated" not in user:
|
|
382
332
|
cmd = [
|
|
@@ -388,70 +338,15 @@ def _test_gcp_credentials(user: str, impersonation_sa: str) -> None:
|
|
|
388
338
|
f"--project={GCP_PROJECT_ID}",
|
|
389
339
|
]
|
|
390
340
|
|
|
391
|
-
|
|
392
|
-
|
|
393
|
-
|
|
394
|
-
|
|
395
|
-
|
|
396
|
-
|
|
397
|
-
)
|
|
398
|
-
if unset_rc != 0:
|
|
399
|
-
print(
|
|
400
|
-
f" {RED}✗ Test Error: Failed to temporarily disable impersonation: {unset_err}{NC}",
|
|
401
|
-
file=sys.stderr,
|
|
402
|
-
)
|
|
403
|
-
# Even if unsetting fails, attempt to restore and continue with impersonation test
|
|
404
|
-
else:
|
|
405
|
-
user_returncode, _, _ = _run_command(
|
|
406
|
-
cmd, suppress_output=True, check=False
|
|
407
|
-
)
|
|
408
|
-
if user_returncode != 0:
|
|
409
|
-
print(
|
|
410
|
-
f" {RED}✗ User Test Failure: Cannot access resources directly as user '{user_short}'. Check roles/project.{NC}",
|
|
411
|
-
file=sys.stderr,
|
|
412
|
-
)
|
|
413
|
-
else:
|
|
414
|
-
print(
|
|
415
|
-
f" {GREEN}✓ User Test ({user_short}): Direct access OK{NC}"
|
|
416
|
-
)
|
|
417
|
-
|
|
418
|
-
# Restore impersonation setting
|
|
419
|
-
set_rc, _, set_err = _gcloud_set_config(
|
|
420
|
-
"auth/impersonate_service_account", orig_sa
|
|
421
|
-
)
|
|
422
|
-
if set_rc != 0:
|
|
423
|
-
print(
|
|
424
|
-
f" {RED}✗ Test Error: Failed to restore impersonation config for {impersonation_short}: {set_err}{NC}",
|
|
425
|
-
file=sys.stderr,
|
|
426
|
-
)
|
|
427
|
-
# If restoring fails, it's a significant issue for the next test
|
|
428
|
-
|
|
429
|
-
# Test 2: Access while impersonating the SA
|
|
430
|
-
print(f" Testing access while impersonating SA ({impersonation_short})...")
|
|
431
|
-
impersonation_returncode, _, _ = _run_command(
|
|
432
|
-
cmd, suppress_output=True, check=False
|
|
341
|
+
print(f" Testing direct access as user ({user_short})...")
|
|
342
|
+
returncode, _, _ = _run_command(cmd, suppress_output=True, check=False)
|
|
343
|
+
if returncode != 0:
|
|
344
|
+
print(
|
|
345
|
+
f" {RED}✗ User Test Failure: Cannot access resources directly as user '{user_short}'. Check roles/project.{NC}",
|
|
346
|
+
file=sys.stderr,
|
|
433
347
|
)
|
|
434
|
-
if impersonation_returncode != 0:
|
|
435
|
-
print(
|
|
436
|
-
f" {RED}✗ Impersonation Test Failure: Cannot access resources impersonating '{impersonation_short}'. Check permissions/config.{NC}",
|
|
437
|
-
file=sys.stderr,
|
|
438
|
-
)
|
|
439
|
-
else:
|
|
440
|
-
print(
|
|
441
|
-
f" {GREEN}✓ Impersonation Test ({impersonation_short}): Access OK{NC}"
|
|
442
|
-
)
|
|
443
|
-
|
|
444
348
|
else:
|
|
445
|
-
|
|
446
|
-
print(f" Testing direct access as user ({user_short})...")
|
|
447
|
-
returncode, _, _ = _run_command(cmd, suppress_output=True, check=False)
|
|
448
|
-
if returncode != 0:
|
|
449
|
-
print(
|
|
450
|
-
f" {RED}✗ User Test Failure: Cannot access resources directly as user '{user_short}'. Check roles/project.{NC}",
|
|
451
|
-
file=sys.stderr,
|
|
452
|
-
)
|
|
453
|
-
else:
|
|
454
|
-
print(f" {GREEN}✓ User Test ({user_short}): Direct access OK{NC}")
|
|
349
|
+
print(f" {GREEN}✓ User Test ({user_short}): Direct access OK{NC}")
|
|
455
350
|
else:
|
|
456
351
|
print(
|
|
457
352
|
f" {YELLOW}User not authenticated, skipping credential access tests.{NC}"
|
|
@@ -561,19 +456,12 @@ aws_app = typer.Typer(
|
|
|
561
456
|
def gcp_status():
|
|
562
457
|
"""Show active GCP credentials for CLI and Libraries/ADC, including staleness."""
|
|
563
458
|
cli_user = _get_current_gcp_user()
|
|
564
|
-
|
|
565
|
-
adc_principal_raw = _get_adc_status() # Raw status string, potentially complex
|
|
459
|
+
adc_principal_raw = _get_adc_status()
|
|
566
460
|
|
|
567
461
|
user_auth_valid = _is_gcp_user_authenticated()
|
|
568
462
|
adc_auth_valid = _is_adc_authenticated()
|
|
569
463
|
|
|
570
|
-
|
|
571
|
-
if cli_impersonation != "None":
|
|
572
|
-
cli_active_short = _get_short_name(cli_impersonation)
|
|
573
|
-
cli_is_impersonating = True
|
|
574
|
-
else:
|
|
575
|
-
cli_active_short = _get_short_name(cli_user)
|
|
576
|
-
cli_is_impersonating = False
|
|
464
|
+
cli_active_short = _get_short_name(cli_user)
|
|
577
465
|
|
|
578
466
|
adc_active_short = _get_short_name(adc_principal_raw)
|
|
579
467
|
|
|
@@ -587,14 +475,6 @@ def gcp_status():
|
|
|
587
475
|
f" └─ Authentication: {RED}STALE/EXPIRED{NC} (Hint: run 'dh gcp login')"
|
|
588
476
|
)
|
|
589
477
|
|
|
590
|
-
if cli_is_impersonating:
|
|
591
|
-
print(
|
|
592
|
-
f" Impersonation ({_get_short_name(cli_impersonation)}): {GREEN}Active{NC}"
|
|
593
|
-
)
|
|
594
|
-
print(f" └─ Access Test: (see results below)")
|
|
595
|
-
else:
|
|
596
|
-
print(f" Impersonation: {YELLOW}Not Active{NC}")
|
|
597
|
-
|
|
598
478
|
print(f"\n{BLUE}--- GCP Library/ADC Credentials ---{NC}")
|
|
599
479
|
print(f" Effective Principal: {GREEN}{adc_active_short}{NC}")
|
|
600
480
|
if adc_principal_raw in ["Not configured", "Error reading", "Invalid format"]:
|
|
@@ -603,227 +483,30 @@ def gcp_status():
|
|
|
603
483
|
print(f" └─ Authentication: {GREEN}VALID{NC}")
|
|
604
484
|
else:
|
|
605
485
|
print(
|
|
606
|
-
f" └─ Authentication: {RED}STALE/EXPIRED{NC} (Hint: run 'dh gcp use
|
|
486
|
+
f" └─ Authentication: {RED}STALE/EXPIRED{NC} (Hint: run 'dh gcp use-user-adc' or 'gcloud auth application-default login ...')"
|
|
607
487
|
)
|
|
608
488
|
|
|
609
489
|
print(f"\n{BLUE}--- GCP Access Tests (for CLI configuration) ---{NC}")
|
|
610
|
-
|
|
611
|
-
_test_gcp_credentials(cli_user, cli_impersonation)
|
|
490
|
+
_test_gcp_credentials(cli_user)
|
|
612
491
|
|
|
613
492
|
|
|
614
493
|
@gcp_app.command("login")
|
|
615
494
|
def gcp_login():
|
|
616
|
-
"""Authenticate user & configure CLI
|
|
495
|
+
"""Authenticate user & configure CLI for GCP access."""
|
|
617
496
|
# Step 1: Authenticate the user (updates ADC for user)
|
|
618
497
|
_run_gcloud_login() # Uses device flow
|
|
619
498
|
|
|
620
|
-
# Step 2: Configure gcloud CLI for devcon SA impersonation
|
|
621
|
-
print(f"\n{BLUE}Configuring gcloud CLI to impersonate {GCP_DEVCON_SA}...{NC}")
|
|
622
|
-
set_sa_rc, _, set_sa_err = _gcloud_set_config(
|
|
623
|
-
"auth/impersonate_service_account", GCP_DEVCON_SA
|
|
624
|
-
)
|
|
625
|
-
if set_sa_rc != 0:
|
|
626
|
-
print(
|
|
627
|
-
f"{RED}Error setting impersonation config: {set_sa_err}{NC}",
|
|
628
|
-
file=sys.stderr,
|
|
629
|
-
)
|
|
630
|
-
print(f"{YELLOW}Warning: CLI impersonation failed to configure.{NC}")
|
|
631
|
-
# Attempt to show status anyway before exiting command
|
|
632
|
-
print("\n{BLUE}Current status:{NC}")
|
|
633
|
-
gcp_status()
|
|
634
|
-
return
|
|
635
|
-
|
|
636
499
|
set_proj_rc, _, set_proj_err = _gcloud_set_config("project", GCP_PROJECT_ID)
|
|
637
500
|
if set_proj_rc != 0:
|
|
638
501
|
print(f"{RED}Error setting project config: {set_proj_err}{NC}", file=sys.stderr)
|
|
639
|
-
# Continue, but warn user
|
|
640
|
-
|
|
641
|
-
# Step 3: Print configuration options
|
|
642
|
-
print(f"\n{GREEN}Login successful. CLI configured for devcon impersonation.{NC}")
|
|
643
|
-
print(f"{BLUE}--- Common Configuration Commands ---\n{NC}")
|
|
644
502
|
|
|
645
|
-
|
|
646
|
-
|
|
647
|
-
print(f" {BLUE}Set CLI to use User:{NC}")
|
|
648
|
-
print(f" {YELLOW}{f'dh gcp use-user':<{cmd_width}}{NC}")
|
|
649
|
-
|
|
650
|
-
print(f" {BLUE}Set CLI to use Devcon SA:{NC}")
|
|
651
|
-
print(
|
|
652
|
-
f" {YELLOW}{f'dh gcp use-devcon':<{cmd_width}}{NC} {GREEN}(Current default after login){NC}"
|
|
653
|
-
)
|
|
654
|
-
|
|
655
|
-
print(f" {BLUE}Set Libraries/Tools (ADC) to use User:{NC}")
|
|
656
|
-
print(f" {YELLOW}{f'dh gcp use-user-adc':<{cmd_width}}{NC}")
|
|
657
|
-
|
|
658
|
-
print(f" {BLUE}Set Libraries/Tools (ADC) to use Devcon SA:{NC}")
|
|
659
|
-
print(f" {YELLOW}{f'dh gcp use-devcon-adc':<{cmd_width}}{NC}")
|
|
503
|
+
print(f"\n{GREEN}Login successful.{NC}")
|
|
660
504
|
|
|
661
505
|
# Step 4: Show current status automatically
|
|
662
506
|
print(f"\n{BLUE}--- Current Status ---{NC}")
|
|
663
507
|
gcp_status()
|
|
664
508
|
|
|
665
509
|
|
|
666
|
-
@gcp_app.command("use-devcon")
|
|
667
|
-
def gcp_use_devcon(
|
|
668
|
-
export: bool = typer.Option(
|
|
669
|
-
False,
|
|
670
|
-
"--export",
|
|
671
|
-
"-x",
|
|
672
|
-
help="Deprecated. Has no effect. Settings are applied directly via gcloud config.",
|
|
673
|
-
hidden=True,
|
|
674
|
-
),
|
|
675
|
-
):
|
|
676
|
-
"""Configure gcloud CLI to impersonate the devcon SA.
|
|
677
|
-
|
|
678
|
-
This command updates gcloud configuration settings directly.
|
|
679
|
-
It DOES NOT modify shell RC files or require `eval`.
|
|
680
|
-
It DOES NOT automatically update Application Default Credentials (ADC) for impersonation.
|
|
681
|
-
|
|
682
|
-
Ensures the primary user login is valid first.
|
|
683
|
-
"""
|
|
684
|
-
if export:
|
|
685
|
-
print(
|
|
686
|
-
f"{YELLOW}Warning: --export/-x is deprecated and has no effect. "
|
|
687
|
-
f"GCP settings are now managed via gcloud config.{NC}",
|
|
688
|
-
file=sys.stderr,
|
|
689
|
-
)
|
|
690
|
-
|
|
691
|
-
if not _is_gcp_user_authenticated():
|
|
692
|
-
print(
|
|
693
|
-
f"{RED}Error: GCP user authentication is invalid or requires interactive login.{NC}",
|
|
694
|
-
file=sys.stderr,
|
|
695
|
-
)
|
|
696
|
-
print(
|
|
697
|
-
f"{YELLOW}Please run 'dh gcp login' interactively first, then try this command again.{NC}",
|
|
698
|
-
file=sys.stderr,
|
|
699
|
-
)
|
|
700
|
-
sys.exit(1)
|
|
701
|
-
|
|
702
|
-
print(f"{BLUE}Configuring gcloud CLI to impersonate {GCP_DEVCON_SA}...{NC}")
|
|
703
|
-
|
|
704
|
-
# Set gcloud CLI impersonation via config
|
|
705
|
-
set_sa_rc, _, set_sa_err = _gcloud_set_config(
|
|
706
|
-
"auth/impersonate_service_account", GCP_DEVCON_SA
|
|
707
|
-
)
|
|
708
|
-
if set_sa_rc != 0:
|
|
709
|
-
print(
|
|
710
|
-
f"{RED}Error setting impersonation config: {set_sa_err}{NC}",
|
|
711
|
-
file=sys.stderr,
|
|
712
|
-
)
|
|
713
|
-
sys.exit(1)
|
|
714
|
-
|
|
715
|
-
set_proj_rc, _, set_proj_err = _gcloud_set_config("project", GCP_PROJECT_ID)
|
|
716
|
-
if set_proj_rc != 0:
|
|
717
|
-
print(f"{RED}Error setting project config: {set_proj_err}{NC}", file=sys.stderr)
|
|
718
|
-
# Continue, but warn user
|
|
719
|
-
|
|
720
|
-
# Check for lingering legacy environment variable
|
|
721
|
-
if _get_env_var("CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT"):
|
|
722
|
-
print(
|
|
723
|
-
f"{YELLOW}Warning: Legacy env var CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT is set.{NC}",
|
|
724
|
-
file=sys.stderr,
|
|
725
|
-
)
|
|
726
|
-
print(
|
|
727
|
-
f"{YELLOW} This may override gcloud config. Consider running:{NC}",
|
|
728
|
-
file=sys.stderr,
|
|
729
|
-
)
|
|
730
|
-
print(
|
|
731
|
-
f"{YELLOW} unset CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT{NC}",
|
|
732
|
-
file=sys.stderr,
|
|
733
|
-
)
|
|
734
|
-
|
|
735
|
-
print(f"\n{GREEN}GCP CLI configured to use devcon SA ({GCP_DEVCON_SA}).{NC}")
|
|
736
|
-
print(f"Project set to: {GCP_PROJECT_ID}")
|
|
737
|
-
print(
|
|
738
|
-
f"{YELLOW}NOTE: If libraries/tools (e.g., for DVC, Terraform) need to use impersonation, update Application Default Credentials (ADC) manually:{NC}"
|
|
739
|
-
)
|
|
740
|
-
print(
|
|
741
|
-
f"{YELLOW} gcloud auth application-default login --impersonate-service-account={GCP_DEVCON_SA}{NC}"
|
|
742
|
-
)
|
|
743
|
-
print(f"Run 'dh gcp status' to verify CLI configuration.")
|
|
744
|
-
|
|
745
|
-
|
|
746
|
-
@gcp_app.command("use-user")
|
|
747
|
-
def gcp_use_user(
|
|
748
|
-
export: bool = typer.Option(
|
|
749
|
-
False,
|
|
750
|
-
"--export",
|
|
751
|
-
"-x",
|
|
752
|
-
help="Deprecated. Has no effect. Settings are applied directly via gcloud config.",
|
|
753
|
-
hidden=True,
|
|
754
|
-
),
|
|
755
|
-
):
|
|
756
|
-
"""Configure gcloud CLI to use the personal user account via gcloud config.
|
|
757
|
-
|
|
758
|
-
This command updates gcloud configuration settings directly.
|
|
759
|
-
It DOES NOT modify shell RC files or require `eval`.
|
|
760
|
-
It DOES NOT automatically update Application Default Credentials (ADC).
|
|
761
|
-
|
|
762
|
-
Ensures the primary user login is valid first.
|
|
763
|
-
"""
|
|
764
|
-
if export:
|
|
765
|
-
print(
|
|
766
|
-
f"{YELLOW}Warning: --export/-x is deprecated and has no effect. "
|
|
767
|
-
f"GCP settings are now managed via gcloud config.{NC}",
|
|
768
|
-
file=sys.stderr,
|
|
769
|
-
)
|
|
770
|
-
|
|
771
|
-
if not _is_gcp_user_authenticated():
|
|
772
|
-
print(
|
|
773
|
-
f"{RED}Error: GCP user authentication is invalid or requires interactive login.{NC}",
|
|
774
|
-
file=sys.stderr,
|
|
775
|
-
)
|
|
776
|
-
print(
|
|
777
|
-
f"{YELLOW}Please run 'dh gcp login' interactively first, then try this command again.{NC}",
|
|
778
|
-
file=sys.stderr,
|
|
779
|
-
)
|
|
780
|
-
sys.exit(1)
|
|
781
|
-
|
|
782
|
-
print(f"{BLUE}Configuring gcloud CLI to use personal user account...{NC}")
|
|
783
|
-
|
|
784
|
-
# Unset gcloud CLI impersonation via config
|
|
785
|
-
unset_sa_rc, _, unset_sa_err = _gcloud_unset_config(
|
|
786
|
-
"auth/impersonate_service_account"
|
|
787
|
-
)
|
|
788
|
-
if unset_sa_rc != 0:
|
|
789
|
-
print(
|
|
790
|
-
f"{RED}Error unsetting impersonation config: {unset_sa_err}{NC}",
|
|
791
|
-
file=sys.stderr,
|
|
792
|
-
)
|
|
793
|
-
# Continue, but warn user
|
|
794
|
-
|
|
795
|
-
set_proj_rc, _, set_proj_err = _gcloud_set_config("project", GCP_PROJECT_ID)
|
|
796
|
-
if set_proj_rc != 0:
|
|
797
|
-
print(f"{RED}Error setting project config: {set_proj_err}{NC}", file=sys.stderr)
|
|
798
|
-
# Continue, but warn user
|
|
799
|
-
|
|
800
|
-
# Check for lingering legacy environment variable
|
|
801
|
-
if _get_env_var("CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT"):
|
|
802
|
-
print(
|
|
803
|
-
f"{YELLOW}Warning: Legacy env var CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT is set.{NC}",
|
|
804
|
-
file=sys.stderr,
|
|
805
|
-
)
|
|
806
|
-
print(
|
|
807
|
-
f"{YELLOW} This may interfere with using your personal account. Consider running:{NC}",
|
|
808
|
-
file=sys.stderr,
|
|
809
|
-
)
|
|
810
|
-
print(
|
|
811
|
-
f"{YELLOW} unset CLOUDSDK_AUTH_IMPERSONATE_SERVICE_ACCOUNT{NC}",
|
|
812
|
-
file=sys.stderr,
|
|
813
|
-
)
|
|
814
|
-
|
|
815
|
-
print(f"\n{GREEN}GCP CLI configured to use personal account.{NC}")
|
|
816
|
-
print(f"Project set to: {GCP_PROJECT_ID}")
|
|
817
|
-
print(
|
|
818
|
-
f"{YELLOW}NOTE: If libraries/tools (e.g., for DVC, Terraform) need to use impersonation, update Application Default Credentials (ADC) manually:{NC}"
|
|
819
|
-
)
|
|
820
|
-
print(f"{YELLOW} gcloud auth application-default login{NC}")
|
|
821
|
-
print(f"Run 'dh gcp status' to verify CLI configuration.")
|
|
822
|
-
|
|
823
|
-
|
|
824
|
-
# === NEW ADC Commands ===
|
|
825
|
-
|
|
826
|
-
|
|
827
510
|
@gcp_app.command("use-user-adc")
|
|
828
511
|
def gcp_use_user_adc():
|
|
829
512
|
"""Configure Libraries/Tools (ADC) to use your PERSONAL account."""
|
|
@@ -863,58 +546,11 @@ def gcp_use_user_adc():
|
|
|
863
546
|
sys.exit(1)
|
|
864
547
|
|
|
865
548
|
|
|
866
|
-
@gcp_app.command("use-devcon-adc")
|
|
867
|
-
def gcp_use_devcon_adc():
|
|
868
|
-
"""Configure Libraries/Tools (ADC) to use the DEVCON service account."""
|
|
869
|
-
if not _is_gcp_user_authenticated():
|
|
870
|
-
print(
|
|
871
|
-
f"{RED}Error: GCP user authentication is invalid or requires interactive login.{NC}",
|
|
872
|
-
file=sys.stderr,
|
|
873
|
-
)
|
|
874
|
-
print(
|
|
875
|
-
f"{YELLOW}Please run 'dh gcp login' interactively first.{NC}",
|
|
876
|
-
file=sys.stderr,
|
|
877
|
-
)
|
|
878
|
-
sys.exit(1)
|
|
879
|
-
|
|
880
|
-
print(f"{BLUE}Attempting to configure ADC for devcon SA ({GCP_DEVCON_SA})...{NC}")
|
|
881
|
-
print(
|
|
882
|
-
f"{YELLOW}This may require you to complete a browser authentication flow.{NC}"
|
|
883
|
-
)
|
|
884
|
-
|
|
885
|
-
gcloud_path = _find_executable("gcloud")
|
|
886
|
-
cmd = [
|
|
887
|
-
gcloud_path,
|
|
888
|
-
"auth",
|
|
889
|
-
"application-default",
|
|
890
|
-
"login",
|
|
891
|
-
f"--impersonate-service-account={GCP_DEVCON_SA}",
|
|
892
|
-
]
|
|
893
|
-
|
|
894
|
-
# Allow interaction, don't capture output
|
|
895
|
-
returncode, _, _ = _run_command(
|
|
896
|
-
cmd, capture=False, check=False, suppress_output=False
|
|
897
|
-
)
|
|
898
|
-
|
|
899
|
-
if returncode == 0:
|
|
900
|
-
print(
|
|
901
|
-
f"\n{GREEN}Successfully configured ADC for devcon SA ({GCP_DEVCON_SA}).{NC}"
|
|
902
|
-
)
|
|
903
|
-
print(f"{BLUE}--- Current Status ---{NC}")
|
|
904
|
-
gcp_status() # Show status after successful change
|
|
905
|
-
else:
|
|
906
|
-
print(
|
|
907
|
-
f"{RED}Failed to configure ADC (Return code: {returncode}). Check messages above.{NC}",
|
|
908
|
-
file=sys.stderr,
|
|
909
|
-
)
|
|
910
|
-
sys.exit(1)
|
|
911
|
-
|
|
912
|
-
|
|
913
549
|
@gcp_app.command("logout")
|
|
914
550
|
def gcp_logout():
|
|
915
551
|
"""Clear all GCP credentials for testing or role switching purposes.
|
|
916
552
|
|
|
917
|
-
This removes the active user's gcloud login
|
|
553
|
+
This removes the active user's gcloud login
|
|
918
554
|
and invalidates Application Default Credentials (ADC).
|
|
919
555
|
"""
|
|
920
556
|
print(f"{BLUE}Clearing all GCP credentials...{NC}")
|
|
@@ -930,20 +566,12 @@ def gcp_logout():
|
|
|
930
566
|
if revoke_code != 0 and revoke_err:
|
|
931
567
|
errors.append(f"Failed to revoke credentials: {revoke_err}")
|
|
932
568
|
|
|
933
|
-
# 2.
|
|
934
|
-
print(f"{BLUE}Disabling service account impersonation...{NC}")
|
|
935
|
-
unset_code, _, unset_err = _gcloud_unset_config(
|
|
936
|
-
"auth/impersonate_service_account"
|
|
937
|
-
)
|
|
938
|
-
if unset_code != 0 and unset_err:
|
|
939
|
-
errors.append(f"Failed to unset impersonation: {unset_err}")
|
|
940
|
-
|
|
941
|
-
# 3. Revoke ADC
|
|
569
|
+
# 2. Revoke ADC
|
|
942
570
|
print(f"{BLUE}Revoking Application Default Credentials (ADC)...{NC}")
|
|
943
571
|
adc_cmd = [gcloud_path, "auth", "application-default", "revoke", "--quiet"]
|
|
944
572
|
adc_code, _, adc_err = _run_command(adc_cmd, capture=True, check=False)
|
|
945
573
|
|
|
946
|
-
#
|
|
574
|
+
# 3. Additionally remove ADC file if it exists (belt-and-suspenders approach)
|
|
947
575
|
adc_file = (
|
|
948
576
|
Path.home() / ".config" / "gcloud" / "application_default_credentials.json"
|
|
949
577
|
)
|
|
@@ -976,9 +604,6 @@ def gcp_logout():
|
|
|
976
604
|
print(f" {YELLOW}gcloud auth application-default revoke{NC}")
|
|
977
605
|
|
|
978
606
|
|
|
979
|
-
# === End NEW ADC Commands ===
|
|
980
|
-
|
|
981
|
-
|
|
982
607
|
# --- AWS Commands ---
|
|
983
608
|
@aws_app.command("status")
|
|
984
609
|
def aws_status(
|
|
@@ -29,7 +29,7 @@ app.command("clean")(delete_local_branch)
|
|
|
29
29
|
app.command("wget")(get_from_warehouse_typer)
|
|
30
30
|
|
|
31
31
|
# Cloud commands
|
|
32
|
-
app.add_typer(gcp_app, name="gcp", help="Manage GCP authentication
|
|
32
|
+
app.add_typer(gcp_app, name="gcp", help="Manage GCP authentication.")
|
|
33
33
|
app.add_typer(aws_app, name="aws", help="Manage AWS SSO authentication.")
|
|
34
34
|
app.add_typer(gh_app, name="gh", help="Manage GitHub authentication.")
|
|
35
35
|
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
"""Structural tests for GCP commands after service account removal.
|
|
2
|
+
|
|
3
|
+
These tests verify the interface contract: which commands exist, which
|
|
4
|
+
constants/functions were removed, and that docstrings no longer reference
|
|
5
|
+
service account impersonation. No GCP credentials required.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import inspect
|
|
9
|
+
|
|
10
|
+
import pytest
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def test_gcp_app_has_exactly_four_commands():
|
|
14
|
+
"""After cleanup, gcp_app should have: login, status, use-user-adc, logout."""
|
|
15
|
+
from dh_cli.cloud_commands import gcp_app
|
|
16
|
+
|
|
17
|
+
registered = {
|
|
18
|
+
cmd_info.name or cmd_info.callback.__name__
|
|
19
|
+
for cmd_info in gcp_app.registered_commands
|
|
20
|
+
}
|
|
21
|
+
expected = {"login", "status", "use-user-adc", "logout"}
|
|
22
|
+
assert registered == expected, (
|
|
23
|
+
f"Expected commands {expected}, got {registered}. "
|
|
24
|
+
f"Extra: {registered - expected}, Missing: {expected - registered}"
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def test_devcon_sa_constant_removed():
|
|
29
|
+
"""GCP_DEVCON_SA constant should no longer exist."""
|
|
30
|
+
import dh_cli.cloud_commands as mod
|
|
31
|
+
|
|
32
|
+
assert not hasattr(mod, "GCP_DEVCON_SA"), "GCP_DEVCON_SA should be removed"
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def test_get_current_gcp_impersonation_removed():
|
|
36
|
+
"""The impersonation helper should no longer exist."""
|
|
37
|
+
import dh_cli.cloud_commands as mod
|
|
38
|
+
|
|
39
|
+
assert not hasattr(mod, "_get_current_gcp_impersonation"), (
|
|
40
|
+
"_get_current_gcp_impersonation should be removed"
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def test_test_gcp_credentials_no_impersonation_param():
|
|
45
|
+
"""_test_gcp_credentials should not accept an impersonation_sa parameter."""
|
|
46
|
+
from dh_cli.cloud_commands import _test_gcp_credentials
|
|
47
|
+
|
|
48
|
+
sig = inspect.signature(_test_gcp_credentials)
|
|
49
|
+
assert "impersonation_sa" not in sig.parameters, (
|
|
50
|
+
"_test_gcp_credentials should not have impersonation_sa parameter"
|
|
51
|
+
)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def test_login_docstring_no_impersonation():
|
|
55
|
+
"""Login command docstring should not mention impersonation or devcon."""
|
|
56
|
+
from dh_cli.cloud_commands import gcp_login
|
|
57
|
+
|
|
58
|
+
doc = (gcp_login.__doc__ or "").lower()
|
|
59
|
+
assert "impersonat" not in doc, f"Docstring still mentions impersonation: {doc}"
|
|
60
|
+
assert "devcon" not in doc, f"Docstring still mentions devcon: {doc}"
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def test_gcp_project_id_still_exists():
|
|
64
|
+
"""GCP_PROJECT_ID should still be present (we still need it)."""
|
|
65
|
+
from dh_cli.cloud_commands import GCP_PROJECT_ID
|
|
66
|
+
|
|
67
|
+
assert GCP_PROJECT_ID == "enzyme-discovery"
|
|
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
|