dh-cli 0.16.1__tar.gz → 0.17.0__tar.gz

{dh_cli-0.16.1 → dh_cli-0.17.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: dh-cli
-Version: 0.16.1
+Version: 0.17.0
 Summary: Dayhoff Labs developer CLI
 Author-email: Dayhoff Labs <dev@dayhofflabs.com>
 License: # PolyForm Noncommercial License 1.0.0

{dh_cli-0.16.1 → dh_cli-0.17.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dh-cli"
-version = "0.16.1"
+version = "0.17.0"
 description = "Dayhoff Labs developer CLI"
 requires-python = ">=3.11"
 readme = "README.md"

{dh_cli-0.16.1 → dh_cli-0.17.0}/src/dh_cli/batch/commands/embed_esmc.py

@@ -27,13 +27,16 @@ DEFAULT_WORKERS = 50
 DEFAULT_SEQS_PER_CHUNK = 5000
 DEFAULT_JOB_DEFINITION = "dayhoff-embed-esmc"
 DEFAULT_IMAGE_URI = "074735440724.dkr.ecr.us-east-1.amazonaws.com/dayhoff:esmc-latest"
+# (vCPUs, MiB) overriding the job def (4 vCPU / 15000 MiB), which 300m/600m use as-is.
+# 6B's loader holds the whole fp32 checkpoint in host RAM: measured peak RSS 34.3 GiB (slice 5), +15%.
+MODEL_RESOURCES = {"6b": (8, 40960)}
 
 
 @click.command()
 @click.argument("input_fasta", type=click.Path(exists=True))
 @click.option(
     "--model",
-    type=click.Choice(["300m", "600m"]),
+    type=click.Choice(["300m", "600m", "6b"]),
     default="600m",
     show_default=True,
     help="ESMC model size",
@@ -114,6 +117,7 @@ def embed_esmc(
     actual_seqs_per_chunk = (num_sequences + num_chunks - 1) // num_chunks
 
     job_id = generate_job_id("esmc")
+    vcpus, memory_mb = MODEL_RESOURCES.get(model, (None, None))
 
     click.echo()
     click.echo(f"Job ID:           {job_id}")
@@ -126,6 +130,8 @@ def embed_esmc(
     click.echo(f"Seqs per chunk:   ~{actual_seqs_per_chunk:,}")
     click.echo(f"Queue:            {queue}")
     click.echo(f"Job definition:   {DEFAULT_JOB_DEFINITION}")
+    if vcpus is not None:
+        click.echo(f"Resources:        {vcpus} vCPU / {memory_mb} MiB (override)")
     click.echo(f"Storage:          {'S3 (remote)' if use_s3 else 'Primordial'}")
 
     if dry_run:
@@ -217,6 +223,8 @@ def embed_esmc(
             retry_attempts=3,
             depends_on=depends_on,
             share_identifier=get_aws_username(),
+            vcpus=vcpus,
+            memory_mb=memory_mb,
         )
     except BatchError as e:
         manifest.status = JobStatus.FAILED

{dh_cli-0.16.1 → dh_cli-0.17.0}/tests/test_embed_esmc_cli.py

@@ -77,10 +77,28 @@ def test_dry_run_prints_plan_no_submit(runner, fasta):
     assert "Dry run" in out
 
 
-def test_model_6b_rejected(runner, fasta):
+def test_model_6b_accepted(runner, fasta):
     result = runner.invoke(esmc_cmd.embed_esmc, [str(fasta), "--dry-run", "--model", "6b"])
-    assert result.exit_code != 0
-    assert "Invalid value" in result.output
+    assert result.exit_code == 0, result.output
+    assert "ESMC-6b" in result.output
+
+
+def test_6b_submit_overrides_resources(runner, fasta):
+    result, client, _, _ = _invoke_submit(runner, fasta, ["--model", "6b"])
+    assert result.exit_code == 0, result.output
+    kwargs = client.submit_job.call_args.kwargs
+    vcpus, memory_mb = esmc_cmd.MODEL_RESOURCES["6b"]
+    assert (kwargs["vcpus"], kwargs["memory_mb"]) == (vcpus, memory_mb)
+    # The loader holds the whole fp32 checkpoint (~25.4 GB) in host RAM.
+    assert memory_mb > 25.4e9 / 2**20
+
+
+def test_600m_submit_uses_job_definition_resources(runner, fasta):
+    result, client, _, _ = _invoke_submit(runner, fasta)
+    assert result.exit_code == 0, result.output
+    kwargs = client.submit_job.call_args.kwargs
+    assert kwargs.get("vcpus") is None
+    assert kwargs.get("memory_mb") is None
 
 
 def test_submit_default_env_s3_mode(runner, fasta):