dh-cli 0.10.0__tar.gz → 0.10.2__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.10.0 → dh_cli-0.10.2}/PKG-INFO +1 -1
- {dh_cli-0.10.0 → dh_cli-0.10.2}/pyproject.toml +1 -1
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/__init__.py +3 -0
- dh_cli-0.10.2/src/dh_cli/batch/commands/esmfold2.py +242 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/finalize.py +36 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/bedrock/cost_report.py +1 -1
- dh_cli-0.10.2/src/dh_cli/bedrock/pricing.yaml +38 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/engine_commands.py +1 -1
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_cost_calc.py +64 -16
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_resolve_base_model.py +6 -0
- dh_cli-0.10.2/tests/test_esmfold2_cli.py +268 -0
- dh_cli-0.10.2/tests/test_finalize_esmfold2.py +204 -0
- dh_cli-0.10.0/src/dh_cli/bedrock/pricing.yaml +0 -28
- {dh_cli-0.10.0 → dh_cli-0.10.2}/.gitignore +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/LICENSE +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/README.md +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/__init__.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/_identity.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/aws_batch.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/__init__.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/boltz.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/cancel.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/clean.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/embed_t5.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/equilibrator.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/list_jobs.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/local.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/logs.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/orca.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/protmpnn.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/protmpnn_to_boltz.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/retry.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/rfd3.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/status.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/submit.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/train.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/wait_for.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/fasta_utils.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/h5_utils.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/job_id.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/manifest.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/s3_transport.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/bedrock/__init__.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/bedrock/commands.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/cloud_commands.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/codeartifact.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/__init__.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/api_client.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/auth.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/progress.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/ssh_config.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/studio_commands.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/github_commands.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/__init__.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/deploy.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/local.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/test.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/tf.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/users.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/main.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/utility_commands.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/warehouse.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/__init__.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/test_aws_batch_resources.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/test_image_override.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/test_submit_cpu_only.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/test_submit_image_validation.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/test_submit_merge.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/conftest.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/A_cache_write.json +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/B_cache_read.json +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/C_plain.json +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/D_cursor_user.json +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/E_service_role.json +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/F_legacy_shared.json +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/G_unknown_model.json +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_build_report.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_classify_arn.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_cli_exit_codes.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_cost_command.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_cur_reconciliation.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_key_command.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_render_formats.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_rotate.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_s3_walker.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/__init__.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/conftest.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_engine_role_cannot_read_github_pat.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_identity.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_login.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_login_error_paths.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_login_security.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_logout.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_rotate.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_status.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/hz/test_init.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/hz/test_suites.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/hz/test_users.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_cloud_gcp.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_equilibrator_command.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_finalize_boltz_tar.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_finalize_equilibrator.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_finalize_protmpnn.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_finalize_rfd3.py +0 -0
- {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_rfd3_command.py +0 -0
|
@@ -15,6 +15,7 @@ from .commands.cancel import cancel
|
|
|
15
15
|
from .commands.clean import clean
|
|
16
16
|
from .commands.embed_t5 import embed_t5
|
|
17
17
|
from .commands.equilibrator import equilibrator
|
|
18
|
+
from .commands.esmfold2 import esmfold2
|
|
18
19
|
from .commands.finalize import finalize
|
|
19
20
|
from .commands.list_jobs import list_jobs
|
|
20
21
|
from .commands.local import local
|
|
@@ -53,6 +54,7 @@ def batch_cli():
|
|
|
53
54
|
\b
|
|
54
55
|
Structure Prediction:
|
|
55
56
|
boltz Predict protein structures with Boltz
|
|
57
|
+
esmfold2 Predict protein structures with ESMFold2
|
|
56
58
|
|
|
57
59
|
\b
|
|
58
60
|
Sequence Design:
|
|
@@ -101,6 +103,7 @@ batch_cli.add_command(wait_for, name="wait-for")
|
|
|
101
103
|
# Register pipeline commands
|
|
102
104
|
batch_cli.add_command(embed_t5, name="embed-t5")
|
|
103
105
|
batch_cli.add_command(equilibrator, name="equilibrator")
|
|
106
|
+
batch_cli.add_command(esmfold2)
|
|
104
107
|
batch_cli.add_command(boltz)
|
|
105
108
|
batch_cli.add_command(protmpnn)
|
|
106
109
|
batch_cli.add_command(orca)
|
|
@@ -0,0 +1,242 @@
|
|
|
1
|
+
"""ESMFold2 protein structure prediction pipeline command."""
|
|
2
|
+
|
|
3
|
+
import math
|
|
4
|
+
import secrets
|
|
5
|
+
import shutil
|
|
6
|
+
from datetime import datetime
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
|
|
9
|
+
import click
|
|
10
|
+
|
|
11
|
+
from dh_batch.esmfold2_config import (
|
|
12
|
+
declares_at_least_one_sequence,
|
|
13
|
+
discover_esmfold2_configs,
|
|
14
|
+
)
|
|
15
|
+
|
|
16
|
+
from ..aws_batch import BatchClient, BatchError, resolve_dependency
|
|
17
|
+
from ..job_id import get_aws_username
|
|
18
|
+
from ..manifest import (
|
|
19
|
+
BATCH_JOBS_BASE,
|
|
20
|
+
BatchConfig,
|
|
21
|
+
InputConfig,
|
|
22
|
+
JobManifest,
|
|
23
|
+
JobStatus,
|
|
24
|
+
OutputConfig,
|
|
25
|
+
create_job_directory,
|
|
26
|
+
save_local_stub,
|
|
27
|
+
save_manifest,
|
|
28
|
+
save_manifest_s3,
|
|
29
|
+
)
|
|
30
|
+
from ..s3_transport import s3_job_prefix, upload_directory
|
|
31
|
+
|
|
32
|
+
DEFAULT_QUEUE = "l40s-1x-spot"
|
|
33
|
+
DEFAULT_JOB_DEFINITION = "dayhoff-esmfold2"
|
|
34
|
+
DEFAULT_IMAGE_URI = "074735440724.dkr.ecr.us-east-1.amazonaws.com/dayhoff:esmfold2-latest"
|
|
35
|
+
MAX_WORKERS = 50
|
|
36
|
+
FILES_PER_WORKER = 1
|
|
37
|
+
RETRY_ATTEMPTS = 5
|
|
38
|
+
TIMEOUT_SECONDS = 2 * 3600
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def _generate_job_id() -> str:
|
|
42
|
+
"""Job ID as ``{username}-esmfold2-{YYYYMMDD}-{random8}``.
|
|
43
|
+
|
|
44
|
+
Built here (not via ``generate_job_id``) so the full ``esmfold2`` pipeline
|
|
45
|
+
name survives — the generic helper truncates to 10 chars.
|
|
46
|
+
"""
|
|
47
|
+
username = get_aws_username()
|
|
48
|
+
date_str = datetime.now().strftime("%Y%m%d")
|
|
49
|
+
return f"{username}-esmfold2-{date_str}-{secrets.token_hex(4)}"
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def _is_primordial_path(path: Path) -> bool:
|
|
53
|
+
"""Check if a path is on Primordial (EFS)."""
|
|
54
|
+
return str(path).startswith("/primordial/")
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
@click.command()
|
|
58
|
+
@click.argument("input_dir", type=click.Path(exists=True))
|
|
59
|
+
@click.option(
|
|
60
|
+
"--workers",
|
|
61
|
+
default=None,
|
|
62
|
+
type=int,
|
|
63
|
+
help="Number of parallel workers [default: ~1 per file]",
|
|
64
|
+
)
|
|
65
|
+
@click.option("--queue", default=DEFAULT_QUEUE, help=f"Batch queue [default: {DEFAULT_QUEUE}]")
|
|
66
|
+
@click.option("--after", "after", multiple=True, help="Job ID(s) to wait for before starting")
|
|
67
|
+
@click.option("--dry-run", is_flag=True, help="Show plan without submitting")
|
|
68
|
+
@click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
|
|
69
|
+
def esmfold2(input_dir, workers, queue, after, dry_run, base_path):
|
|
70
|
+
"""Predict protein structures with ESMFold2.
|
|
71
|
+
|
|
72
|
+
Processes a directory of per-input YAML/JSON configs, each declaring the
|
|
73
|
+
entities to fold, and runs them in parallel using AWS Batch array jobs.
|
|
74
|
+
|
|
75
|
+
\b
|
|
76
|
+
Examples:
|
|
77
|
+
# Submit to AWS Batch (auto-sized workers)
|
|
78
|
+
dh batch esmfold2 input/
|
|
79
|
+
|
|
80
|
+
# Force a specific worker count
|
|
81
|
+
dh batch esmfold2 input/ --workers 20
|
|
82
|
+
|
|
83
|
+
# Preview what would run
|
|
84
|
+
dh batch esmfold2 input/ --dry-run
|
|
85
|
+
"""
|
|
86
|
+
input_path = Path(input_dir).resolve()
|
|
87
|
+
|
|
88
|
+
click.echo(f"Scanning {input_path} for configs...")
|
|
89
|
+
try:
|
|
90
|
+
configs = discover_esmfold2_configs(input_path)
|
|
91
|
+
except ValueError as e:
|
|
92
|
+
raise click.ClickException(str(e))
|
|
93
|
+
|
|
94
|
+
for cfg in configs:
|
|
95
|
+
try:
|
|
96
|
+
declares_at_least_one_sequence(cfg)
|
|
97
|
+
except ValueError as e:
|
|
98
|
+
raise click.ClickException(f"Invalid config {cfg.name}: {e}")
|
|
99
|
+
|
|
100
|
+
num_files = len(configs)
|
|
101
|
+
click.echo(f"Found {num_files} config(s) to process")
|
|
102
|
+
|
|
103
|
+
if workers is None:
|
|
104
|
+
workers = max(1, min(math.ceil(num_files / FILES_PER_WORKER), MAX_WORKERS))
|
|
105
|
+
array_size = max(1, min(workers, num_files))
|
|
106
|
+
|
|
107
|
+
use_s3 = not _is_primordial_path(input_path)
|
|
108
|
+
|
|
109
|
+
job_id = _generate_job_id()
|
|
110
|
+
|
|
111
|
+
click.echo()
|
|
112
|
+
click.echo(f"Job ID: {job_id}")
|
|
113
|
+
click.echo(f"Input: {input_path}")
|
|
114
|
+
click.echo(f"Configs: {num_files}")
|
|
115
|
+
click.echo(f"Workers: {array_size}")
|
|
116
|
+
click.echo(f"Queue: {queue}")
|
|
117
|
+
click.echo(f"Job definition: {DEFAULT_JOB_DEFINITION}")
|
|
118
|
+
click.echo(f"Image: {DEFAULT_IMAGE_URI}")
|
|
119
|
+
click.echo(f"Storage: {'S3 (remote)' if use_s3 else 'Primordial'}")
|
|
120
|
+
|
|
121
|
+
if dry_run:
|
|
122
|
+
click.echo()
|
|
123
|
+
click.echo(click.style("Dry run - job not submitted", fg="yellow"))
|
|
124
|
+
return
|
|
125
|
+
|
|
126
|
+
if not click.confirm("\nSubmit job?", default=True):
|
|
127
|
+
click.echo("Cancelled.")
|
|
128
|
+
raise SystemExit(0)
|
|
129
|
+
click.echo()
|
|
130
|
+
|
|
131
|
+
client = BatchClient()
|
|
132
|
+
|
|
133
|
+
response = client.batch.describe_job_definitions(jobDefinitionName=DEFAULT_JOB_DEFINITION, status="ACTIVE")
|
|
134
|
+
if not response.get("jobDefinitions"):
|
|
135
|
+
raise click.ClickException(
|
|
136
|
+
"esmfold2 job definition not yet deployed; apply Terraform in environments/dev first"
|
|
137
|
+
)
|
|
138
|
+
|
|
139
|
+
s3_prefix = s3_job_prefix(job_id) if use_s3 else None
|
|
140
|
+
job_dir_path = Path(base_path) / job_id
|
|
141
|
+
|
|
142
|
+
if use_s3:
|
|
143
|
+
click.echo("Uploading configs to S3...")
|
|
144
|
+
upload_directory(input_path, f"{s3_prefix}input/", glob="*.yaml")
|
|
145
|
+
upload_directory(input_path, f"{s3_prefix}input/", glob="*.json")
|
|
146
|
+
click.echo(f"Uploaded {num_files} config(s) to S3")
|
|
147
|
+
else:
|
|
148
|
+
job_dir_path = create_job_directory(job_id, base_path)
|
|
149
|
+
click.echo(f"Created job directory: {job_dir_path}")
|
|
150
|
+
dest = job_dir_path / "input"
|
|
151
|
+
dest.mkdir(parents=True, exist_ok=True)
|
|
152
|
+
for cfg in configs:
|
|
153
|
+
shutil.copy2(cfg, dest / cfg.name)
|
|
154
|
+
click.echo(f"Copied {num_files} config(s)")
|
|
155
|
+
|
|
156
|
+
manifest = JobManifest(
|
|
157
|
+
job_id=job_id,
|
|
158
|
+
user=job_id.split("-")[0],
|
|
159
|
+
pipeline="esmfold2",
|
|
160
|
+
status=JobStatus.PENDING,
|
|
161
|
+
image_uri=DEFAULT_IMAGE_URI,
|
|
162
|
+
storage_mode="s3" if use_s3 else "primordial",
|
|
163
|
+
s3_prefix=s3_prefix,
|
|
164
|
+
input=InputConfig(
|
|
165
|
+
source=str(input_path),
|
|
166
|
+
num_sequences=num_files,
|
|
167
|
+
num_chunks=array_size,
|
|
168
|
+
),
|
|
169
|
+
batch=BatchConfig(
|
|
170
|
+
queue=queue,
|
|
171
|
+
job_definition=DEFAULT_JOB_DEFINITION,
|
|
172
|
+
array_size=array_size,
|
|
173
|
+
),
|
|
174
|
+
output=OutputConfig(
|
|
175
|
+
destination=None,
|
|
176
|
+
finalized=False,
|
|
177
|
+
),
|
|
178
|
+
depends_on=list(after) if after else None,
|
|
179
|
+
)
|
|
180
|
+
|
|
181
|
+
if use_s3:
|
|
182
|
+
save_manifest_s3(manifest)
|
|
183
|
+
save_local_stub(job_id, s3_prefix)
|
|
184
|
+
else:
|
|
185
|
+
save_manifest(manifest, base_path)
|
|
186
|
+
|
|
187
|
+
try:
|
|
188
|
+
resolved = [resolve_dependency(jid, base_path) for jid in after]
|
|
189
|
+
depends_on = [{"jobId": aws_id} for aws_id in resolved if aws_id is not None] or None
|
|
190
|
+
|
|
191
|
+
environment = {
|
|
192
|
+
"JOB_DIR": str(job_dir_path),
|
|
193
|
+
"JOB_ID": job_id,
|
|
194
|
+
"BATCH_ARRAY_SIZE": str(array_size),
|
|
195
|
+
}
|
|
196
|
+
if use_s3:
|
|
197
|
+
environment["STORAGE_MODE"] = "s3"
|
|
198
|
+
environment["S3_JOB_PREFIX"] = s3_prefix
|
|
199
|
+
|
|
200
|
+
batch_job_id = client.submit_job(
|
|
201
|
+
job_name=job_id,
|
|
202
|
+
job_definition=DEFAULT_JOB_DEFINITION,
|
|
203
|
+
job_queue=queue,
|
|
204
|
+
array_size=array_size,
|
|
205
|
+
environment=environment,
|
|
206
|
+
timeout_seconds=TIMEOUT_SECONDS,
|
|
207
|
+
retry_attempts=RETRY_ATTEMPTS,
|
|
208
|
+
depends_on=depends_on,
|
|
209
|
+
share_identifier=get_aws_username(),
|
|
210
|
+
)
|
|
211
|
+
|
|
212
|
+
manifest.status = JobStatus.SUBMITTED
|
|
213
|
+
manifest.batch.job_id = batch_job_id
|
|
214
|
+
if use_s3:
|
|
215
|
+
save_manifest_s3(manifest)
|
|
216
|
+
else:
|
|
217
|
+
save_manifest(manifest, base_path)
|
|
218
|
+
|
|
219
|
+
click.echo()
|
|
220
|
+
click.echo(click.style("Job submitted successfully!", fg="green"))
|
|
221
|
+
click.echo()
|
|
222
|
+
click.echo(f"AWS Batch Job ID: {batch_job_id}")
|
|
223
|
+
if depends_on:
|
|
224
|
+
click.echo(f"Waiting on: {', '.join(after)}")
|
|
225
|
+
click.echo()
|
|
226
|
+
click.echo("Next steps:")
|
|
227
|
+
click.echo(f" Check status: dh batch status {job_id}")
|
|
228
|
+
click.echo(f" View logs: dh batch logs {job_id}")
|
|
229
|
+
click.echo(f" Cancel: dh batch cancel {job_id}")
|
|
230
|
+
click.echo()
|
|
231
|
+
click.echo("After completion:")
|
|
232
|
+
click.echo(f" Finalize: dh batch finalize {job_id} --output ./results/")
|
|
233
|
+
|
|
234
|
+
except BatchError as e:
|
|
235
|
+
manifest.status = JobStatus.FAILED
|
|
236
|
+
manifest.error_message = str(e)
|
|
237
|
+
if use_s3:
|
|
238
|
+
save_manifest_s3(manifest)
|
|
239
|
+
else:
|
|
240
|
+
save_manifest(manifest, base_path)
|
|
241
|
+
click.echo(click.style(f"Failed to submit job: {e}", fg="red"), err=True)
|
|
242
|
+
raise SystemExit(1)
|
|
@@ -146,6 +146,8 @@ def finalize(job_id, output, force, keep_intermediates, full_output, skip_dedup,
|
|
|
146
146
|
_finalize_protmpnn(output_dir, output_path)
|
|
147
147
|
elif manifest.pipeline == "rfd3":
|
|
148
148
|
_finalize_rfd3(output_dir, output_path)
|
|
149
|
+
elif manifest.pipeline == "esmfold2":
|
|
150
|
+
_finalize_esmfold2(output_dir, output_path)
|
|
149
151
|
elif manifest.pipeline == "equilibrator":
|
|
150
152
|
_finalize_equilibrator(output_dir, output_path, job_id=job_id)
|
|
151
153
|
else:
|
|
@@ -585,6 +587,40 @@ def _finalize_rfd3(output_dir: Path, output_path: Path):
|
|
|
585
587
|
click.echo(click.style(f"Results: {output_path}", fg="green"))
|
|
586
588
|
|
|
587
589
|
|
|
590
|
+
def _finalize_esmfold2(output_dir: Path, output_path: Path):
|
|
591
|
+
"""Collate completed ESMFold2 per-config prediction dirs into the final output.
|
|
592
|
+
|
|
593
|
+
The worker promotes each config's structure + confidence files into
|
|
594
|
+
`output/<stem>/` and drops a `output/esmfold2_<stem>.done` marker only
|
|
595
|
+
when the prediction succeeded. Finalize copies each completed config's
|
|
596
|
+
directory into `output_path/<stem>/`, skipping configs without a `.done`
|
|
597
|
+
marker and the loose worker markers.
|
|
598
|
+
"""
|
|
599
|
+
completed = []
|
|
600
|
+
for marker in sorted(output_dir.glob("esmfold2_*.done")):
|
|
601
|
+
stem = marker.name[len("esmfold2_") : -len(".done")]
|
|
602
|
+
if (output_dir / stem).is_dir():
|
|
603
|
+
completed.append(stem)
|
|
604
|
+
|
|
605
|
+
if not completed:
|
|
606
|
+
click.echo("No completed ESMFold2 configs found in output directory.", err=True)
|
|
607
|
+
raise SystemExit(1)
|
|
608
|
+
|
|
609
|
+
output_path.mkdir(parents=True, exist_ok=True)
|
|
610
|
+
|
|
611
|
+
copied_count = 0
|
|
612
|
+
for stem in completed:
|
|
613
|
+
dest = output_path / stem
|
|
614
|
+
if dest.exists():
|
|
615
|
+
click.echo(f" Skipping {stem} (already exists)")
|
|
616
|
+
continue
|
|
617
|
+
shutil.copytree(output_dir / stem, dest)
|
|
618
|
+
copied_count += 1
|
|
619
|
+
|
|
620
|
+
click.echo(f"Collated {copied_count} ESMFold2 config(s)")
|
|
621
|
+
click.echo(click.style(f"Results: {output_path}", fg="green"))
|
|
622
|
+
|
|
623
|
+
|
|
588
624
|
def _finalize_equilibrator(output_dir: Path, output_path: Path, job_id: str | None = None) -> None:
|
|
589
625
|
"""Combine per-chunk JSONL worker outputs into one sorted JSONL file.
|
|
590
626
|
|
|
@@ -103,7 +103,7 @@ def classify_arn(arn: str) -> Principal:
|
|
|
103
103
|
return Principal("unknown", arn)
|
|
104
104
|
|
|
105
105
|
|
|
106
|
-
_BASE_MODELS = ("opus-4-7", "opus-4-6", "sonnet-4-6", "sonnet-4-5", "haiku-4-5")
|
|
106
|
+
_BASE_MODELS = ("opus-4-8", "opus-4-7", "opus-4-6", "sonnet-4-6", "sonnet-4-5", "haiku-4-5", "fable-5")
|
|
107
107
|
|
|
108
108
|
|
|
109
109
|
def resolve_base_model(model_id: str) -> str:
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
inputTokenCount_inclusive: false
|
|
2
|
+
|
|
3
|
+
models:
|
|
4
|
+
opus-4-8:
|
|
5
|
+
input_per_1m: 5.00
|
|
6
|
+
output_per_1m: 25.00
|
|
7
|
+
cache_write_per_1m: 6.25
|
|
8
|
+
cache_read_per_1m: 0.50
|
|
9
|
+
opus-4-7:
|
|
10
|
+
input_per_1m: 5.00
|
|
11
|
+
output_per_1m: 25.00
|
|
12
|
+
cache_write_per_1m: 6.25
|
|
13
|
+
cache_read_per_1m: 0.50
|
|
14
|
+
opus-4-6:
|
|
15
|
+
input_per_1m: 5.00
|
|
16
|
+
output_per_1m: 25.00
|
|
17
|
+
cache_write_per_1m: 6.25
|
|
18
|
+
cache_read_per_1m: 0.50
|
|
19
|
+
fable-5:
|
|
20
|
+
input_per_1m: 10.00
|
|
21
|
+
output_per_1m: 50.00
|
|
22
|
+
cache_write_per_1m: 12.50
|
|
23
|
+
cache_read_per_1m: 1.00
|
|
24
|
+
sonnet-4-6:
|
|
25
|
+
input_per_1m: 3.00
|
|
26
|
+
output_per_1m: 15.00
|
|
27
|
+
cache_write_per_1m: 3.75
|
|
28
|
+
cache_read_per_1m: 0.30
|
|
29
|
+
sonnet-4-5:
|
|
30
|
+
input_per_1m: 3.00
|
|
31
|
+
output_per_1m: 15.00
|
|
32
|
+
cache_write_per_1m: 3.75
|
|
33
|
+
cache_read_per_1m: 0.30
|
|
34
|
+
haiku-4-5:
|
|
35
|
+
input_per_1m: 0.80
|
|
36
|
+
output_per_1m: 4.00
|
|
37
|
+
cache_write_per_1m: 1.00
|
|
38
|
+
cache_read_per_1m: 0.08
|
|
@@ -28,7 +28,7 @@ def engine_cli():
|
|
|
28
28
|
"--type",
|
|
29
29
|
"engine_type",
|
|
30
30
|
required=True,
|
|
31
|
-
type=click.Choice(["cpu", "cpumax", "t4", "a10g", "a100", "4_t4", "8_t4", "4_a10g", "8_a10g"]),
|
|
31
|
+
type=click.Choice(["cpu", "cpumax", "t4", "a10g", "l40s", "a100", "4_t4", "8_t4", "4_a10g", "8_a10g"]),
|
|
32
32
|
)
|
|
33
33
|
@click.option("--size", "boot_disk_size", type=int, help="Boot disk size in GB")
|
|
34
34
|
@click.option(
|
|
@@ -25,11 +25,11 @@ def test_cost_fixture_A_cache_write(load_fixture, pricing):
|
|
|
25
25
|
from dh_cli.bedrock.cost_report import cost_of
|
|
26
26
|
|
|
27
27
|
record = load_fixture("A_cache_write.json")
|
|
28
|
-
# 100 *
|
|
29
|
-
# 2000 *
|
|
30
|
-
# 500 *
|
|
31
|
-
# total = 0.
|
|
32
|
-
expected = 0.
|
|
28
|
+
# 100 * 5.00 / 1e6 = 0.0005
|
|
29
|
+
# 2000 * 6.25 / 1e6 = 0.0125
|
|
30
|
+
# 500 * 25.00 / 1e6 = 0.0125
|
|
31
|
+
# total = 0.0255
|
|
32
|
+
expected = 0.0005 + 0.0125 + 0.0125
|
|
33
33
|
assert cost_of(record, pricing) == _approx(expected)
|
|
34
34
|
|
|
35
35
|
|
|
@@ -38,10 +38,10 @@ def test_cost_fixture_B_cache_read(load_fixture, pricing):
|
|
|
38
38
|
from dh_cli.bedrock.cost_report import cost_of
|
|
39
39
|
|
|
40
40
|
record = load_fixture("B_cache_read.json")
|
|
41
|
-
# 1 *
|
|
42
|
-
# 2000 *
|
|
43
|
-
# 500 *
|
|
44
|
-
expected = 0.
|
|
41
|
+
# 1 * 5.00 / 1e6 = 0.000005
|
|
42
|
+
# 2000 * 0.50 / 1e6 = 0.001
|
|
43
|
+
# 500 * 25.00 / 1e6 = 0.0125
|
|
44
|
+
expected = 0.000005 + 0.001 + 0.0125
|
|
45
45
|
assert cost_of(record, pricing) == _approx(expected)
|
|
46
46
|
|
|
47
47
|
|
|
@@ -50,9 +50,9 @@ def test_cost_fixture_C_plain(load_fixture, pricing):
|
|
|
50
50
|
from dh_cli.bedrock.cost_report import cost_of
|
|
51
51
|
|
|
52
52
|
record = load_fixture("C_plain.json")
|
|
53
|
-
# 100 *
|
|
54
|
-
# 500 *
|
|
55
|
-
expected = 0.
|
|
53
|
+
# 100 * 5.00 / 1e6 = 0.0005
|
|
54
|
+
# 500 * 25.00 / 1e6 = 0.0125
|
|
55
|
+
expected = 0.0005 + 0.0125
|
|
56
56
|
assert cost_of(record, pricing) == _approx(expected)
|
|
57
57
|
|
|
58
58
|
|
|
@@ -103,11 +103,59 @@ def test_cost_exclusive_semantics_sanity(pricing):
|
|
|
103
103
|
"output": {"outputTokenCount": 0},
|
|
104
104
|
}
|
|
105
105
|
# Exclusive (pricing as-loaded, flag = false):
|
|
106
|
-
# 100 *
|
|
106
|
+
# 100 * 5 / 1e6 + 40 * 0.50 / 1e6
|
|
107
107
|
exclusive = cost_of(record, pricing)
|
|
108
108
|
inclusive_pricing = {**pricing, "inputTokenCount_inclusive": True}
|
|
109
|
-
# (100 - 40) *
|
|
109
|
+
# (100 - 40) * 5 / 1e6 + 40 * 0.50 / 1e6
|
|
110
110
|
inclusive = cost_of(record, inclusive_pricing)
|
|
111
111
|
assert exclusive > inclusive
|
|
112
|
-
assert math.isclose(inclusive, (60 *
|
|
113
|
-
assert math.isclose(exclusive, (100 *
|
|
112
|
+
assert math.isclose(inclusive, (60 * 5 + 40 * 0.50) / 1_000_000, rel_tol=1e-9)
|
|
113
|
+
assert math.isclose(exclusive, (100 * 5 + 40 * 0.50) / 1_000_000, rel_tol=1e-9)
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
def test_cost_opus_4_8_prices_nonzero(pricing):
|
|
117
|
+
"""Opus 4.8 (global inference profile) must price at the Bedrock
|
|
118
|
+
Opus tier: input 5 / output 25 / cache_write 6.25 / cache_read 0.50."""
|
|
119
|
+
from dh_cli.bedrock.cost_report import cost_of
|
|
120
|
+
|
|
121
|
+
record = {
|
|
122
|
+
"modelId": "global.anthropic.claude-opus-4-8",
|
|
123
|
+
"input": {
|
|
124
|
+
"inputTokenCount": 100,
|
|
125
|
+
"cacheReadInputTokenCount": 2000,
|
|
126
|
+
"cacheWriteInputTokenCount": 1000,
|
|
127
|
+
},
|
|
128
|
+
"output": {"outputTokenCount": 500},
|
|
129
|
+
}
|
|
130
|
+
# 100 * 5.00 / 1e6 = 0.0005
|
|
131
|
+
# 500 * 25.00 / 1e6 = 0.0125
|
|
132
|
+
# 2000 * 0.50 / 1e6 = 0.001
|
|
133
|
+
# 1000 * 6.25 / 1e6 = 0.00625
|
|
134
|
+
expected = 0.0005 + 0.0125 + 0.001 + 0.00625
|
|
135
|
+
result = cost_of(record, pricing)
|
|
136
|
+
assert result > 0
|
|
137
|
+
assert result == _approx(expected)
|
|
138
|
+
|
|
139
|
+
|
|
140
|
+
def test_cost_fable_5_prices_nonzero(pricing):
|
|
141
|
+
"""Fable 5 (global inference profile) must price at the Fable tier:
|
|
142
|
+
input 10 / output 50 / cache_write 12.50 / cache_read 1.00."""
|
|
143
|
+
from dh_cli.bedrock.cost_report import cost_of
|
|
144
|
+
|
|
145
|
+
record = {
|
|
146
|
+
"modelId": "global.anthropic.claude-fable-5",
|
|
147
|
+
"input": {
|
|
148
|
+
"inputTokenCount": 100,
|
|
149
|
+
"cacheReadInputTokenCount": 2000,
|
|
150
|
+
"cacheWriteInputTokenCount": 1000,
|
|
151
|
+
},
|
|
152
|
+
"output": {"outputTokenCount": 500},
|
|
153
|
+
}
|
|
154
|
+
# 100 * 10.00 / 1e6 = 0.001
|
|
155
|
+
# 500 * 50.00 / 1e6 = 0.025
|
|
156
|
+
# 2000 * 1.00 / 1e6 = 0.002
|
|
157
|
+
# 1000 * 12.50 / 1e6 = 0.0125
|
|
158
|
+
expected = 0.001 + 0.025 + 0.002 + 0.0125
|
|
159
|
+
result = cost_of(record, pricing)
|
|
160
|
+
assert result > 0
|
|
161
|
+
assert result == _approx(expected)
|
|
@@ -20,6 +20,9 @@ import pytest
|
|
|
20
20
|
("anthropic.claude-opus-4-7", "opus-4-7"),
|
|
21
21
|
("us.anthropic.claude-opus-4-7", "opus-4-7"),
|
|
22
22
|
("global.anthropic.claude-opus-4-7", "opus-4-7"),
|
|
23
|
+
("anthropic.claude-opus-4-8", "opus-4-8"),
|
|
24
|
+
("us.anthropic.claude-opus-4-8", "opus-4-8"),
|
|
25
|
+
("global.anthropic.claude-opus-4-8", "opus-4-8"),
|
|
23
26
|
("anthropic.claude-opus-4-6-v1", "opus-4-6"),
|
|
24
27
|
("us.anthropic.claude-opus-4-6-v1", "opus-4-6"),
|
|
25
28
|
("global.anthropic.claude-opus-4-6-v1", "opus-4-6"),
|
|
@@ -29,6 +32,9 @@ import pytest
|
|
|
29
32
|
("anthropic.claude-haiku-4-5-20251001-v1:0", "haiku-4-5"),
|
|
30
33
|
("us.anthropic.claude-haiku-4-5-20251001-v1:0", "haiku-4-5"),
|
|
31
34
|
("global.anthropic.claude-haiku-4-5-20251001-v1:0", "haiku-4-5"),
|
|
35
|
+
("anthropic.claude-fable-5", "fable-5"),
|
|
36
|
+
("us.anthropic.claude-fable-5", "fable-5"),
|
|
37
|
+
("global.anthropic.claude-fable-5", "fable-5"),
|
|
32
38
|
],
|
|
33
39
|
)
|
|
34
40
|
def test_resolve_base_model_known(model_id: str, expected_base: str):
|