dh-cli 0.9.0__tar.gz → 0.10.0__tar.gz

{dh_cli-0.9.0 → dh_cli-0.10.0}/PKG-INFO

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

{dh_cli-0.9.0 → dh_cli-0.10.0}/pyproject.toml

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

{dh_cli-0.9.0 → dh_cli-0.10.0}/src/dh_cli/batch/__init__.py

@@ -14,6 +14,7 @@ from .commands.boltz import boltz
 from .commands.cancel import cancel
 from .commands.clean import clean
 from .commands.embed_t5 import embed_t5
+from .commands.equilibrator import equilibrator
 from .commands.finalize import finalize
 from .commands.list_jobs import list_jobs
 from .commands.local import local
@@ -99,6 +100,7 @@ batch_cli.add_command(wait_for, name="wait-for")
 
 # Register pipeline commands
 batch_cli.add_command(embed_t5, name="embed-t5")
+batch_cli.add_command(equilibrator, name="equilibrator")
 batch_cli.add_command(boltz)
 batch_cli.add_command(protmpnn)
 batch_cli.add_command(orca)

{dh_cli-0.9.0 → dh_cli-0.10.0}/src/dh_cli/batch/commands/__init__.py

@@ -3,7 +3,6 @@
 from .boltz import boltz
 from .cancel import cancel
 from .embed_t5 import embed_t5
-from .finalize import finalize
 from .list_jobs import list_jobs
 from .local import local
 from .logs import logs
