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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dh-cli
-Version: 0.8.8
+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.8.8 → dh_cli-0.10.0}/pyproject.toml

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

{dh_cli-0.8.8 → 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
@@ -22,6 +23,7 @@ from .commands.orca import orca
 from .commands.protmpnn import protmpnn
 from .commands.protmpnn_to_boltz import protmpnn_to_boltz
 from .commands.retry import retry
+from .commands.rfd3 import rfd3
 from .commands.status import status
 from .commands.submit import submit
 from .commands.train import train
@@ -98,10 +100,12 @@ 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)
 batch_cli.add_command(protmpnn_to_boltz, name="protmpnn-to-boltz")
+batch_cli.add_command(rfd3)
 batch_cli.add_command(train)
 
 __all__ = ["batch_cli"]

{dh_cli-0.8.8 → 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.8.8 → 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
@@ -143,6 +144,10 @@ def finalize(job_id, output, force, keep_intermediates, full_output, skip_dedup,
         _finalize_boltz(output_dir, output_path, full_output=full_output)
     elif manifest.pipeline == "protmpnn":
         _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)
 
@@ -544,3 +549,65 @@ def _finalize_protmpnn(output_dir: Path, output_path: Path):
 
     click.echo(f"{num_variants} variants from {num_configs} config(s), top confidence: {top_conf}")
     click.echo(click.style(f"Results: {output_path / 'results.csv'}", fg="green"))
+
+
+def _finalize_rfd3(output_dir: Path, output_path: Path):
+    """Collate completed RFD3 per-config backbone dirs into the final output.
+
+    The worker promotes each config's flat `.cif.gz` + `.json` design files
+    into `output/<stem>/` and drops a `output/rfd3_<stem>.done` marker only
+    when validation passed. Finalize copies each completed config's directory
+    into `output_path/<stem>/`, skipping configs without a `.done` marker and
+    the loose worker-index markers.
+    """
+    completed = []
+    for marker in sorted(output_dir.glob("rfd3_*.done")):
+        stem = marker.name[len("rfd3_") : -len(".done")]
+        if (output_dir / stem).is_dir():
+            completed.append(stem)
+
+    if not completed:
+        click.echo("No completed RFD3 configs found in output directory.", err=True)
+        raise SystemExit(1)
+
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    copied_count = 0
+    for stem in completed:
+        dest = output_path / stem
+        if dest.exists():
+            click.echo(f"  Skipping {stem} (already exists)")
+            continue
+        shutil.copytree(output_dir / stem, dest)
+        copied_count += 1
+
+    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}")