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.
- {dh_cli-0.15.0 → dh_cli-0.16.0}/PKG-INFO +1 -1
- {dh_cli-0.15.0 → dh_cli-0.16.0}/pyproject.toml +1 -1
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/__init__.py +3 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/__init__.py +2 -0
- dh_cli-0.16.0/src/dh_cli/batch/commands/embed_esmc.py +244 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/finalize.py +2 -2
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/retry.py +13 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/bedrock/cost_report.py +1 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/bedrock/pricing.yaml +5 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_resolve_base_model.py +26 -0
- dh_cli-0.16.0/tests/test_embed_esmc_cli.py +195 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/.gitignore +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/LICENSE +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/README.md +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/__init__.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/_identity.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/aws_batch.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/boltz.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/cancel.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/clean.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/cost.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/disco.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/embed_t5.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/equilibrator.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/esmfold2.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/list_jobs.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/local.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/logs.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/orca.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/protmpnn.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/protmpnn_to_boltz.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/rfd3.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/status.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/submit.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/commands/wait_for.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/cost_report.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/fasta_utils.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/h5_utils.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/job_id.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/manifest.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/batch/s3_transport.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/bedrock/__init__.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/bedrock/commands.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/cloud_commands.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/codeartifact.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/__init__.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/api_client.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/auth.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/engine_commands.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/progress.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/retry.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/ssh_config.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/studio_commands.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/engines_studios/studio_resolve.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/github_commands.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/__init__.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/auth.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/deploy.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/local.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/mcp_server.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/test.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/tf.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/hz/users.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/main.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/sweep.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/src/dh_cli/warehouse.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/__init__.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/cur_fixtures.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_aws_batch_resources.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_batch_cost_command.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_batch_cost_filters.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_batch_cost_reader.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_batch_cost_reconcile.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_batch_cost_report.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_image_override.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_no_train_command.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_submit_cpu_only.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_submit_image_validation.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_submit_job_tags.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_submit_merge.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/batch/test_submit_mlflow_env.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/conftest.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/A_cache_write.json +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/B_cache_read.json +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/C_plain.json +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/D_cursor_user.json +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/E_service_role.json +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/F_legacy_shared.json +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/fixtures/G_unknown_model.json +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_build_report.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_classify_arn.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_cli_exit_codes.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_cost_calc.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_cost_command.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_cur_reconciliation.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_key_command.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_render_formats.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_rotate.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/bedrock/test_s3_walker.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/__init__.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/test_engine_retry_commands.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/test_retry.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/test_ssh_config.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/test_studio_resolution_commands.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/engines_studios/test_studio_resolve.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/__init__.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/conftest.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_engine_role_cannot_read_github_pat.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_identity.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_login.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_login_error_paths.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_login_security.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_logout.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_rotate.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/github/test_status.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/live_smoke.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/test_auth.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/test_init.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/test_mcp_server.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/test_suites.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/hz/test_users.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_aws_status.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_cloud_gcp.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_disco_command.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_disco_jobdef.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_equilibrator_command.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_esmfold2_cli.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_boltz_tar.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_disco.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_disco_dispatch.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_equilibrator.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_esmfold2.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_protmpnn.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_finalize_rfd3.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_rfd3_command.py +0 -0
- {dh_cli-0.15.0 → dh_cli-0.16.0}/tests/test_sweep.py +0 -0
|
@@ -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
|
|
@@ -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
|
|
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
|
|
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
|
|
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
|