dh-cli 0.15.0__tar.gz → 0.16.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (136) hide show
  1. {dh_cli-0.15.0 → dh_cli-0.16.0}/PKG-INFO +1 -1
  2. {dh_cli-0.15.0 → dh_cli-0.16.0}/pyproject.toml +1 -1
  3. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/__init__.py +3 -0
  4. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/__init__.py +2 -0
  5. dh_cli-0.16.0/src/dh_cli/batch/commands/embed_esmc.py +244 -0
  6. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/finalize.py +2 -2
  7. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/retry.py +13 -0
  8. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/bedrock/cost_report.py +1 -0
  9. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/bedrock/pricing.yaml +5 -0
  10. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_resolve_base_model.py +26 -0
  11. dh_cli-0.16.0/tests/test_embed_esmc_cli.py +195 -0
  12. {dh_cli-0.15.0 → dh_cli-0.16.0}/.gitignore +0 -0
  13. {dh_cli-0.15.0 → dh_cli-0.16.0}/LICENSE +0 -0
  14. {dh_cli-0.15.0 → dh_cli-0.16.0}/README.md +0 -0
  15. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/__init__.py +0 -0
  16. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/_identity.py +0 -0
  17. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/aws_batch.py +0 -0
  18. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/boltz.py +0 -0
  19. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/cancel.py +0 -0
  20. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/clean.py +0 -0
  21. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/cost.py +0 -0
  22. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/disco.py +0 -0
  23. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/embed_t5.py +0 -0
  24. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/equilibrator.py +0 -0
  25. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/esmfold2.py +0 -0
  26. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/list_jobs.py +0 -0
  27. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/local.py +0 -0
  28. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/logs.py +0 -0
  29. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/orca.py +0 -0
  30. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/protmpnn.py +0 -0
  31. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/protmpnn_to_boltz.py +0 -0
  32. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/rfd3.py +0 -0
  33. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/status.py +0 -0
  34. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/submit.py +0 -0
  35. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/wait_for.py +0 -0
  36. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/cost_report.py +0 -0
  37. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/fasta_utils.py +0 -0
  38. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/h5_utils.py +0 -0
  39. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/job_id.py +0 -0
  40. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/manifest.py +0 -0
  41. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/s3_transport.py +0 -0
  42. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/bedrock/__init__.py +0 -0
  43. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/bedrock/commands.py +0 -0
  44. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/cloud_commands.py +0 -0
  45. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/codeartifact.py +0 -0
  46. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/__init__.py +0 -0
  47. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/api_client.py +0 -0
  48. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/auth.py +0 -0
  49. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/engine_commands.py +0 -0
  50. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/progress.py +0 -0
  51. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/retry.py +0 -0
  52. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/ssh_config.py +0 -0
  53. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/studio_commands.py +0 -0
  54. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/studio_resolve.py +0 -0
  55. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/github_commands.py +0 -0
  56. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/__init__.py +0 -0
  57. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/auth.py +0 -0
  58. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/deploy.py +0 -0
  59. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/local.py +0 -0
  60. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/mcp_server.py +0 -0
  61. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/test.py +0 -0
  62. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/tf.py +0 -0
  63. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/users.py +0 -0
  64. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/main.py +0 -0
  65. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/sweep.py +0 -0
  66. {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/warehouse.py +0 -0
  67. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/__init__.py +0 -0
  68. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/cur_fixtures.py +0 -0
  69. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_aws_batch_resources.py +0 -0
  70. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_batch_cost_command.py +0 -0
  71. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_batch_cost_filters.py +0 -0
  72. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_batch_cost_reader.py +0 -0
  73. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_batch_cost_reconcile.py +0 -0
  74. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_batch_cost_report.py +0 -0
  75. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_image_override.py +0 -0
  76. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_no_train_command.py +0 -0
  77. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_submit_cpu_only.py +0 -0
  78. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_submit_image_validation.py +0 -0
  79. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_submit_job_tags.py +0 -0
  80. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_submit_merge.py +0 -0
  81. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_submit_mlflow_env.py +0 -0
  82. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/conftest.py +0 -0
  83. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/A_cache_write.json +0 -0
  84. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/B_cache_read.json +0 -0
  85. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/C_plain.json +0 -0
  86. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/D_cursor_user.json +0 -0
  87. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/E_service_role.json +0 -0
  88. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/F_legacy_shared.json +0 -0
  89. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/G_unknown_model.json +0 -0
  90. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_build_report.py +0 -0
  91. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_classify_arn.py +0 -0
  92. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_cli_exit_codes.py +0 -0
  93. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_cost_calc.py +0 -0
  94. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_cost_command.py +0 -0
  95. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_cur_reconciliation.py +0 -0
  96. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_key_command.py +0 -0
  97. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_render_formats.py +0 -0
  98. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_rotate.py +0 -0
  99. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_s3_walker.py +0 -0
  100. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/__init__.py +0 -0
  101. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/test_engine_retry_commands.py +0 -0
  102. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/test_retry.py +0 -0
  103. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/test_ssh_config.py +0 -0
  104. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/test_studio_resolution_commands.py +0 -0
  105. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/test_studio_resolve.py +0 -0
  106. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/__init__.py +0 -0
  107. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/conftest.py +0 -0
  108. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_engine_role_cannot_read_github_pat.py +0 -0
  109. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_identity.py +0 -0
  110. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_login.py +0 -0
  111. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_login_error_paths.py +0 -0
  112. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_login_security.py +0 -0
  113. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_logout.py +0 -0
  114. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_rotate.py +0 -0
  115. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_status.py +0 -0
  116. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/live_smoke.py +0 -0
  117. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/test_auth.py +0 -0
  118. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/test_init.py +0 -0
  119. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/test_mcp_server.py +0 -0
  120. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/test_suites.py +0 -0
  121. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/test_users.py +0 -0
  122. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_aws_status.py +0 -0
  123. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_cloud_gcp.py +0 -0
  124. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_disco_command.py +0 -0
  125. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_disco_jobdef.py +0 -0
  126. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_equilibrator_command.py +0 -0
  127. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_esmfold2_cli.py +0 -0
  128. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_boltz_tar.py +0 -0
  129. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_disco.py +0 -0
  130. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_disco_dispatch.py +0 -0
  131. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_equilibrator.py +0 -0
  132. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_esmfold2.py +0 -0
  133. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_protmpnn.py +0 -0
  134. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_rfd3.py +0 -0
  135. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_rfd3_command.py +0 -0
  136. {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_sweep.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: dh-cli
3
- Version: 0.15.0
3
+ Version: 0.16.0
4
4
  Summary: Dayhoff Labs developer CLI
5
5
  Author-email: Dayhoff Labs <dev@dayhofflabs.com>
6
6
  License: # PolyForm Noncommercial License 1.0.0
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "dh-cli"
7
- version = "0.15.0"
7
+ version = "0.16.0"
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.cost import cost
17
17
  from .commands.disco import disco
18
+ from .commands.embed_esmc import embed_esmc
18
19
  from .commands.embed_t5 import embed_t5
19
20
  from .commands.equilibrator import equilibrator
20
21
  from .commands.esmfold2 import esmfold2
@@ -55,6 +56,7 @@ def batch_cli():
55
56
  \b
56
57
  Embedding Pipelines:
57
58
  embed-t5 Generate T5 protein embeddings
59
+ embed-esmc Generate ESMC protein embeddings
58
60
 
59
61
  \b
60
62
  Structure Prediction:
@@ -106,6 +108,7 @@ batch_cli.add_command(cost)
106
108
 
107
109
  # Register pipeline commands
108
110
  batch_cli.add_command(embed_t5, name="embed-t5")
111
+ batch_cli.add_command(embed_esmc, name="embed-esmc")
109
112
  batch_cli.add_command(equilibrator, name="equilibrator")
110
113
  batch_cli.add_command(esmfold2)
111
114
  batch_cli.add_command(boltz)
@@ -3,6 +3,7 @@
3
3
  from .boltz import boltz
4
4
  from .cancel import cancel
5
5
  from .cost import cost
6
+ from .embed_esmc import embed_esmc
6
7
  from .embed_t5 import embed_t5
7
8
  from .list_jobs import list_jobs
8
9
  from .local import local
@@ -15,6 +16,7 @@ __all__ = [
15
16
  "boltz",
16
17
  "cancel",
17
18
  "cost",
19
+ "embed_esmc",
18
20
  "embed_t5",
19
21
  "list_jobs",
20
22
  "local",
@@ -0,0 +1,244 @@
1
+ """ESMC embedding pipeline command."""
2
+
3
+ from pathlib import Path
4
+
5
+ import click
6
+
7
+ from ..aws_batch import BatchClient, BatchError, resolve_dependency
8
+ from ..job_id import generate_job_id, get_aws_username
9
+ from ..manifest import (
10
+ BATCH_JOBS_BASE,
11
+ BatchConfig,
12
+ InputConfig,
13
+ JobManifest,
14
+ JobStatus,
15
+ OutputConfig,
16
+ create_job_directory,
17
+ save_local_stub,
18
+ save_manifest,
19
+ save_manifest_s3,
20
+ )
21
+ from ..s3_transport import s3_job_prefix, upload_directory
22
+ from .embed_t5 import _count_sequences, _is_primordial_path, _split_fasta
23
+
24
+ # A10G: ~79 seq/s vs T4 ~6 seq/s at 600M bf16 (T4 emulates bf16); ~$2.9 vs ~$13.7 per 1M seqs on spot.
25
+ DEFAULT_QUEUE = "a10g-1x-spot"
26
+ DEFAULT_WORKERS = 50
27
+ DEFAULT_SEQS_PER_CHUNK = 5000
28
+ DEFAULT_JOB_DEFINITION = "dayhoff-embed-esmc"
29
+ DEFAULT_IMAGE_URI = "074735440724.dkr.ecr.us-east-1.amazonaws.com/dayhoff:esmc-latest"
30
+
31
+
32
+ @click.command()
33
+ @click.argument("input_fasta", type=click.Path(exists=True))
34
+ @click.option(
35
+ "--model",
36
+ type=click.Choice(["300m", "600m"]),
37
+ default="600m",
38
+ show_default=True,
39
+ help="ESMC model size",
40
+ )
41
+ @click.option(
42
+ "--layer",
43
+ type=int,
44
+ default=-1,
45
+ show_default=True,
46
+ help="Hidden-state layer to mean-pool (0 = embeddings, -1 = final output)",
47
+ )
48
+ @click.option(
49
+ "--keep-alignment-tokens",
50
+ default="-",
51
+ show_default=True,
52
+ help='Which of "-.|" to keep as ESMC tokens; the rest are stripped ("" strips all)',
53
+ )
54
+ @click.option(
55
+ "--workers",
56
+ default=DEFAULT_WORKERS,
57
+ type=int,
58
+ help=f"Number of parallel workers [default: {DEFAULT_WORKERS}]",
59
+ )
60
+ @click.option("--queue", default=DEFAULT_QUEUE, help=f"Batch queue [default: {DEFAULT_QUEUE}]")
61
+ @click.option(
62
+ "--seqs-per-chunk",
63
+ default=DEFAULT_SEQS_PER_CHUNK,
64
+ type=int,
65
+ help=f"Sequences per chunk [default: {DEFAULT_SEQS_PER_CHUNK}]",
66
+ )
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
+ @click.option("--after", "after", multiple=True, help="Job ID(s) to wait for before starting")
70
+ def embed_esmc(
71
+ input_fasta,
72
+ model,
73
+ layer,
74
+ keep_alignment_tokens,
75
+ workers,
76
+ queue,
77
+ seqs_per_chunk,
78
+ dry_run,
79
+ base_path,
80
+ after,
81
+ ):
82
+ """Generate ESMC mean-pooled protein embeddings for a FASTA file.
83
+
84
+ Splits the input FASTA into chunks and embeds them in parallel as an
85
+ AWS Batch array job. Each worker writes an `ids`/`vectors` H5 file.
86
+
87
+ \b
88
+ Examples:
89
+ dh batch embed-esmc /primordial/proteins.fasta
90
+ dh batch embed-esmc proteins.fasta --model 300m --layer 20
91
+ dh batch embed-esmc msa.fasta --keep-alignment-tokens "-."
92
+
93
+ \b
94
+ After job completes:
95
+ dh batch status <job-id>
96
+ dh batch finalize <job-id> --output out.h5
97
+
98
+ `dh batch retry` does not carry the ESMC settings and refuses these jobs;
99
+ resubmit the failed chunks' sequences with `dh batch embed-esmc` instead.
100
+ """
101
+ input_path = Path(input_fasta).resolve()
102
+
103
+ click.echo(f"Counting sequences in {input_path}...")
104
+ num_sequences = _count_sequences(input_path)
105
+ click.echo(f"Found {num_sequences:,} sequences")
106
+
107
+ if num_sequences == 0:
108
+ click.echo(click.style("Error: No sequences found in input file", fg="red"), err=True)
109
+ raise SystemExit(1)
110
+
111
+ use_s3 = not _is_primordial_path(input_path)
112
+
113
+ num_chunks = min((num_sequences + seqs_per_chunk - 1) // seqs_per_chunk, workers)
114
+ actual_seqs_per_chunk = (num_sequences + num_chunks - 1) // num_chunks
115
+
116
+ job_id = generate_job_id("esmc")
117
+
118
+ click.echo()
119
+ click.echo(f"Job ID: {job_id}")
120
+ click.echo(f"Input: {input_path}")
121
+ click.echo(f"Sequences: {num_sequences:,}")
122
+ click.echo(f"Model: ESMC-{model}")
123
+ click.echo(f"Layer: {layer}")
124
+ click.echo(f"Keep tokens: {keep_alignment_tokens!r}")
125
+ click.echo(f"Chunks: {num_chunks}")
126
+ click.echo(f"Seqs per chunk: ~{actual_seqs_per_chunk:,}")
127
+ click.echo(f"Queue: {queue}")
128
+ click.echo(f"Job definition: {DEFAULT_JOB_DEFINITION}")
129
+ click.echo(f"Storage: {'S3 (remote)' if use_s3 else 'Primordial'}")
130
+
131
+ if dry_run:
132
+ click.echo()
133
+ click.echo(click.style("Dry run - job not submitted", fg="yellow"))
134
+ return
135
+
136
+ if not click.confirm("\nSubmit job?", default=True):
137
+ click.echo("Cancelled.")
138
+ raise SystemExit(0)
139
+ click.echo()
140
+
141
+ s3_prefix = s3_job_prefix(job_id) if use_s3 else None
142
+ job_dir_path = Path(base_path) / job_id
143
+
144
+ if use_s3:
145
+ import tempfile
146
+
147
+ with tempfile.TemporaryDirectory() as tmpdir:
148
+ tmp_input_dir = Path(tmpdir) / "input"
149
+ tmp_input_dir.mkdir()
150
+ click.echo("Splitting FASTA into chunks...")
151
+ actual_chunks = _split_fasta(input_path, tmp_input_dir, actual_seqs_per_chunk)
152
+ click.echo(f"Created {actual_chunks} chunks")
153
+ click.echo("Uploading chunks to S3...")
154
+ upload_directory(tmp_input_dir, f"{s3_prefix}input/", glob="*.fasta")
155
+ else:
156
+ job_dir_path = create_job_directory(job_id, base_path)
157
+ click.echo(f"Created job directory: {job_dir_path}")
158
+ click.echo("Splitting FASTA into chunks...")
159
+ actual_chunks = _split_fasta(input_path, job_dir_path / "input", actual_seqs_per_chunk)
160
+ click.echo(f"Created {actual_chunks} chunks")
161
+
162
+ manifest = JobManifest(
163
+ job_id=job_id,
164
+ user=job_id.split("-")[0],
165
+ pipeline="embed-esmc",
166
+ status=JobStatus.PENDING,
167
+ image_uri=DEFAULT_IMAGE_URI,
168
+ storage_mode="s3" if use_s3 else "primordial",
169
+ s3_prefix=s3_prefix,
170
+ input=InputConfig(
171
+ source=str(input_path),
172
+ num_sequences=num_sequences,
173
+ num_chunks=actual_chunks,
174
+ sequences_per_chunk=actual_seqs_per_chunk,
175
+ ),
176
+ batch=BatchConfig(
177
+ queue=queue,
178
+ job_definition=DEFAULT_JOB_DEFINITION,
179
+ array_size=actual_chunks,
180
+ ),
181
+ output=OutputConfig(destination=None, finalized=False),
182
+ depends_on=list(after) if after else None,
183
+ )
184
+
185
+ def _save():
186
+ if use_s3:
187
+ save_manifest_s3(manifest)
188
+ else:
189
+ save_manifest(manifest, base_path)
190
+
191
+ _save()
192
+ if use_s3:
193
+ save_local_stub(job_id, s3_prefix)
194
+
195
+ try:
196
+ resolved = [resolve_dependency(jid, base_path) for jid in after]
197
+ depends_on = [{"jobId": aws_id} for aws_id in resolved if aws_id is not None] or None
198
+
199
+ environment = {
200
+ "JOB_DIR": str(job_dir_path),
201
+ "JOB_ID": job_id,
202
+ "ESMC_MODEL": model,
203
+ "ESMC_LAYER": str(layer),
204
+ "ESMC_KEEP_ALIGNMENT_TOKENS": keep_alignment_tokens,
205
+ }
206
+ if use_s3:
207
+ environment["STORAGE_MODE"] = "s3"
208
+ environment["S3_JOB_PREFIX"] = s3_prefix
209
+
210
+ batch_job_id = BatchClient().submit_job(
211
+ job_name=job_id,
212
+ job_definition=DEFAULT_JOB_DEFINITION,
213
+ job_queue=queue,
214
+ array_size=actual_chunks,
215
+ environment=environment,
216
+ timeout_seconds=6 * 3600,
217
+ retry_attempts=3,
218
+ depends_on=depends_on,
219
+ share_identifier=get_aws_username(),
220
+ )
221
+ except BatchError as e:
222
+ manifest.status = JobStatus.FAILED
223
+ manifest.error_message = str(e)
224
+ _save()
225
+ click.echo(click.style(f"✗ Failed to submit job: {e}", fg="red"), err=True)
226
+ raise SystemExit(1)
227
+
228
+ manifest.status = JobStatus.SUBMITTED
229
+ manifest.batch.job_id = batch_job_id
230
+ _save()
231
+
232
+ click.echo()
233
+ click.echo(click.style("✓ Job submitted successfully!", fg="green"))
234
+ click.echo()
235
+ click.echo(f"AWS Batch Job ID: {batch_job_id}")
236
+ click.echo()
237
+ click.echo("Next steps:")
238
+ click.echo(f" Check status: dh batch status {job_id}")
239
+ click.echo(f" View logs: dh batch logs {job_id}")
240
+ click.echo(f" Cancel: dh batch cancel {job_id}")
241
+ click.echo()
242
+ click.echo("After completion:")
243
+ finalize_output = "./embeddings.h5" if use_s3 else "/primordial/embeddings.h5"
244
+ click.echo(f" Finalize: dh batch finalize {job_id} --output {finalize_output}")
@@ -85,7 +85,7 @@ def finalize(job_id, output, force, keep_intermediates, full_output, skip_dedup,
85
85
  output_path = Path(output).resolve()
86
86
 
87
87
  # For embed jobs: if output is an existing directory, append a filename
88
- if manifest.pipeline in ("embed-t5", "embed") and output_path.is_dir():
88
+ if manifest.pipeline in ("embed-t5", "embed", "embed-esmc") and output_path.is_dir():
89
89
  output_path = output_path / f"{job_id}.h5"
90
90
  click.echo(f"Output is a directory, writing to: {output_path}")
91
91
 
@@ -127,7 +127,7 @@ def finalize(job_id, output, force, keep_intermediates, full_output, skip_dedup,
127
127
 
128
128
  # Finalize based on pipeline type
129
129
  click.echo()
130
- if manifest.pipeline in ("embed-t5", "embed"):
130
+ if manifest.pipeline in ("embed-t5", "embed", "embed-esmc"):
131
131
  _finalize_embeddings(output_dir, output_path, skip_dedup=skip_dedup)
132
132
  elif manifest.pipeline == "boltz":
133
133
  # Check if S3 essential mode was used — output is already filtered
@@ -58,6 +58,19 @@ def retry(job_id, indices, queue, reslice, dry_run, base_path):
58
58
  click.echo(f"Job not found: {job_id}", err=True)
59
59
  raise SystemExit(1)
60
60
 
61
+ if manifest.pipeline == "embed-esmc":
62
+ click.echo(
63
+ click.style(
64
+ "Refusing to retry an embed-esmc job: retry does not carry the ESMC config "
65
+ "(ESMC_MODEL/ESMC_LAYER/ESMC_KEEP_ALIGNMENT_TOKENS), so retried chunks would "
66
+ "silently use the worker defaults. Resubmit the job, or run `dh batch embed-esmc` "
67
+ "on the failed chunks' sequences with the original options.",
68
+ fg="red",
69
+ ),
70
+ err=True,
71
+ )
72
+ raise SystemExit(1)
73
+
61
74
  # Get failed indices
62
75
  if indices:
63
76
  # User specified indices
@@ -104,6 +104,7 @@ def classify_arn(arn: str) -> Principal:
104
104
 
105
105
 
106
106
  _BASE_MODELS = (
107
+ "opus-5-5",
107
108
  "opus-5",
108
109
  "opus-4-8",
109
110
  "opus-4-7",
@@ -1,6 +1,11 @@
1
1
  inputTokenCount_inclusive: false
2
2
 
3
3
  models:
4
+ opus-5-5:
5
+ input_per_1m: 4.00
6
+ output_per_1m: 20.00
7
+ cache_write_per_1m: 5.00
8
+ cache_read_per_1m: 0.20
4
9
  opus-5:
5
10
  input_per_1m: 5.00
6
11
  output_per_1m: 25.00
@@ -43,6 +43,15 @@ import pytest
43
43
  ("anthropic.claude-fable-5-1", "fable-5-1"),
44
44
  ("us.anthropic.claude-fable-5-1", "fable-5-1"),
45
45
  ("global.anthropic.claude-fable-5-1", "fable-5-1"),
46
+ # Same trap, second occurrence: `opus-5` is a substring of
47
+ # `claude-opus-5-5`. Here the newer model is *cheaper* ($4/$20 vs
48
+ # $5/$25), so the silent failure over-reports 5.5 spend by ~25%.
49
+ ("anthropic.claude-opus-5", "opus-5"),
50
+ ("us.anthropic.claude-opus-5", "opus-5"),
51
+ ("global.anthropic.claude-opus-5", "opus-5"),
52
+ ("anthropic.claude-opus-5-5", "opus-5-5"),
53
+ ("us.anthropic.claude-opus-5-5", "opus-5-5"),
54
+ ("global.anthropic.claude-opus-5-5", "opus-5-5"),
46
55
  ],
47
56
  )
48
57
  def test_resolve_base_model_known(model_id: str, expected_base: str):
@@ -63,6 +72,8 @@ def test_resolve_prefers_the_most_specific_base():
63
72
  cost_report._BASE_MODELS = tuple(reversed(original))
64
73
  assert cost_report.resolve_base_model("global.anthropic.claude-fable-5-1") == "fable-5-1"
65
74
  assert cost_report.resolve_base_model("global.anthropic.claude-fable-5") == "fable-5"
75
+ assert cost_report.resolve_base_model("global.anthropic.claude-opus-5-5") == "opus-5-5"
76
+ assert cost_report.resolve_base_model("global.anthropic.claude-opus-5") == "opus-5"
66
77
  finally:
67
78
  cost_report._BASE_MODELS = original
68
79
 
@@ -91,6 +102,21 @@ def test_fable_51_cache_reads_are_priced_below_fable_5():
91
102
  assert pricing["fable-5-1"][rate] == pricing["fable-5"][rate]
92
103
 
93
104
 
105
+ def test_opus_55_is_priced_below_opus_5():
106
+ """5.5 undercuts 5 on every dimension; a substring collapse hides it."""
107
+ from dh_cli.bedrock.cost_report import default_pricing_path, load_pricing
108
+
109
+ pricing = load_pricing(default_pricing_path())
110
+ assert pricing["opus-5-5"] == {
111
+ "input_per_1m": 4.00,
112
+ "output_per_1m": 20.00,
113
+ "cache_write_per_1m": 5.00,
114
+ "cache_read_per_1m": 0.20,
115
+ }
116
+ for rate in ("input_per_1m", "output_per_1m", "cache_write_per_1m", "cache_read_per_1m"):
117
+ assert pricing["opus-5-5"][rate] < pricing["opus-5"][rate]
118
+
119
+
94
120
  def test_resolve_base_model_unknown_raises():
95
121
  from dh_cli.bedrock.cost_report import UnknownModel, resolve_base_model
96
122
 
@@ -0,0 +1,195 @@
1
+ """Tests for the `dh batch embed-esmc` CLI command, plus its finalize/retry hooks.
2
+
3
+ All AWS + S3 interactions are mocked; no real Batch or S3 calls.
4
+ """
5
+
6
+ import importlib
7
+ import re
8
+ import tempfile
9
+ from pathlib import Path
10
+ from unittest.mock import MagicMock, patch
11
+
12
+ import h5py
13
+ import numpy as np
14
+ import pytest
15
+ from click.testing import CliRunner
16
+
17
+ from dh_cli.batch.manifest import BatchConfig, JobManifest, JobStatus, OutputConfig
18
+
19
+ # `commands/__init__.py` re-exports Click commands under the module names, so
20
+ # `from dh_cli.batch.commands import retry` would return the command, not the module.
21
+ esmc_cmd = importlib.import_module("dh_cli.batch.commands.embed_esmc")
22
+ fin = importlib.import_module("dh_cli.batch.commands.finalize")
23
+ retry_cmd = importlib.import_module("dh_cli.batch.commands.retry")
24
+
25
+ JOB_ID_RE = re.compile(r"^[a-z0-9]+-esmc-\d{8}-[0-9a-f]{8}$")
26
+
27
+
28
+ @pytest.fixture
29
+ def runner():
30
+ return CliRunner()
31
+
32
+
33
+ @pytest.fixture
34
+ def fasta():
35
+ with tempfile.TemporaryDirectory() as tmpdir:
36
+ path = Path(tmpdir) / "in.fasta"
37
+ path.write_text(">a\nMKTAYIAKQR\n>b\nMKT-AYIAK\n>c\nMSTNPKPQRK\n")
38
+ yield path
39
+
40
+
41
+ def _invoke_submit(runner, input_path, args=(), primordial=False):
42
+ client = MagicMock()
43
+ client.submit_job.return_value = "aws-job-123"
44
+ captured = {}
45
+ with (
46
+ patch.object(esmc_cmd, "BatchClient", return_value=client),
47
+ patch.object(esmc_cmd, "resolve_dependency", return_value="dep-aws-id"),
48
+ patch.object(esmc_cmd, "upload_directory") as up,
49
+ patch.object(esmc_cmd, "save_manifest_s3", side_effect=lambda m: captured.update(manifest=m)),
50
+ patch.object(esmc_cmd, "save_manifest", side_effect=lambda m, b: captured.update(manifest=m)),
51
+ patch.object(esmc_cmd, "save_local_stub"),
52
+ patch.object(esmc_cmd, "get_aws_username", return_value="tester"),
53
+ patch.object(esmc_cmd, "_is_primordial_path", return_value=primordial),
54
+ ):
55
+ result = runner.invoke(
56
+ esmc_cmd.embed_esmc,
57
+ [str(input_path), "--base-path", str(Path(input_path).parent / "_jobs"), *args],
58
+ input="y\n",
59
+ )
60
+ return result, client, up, captured.get("manifest")
61
+
62
+
63
+ def test_dry_run_prints_plan_no_submit(runner, fasta):
64
+ client = MagicMock()
65
+ with patch.object(esmc_cmd, "BatchClient", return_value=client):
66
+ result = runner.invoke(
67
+ esmc_cmd.embed_esmc,
68
+ [str(fasta), "--dry-run", "--model", "300m", "--layer", "12", "--keep-alignment-tokens", ".|"],
69
+ )
70
+ assert result.exit_code == 0, result.output
71
+ client.submit_job.assert_not_called()
72
+ out = result.output
73
+ assert "300m" in out
74
+ assert "12" in out
75
+ assert ".|" in out
76
+ assert "dayhoff-embed-esmc" in out
77
+ assert "Dry run" in out
78
+
79
+
80
+ def test_model_6b_rejected(runner, fasta):
81
+ result = runner.invoke(esmc_cmd.embed_esmc, [str(fasta), "--dry-run", "--model", "6b"])
82
+ assert result.exit_code != 0
83
+ assert "Invalid value" in result.output
84
+
85
+
86
+ def test_submit_default_env_s3_mode(runner, fasta):
87
+ result, client, up, _ = _invoke_submit(runner, fasta)
88
+ assert result.exit_code == 0, result.output
89
+ kwargs = client.submit_job.call_args.kwargs
90
+ assert kwargs["job_definition"] == "dayhoff-embed-esmc"
91
+ env = kwargs["environment"]
92
+ assert env["ESMC_MODEL"] == "600m"
93
+ assert env["ESMC_LAYER"] == "-1"
94
+ assert env["ESMC_KEEP_ALIGNMENT_TOKENS"] == "-"
95
+ assert env["STORAGE_MODE"] == "s3"
96
+ assert env["S3_JOB_PREFIX"]
97
+ assert env["JOB_ID"] == kwargs["job_name"]
98
+ up.assert_called()
99
+
100
+
101
+ def test_submit_overrides_env_primordial_mode(runner, fasta):
102
+ result, client, up, _ = _invoke_submit(
103
+ runner,
104
+ fasta,
105
+ ["--model", "300m", "--layer", "20", "--keep-alignment-tokens", ""],
106
+ primordial=True,
107
+ )
108
+ assert result.exit_code == 0, result.output
109
+ env = client.submit_job.call_args.kwargs["environment"]
110
+ assert env["ESMC_MODEL"] == "300m"
111
+ assert env["ESMC_LAYER"] == "20"
112
+ assert env["ESMC_KEEP_ALIGNMENT_TOKENS"] == ""
113
+ assert "STORAGE_MODE" not in env
114
+ assert "S3_JOB_PREFIX" not in env
115
+ up.assert_not_called()
116
+
117
+
118
+ def test_manifest_pipeline_and_job_id(runner, fasta):
119
+ result, client, _, manifest = _invoke_submit(runner, fasta)
120
+ assert result.exit_code == 0, result.output
121
+ assert manifest.pipeline == "embed-esmc"
122
+ assert manifest.batch.job_definition == "dayhoff-embed-esmc"
123
+ assert manifest.image_uri.endswith("dayhoff:esmc-latest")
124
+ assert JOB_ID_RE.match(client.submit_job.call_args.kwargs["job_name"])
125
+
126
+
127
+ def test_default_queue_is_a10g_spot(runner, fasta):
128
+ # Set from the slice-4 smoke runs: A10G ~79 seq/s vs T4 ~6 seq/s at 600M bf16.
129
+ result, client, _, _ = _invoke_submit(runner, fasta)
130
+ assert result.exit_code == 0, result.output
131
+ assert client.submit_job.call_args.kwargs["job_queue"] == "a10g-1x-spot"
132
+
133
+
134
+ def test_registered_in_batch_cli():
135
+ from dh_cli.batch import batch_cli
136
+
137
+ assert "embed-esmc" in batch_cli.commands
138
+ result = CliRunner().invoke(batch_cli, ["--help"])
139
+ assert "embed-esmc" in result.output
140
+
141
+
142
+ def _manifest(pipeline="embed-esmc"):
143
+ return JobManifest(
144
+ job_id="tester-esmc-20260924-deadbeef",
145
+ user="tester",
146
+ pipeline=pipeline,
147
+ status=JobStatus.SUCCEEDED,
148
+ batch=BatchConfig(queue="t4-1x-spot", job_definition="dayhoff-embed-esmc", array_size=2),
149
+ output=OutputConfig(destination=None, finalized=False),
150
+ )
151
+
152
+
153
+ def _write_h5(path: Path, ids: list[str], dim: int = 4):
154
+ with h5py.File(path, "w") as f:
155
+ f.create_dataset("ids", data=np.array(ids, dtype="S"))
156
+ f.create_dataset("vectors", data=np.random.rand(len(ids), dim).astype(np.float32))
157
+
158
+
159
+ def test_finalize_merges_embed_h5(runner, tmp_path):
160
+ job_dir = tmp_path / "tester-esmc-20260924-deadbeef"
161
+ out_dir = job_dir / "output"
162
+ out_dir.mkdir(parents=True)
163
+ _write_h5(out_dir / "embed_000.h5", ["a", "b"])
164
+ _write_h5(out_dir / "embed_001.h5", ["c"])
165
+ for i in range(2):
166
+ (out_dir / f"embed_{i:03d}.done").touch()
167
+ dest = tmp_path / "merged.h5"
168
+ with (
169
+ patch.object(fin, "load_manifest", return_value=_manifest()),
170
+ patch.object(fin, "save_manifest"),
171
+ patch.object(fin, "get_job_dir", return_value=job_dir),
172
+ patch.object(fin, "_check_completion", return_value=[]),
173
+ ):
174
+ result = runner.invoke(
175
+ fin.finalize,
176
+ [job_dir.name, "--output", str(dest), "--keep-intermediates", "--base-path", str(tmp_path)],
177
+ )
178
+ assert result.exit_code == 0, result.output
179
+ with h5py.File(dest) as f:
180
+ assert sorted(x.decode() for x in f["ids"][:]) == ["a", "b", "c"]
181
+ assert f["vectors"].shape == (3, 4)
182
+
183
+
184
+ def test_retry_refuses_embed_esmc(runner):
185
+ client = MagicMock()
186
+ with (
187
+ patch.object(retry_cmd, "load_manifest", return_value=_manifest()),
188
+ patch.object(retry_cmd, "_find_incomplete_chunks", return_value=[1]),
189
+ patch.object(retry_cmd, "BatchClient", return_value=client),
190
+ ):
191
+ result = runner.invoke(retry_cmd.retry, ["tester-esmc-20260924-deadbeef", "--indices", "1"])
192
+ assert result.exit_code != 0
193
+ assert "embed-esmc" in result.output
194
+ client.submit_job.assert_not_called()
195
+ client.submit_array_job_with_indices.assert_not_called()
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