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

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

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

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

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

@@ -22,6 +22,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
@@ -102,6 +103,7 @@ 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.9.0}/src/dh_cli/batch/commands/finalize.py

@@ -143,6 +143,8 @@ 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)
     else:
         _finalize_generic(output_dir, output_path)
 
@@ -544,3 +546,37 @@ 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"))

dh_cli-0.9.0/src/dh_cli/batch/commands/rfd3.py

@@ -0,0 +1,402 @@
+"""RFD3 protein backbone design pipeline command."""
+
+import math
+import platform
+import shutil
+import subprocess
+from pathlib import Path
+
+import click
+import yaml
+
+from ..aws_batch import BatchClient, BatchError, resolve_dependency
+from ..job_id import generate_job_id, get_aws_username
+from ..manifest import (
+    BATCH_JOBS_BASE,
+    BatchConfig,
+    InputConfig,
+    JobManifest,
+    JobStatus,
+    OutputConfig,
+    create_job_directory,
+    save_manifest,
+)
+
+DEFAULT_QUEUE = "a10g-1x-spot"
+MAX_WORKERS = 50
+FILES_PER_WORKER = 2
+DEFAULT_JOB_DEFINITION = "dayhoff-rfd3"
+DEFAULT_IMAGE_URI = "074735440724.dkr.ecr.us-east-1.amazonaws.com/dayhoff:rfd3-latest"
+OFFLINE_ENV_VAR = "HF_HUB_OFFLINE"
+
+CHECKPOINT_SUBTREE = "/primordial/.cache/foundry/checkpoints"
+PRIMORDIAL_CHECKPOINT_DIR = Path(CHECKPOINT_SUBTREE)
+
+
+@click.command()
+@click.argument("input_dir", type=click.Path(exists=True))
+@click.option(
+    "--workers",
+    default=None,
+    type=int,
+    help="Number of parallel workers [default: ~1 per 2 files]",
+)
+@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(
+    "--local",
+    "run_local",
+    is_flag=True,
+    help="Force local execution via Docker",
+)
+@click.option(
+    "--remote",
+    "run_remote",
+    is_flag=True,
+    help="Force Batch submission (override auto-detect)",
+)
+@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")
+@click.option("--num-designs", default=1, type=int, help="Number of designs per config")
+def rfd3(
+    input_dir,
+    workers,
+    queue,
+    dry_run,
+    run_local,
+    run_remote,
+    base_path,
+    after,
+    num_designs,
+):
+    """Design protein backbones with RFD3.
+
+    Processes a directory of YAML config files, each specifying design
+    parameters and an optional motif structure.
+
+    \b
+    Examples:
+      # Run on a GPU engine (auto-detects GPU, runs locally)
+      dh batch rfd3 input/
+
+      # Force remote Batch submission
+      dh batch rfd3 input/ --remote
+
+      # Preview what would run
+      dh batch rfd3 input/ --dry-run
+
+      # Run locally via Docker
+      dh batch rfd3 input/ --local
+    """
+    input_path = Path(input_dir).resolve()
+
+    if run_local:
+        _run_local_mode(input_path, base_path)
+        return
+
+    if not run_remote and not dry_run:
+        if _has_local_gpu():
+            click.echo("GPU detected — running locally (use --remote to override)")
+            _run_local_mode(input_path, base_path)
+            return
+
+    _submit_batch_job(input_path, workers, queue, dry_run, base_path, after, num_designs)
+
+
+def _has_local_gpu() -> bool:
+    """Check if a local NVIDIA GPU is available."""
+    try:
+        result = subprocess.run(
+            ["nvidia-smi"],
+            capture_output=True,
+            timeout=5,
+        )
+        return result.returncode == 0
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        return False
+
+
+def _has_nvidia_runtime() -> bool:
+    """Check if the NVIDIA container runtime is available to Docker."""
+    try:
+        result = subprocess.run(
+            ["docker", "info", "--format", "{{json .Runtimes}}"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        return result.returncode == 0 and "nvidia" in result.stdout
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        return False
+
+
+def _discover_configs(input_path: Path) -> list[Path]:
+    return sorted(input_path.glob("*.yaml"))
+
+
+def _validate_motif_paths(yaml_path: Path, input_root: Path):
+    """Reject absolute, escaping, or missing motif paths referenced by a config."""
+    with open(yaml_path) as f:
+        data = yaml.safe_load(f)
+
+    if not isinstance(data, dict):
+        return
+
+    for entry in data.values():
+        if not isinstance(entry, dict) or "input" not in entry:
+            continue
+        motif = entry["input"]
+
+        if Path(motif).is_absolute():
+            raise click.ClickException(f"Motif path must be relative, got absolute: {motif}")
+
+        resolved = (input_root / motif).resolve()
+        if not resolved.is_relative_to(input_root):
+            raise click.ClickException(f"Motif path escapes input dir: {motif}")
+
+        if not resolved.exists():
+            raise click.ClickException(f"Motif file not found: {motif}")
+
+
+def _validate_inputs(input_path: Path) -> list[Path]:
+    """Discover and validate configs; raise a clean CLI error on any violation."""
+    configs = _discover_configs(input_path)
+    if not configs:
+        raise click.ClickException(f"No YAML config files found in {input_path}")
+
+    for config in configs:
+        _validate_motif_paths(config, input_path)
+
+    return configs
+
+
+def _referenced_motifs(config: Path) -> list[str]:
+    """Return the relative motif ``input:`` paths referenced by a config."""
+    with open(config) as f:
+        data = yaml.safe_load(f)
+
+    if not isinstance(data, dict):
+        return []
+
+    motifs = []
+    for entry in data.values():
+        if isinstance(entry, dict) and "input" in entry:
+            motifs.append(entry["input"])
+    return motifs
+
+
+def _copy_inputs_to_job_dir(input_path: Path, job_dir: Path) -> int:
+    """Copy configs and their referenced motif files into the job directory.
+
+    Motifs are copied at their relative subpath (e.g. ``motifs/x.pdb``) so
+    RFD3 resolves the config's relative ``input:`` path from the copied dir.
+    """
+    dest = job_dir / "input"
+    dest.mkdir(parents=True, exist_ok=True)
+
+    count = 0
+    for yaml_file in sorted(input_path.glob("*.yaml")):
+        shutil.copy2(yaml_file, dest / yaml_file.name)
+        count += 1
+        for motif in _referenced_motifs(yaml_file):
+            src = input_path / motif
+            target = dest / motif
+            target.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(src, target)
+
+    return count
+
+
+def _submit_batch_job(
+    input_path: Path,
+    workers: int | None,
+    queue: str,
+    dry_run: bool,
+    base_path: str,
+    after: tuple[str, ...],
+    num_designs: int,
+):
+    """Submit an RFD3 job to AWS Batch."""
+    click.echo(f"Scanning {input_path} for YAML files...")
+    configs = _validate_inputs(input_path)
+    num_files = len(configs)
+    click.echo(f"Found {num_files} config(s) to process")
+
+    if workers is None:
+        workers = max(1, min(math.ceil(num_files / FILES_PER_WORKER), MAX_WORKERS))
+    array_size = min(num_files, workers)
+
+    job_id = generate_job_id("rfd3")
+
+    click.echo()
+    click.echo(f"Job ID:           {job_id}")
+    click.echo(f"Input:            {input_path}")
+    click.echo(f"Configs:          {num_files}")
+    click.echo(f"Workers:          {array_size}")
+    click.echo(f"Num designs:      {num_designs}")
+    click.echo(f"Queue:            {DEFAULT_QUEUE}")
+    click.echo(f"Job definition:   {DEFAULT_JOB_DEFINITION}")
+    click.echo(f"Image:            {DEFAULT_IMAGE_URI}")
+
+    if dry_run:
+        click.echo()
+        click.echo(click.style("Dry run - job not submitted (REMOTE plan)", 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("rfd3 job definition not yet deployed; apply Terraform in environments/dev first")
+
+    job_dir = create_job_directory(job_id, base_path)
+    click.echo(f"Created job directory: {job_dir}")
+
+    click.echo("Copying input files...")
+    copied = _copy_inputs_to_job_dir(input_path, job_dir)
+    click.echo(f"Copied {copied} YAML files")
+
+    manifest = JobManifest(
+        job_id=job_id,
+        user=job_id.split("-")[0],
+        pipeline="rfd3",
+        status=JobStatus.PENDING,
+        image_uri=DEFAULT_IMAGE_URI,
+        input=InputConfig(
+            source=str(input_path),
+            num_sequences=num_files,
+            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,
+    )
+
+    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),
+            "JOB_ID": job_id,
+            "BATCH_ARRAY_SIZE": str(array_size),
+            "BATCH_NUM_FILES": str(num_files),
+            "FOUNDRY_CHECKPOINT_DIRS": CHECKPOINT_SUBTREE,
+            OFFLINE_ENV_VAR: "1",
+            "NUM_DESIGNS": str(num_designs),
+        }
+
+        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=12 * 3600,
+            retry_attempts=2,
+            depends_on=depends_on,
+            share_identifier=get_aws_username(),
+        )
+
+        manifest.status = JobStatus.SUBMITTED
+        manifest.batch.job_id = batch_job_id
+        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)
+        save_manifest(manifest, base_path)
+        click.echo(click.style(f"Failed to submit job: {e}", fg="red"), err=True)
+        raise SystemExit(1)
+
+
+def _run_local_mode(input_path: Path, base_path: str = BATCH_JOBS_BASE):
+    """Run RFD3 locally in a Docker container after host-side fast-fail checks."""
+    machine = platform.machine()
+    if machine not in {"x86_64", "amd64"}:
+        raise click.ClickException(
+            f"Local RFD3 requires an x86_64 host (got {machine}); "
+            "no local image exists for this architecture. Use --remote instead."
+        )
+
+    if not PRIMORDIAL_CHECKPOINT_DIR.exists() or not any(PRIMORDIAL_CHECKPOINT_DIR.iterdir()):
+        raise click.ClickException(
+            f"Foundry checkpoints not found at {PRIMORDIAL_CHECKPOINT_DIR}; "
+            "this host cannot run RFD3 locally. Use --remote instead."
+        )
+
+    if not _has_nvidia_runtime():
+        raise click.ClickException(
+            "NVIDIA container runtime not available to Docker; cannot run RFD3 locally. Use --remote instead."
+        )
+
+    configs = _validate_inputs(input_path)
+
+    temp_job_dir = input_path / ".local_rfd3_job"
+    if temp_job_dir.exists():
+        shutil.rmtree(temp_job_dir)
+    (temp_job_dir / "output").mkdir(parents=True)
+    _copy_inputs_to_job_dir(input_path, temp_job_dir)
+
+    click.echo(f"Running RFD3 locally on {len(configs)} config(s)...")
+
+    cmd = [
+        "docker",
+        "run",
+        "--rm",
+        "--gpus",
+        "all",
+        "-v",
+        "/primordial:/primordial",
+        "-v",
+        f"{temp_job_dir}:{temp_job_dir}",
+        "-e",
+        f"JOB_DIR={temp_job_dir}",
+        "-e",
+        "AWS_BATCH_JOB_ARRAY_INDEX=0",
+        "-e",
+        "BATCH_ARRAY_SIZE=1",
+        "-e",
+        f"BATCH_NUM_FILES={len(configs)}",
+        "-e",
+        f"FOUNDRY_CHECKPOINT_DIRS={CHECKPOINT_SUBTREE}",
+        "-e",
+        f"{OFFLINE_ENV_VAR}=1",
+        DEFAULT_IMAGE_URI,
+    ]
+
+    click.echo(f"Running: {' '.join(cmd)}")
+    subprocess.run(cmd)

{dh_cli-0.8.8 → dh_cli-0.9.0}/src/dh_cli/bedrock/commands.py

@@ -231,14 +231,11 @@ def bedrock_rotate(handle: str, region: str):
         )
 
     prior_active_ids = [
-        meta["AccessKeyId"]
-        for meta in resp.get("AccessKeyMetadata", [])
-        if meta.get("Status") == "Active"
+        meta["AccessKeyId"] for meta in resp.get("AccessKeyMetadata", []) if meta.get("Status") == "Active"
     ]
     if len(prior_active_ids) >= 2:
         raise click.ClickException(
-            f"`{user_name}` already has 2 active keys (the IAM per-user limit). "
-            f"Delete one before rotating."
+            f"`{user_name}` already has 2 active keys (the IAM per-user limit). Delete one before rotating."
         )
 
     # Transaction: create → put → delete prior.

