dh-cli 0.10.2__tar.gz → 0.11.0__tar.gz

{dh_cli-0.10.2 → dh_cli-0.11.0}/.gitignore

@@ -28,6 +28,11 @@
 # nohup output
 **/nohup.out
 
+# Container build contexts stage multi-GB upstream source out-of-band before
+# `docker build`; never commit these vendored trees.
+containers/disco/DISCO/
+containers/boltz/thunderboltz/
+
 # Files with sensitive information
 key.json
 .env.secret

{dh_cli-0.10.2 → dh_cli-0.11.0}/PKG-INFO

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

{dh_cli-0.10.2 → dh_cli-0.11.0}/pyproject.toml

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

{dh_cli-0.10.2 → dh_cli-0.11.0}/src/dh_cli/batch/__init__.py

@@ -13,6 +13,7 @@ import click
 from .commands.boltz import boltz
 from .commands.cancel import cancel
 from .commands.clean import clean
+from .commands.disco import disco
 from .commands.embed_t5 import embed_t5
 from .commands.equilibrator import equilibrator
 from .commands.esmfold2 import esmfold2
@@ -109,6 +110,7 @@ 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(disco)
 batch_cli.add_command(train)
 
 __all__ = ["batch_cli"]

dh_cli-0.11.0/src/dh_cli/batch/commands/disco.py