@@ -15,7 +14,6 @@ __all__ = [
     "boltz",
     "cancel",
     "embed_t5",
-    "finalize",
     "list_jobs",
     "local",
     "logs",

dh_cli-0.10.0/src/dh_cli/batch/commands/equilibrator.py

@@ -0,0 +1,319 @@
+"""eQuilibrator reaction ΔG°' pipeline command."""
+
+import math
+import secrets
+import subprocess
+from datetime import datetime
+from pathlib import Path
+
+import click
+
+from ..aws_batch import BatchClient, BatchError, resolve_dependency
+from ..job_id import get_aws_username
+from ..manifest import (
+    BATCH_JOBS_BASE,
+    BatchConfig,
+    InputConfig,
+    JobManifest,
+    JobStatus,
+    OutputConfig,
+    create_job_directory,
+    save_local_stub,
+    save_manifest,
+    save_manifest_s3,
+)
+from ..s3_transport import s3_job_prefix, upload_directory
+
+DEFAULT_QUEUE = "cpu-spot"
+DEFAULT_JOB_DEFINITION = "dayhoff-equilibrator"
+DEFAULT_IMAGE_URI = "074735440724.dkr.ecr.us-east-1.amazonaws.com/dayhoff:equilibrator-latest"
+REACTIONS_PER_WORKER = 1000
+MAX_WORKERS = 50
+
+ASSETS_REPO_PATH = "/primordial/equilibrator-assets"
+# ~4s/reaction * REACTIONS_PER_WORKER + sqlite copy + headroom.
+TIMEOUT_SECONDS = 3 * 3600
+
+
+def _generate_job_id() -> str:
+    """Job ID as ``{username}-equilibrator-{YYYYMMDD}-{random8}``.
+
+    Built here (not via ``generate_job_id``) so the full ``equilibrator``
+    pipeline name survives — the generic helper truncates to 10 chars.
+    """
+    username = get_aws_username()
+    date_str = datetime.now().strftime("%Y%m%d")
+    return f"{username}-equilibrator-{date_str}-{secrets.token_hex(4)}"
+
+
+def _is_primordial_path(path: Path) -> bool:
+    """Check if a path is on Primordial (EFS)."""
+    return str(path).startswith("/primordial/")
+
+
+def _resolve_image_digest() -> str:
+    """Best-effort resolution of the container image digest via ECR.
+
+    Patched out in tests. Never crashes the command.
+    """
+    try:
+        result = subprocess.run(
+            [
+                "aws",
+                "ecr",
+                "describe-images",
+                "--repository-name",
+                "dayhoff",
+                "--image-ids",
+                "imageTag=equilibrator-latest",
+                "--query",
+                "imageDetails[0].imageDigest",
+                "--output",
+                "text",
+            ],
+            capture_output=True,
+            text=True,
+            timeout=15,
+        )
+        digest = result.stdout.strip()
+        if result.returncode == 0 and digest and digest != "None":
+            return digest
+    except Exception:
+        pass
+    return DEFAULT_IMAGE_URI
+
+
+def _resolve_assets_commit() -> str:
+    """Best-effort current equilibrator-assets git commit short SHA.
+
+    Patched out in tests. Returns "unknown" on failure.
+    """
+    try:
+        result = subprocess.run(
+            ["git", "-C", ASSETS_REPO_PATH, "rev-parse", "--short", "HEAD"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        sha = result.stdout.strip()
+        if result.returncode == 0 and sha:
+            return sha
+    except Exception:
+        pass
+    return "unknown"
+
+
+def _read_reactions(input_path: Path) -> list[str]:
+    """Return non-blank reaction lines from the input file, in order."""
+    return [line for line in input_path.read_text().splitlines() if line.strip()]
+
+
+def _chunk_reactions(reactions: list[str], array_size: int) -> list[list[str]]:
+    """Split reactions into ``array_size`` contiguous, balanced chunks.
+
+    First ``n % k`` chunks get ceil(n/k) lines, the rest get floor(n/k).
+    """
+    n = len(reactions)
+    base = n // array_size
+    remainder = n % array_size
+    chunks: list[list[str]] = []
+    start = 0
+    for i in range(array_size):
+        size = base + (1 if i < remainder else 0)
+        chunks.append(reactions[start : start + size])
+        start += size
+    return chunks
+
+
+def _write_chunks(input_dir: Path, chunks: list[list[str]]) -> None:
+    """Write each chunk to ``input_dir/chunk_{i:03d}.txt``."""
+    input_dir.mkdir(parents=True, exist_ok=True)
+    for i, chunk in enumerate(chunks):
+        (input_dir / f"chunk_{i:03d}.txt").write_text("\n".join(chunk) + "\n")
+
+
+@click.command()
+@click.argument("input_file", type=click.Path(exists=True))
+@click.option(
+    "--workers",
+    default=None,
+    type=int,
+    help=f"Number of parallel workers [default: ~1 per {REACTIONS_PER_WORKER} reactions]",
+)
+@click.option("--queue", default=DEFAULT_QUEUE, help=f"Batch queue [default: {DEFAULT_QUEUE}]")
+@click.option("--dry-run", is_flag=True, help="Show plan without submitting")
+@click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
+@click.option("--after", "after", multiple=True, help="Job ID(s) to wait for before starting")
+def equilibrator(input_file, workers, queue, dry_run, base_path, after):
+    """Compute reaction ΔG°' with eQuilibrator.
+
+    Splits an input file of reaction SMILES (one per line) into chunks and
+    processes them in parallel using AWS Batch array jobs.
+
+    \b
+    Examples:
+      # Submit to AWS Batch (auto-sized workers)
+      dh batch equilibrator /primordial/reactions.txt
+
+      # Force a specific worker count
+      dh batch equilibrator reactions.txt --workers 20
+
+      # Preview what would run
+      dh batch equilibrator reactions.txt --dry-run
+    """
+    input_path = Path(input_file).resolve()
+
+    click.echo(f"Counting reactions in {input_path}...")
+    reactions = _read_reactions(input_path)
+    num_reactions = len(reactions)
+    click.echo(f"Found {num_reactions} reaction(s)")
+
+    if num_reactions == 0:
+        click.echo(click.style("Error: No reactions found in input file", fg="red"), err=True)
+        raise SystemExit(1)
+
+    if workers is None:
+        workers = min(math.ceil(num_reactions / REACTIONS_PER_WORKER), MAX_WORKERS)
+    array_size = min(workers, num_reactions)
+
+    use_s3 = not _is_primordial_path(input_path) and str(base_path) == BATCH_JOBS_BASE
+
+    job_id = _generate_job_id()
+
+    click.echo()
+    click.echo(f"Job ID:           {job_id}")
+    click.echo(f"Input:            {input_path}")
+    click.echo(f"Reactions:        {num_reactions}")
+    click.echo(f"Workers:          {array_size}")
+    click.echo(f"Queue:            {queue}")
+    click.echo(f"Job definition:   {DEFAULT_JOB_DEFINITION}")
+    click.echo(f"Storage:          {'S3 (remote)' if use_s3 else 'Primordial'}")
+
+    if dry_run:
+        click.echo()
+        click.echo(click.style("Dry run - job not submitted", fg="yellow"))
+        return
+
+    if not click.confirm("\nSubmit job?", default=True):
+        click.echo("Cancelled.")
+        raise SystemExit(0)
+    click.echo()
+
+    client = BatchClient()
+
+    response = client.batch.describe_job_definitions(jobDefinitionName=DEFAULT_JOB_DEFINITION, status="ACTIVE")
+    if not response.get("jobDefinitions"):
+        raise click.ClickException(
+            "equilibrator job definition not yet deployed; apply Terraform in environments/dev first"
+        )
+
+    chunks = _chunk_reactions(reactions, array_size)
+    image_digest = _resolve_image_digest()
+    source_commit = _resolve_assets_commit()
+
+    s3_prefix = s3_job_prefix(job_id) if use_s3 else None
+    job_dir_path = Path(base_path) / job_id
+
+    if use_s3:
+        import tempfile as _tempfile
+
+        with _tempfile.TemporaryDirectory() as tmpdir:
+            tmp_input_dir = Path(tmpdir) / "input"
+            _write_chunks(tmp_input_dir, chunks)
+            click.echo("Uploading chunks to S3...")
+            upload_directory(tmp_input_dir, f"{s3_prefix}input/", glob="*.txt")
+            click.echo(f"Uploaded {array_size} chunks to S3")
+    else:
+        job_dir_path = create_job_directory(job_id, base_path)
+        click.echo(f"Created job directory: {job_dir_path}")
+        _write_chunks(job_dir_path / "input", chunks)
+        click.echo(f"Created {array_size} chunks")
+
+    manifest = JobManifest(
+        job_id=job_id,
+        user=job_id.split("-")[0],
+        pipeline="equilibrator",
+        status=JobStatus.PENDING,
+        image_uri=DEFAULT_IMAGE_URI,
+        image_digest=image_digest,
+        source_commit=source_commit,
+        storage_mode="s3" if use_s3 else "primordial",
+        s3_prefix=s3_prefix,
+        input=InputConfig(
+            source=str(input_path),
+            num_sequences=num_reactions,
+            num_chunks=array_size,
+        ),
+        batch=BatchConfig(
+            queue=queue,
+            job_definition=DEFAULT_JOB_DEFINITION,
+            array_size=array_size,
+        ),
+        output=OutputConfig(
+            destination=None,
+            finalized=False,
+        ),
+        depends_on=list(after) if after else None,
+    )
+
+    if use_s3:
+        save_manifest_s3(manifest)
+        save_local_stub(job_id, s3_prefix)
+    else:
+        save_manifest(manifest, base_path)
+
+    try:
+        resolved = [resolve_dependency(jid, base_path) for jid in after]
+        depends_on = [{"jobId": aws_id} for aws_id in resolved if aws_id is not None] or None
+
+        environment = {
+            "JOB_DIR": str(job_dir_path),
+            "JOB_ID": job_id,
+        }
+        if use_s3:
+            environment["STORAGE_MODE"] = "s3"
+            environment["S3_JOB_PREFIX"] = s3_prefix
+
+        batch_job_id = client.submit_job(
+            job_name=job_id,
+            job_definition=DEFAULT_JOB_DEFINITION,
+            job_queue=queue,
+            array_size=array_size,
+            environment=environment,
+            timeout_seconds=TIMEOUT_SECONDS,
+            retry_attempts=5,
+            depends_on=depends_on,
+            share_identifier=get_aws_username(),
+        )
+
+        manifest.status = JobStatus.SUBMITTED
+        manifest.batch.job_id = batch_job_id
+        if use_s3:
+            save_manifest_s3(manifest)
+        else:
+            save_manifest(manifest, base_path)
+
+        click.echo()
+        click.echo(click.style("Job submitted successfully!", fg="green"))
+        click.echo()
+        click.echo(f"AWS Batch Job ID: {batch_job_id}")
+        if depends_on:
+            click.echo(f"Waiting on:       {', '.join(after)}")
+        click.echo()
+        click.echo("Next steps:")
+        click.echo(f"  Check status:  dh batch status {job_id}")
+        click.echo(f"  View logs:     dh batch logs {job_id}")
+        click.echo(f"  Cancel:        dh batch cancel {job_id}")
+        click.echo()
+        click.echo("After completion:")
+        click.echo(f"  Finalize:      dh batch finalize {job_id} --output ./results/")
+
+    except BatchError as e:
+        manifest.status = JobStatus.FAILED
+        manifest.error_message = str(e)
+        if use_s3:
+            save_manifest_s3(manifest)
+        else:
+            save_manifest(manifest, base_path)
+        click.echo(click.style(f"Failed to submit job: {e}", fg="red"), err=True)
+        raise SystemExit(1)

{dh_cli-0.9.0 → dh_cli-0.10.0}/src/dh_cli/batch/commands/finalize.py

@@ -1,5 +1,6 @@
 """Finalize command for combining results and cleaning up."""
 
+import json
 import shutil
 import tarfile
 import tempfile
@@ -145,6 +146,8 @@ def finalize(job_id, output, force, keep_intermediates, full_output, skip_dedup,
         _finalize_protmpnn(output_dir, output_path)
     elif manifest.pipeline == "rfd3":
         _finalize_rfd3(output_dir, output_path)
+    elif manifest.pipeline == "equilibrator":
+        _finalize_equilibrator(output_dir, output_path, job_id=job_id)
     else:
         _finalize_generic(output_dir, output_path)
 
@@ -580,3 +583,31 @@ def _finalize_rfd3(output_dir: Path, output_path: Path):
 
     click.echo(f"Collated {copied_count} RFD3 config(s)")
     click.echo(click.style(f"Results: {output_path}", fg="green"))
+
+
+def _finalize_equilibrator(output_dir: Path, output_path: Path, job_id: str | None = None) -> None:
+    """Combine per-chunk JSONL worker outputs into one sorted JSONL file.
+
+    Reads every ``chunk_*.jsonl`` under ``output_dir``, sorts all records by
+    ``input_index``, and writes them to the destination. A missing chunk
+    (gap in indices) is tolerated. If ``output_path`` is an existing
+    directory, the combined file is named ``{job_id}.jsonl`` inside it.
+    """
+    if output_path.is_dir():
+        output_path = output_path / f"{job_id}.jsonl"
+
+    records: list[dict] = []
+    for chunk_file in sorted(output_dir.glob("chunk_*.jsonl")):
+        for line in chunk_file.read_text().splitlines():
+            if line.strip():
+                records.append(json.loads(line))
+
+    records.sort(key=lambda r: r["input_index"])
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, "w") as f:
+        for rec in records:
+            f.write(json.dumps(rec) + "\n")
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+    click.echo(f"Combined {len(records)} records into {output_path}")

{dh_cli-0.9.0 → dh_cli-0.10.0}/src/dh_cli/batch/manifest.py

@@ -90,6 +90,8 @@ class JobManifest(BaseModel):
 
     # Additional metadata
     image_uri: str | None = Field(None, description="Container image URI")
+    image_digest: str | None = Field(None, description="Resolved container image digest")
+    source_commit: str | None = Field(None, description="Assets/source repo commit for reproducibility")
     command: str | None = Field(None, description="Command to run")
     error_message: str | None = Field(None, description="Error message if failed")
 

dh_cli-0.10.0/tests/test_equilibrator_command.py

@@ -0,0 +1,229 @@
+"""Tests for the `dh batch equilibrator` CLI command.
+
+All AWS interactions are mocked; no real Batch calls. Uses Click's CliRunner
+and unittest.mock, mirroring the existing dh-cli test conventions
+(see test_rfd3_command.py).
+"""
+
+import math
+import re
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+from click.testing import CliRunner
+
+from dh_cli.batch.commands import equilibrator as eq_cmd
+
+JOB_ID_RE = re.compile(r"^[a-z0-9]+-equilibrator-\d{8}-[0-9a-f]{8}$")
+
+
+@pytest.fixture
+def runner():
+    return CliRunner()
+
+
+@pytest.fixture
+def input_dir():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        yield Path(tmpdir)
+
+
+def _write_reactions(input_dir: Path, n: int, name: str = "reactions.txt") -> Path:
+    """Write ``n`` distinct reaction SMILES lines to an input file."""
+    lines = [f"CC{'C' * i}O.O>>CC{'C' * i}(=O)O" for i in range(n)]
+    p = input_dir / name
+    p.write_text("\n".join(lines) + "\n")
+    return p
+
+
+def _submit_mock():
+    """A BatchClient whose submit_job returns a fake AWS job id and whose
+    job-definition preflight reports an ACTIVE dayhoff-equilibrator."""
+    client = MagicMock()
+    client.submit_job.return_value = "aws-job-123"
+    client.batch.describe_job_definitions.return_value = {
+        "jobDefinitions": [{"jobDefinitionName": "dayhoff-equilibrator", "status": "ACTIVE"}]
+    }
+    return client
+
+
+def _invoke_submit(runner, input_file, args=(), client=None, base=None):
+    """Invoke the command on the submit path with AWS + metadata resolution mocked."""
+    client = client or _submit_mock()
+    if base is None:
+        base = Path(input_file).parent / "_jobs"
+    with (
+        patch.object(eq_cmd, "BatchClient", return_value=client),
+        patch.object(eq_cmd, "resolve_dependency", return_value="dep-aws-id"),
+        patch.object(eq_cmd, "_resolve_image_digest", return_value="sha256:deadbeef"),
+        patch.object(eq_cmd, "_resolve_assets_commit", return_value="abc1234"),
+    ):
+        result = runner.invoke(
+            eq_cmd.equilibrator,
+            [str(input_file), "--base-path", str(base), *args],
+            input="y\n",
+        )
+    return result, client, base
+
+
+class TestJobIdAndDefaults:
+    def test_job_id_format(self, runner, input_dir):
+        rxns = _write_reactions(input_dir, 3)
+        result, client, _ = _invoke_submit(runner, rxns)
+        assert result.exit_code == 0, result.output
+        job_name = client.submit_job.call_args.kwargs["job_name"]
+        assert JOB_ID_RE.match(job_name), job_name
+
+    def test_default_queue_is_cpu(self, runner, input_dir):
+        rxns = _write_reactions(input_dir, 3)
+        result, client, _ = _invoke_submit(runner, rxns)
+        assert result.exit_code == 0, result.output
+        assert client.submit_job.call_args.kwargs["job_queue"] == "cpu-spot"
+        assert client.submit_job.call_args.kwargs["job_definition"] == "dayhoff-equilibrator"
+        assert eq_cmd.DEFAULT_QUEUE == "cpu-spot"
+        assert eq_cmd.DEFAULT_JOB_DEFINITION == "dayhoff-equilibrator"
+
+
+class TestChunking:
+    def test_one_chunk_file_per_index(self, runner, input_dir):
+        n = 7
+        rxns = _write_reactions(input_dir, n)
+        base = input_dir / "_jobs"
+        result, client, base = _invoke_submit(runner, rxns, ["--workers", str(n)], base=base)
+        assert result.exit_code == 0, result.output
+
+        array_size = client.submit_job.call_args.kwargs["array_size"]
+        chunk_files = sorted(Path(base).rglob("input/chunk_*.txt"))
+        assert len(chunk_files) == array_size
+
+        # Index N's file is chunk_{N:03d}.txt.
+        names = {p.name for p in chunk_files}
+        for i in range(array_size):
+            assert f"chunk_{i:03d}.txt" in names
+
+        # Union of chunk lines equals the input reactions, in order.
+        original = [ln for ln in rxns.read_text().splitlines() if ln.strip()]
+        combined: list[str] = []
+        for p in chunk_files:
+            combined.extend(ln for ln in p.read_text().splitlines() if ln.strip())
+        assert combined == original
+
+    def test_chunking_default(self, runner, input_dir):
+        n = 5
+        rxns = _write_reactions(input_dir, n)
+        result, client, _ = _invoke_submit(runner, rxns)
+        assert result.exit_code == 0, result.output
+        expected = min(math.ceil(n / eq_cmd.REACTIONS_PER_WORKER), eq_cmd.MAX_WORKERS)
+        assert client.submit_job.call_args.kwargs["array_size"] == expected
+
+    def test_workers_override(self, runner, input_dir):
+        n = eq_cmd.REACTIONS_PER_WORKER * 10  # large input
+        rxns = _write_reactions(input_dir, n)
+        base = input_dir / "_jobs"
+        result, client, base = _invoke_submit(runner, rxns, ["--workers", "4"], base=base)
+        assert result.exit_code == 0, result.output
+        assert client.submit_job.call_args.kwargs["array_size"] == 4
+
+        chunk_files = sorted(Path(base).rglob("input/chunk_*.txt"))
+        assert len(chunk_files) == 4
+        counts = [len([ln for ln in p.read_text().splitlines() if ln.strip()]) for p in chunk_files]
+        # Balanced: chunk sizes differ by at most 1.
+        assert max(counts) - min(counts) <= 1
+        assert sum(counts) == n
+
+    def test_blank_lines_skipped(self, runner, input_dir):
+        p = input_dir / "reactions.txt"
+        p.write_text("CCO.O>>CC(=O)O\n\n   \n CCCO.O>>CCC(=O)O \n\n")
+        result, client, _ = _invoke_submit(runner, p)
+        assert result.exit_code == 0, result.output
+        assert client.submit_job.call_args.kwargs["array_size"] == min(
+            math.ceil(2 / eq_cmd.REACTIONS_PER_WORKER), eq_cmd.MAX_WORKERS
+        )
+        # Only 2 non-blank reactions counted.
+        assert "2" in result.output
+
+
+class TestDryRun:
+    def test_dry_run_no_submit(self, runner, input_dir):
+        rxns = _write_reactions(input_dir, 4)
+        base = input_dir / "_jobs"
+        client = _submit_mock()
+        with (
+            patch.object(eq_cmd, "BatchClient", return_value=client),
+            patch.object(eq_cmd, "_resolve_image_digest", return_value="sha256:deadbeef"),
+            patch.object(eq_cmd, "_resolve_assets_commit", return_value="abc1234"),
+        ):
+            result = runner.invoke(eq_cmd.equilibrator, [str(rxns), "--base-path", str(base), "--dry-run"])
+        assert result.exit_code == 0, result.output
+        client.submit_job.assert_not_called()
+        # No job directory created.
+        assert not list(Path(base).rglob("chunk_*.txt")) if base.exists() else True
+        # Plan is printed.
+        out = result.output
+        assert "4" in out  # reaction count
+        assert eq_cmd.DEFAULT_QUEUE in out
+        assert eq_cmd.DEFAULT_JOB_DEFINITION in out
+
+
+class TestEmptyInput:
+    def test_empty_input_errors(self, runner, input_dir):
+        p = input_dir / "empty.txt"
+        p.write_text("\n  \n\n")
+        result, client, _ = _invoke_submit(runner, p)
+        assert result.exit_code != 0
+        client.submit_job.assert_not_called()
+
+
+class TestManifest:
+    def test_manifest_written(self, runner, input_dir):
+        import json
+
+        rxns = _write_reactions(input_dir, 3)
+        base = input_dir / "_jobs"
+        result, client, base = _invoke_submit(runner, rxns, base=base)
+        assert result.exit_code == 0, result.output
+
+        manifests = list(Path(base).rglob("manifest.json"))
+        assert len(manifests) == 1
+        data = json.loads(manifests[0].read_text())
+        assert data["pipeline"] == "equilibrator"
+        assert data["batch"]["queue"] == "cpu-spot"
+        assert data["batch"]["job_definition"] == "dayhoff-equilibrator"
+        assert data["batch"]["array_size"] >= 1
+        assert data["input"]["num_chunks"] == data["batch"]["array_size"]
+        # D-14: reproducibility — image digest + equilibrator-assets commit.
+        blob = json.dumps(data)
+        assert "sha256:deadbeef" in blob
+        assert "abc1234" in blob
+
+
+class TestSubmissionEnv:
+    def test_env_vars_passed(self, runner, input_dir):
+        rxns = _write_reactions(input_dir, 3)
+        result, client, _ = _invoke_submit(runner, rxns)
+        assert result.exit_code == 0, result.output
+        env = client.submit_job.call_args.kwargs["environment"]
+        assert "JOB_DIR" in env
+        assert "JOB_ID" in env
+        # D-10: worker derives its file from the array index, not a file count.
+        assert "BATCH_NUM_FILES" not in env
+
+
+class TestSubmissionParams:
+    def test_retry_and_share_identifier(self, runner, input_dir):
+        rxns = _write_reactions(input_dir, 3)
+        with patch.object(eq_cmd, "get_aws_username", return_value="tester"):
+            result, client, _ = _invoke_submit(runner, rxns)
+        assert result.exit_code == 0, result.output
+        kwargs = client.submit_job.call_args.kwargs
+        assert kwargs["retry_attempts"] == 5
+        assert kwargs["timeout_seconds"] > 0
+
+
+class TestRegistration:
+    def test_equilibrator_registered_on_group(self):
+        from dh_cli.batch import batch_cli
+
+        assert "equilibrator" in batch_cli.commands

dh_cli-0.10.0/tests/test_finalize_equilibrator.py

@@ -0,0 +1,191 @@
+"""Tests for eQuilibrator JSON-lines chunk combining in finalize.
+
+Covers the pure combiner (``_finalize_equilibrator``) and the S3
+download-then-delete path via the finalize command dispatch.
+"""
+
+import json
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from dh_cli.batch.commands.finalize import _finalize_equilibrator
+
+
+@pytest.fixture
+def temp_dir():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        yield Path(tmpdir)
+
+
+def _record(input_index: int, dg: float | None = -10.5, error: str | None = None) -> dict:
+    """A worker output record with the full fixed schema."""
+    return {
+        "reaction": f"R{input_index}",
+        "equation": f"eq{input_index}",
+        "updated_equation": f"ueq{input_index}",
+        "dg_prime": dg,
+        "dg_prime_stddev": 1.0 if error is None else None,
+        "units": "kJ/mol",
+        "compound_map": {"a": 1, "b": 2},
+        "compound_refs": {"a": {"id": "MNXM1"}},
+        "error": error,
+        "input_index": input_index,
+    }
+
+
+def _write_chunk(output_dir: Path, index: int, records: list[dict], done: bool = True) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    path = output_dir / f"chunk_{index:03d}.jsonl"
+    path.write_text("".join(json.dumps(r) + "\n" for r in records))
+    if done:
+        (output_dir / f"chunk_{index:03d}.done").touch()
+
+
+def _read_jsonl(path: Path) -> list[dict]:
+    return [json.loads(ln) for ln in path.read_text().splitlines() if ln.strip()]
+
+
+class TestCombine:
+    def test_combine_multiple_chunks(self, temp_dir):
+        output_dir = temp_dir / "output"
+        _write_chunk(output_dir, 0, [_record(0), _record(1)])
+        _write_chunk(output_dir, 1, [_record(2), _record(3)])
+        _write_chunk(output_dir, 2, [_record(4)])
+
+        out = temp_dir / "combined.jsonl"
+        _finalize_equilibrator(output_dir, out)
+
+        recs = _read_jsonl(out)
+        assert len(recs) == 5
+        assert [r["input_index"] for r in recs] == [0, 1, 2, 3, 4]
+
+    def test_records_passthrough(self, temp_dir):
+        output_dir = temp_dir / "output"
+        rec = _record(0)
+        _write_chunk(output_dir, 0, [rec])
+
+        out = temp_dir / "combined.jsonl"
+        _finalize_equilibrator(output_dir, out)
+
+        got = _read_jsonl(out)[0]
+        assert got == rec
+        assert got["compound_map"] == {"a": 1, "b": 2}
+        assert got["compound_refs"] == {"a": {"id": "MNXM1"}}
+
+    def test_error_records_included(self, temp_dir):
+        output_dir = temp_dir / "output"
+        _write_chunk(output_dir, 0, [_record(0), _record(1, dg=None, error="bad reaction")])
+
+        out = temp_dir / "combined.jsonl"
+        _finalize_equilibrator(output_dir, out)
+
+        recs = _read_jsonl(out)
+        assert len(recs) == 2
+        errored = [r for r in recs if r["error"] is not None]
+        assert len(errored) == 1
+        assert errored[0]["error"] == "bad reaction"
+
+    def test_done_markers_dropped(self, temp_dir):
+        output_dir = temp_dir / "output"
+        _write_chunk(output_dir, 0, [_record(0)])
+        _write_chunk(output_dir, 1, [_record(1)])
+
+        out = temp_dir / "combined.jsonl"
+        _finalize_equilibrator(output_dir, out)
+
+        # The finalized artifact is exactly the combined file — no .done, no chunks.
+        siblings = list(out.parent.iterdir())
+        assert siblings == [out]
+        assert not list(out.parent.glob("*.done"))
+        assert not list(out.parent.glob("chunk_*.jsonl"))
+
+    def test_missing_chunk_tolerated(self, temp_dir):
+        output_dir = temp_dir / "output"
+        _write_chunk(output_dir, 0, [_record(0)])
+        # index 1 produced no output (missing chunk).
+        _write_chunk(output_dir, 2, [_record(2)])
+
+        out = temp_dir / "combined.jsonl"
+        _finalize_equilibrator(output_dir, out)  # must not raise
+
+        recs = _read_jsonl(out)
+        assert [r["input_index"] for r in recs] == [0, 2]
+
+
+class TestOutputAutoname:
+    def test_output_dir_autoname(self, temp_dir):
+        output_dir = temp_dir / "output"
+        _write_chunk(output_dir, 0, [_record(0)])
+
+        dest_dir = temp_dir / "results"
+        dest_dir.mkdir()
+        job_id = "tester-equilibrator-20260707-abcd1234"
+
+        _finalize_equilibrator(output_dir, dest_dir, job_id=job_id)
+
+        expected = dest_dir / f"{job_id}.jsonl"
+        assert expected.exists()
+        assert len(_read_jsonl(expected)) == 1
+
+
+class TestS3Finalize:
+    def test_s3_download_before_combine(self, temp_dir):
+        """For an S3-mode job, finalize downloads the output prefix, combines,
+        then deletes the S3 prefix (existing convention)."""
+        from dh_cli.batch.commands import finalize as fin
+
+        base = temp_dir / "base"
+        base.mkdir()
+        # Job dir with an input/ so completion checks are inert.
+        job_id = "tester-equilibrator-20260707-abcd1234"
+        (base / job_id / "input").mkdir(parents=True)
+
+        manifest = MagicMock()
+        manifest.status = "submitted"
+        manifest.pipeline = "equilibrator"
+        manifest.storage_mode = "s3"
+        manifest.s3_prefix = "s3://dayhoff-atlas/ephemeral/batch-jobs/job/"
+        manifest.output = MagicMock()
+
+        downloaded = {}
+
+        def fake_download(prefix, local_dir):
+            # Simulate S3 output containing two chunks.
+            _write_chunk(Path(local_dir), 0, [_record(0)])
+            _write_chunk(Path(local_dir), 1, [_record(1)])
+            downloaded["prefix"] = prefix
+
+        deleted = {}
+
+        def fake_delete(prefix):
+            deleted["prefix"] = prefix
+
+        dest = temp_dir / "out"
+        dest.mkdir()
+
+        with (
+            patch.object(fin, "load_manifest", return_value=manifest),
+            patch.object(fin, "save_manifest_s3"),
+            patch("dh_cli.batch.s3_transport.download_directory", side_effect=fake_download),
+            patch("dh_cli.batch.s3_transport.delete_prefix", side_effect=fake_delete),
+        ):
+            from click.testing import CliRunner
+
+            runner = CliRunner()
+            result = runner.invoke(
+                fin.finalize,
+                [job_id, "--output", str(dest), "--base-path", str(base)],
+            )
+
+        assert result.exit_code == 0, result.output
+        # Downloaded from the manifest's output prefix.
+        assert downloaded["prefix"].startswith("s3://dayhoff-atlas/ephemeral/batch-jobs/job/")
+        # Deleted the S3 prefix after combining.
+        assert deleted["prefix"] == manifest.s3_prefix
+        # Combined output written.
+        out_file = dest / f"{job_id}.jsonl"
+        recs = _read_jsonl(out_file)
+        assert [r["input_index"] for r in recs] == [0, 1]