{dh_cli-0.8.8 → dh_cli-0.9.0}/tests/bedrock/test_rotate.py

@@ -59,12 +59,8 @@ def _build_iam_stub(
         iam.list_access_keys.side_effect = list_keys_error
     else:
         ids = existing_key_ids if existing_key_ids is not None else []
-        iam.list_access_keys.return_value = {
-            "AccessKeyMetadata": [{"AccessKeyId": k, "Status": "Active"} for k in ids]
-        }
-    iam.create_access_key.return_value = {
-        "AccessKey": {"AccessKeyId": new_akid, "SecretAccessKey": new_secret}
-    }
+        iam.list_access_keys.return_value = {"AccessKeyMetadata": [{"AccessKeyId": k, "Status": "Active"} for k in ids]}
+    iam.create_access_key.return_value = {"AccessKey": {"AccessKeyId": new_akid, "SecretAccessKey": new_secret}}
     return iam
 
 

{dh_cli-0.8.8 → dh_cli-0.9.0}/tests/github/conftest.py

@@ -30,7 +30,7 @@ import pytest
 DEV_ARN = "arn:aws:sts::074735440724:assumed-role/AWSReservedSSO_DeveloperAccess_abcdef1234567890/dma"
 HANDLE = "dma"
 PAT_VALUE = "ghp_REALTOKEN1234567890ABCDEF"
-PAT_CREATED = dt.datetime(2026, 4, 1, 12, 0, tzinfo=dt.timezone.utc)
+PAT_CREATED = dt.datetime.now(dt.timezone.utc) - dt.timedelta(days=10)
 
 
 def _build_sm_stub(

dh_cli-0.9.0/tests/test_finalize_rfd3.py

@@ -0,0 +1,130 @@
+"""Tests for RFD3 backbone-dir collation in _finalize_rfd3()."""
+
+import gzip
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from dh_cli.batch.commands.finalize import _finalize_rfd3
+
+
+@pytest.fixture
+def temp_dir():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        yield Path(tmpdir)
+
+
+def _create_rfd3_output(
+    output_dir: Path,
+    completed: list[str],
+    incomplete: list[str] | None = None,
+    num_designs: int = 2,
+):
+    """Build a mock RFD3 worker output directory.
+
+    For each completed config ``stem``, creates:
+      - ``output/<stem>/`` holding ``num_designs`` ``.cif.gz`` + ``.json`` files
+        (the flat per-design files RFD3 emits, promoted by the worker).
+      - a loose ``output/rfd3_<stem>.done`` marker.
+
+    Incomplete configs get their ``<stem>/`` dir but NO ``.done`` marker,
+    mirroring a config that failed output validation and was never promoted.
+    Also drops a worker-level ``rfd3_0.done`` marker (array-index marker),
+    which finalize must not treat as a config.
+    """
+    incomplete = incomplete or []
+
+    def _write_designs(stem: str):
+        config_dir = output_dir / stem
+        config_dir.mkdir(parents=True)
+        for i in range(num_designs):
+            cif = config_dir / f"{stem}_motif_0_model_{i}.cif.gz"
+            with gzip.open(cif, "wt") as f:
+                f.write("data_design\n_cell.length_a 1.0\n")
+            (config_dir / f"{stem}_motif_0_model_{i}.json").write_text('{"metrics": {"num_residues": 150}}')
+
+    for stem in completed:
+        _write_designs(stem)
+        (output_dir / f"rfd3_{stem}.done").touch()
+
+    for stem in incomplete:
+        _write_designs(stem)
+
+    # Worker array-index marker (not a config); finalize must ignore it.
+    (output_dir / "rfd3_0.done").touch()
+
+
+class TestFinalizeRfd3CollatesBackboneDirs:
+    def test_finalize_collates_backbone_dirs(self, temp_dir):
+        output_dir = temp_dir / "output"
+        output_dir.mkdir()
+        dest = temp_dir / "final"
+
+        _create_rfd3_output(
+            output_dir,
+            completed=["sortase_a", "sortase_b"],
+            num_designs=3,
+        )
+
+        _finalize_rfd3(output_dir, dest)
+
+        # Each completed config becomes a per-config backbone dir in the output.
+        assert (dest / "sortase_a").is_dir()
+        assert (dest / "sortase_b").is_dir()
+
+        # The flat per-design .cif.gz + .json files are carried through.
+        a_cifs = sorted((dest / "sortase_a").glob("*.cif.gz"))
+        a_jsons = sorted((dest / "sortase_a").glob("*.json"))
+        assert len(a_cifs) == 3
+        assert len(a_jsons) == 3
+
+        # Loose .done markers are not copied into the finalized output.
+        assert not list(dest.glob("*.done"))
+
+    def test_no_done_markers_in_finalized_output(self, temp_dir):
+        output_dir = temp_dir / "output"
+        output_dir.mkdir()
+        dest = temp_dir / "final"
+
+        _create_rfd3_output(output_dir, completed=["cfg"], num_designs=1)
+
+        _finalize_rfd3(output_dir, dest)
+
+        assert not (dest / "rfd3_cfg.done").exists()
+        assert not (dest / "rfd3_0.done").exists()
+        assert (dest / "cfg").is_dir()
+
+
+class TestFinalizeRfd3SkipsIncomplete:
+    def test_finalize_skips_incomplete(self, temp_dir):
+        output_dir = temp_dir / "output"
+        output_dir.mkdir()
+        dest = temp_dir / "final"
+
+        _create_rfd3_output(
+            output_dir,
+            completed=["good"],
+            incomplete=["bad"],
+            num_designs=2,
+        )
+
+        _finalize_rfd3(output_dir, dest)
+
+        # The config with a .done marker is collated.
+        assert (dest / "good").is_dir()
+        # The config missing its .done marker is excluded.
+        assert not (dest / "bad").exists()
+
+
+class TestFinalizeRfd3Empty:
+    def test_no_completed_configs_errors(self, temp_dir):
+        output_dir = temp_dir / "output"
+        output_dir.mkdir()
+        dest = temp_dir / "final"
+
+        # Only a worker-index marker, no per-config .done markers.
+        (output_dir / "rfd3_0.done").touch()
+
+        with pytest.raises(SystemExit):
+            _finalize_rfd3(output_dir, dest)