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

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

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

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

{dh_cli-0.8.7 → 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.7 → 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.7 → dh_cli-0.9.0}/src/dh_cli/bedrock/commands.py

@@ -17,15 +17,30 @@ from __future__ import annotations
 
 import datetime as dt
 import json
+import re
 import sys
-from typing import Optional
+from typing import Optional, TypedDict
+from uuid import uuid4
 
 import click
 
 from dh_cli._identity import HandleResolutionError, resolve_handle_from_session
+from dh_cli.github_commands import _scrub_token
 
 from . import cost_report as cr
 
+
+class BedrockKeyPayload(TypedDict):
+    """The JSON shape stored in `cursor-bedrock/<handle>`."""
+
+    access_key_id: str
+    secret_access_key: str
+    region: str
+
+
+# Matches the admin SSO role within an assumed-role ARN.
+_ADMIN_ROLE_RE = re.compile(r"AWSReservedSSO_DevAdminAccess_")
+
 # Hard-coded for the dev account. Kept here (not as user-facing flags)
 # because they never change in practice and making every dev remember
 # them is the whole reason this command exists.
@@ -124,9 +139,18 @@ def bedrock_key(handle: Optional[str], region: str, mode: str):
     resolved = handle or _resolve_handle()
     secret_id = f"{_SECRET_PREFIX}{resolved}"
 
+    sm = boto3.client("secretsmanager", region_name=region)
+    not_found = getattr(sm.exceptions, "ResourceNotFoundException", None)
+    if not (isinstance(not_found, type) and issubclass(not_found, BaseException)):
+        not_found = ()
     try:
-        sm = boto3.client("secretsmanager", region_name=region)
         resp = sm.get_secret_value(SecretId=secret_id)
+    except not_found:
+        raise click.ClickException(
+            f"No key yet for `{resolved}` (secret `{secret_id}` has no current "
+            f"version). Ask an admin to run `dh bedrock rotate --handle {resolved}` "
+            "to mint your first key."
+        )
     except Exception as exc:
         raise click.ClickException(
             f"Could not read secret `{secret_id}`: {exc}\n"
@@ -151,6 +175,104 @@ def bedrock_key(handle: Optional[str], region: str, mode: str):
         click.echo(json.dumps(payload, indent=2))
 
 
+# --------------------------------------------------------------------------
+# `dh bedrock rotate`
+# --------------------------------------------------------------------------
+
+
+@bedrock.command("rotate")
+@click.option(
+    "--handle",
+    "-u",
+    required=True,
+    help="Developer handle (e.g. 'dma'). Mints a key for `dev-cursor-<handle>`.",
+)
+@click.option(
+    "--region",
+    default=_DEFAULT_REGION,
+    show_default=True,
+    help="AWS region for the Secrets Manager secret.",
+)
+def bedrock_rotate(handle: str, region: str):
+    """Admin: mint a fresh Cursor Bedrock key for a developer.
+
+    Runs as `dev-admin`. Creates a new IAM access key for
+    `dev-cursor-<handle>`, writes it to `cursor-bedrock/<handle>`, then
+    deletes the prior key (delete-now).
+
+    \b
+    Examples:
+      dh bedrock rotate --handle dma
+    """
+    import boto3
+
+    user_name = f"dev-cursor-{handle}"
+    secret_id = f"{_SECRET_PREFIX}{handle}"
+
+    # Admin pre-check (defense-in-depth) — before any AWS mutation.
+    sts = boto3.Session().client("sts")
+    caller_arn = sts.get_caller_identity()["Arn"]
+    if not _ADMIN_ROLE_RE.search(caller_arn):
+        raise click.ClickException(
+            f"This command must be run as dev-admin (`awslogin dev-admin`). "
+            f"Current caller is not a DevAdminAccess role: {caller_arn}"
+        )
+
+    iam = boto3.client("iam")
+    sm = boto3.client("secretsmanager", region_name=region)
+
+    # Preflight — before create_access_key.
+    try:
+        resp = iam.list_access_keys(UserName=user_name)
+    except Exception as exc:
+        raise click.ClickException(
+            f"Could not list access keys for `{user_name}`: {exc}\n"
+            f"The IAM user may not exist yet — run `tfa` to provision it first."
+        )
+
+    prior_active_ids = [
+        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). Delete one before rotating."
+        )
+
+    # Transaction: create → put → delete prior.
+    create_resp = iam.create_access_key(UserName=user_name)
+    new_akid = create_resp["AccessKey"]["AccessKeyId"]
+    new_secret = create_resp["AccessKey"]["SecretAccessKey"]
+
+    payload: BedrockKeyPayload = {
+        "access_key_id": new_akid,
+        "secret_access_key": new_secret,
+        "region": region,
+    }
+    try:
+        sm.put_secret_value(
+            SecretId=secret_id,
+            SecretString=json.dumps(payload),
+            ClientRequestToken=uuid4().hex,
+        )
+    except Exception as exc:
+        # Roll back the just-created key. Scrub the secret BEFORE building
+        # the message so the ClickException output can never leak it.
+        iam.delete_access_key(UserName=user_name, AccessKeyId=new_akid)
+        scrubbed = _scrub_token(str(exc), new_secret)
+        raise click.ClickException(
+            f"PutSecretValue failed for `{secret_id}`; rolled back new key "
+            f"{new_akid}. No prior key was deleted. Error: {scrubbed}"
+        )
+
+    # Only after put succeeds: delete prior active keys (not the new one).
+    for old_id in prior_active_ids:
+        if old_id == new_akid:
+            continue
+        iam.delete_access_key(UserName=user_name, AccessKeyId=old_id)
+
+    click.echo(f"Rotated `{secret_id}`. New access_key_id: {new_akid}")
+
+
 # --------------------------------------------------------------------------
 # `dh bedrock cost`
 # --------------------------------------------------------------------------