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.
Files changed (105) hide show
  1. {dh_cli-0.10.0 → dh_cli-0.10.2}/PKG-INFO +1 -1
  2. {dh_cli-0.10.0 → dh_cli-0.10.2}/pyproject.toml +1 -1
  3. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/__init__.py +3 -0
  4. dh_cli-0.10.2/src/dh_cli/batch/commands/esmfold2.py +242 -0
  5. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/finalize.py +36 -0
  6. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/bedrock/cost_report.py +1 -1
  7. dh_cli-0.10.2/src/dh_cli/bedrock/pricing.yaml +38 -0
  8. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/engine_commands.py +1 -1
  9. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_cost_calc.py +64 -16
  10. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_resolve_base_model.py +6 -0
  11. dh_cli-0.10.2/tests/test_esmfold2_cli.py +268 -0
  12. dh_cli-0.10.2/tests/test_finalize_esmfold2.py +204 -0
  13. dh_cli-0.10.0/src/dh_cli/bedrock/pricing.yaml +0 -28
  14. {dh_cli-0.10.0 → dh_cli-0.10.2}/.gitignore +0 -0
  15. {dh_cli-0.10.0 → dh_cli-0.10.2}/LICENSE +0 -0
  16. {dh_cli-0.10.0 → dh_cli-0.10.2}/README.md +0 -0
  17. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/__init__.py +0 -0
  18. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/_identity.py +0 -0
  19. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/aws_batch.py +0 -0
  20. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/__init__.py +0 -0
  21. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/boltz.py +0 -0
  22. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/cancel.py +0 -0
  23. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/clean.py +0 -0
  24. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/embed_t5.py +0 -0
  25. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/equilibrator.py +0 -0
  26. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/list_jobs.py +0 -0
  27. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/local.py +0 -0
  28. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/logs.py +0 -0
  29. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/orca.py +0 -0
  30. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/protmpnn.py +0 -0
  31. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/protmpnn_to_boltz.py +0 -0
  32. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/retry.py +0 -0
  33. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/rfd3.py +0 -0
  34. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/status.py +0 -0
  35. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/submit.py +0 -0
  36. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/train.py +0 -0
  37. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/commands/wait_for.py +0 -0
  38. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/fasta_utils.py +0 -0
  39. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/h5_utils.py +0 -0
  40. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/job_id.py +0 -0
  41. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/manifest.py +0 -0
  42. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/batch/s3_transport.py +0 -0
  43. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/bedrock/__init__.py +0 -0
  44. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/bedrock/commands.py +0 -0
  45. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/cloud_commands.py +0 -0
  46. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/codeartifact.py +0 -0
  47. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/__init__.py +0 -0
  48. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/api_client.py +0 -0
  49. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/auth.py +0 -0
  50. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/progress.py +0 -0
  51. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/ssh_config.py +0 -0
  52. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/engines_studios/studio_commands.py +0 -0
  53. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/github_commands.py +0 -0
  54. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/__init__.py +0 -0
  55. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/deploy.py +0 -0
  56. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/local.py +0 -0
  57. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/test.py +0 -0
  58. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/tf.py +0 -0
  59. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/hz/users.py +0 -0
  60. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/main.py +0 -0
  61. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/utility_commands.py +0 -0
  62. {dh_cli-0.10.0 → dh_cli-0.10.2}/src/dh_cli/warehouse.py +0 -0
  63. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/__init__.py +0 -0
  64. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/test_aws_batch_resources.py +0 -0
  65. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/test_image_override.py +0 -0
  66. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/test_submit_cpu_only.py +0 -0
  67. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/test_submit_image_validation.py +0 -0
  68. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/batch/test_submit_merge.py +0 -0
  69. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/conftest.py +0 -0
  70. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/A_cache_write.json +0 -0
  71. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/B_cache_read.json +0 -0
  72. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/C_plain.json +0 -0
  73. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/D_cursor_user.json +0 -0
  74. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/E_service_role.json +0 -0
  75. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/F_legacy_shared.json +0 -0
  76. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/fixtures/G_unknown_model.json +0 -0
  77. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_build_report.py +0 -0
  78. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_classify_arn.py +0 -0
  79. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_cli_exit_codes.py +0 -0
  80. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_cost_command.py +0 -0
  81. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_cur_reconciliation.py +0 -0
  82. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_key_command.py +0 -0
  83. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_render_formats.py +0 -0
  84. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_rotate.py +0 -0
  85. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/bedrock/test_s3_walker.py +0 -0
  86. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/__init__.py +0 -0
  87. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/conftest.py +0 -0
  88. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_engine_role_cannot_read_github_pat.py +0 -0
  89. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_identity.py +0 -0
  90. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_login.py +0 -0
  91. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_login_error_paths.py +0 -0
  92. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_login_security.py +0 -0
  93. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_logout.py +0 -0
  94. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_rotate.py +0 -0
  95. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/github/test_status.py +0 -0
  96. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/hz/test_init.py +0 -0
  97. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/hz/test_suites.py +0 -0
  98. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/hz/test_users.py +0 -0
  99. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_cloud_gcp.py +0 -0
  100. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_equilibrator_command.py +0 -0
  101. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_finalize_boltz_tar.py +0 -0
  102. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_finalize_equilibrator.py +0 -0
  103. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_finalize_protmpnn.py +0 -0
  104. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_finalize_rfd3.py +0 -0
  105. {dh_cli-0.10.0 → dh_cli-0.10.2}/tests/test_rfd3_command.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dh-cli
3
- Version: 0.10.0
3
+ Version: 0.10.2
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.10.0"
7
+ version = "0.10.2"
8
8
  description = "Dayhoff Labs developer CLI"
9
9
  requires-python = ">=3.11"
10
10
  readme = "README.md"
@@ -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 * 15.00 / 1e6 = 0.0015
29
- # 2000 * 18.75 / 1e6 = 0.0375
30
- # 500 * 75.00 / 1e6 = 0.0375
31
- # total = 0.0765
32
- expected = 0.0015 + 0.0375 + 0.0375
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 * 15.00 / 1e6 = 0.000015
42
- # 2000 * 1.50 / 1e6 = 0.003
43
- # 500 * 75.00 / 1e6 = 0.0375
44
- expected = 0.000015 + 0.003 + 0.0375
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 * 15.00 / 1e6 = 0.0015
54
- # 500 * 75.00 / 1e6 = 0.0375
55
- expected = 0.0015 + 0.0375
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 * 15 / 1e6 + 40 * 1.50 / 1e6
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) * 15 / 1e6 + 40 * 1.50 / 1e6
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 * 15 + 40 * 1.50) / 1_000_000, rel_tol=1e-9)
113
- assert math.isclose(exclusive, (100 * 15 + 40 * 1.50) / 1_000_000, rel_tol=1e-9)
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):