@@ -0,0 +1,491 @@
+"""DISCO structure/design pipeline command."""
+
+import json
+import math
+import platform
+import re
+import shutil
+import subprocess
+from pathlib import Path
+
+import click
+
+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-disco"
+DEFAULT_IMAGE_URI = "074735440724.dkr.ecr.us-east-1.amazonaws.com/dayhoff:disco-latest"
+
+DISCO_TIMEOUT_SECONDS = 5400
+DISCO_RETRY_ATTEMPTS = 2
+
+# Task-01 spike measured VRAM fitting 24 GB A10G up to 250 aa (16.8 GB peak),
+# with cost ~quadratic in tokens. Beyond this the A10G ceiling is unproven, so
+# warn and point at the L40S fallback (D-15) rather than silently accepting it.
+A10G_SAFE_PROTEIN_LENGTH = 250
+
+HF_HUB_CACHE_DIR = "/primordial/.cache/huggingface"
+CHECKPOINT_SUBTREE = "/primordial/.cache/foundry/checkpoints/disco"
+PRIMORDIAL_CHECKPOINT_DIR = Path(CHECKPOINT_SUBTREE)
+DISCO_CHECKPOINT_PATH = f"{CHECKPOINT_SUBTREE}/DISCO.pt"
+CUTLASS_PATH = "/opt/cutlass"
+
+_JOB_ID_PATTERN = re.compile(r"^[a-z0-9-]+-[a-z0-9]+-\d{8}-[0-9a-f]{8}$")
+_ENTITY_KEYS = {"proteinChain", "ligand", "dnaSequence", "rnaSequence", "ion"}
+
+
+@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")
+def disco(
+    input_dir,
+    workers,
+    queue,
+    dry_run,
+    run_local,
+    run_remote,
+    base_path,
+    after,
+):
+    """Design structures with DISCO.
+
+    Processes a directory of native DISCO JSON config files, each a list of
+    jobs with sequences and optional ligand structure files.
+
+    \b
+    Examples:
+      # Run on a GPU engine (auto-detects GPU, runs locally)
+      dh batch disco input/
+
+      # Force remote Batch submission
+      dh batch disco input/ --remote
+
+      # Preview what would run
+      dh batch disco input/ --dry-run
+
+      # Run locally via Docker
+      dh batch disco 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)
+
+
+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(p for p in input_path.glob("*.json") if not p.name.endswith(".disco-meta.json"))
+
+
+def _referenced_file_ligands(config: Path) -> list[str]:
+    """Return relative ligand paths referenced via a ``FILE_`` prefix."""
+    with open(config) as f:
+        jobs = json.load(f)
+
+    relpaths = []
+    for job in jobs:
+        for entry in job.get("sequences", []):
+            if not isinstance(entry, dict):
+                continue
+            ligand = entry.get("ligand")
+            if not isinstance(ligand, dict):
+                continue
+            ref = ligand.get("ligand")
+            if isinstance(ref, str) and ref.startswith("FILE_"):
+                relpaths.append(ref[len("FILE_") :])
+    return relpaths
+
+
+def _warn_if_over_a10g_headroom(config: Path, name: str, sequences: list) -> None:
+    """Warn (non-fatal) when a protein chain exceeds the A10G-measured length.
+
+    The spike only proved the 24 GB fit up to 250 aa (D-15); longer chains may
+    approach the ceiling, so point at the L40S fallback rather than silently
+    submitting.
+    """
+    for entry in sequences:
+        if not isinstance(entry, dict):
+            continue
+        chain = entry.get("proteinChain")
+        if not isinstance(chain, dict):
+            continue
+        seq = chain.get("sequence")
+        if isinstance(seq, str) and len(seq) > A10G_SAFE_PROTEIN_LENGTH:
+            click.echo(
+                click.style(
+                    f"Warning: {config.name} job {name!r} has a protein chain of "
+                    f"{len(seq)} aa (> {A10G_SAFE_PROTEIN_LENGTH} aa measured-safe "
+                    f"on A10G {DEFAULT_QUEUE}). VRAM may approach the 24 GB ceiling; "
+                    "consider the l40s-1x-spot fallback (--queue).",
+                    fg="yellow",
+                ),
+                err=True,
+            )
+
+
+def _validate_config(config: Path, input_root: Path):
+    """Validate a single DISCO config; raise a clean CLI error on any violation."""
+    try:
+        jobs = json.loads(config.read_text())
+    except json.JSONDecodeError:
+        raise click.ClickException(f"Config is not valid JSON: {config.name}")
+
+    if not isinstance(jobs, list) or not jobs:
+        raise click.ClickException(f"Config must be a non-empty list of jobs: {config.name}")
+
+    seen_names = set()
+    for job in jobs:
+        if not isinstance(job, dict):
+            raise click.ClickException(f"Job must be an object in {config.name}")
+
+        name = job.get("name")
+        if not name or not isinstance(name, str):
+            raise click.ClickException(f"Job missing 'name' in {config.name}")
+
+        if "/" in name or "\\" in name or ".." in name or name != Path(name).name:
+            raise click.ClickException(f"Invalid job name (path separators): {name!r}")
+
+        if name in seen_names:
+            raise click.ClickException(f"Duplicate job name in {config.name}: {name!r}")
+        seen_names.add(name)
+
+        sequences = job.get("sequences")
+        if not isinstance(sequences, list) or not sequences:
+            raise click.ClickException(f"Job {name!r} has no entities in {config.name}")
+
+        _warn_if_over_a10g_headroom(config, name, sequences)
+
+    _validate_file_ligands(config, input_root)
+
+    meta = config.with_name(f"{config.stem}.disco-meta.json")
+    if meta.exists():
+        try:
+            json.loads(meta.read_text())
+        except json.JSONDecodeError:
+            raise click.ClickException(f"Malformed sidecar: {meta.name}")
+
+
+def _validate_file_ligands(config: Path, input_root: Path):
+    """Reject absolute, escaping, symlink-escaping, or missing FILE_ ligands."""
+    input_root = input_root.resolve()
+    for relpath in _referenced_file_ligands(config):
+        if Path(relpath).is_absolute():
+            raise click.ClickException(f"FILE_ ligand path must be relative: {relpath}")
+
+        resolved = (input_root / relpath).resolve()
+        if not resolved.is_relative_to(input_root):
+            raise click.ClickException(f"FILE_ ligand path escapes input dir: {relpath}")
+
+        if not resolved.exists():
+            raise click.ClickException(f"FILE_ ligand file not found: {relpath}")
+
+
+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 DISCO JSON config files found in {input_path}")
+
+    for config in configs:
+        _validate_config(config, input_path)
+
+    return configs
+
+
+def _copy_inputs_to_job_dir(input_path: Path, job_dir: Path) -> int:
+    """Copy configs, their FILE_ ligands, and meta sidecars into the job dir."""
+    dest = job_dir / "input"
+    dest.mkdir(parents=True, exist_ok=True)
+
+    count = 0
+    for config in _discover_configs(input_path):
+        shutil.copy2(config, dest / config.name)
+        count += 1
+
+        for relpath in _referenced_file_ligands(config):
+            src = input_path / relpath
+            target = dest / relpath
+            target.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(src, target)
+
+        meta = config.with_name(f"{config.stem}.disco-meta.json")
+        if meta.exists():
+            shutil.copy2(meta, dest / meta.name)
+
+    return count
+
+
+def _submit_batch_job(
+    input_path: Path,
+    workers: int | None,
+    queue: str,
+    dry_run: bool,
+    base_path: str,
+    after: tuple[str, ...],
+):
+    """Submit a DISCO job to AWS Batch."""
+    click.echo(f"Scanning {input_path} for JSON 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("disco")
+
+    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"Queue:            {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
+
+    for jid in after:
+        if not _JOB_ID_PATTERN.match(jid):
+            raise click.ClickException(f"Invalid --after job ID: {jid!r}")
+
+    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(
+            "dayhoff-disco 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} config file(s)")
+
+    manifest = JobManifest(
+        job_id=job_id,
+        user=job_id.split("-")[0],
+        pipeline="disco",
+        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),
+            "HF_HUB_OFFLINE": "1",
+            "TRANSFORMERS_OFFLINE": "1",
+            "HF_HOME": HF_HUB_CACHE_DIR,
+            "DISCO_CHECKPOINT_PATH": DISCO_CHECKPOINT_PATH,
+            "load_checkpoint_path": DISCO_CHECKPOINT_PATH,
+        }
+
+        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=DISCO_TIMEOUT_SECONDS,
+            retry_attempts=DISCO_RETRY_ATTEMPTS,
+            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 DISCO 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 DISCO 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"DISCO checkpoints not found at {PRIMORDIAL_CHECKPOINT_DIR}; "
+            "this host cannot run DISCO locally. Use --remote instead."
+        )
+
+    if not _has_nvidia_runtime():
+        raise click.ClickException(
+            "NVIDIA container runtime not available to Docker; cannot run DISCO locally. Use --remote instead."
+        )
+
+    configs = _validate_inputs(input_path)
+
+    temp_job_dir = input_path / ".local_disco_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 DISCO 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",
+        "HF_HUB_OFFLINE=1",
+        "-e",
+        "TRANSFORMERS_OFFLINE=1",
+        "-e",
+        f"HF_HOME={HF_HUB_CACHE_DIR}",
+        "-e",
+        f"CUTLASS_PATH={CUTLASS_PATH}",
+        "-e",
+        f"DISCO_CHECKPOINT_PATH={DISCO_CHECKPOINT_PATH}",
+        DEFAULT_IMAGE_URI,
+    ]
+
+    click.echo(f"Running: {' '.join(cmd)}")
+    subprocess.run(cmd)

{dh_cli-0.10.2 → dh_cli-0.11.0}/src/dh_cli/batch/commands/finalize.py

@@ -148,6 +148,8 @@ def finalize(job_id, output, force, keep_intermediates, full_output, skip_dedup,
         _finalize_rfd3(output_dir, output_path)
     elif manifest.pipeline == "esmfold2":
         _finalize_esmfold2(output_dir, output_path)
+    elif manifest.pipeline == "disco":
+        _finalize_disco(output_dir, output_path)
     elif manifest.pipeline == "equilibrator":
         _finalize_equilibrator(output_dir, output_path, job_id=job_id)
     else:
@@ -621,6 +623,43 @@ def _finalize_esmfold2(output_dir: Path, output_path: Path):
     click.echo(click.style(f"Results: {output_path}", fg="green"))
 
 
+def _finalize_disco(output_dir: Path, output_path: Path):
+    """Collate completed DISCO per-config output dirs into the final output.
+
+    The worker promotes each config's `pdbs/`, `sequences/`, and `ERR/`
+    subtrees into `output/<stem>/` and drops a `output/disco_<stem>.done`
+    marker only when the config succeeded. Finalize copies each completed
+    config's directory into `output_path/<stem>/`, skipping configs without a
+    `.done` marker and the loose `disco_worker_<NNN>.done` per-worker-index
+    markers, which live in a distinct namespace.
+    """
+    completed = []
+    for marker in sorted(output_dir.glob("disco_*.done")):
+        if marker.name.startswith("disco_worker_"):
+            continue
+        stem = marker.name[len("disco_") : -len(".done")]
+        if (output_dir / stem).is_dir():
+            completed.append(stem)
+
+    if not completed:
+        click.echo("No completed DISCO 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} DISCO 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.