dayhoff-tools 1.14.1__py3-none-any.whl → 1.14.2__py3-none-any.whl

dayhoff_tools/cli/batch/commands/embed_t5.py

@@ -0,0 +1,303 @@
+"""T5 embedding pipeline command."""
+
+import os
+from pathlib import Path
+
+import click
+
+from ..aws_batch import BatchClient, BatchError
+from ..job_id import generate_job_id
+from ..manifest import (
+    BATCH_JOBS_BASE,
+    BatchConfig,
+    InputConfig,
+    JobManifest,
+    JobStatus,
+    OutputConfig,
+    create_job_directory,
+    get_job_dir,
+    save_manifest,
+)
+
+
+# Default settings for T5 embedding
+DEFAULT_QUEUE = "t4-1x-spot"
+DEFAULT_WORKERS = 50
+DEFAULT_SEQS_PER_CHUNK = 5000
+DEFAULT_JOB_DEFINITION = "dayhoff-embed-t5"
+DEFAULT_IMAGE_URI = "074735440724.dkr.ecr.us-east-1.amazonaws.com/dayhoff:embed-latest"
+
+
+@click.command()
+@click.argument("input_fasta", type=click.Path(exists=True))
+@click.option("--workers", default=DEFAULT_WORKERS, type=int, help=f"Number of parallel workers [default: {DEFAULT_WORKERS}]")
+@click.option("--queue", default=DEFAULT_QUEUE, help=f"Batch queue [default: {DEFAULT_QUEUE}]")
+@click.option("--seqs-per-chunk", default=DEFAULT_SEQS_PER_CHUNK, type=int, help=f"Sequences per chunk [default: {DEFAULT_SEQS_PER_CHUNK}]")
+@click.option("--local", "run_local", is_flag=True, help="Run single chunk locally instead of Batch")
+@click.option("--shell", "run_shell", is_flag=True, help="Drop into container shell for debugging")
+@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")
+def embed_t5(input_fasta, workers, queue, seqs_per_chunk, run_local, run_shell, dry_run, base_path):
+    """Generate T5 protein embeddings for a FASTA file.
+
+    Splits the input FASTA into chunks and processes them in parallel using
+    AWS Batch array jobs. Each worker generates embeddings for its chunk
+    and writes an H5 file.
+
+    \b
+    Examples:
+      # Submit to AWS Batch with 50 workers
+      dh batch embed-t5 /primordial/proteins.fasta --workers 50
+
+      # Use a faster queue with more workers
+      dh batch embed-t5 /primordial/big.fasta --workers 100 --queue a10g-1x-spot
+
+      # Test locally with a single chunk
+      dh batch embed-t5 /primordial/test.fasta --local
+
+      # Debug by dropping into container shell
+      dh batch embed-t5 /primordial/test.fasta --shell
+
+    \b
+    After job completes:
+      dh batch status <job-id>                      # Check status
+      dh batch finalize <job-id> --output out.h5   # Combine results
+    """
+    input_path = Path(input_fasta).resolve()
+
+    if run_shell:
+        _run_shell_mode(input_path)
+        return
+
+    if run_local:
+        _run_local_mode(input_path)
+        return
+
+    # Batch submission mode
+    _submit_batch_job(input_path, workers, queue, seqs_per_chunk, dry_run, base_path)
+
+
+def _count_sequences(fasta_path: Path) -> int:
+    """Count sequences in a FASTA file (fast, just counts > lines)."""
+    count = 0
+    with open(fasta_path) as f:
+        for line in f:
+            if line.startswith(">"):
+                count += 1
+    return count
+
+
+def _split_fasta(input_path: Path, output_dir: Path, seqs_per_chunk: int) -> int:
+    """Split FASTA file into chunks.
+
+    Returns:
+        Number of chunks created
+    """
+    from dayhoff_tools.fasta import split_fasta
+
+    num_chunks = split_fasta(
+        fasta_file=str(input_path),
+        target_folder=str(output_dir),
+        base_name="chunk",
+        sequences_per_file=seqs_per_chunk,
+        show_progress=True,
+    )
+
+    # Rename files to use zero-padded indices (chunk_000.fasta, etc.)
+    for i in range(1, num_chunks + 1):
+        old_name = output_dir / f"chunk_{i}.fasta"
+        new_name = output_dir / f"chunk_{i-1:03d}.fasta"
+        if old_name.exists():
+            old_name.rename(new_name)
+
+    return num_chunks
+
+
+def _submit_batch_job(input_path: Path, workers: int, queue: str, seqs_per_chunk: int, dry_run: bool, base_path: str):
+    """Submit embedding job to AWS Batch."""
+    # Count sequences
+    click.echo(f"Counting sequences in {input_path}...")
+    num_sequences = _count_sequences(input_path)
+    click.echo(f"Found {num_sequences:,} sequences")
+
+    if num_sequences == 0:
+        click.echo(click.style("Error: No sequences found in input file", fg="red"), err=True)
+        raise SystemExit(1)
+
+    # Calculate chunks
+    num_chunks = min((num_sequences + seqs_per_chunk - 1) // seqs_per_chunk, workers)
+    actual_seqs_per_chunk = (num_sequences + num_chunks - 1) // num_chunks
+
+    # Generate job ID
+    job_id = generate_job_id("embed")
+
+    # Show plan
+    click.echo()
+    click.echo(f"Job ID:           {job_id}")
+    click.echo(f"Input:            {input_path}")
+    click.echo(f"Sequences:        {num_sequences:,}")
+    click.echo(f"Chunks:           {num_chunks}")
+    click.echo(f"Seqs per chunk:   ~{actual_seqs_per_chunk:,}")
+    click.echo(f"Queue:            {queue}")
+    click.echo(f"Job definition:   {DEFAULT_JOB_DEFINITION}")
+
+    if dry_run:
+        click.echo()
+        click.echo(click.style("Dry run - job not submitted", fg="yellow"))
+        return
+
+    click.echo()
+
+    # Create job directory
+    job_dir = create_job_directory(job_id, base_path)
+    input_dir = job_dir / "input"
+    output_dir = job_dir / "output"
+
+    click.echo(f"Created job directory: {job_dir}")
+
+    # Split FASTA into chunks
+    click.echo("Splitting FASTA into chunks...")
+    actual_chunks = _split_fasta(input_path, input_dir, actual_seqs_per_chunk)
+    click.echo(f"Created {actual_chunks} chunks")
+
+    # Create manifest
+    manifest = JobManifest(
+        job_id=job_id,
+        user=job_id.split("-")[0],
+        pipeline="embed-t5",
+        status=JobStatus.PENDING,
+        image_uri=DEFAULT_IMAGE_URI,
+        input=InputConfig(
+            source=str(input_path),
+            num_sequences=num_sequences,
+            num_chunks=actual_chunks,
+            sequences_per_chunk=actual_seqs_per_chunk,
+        ),
+        batch=BatchConfig(
+            queue=queue,
+            job_definition=DEFAULT_JOB_DEFINITION,
+            array_size=actual_chunks,
+        ),
+        output=OutputConfig(
+            destination=None,
+            finalized=False,
+        ),
+    )
+
+    save_manifest(manifest, base_path)
+
+    # Submit to AWS Batch
+    try:
+        client = BatchClient()
+
+        environment = {
+            "JOB_DIR": str(job_dir),
+            "JOB_ID": job_id,
+        }
+
+        batch_job_id = client.submit_job(
+            job_name=job_id,
+            job_definition=DEFAULT_JOB_DEFINITION,
+            job_queue=queue,
+            array_size=actual_chunks,
+            environment=environment,
+            timeout_seconds=6 * 3600,  # 6 hours
+            retry_attempts=3,
+        )
+
+        # Update manifest
+        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}")
+        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 /primordial/embeddings.h5")
+
+    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):
+    """Run embedding locally for a single chunk."""
+    import subprocess
+
+    click.echo("Running T5 embedding locally...")
+    click.echo(f"Input: {input_path}")
+
+    # Check if we have the embedder available
+    try:
+        from dayhoff_tools.embedders import T5Embedder
+
+        output_file = input_path.with_suffix(".h5")
+        click.echo(f"Output: {output_file}")
+        click.echo()
+
+        embedder = T5Embedder(
+            max_seq_length=4500,
+            large_protein_threshold=2500,
+            batch_residue_limit=4500,
+        )
+        embedder.run(str(input_path), str(output_file))
+
+        click.echo()
+        click.echo(click.style("✓ Embedding complete!", fg="green"))
+        click.echo(f"Output: {output_file}")
+
+    except ImportError:
+        click.echo(
+            click.style(
+                "Error: T5Embedder requires 'embedders' extra. "
+                "Install with: pip install 'dayhoff-tools[embedders]'",
+                fg="red",
+            ),
+            err=True,
+        )
+        raise SystemExit(1)
+
+
+def _run_shell_mode(input_path: Path):
+    """Drop into container shell for debugging."""
+    import subprocess
+
+    click.echo("Dropping into container shell...")
+    click.echo(f"Input will be available at: /input/{input_path.name}")
+    click.echo()
+
+    input_dir = input_path.parent
+
+    cmd = [
+        "docker", "run", "--rm", "-it",
+        "--gpus", "all",
+        "-v", "/primordial:/primordial",
+        "-v", f"{input_dir}:/input",
+        "-e", "JOB_DIR=/input",
+        "-e", "AWS_BATCH_JOB_ARRAY_INDEX=0",
+        "--entrypoint", "/bin/bash",
+        DEFAULT_IMAGE_URI,
+    ]
+
+    click.echo(f"Running: {' '.join(cmd)}")
+    click.echo()
+
+    try:
+        subprocess.run(cmd)
+    except FileNotFoundError:
+        click.echo(
+            click.style("Error: Docker not found. Is Docker installed and running?", fg="red"),
+            err=True,
+        )
+        raise SystemExit(1)

dayhoff_tools/cli/batch/commands/finalize.py

@@ -0,0 +1,206 @@
+"""Finalize command for combining results and cleaning up."""
+
+import shutil
+from pathlib import Path
+
+import click
+
+from ..manifest import (
+    BATCH_JOBS_BASE,
+    JobStatus,
+    delete_job_directory,
+    get_job_dir,
+    load_manifest,
+    save_manifest,
+)
+
+
+@click.command()
+@click.argument("job_id")
+@click.option("--output", required=True, type=click.Path(), help="Output path for combined results")
+@click.option("--force", is_flag=True, help="Finalize even if some chunks failed")
+@click.option("--keep-intermediates", is_flag=True, help="Don't delete job directory after finalizing")
+@click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
+def finalize(job_id, output, force, keep_intermediates, base_path):
+    """Combine results and clean up job intermediates.
+
+    For embedding jobs, combines H5 files into a single output file.
+    For structure prediction, moves outputs to the destination.
+
+    \b
+    Examples:
+      dh batch finalize dma-embed-20260109-a3f2 --output /primordial/embeddings.h5
+      dh batch finalize dma-embed-20260109-a3f2 --output /primordial/embeddings.h5 --force
+      dh batch finalize dma-embed-20260109-a3f2 --output /primordial/out.h5 --keep-intermediates
+    """
+    # Load manifest
+    try:
+        manifest = load_manifest(job_id, base_path)
+    except FileNotFoundError:
+        click.echo(f"Job not found: {job_id}", err=True)
+        raise SystemExit(1)
+
+    # Check job status
+    if manifest.status == JobStatus.FINALIZED:
+        click.echo(f"Job {job_id} is already finalized.", err=True)
+        raise SystemExit(1)
+
+    job_dir = get_job_dir(job_id, base_path)
+    output_dir = job_dir / "output"
+    output_path = Path(output).resolve()
+
+    # Check completion status
+    incomplete = _check_completion(job_id, base_path)
+    if incomplete:
+        click.echo(f"Found {len(incomplete)} incomplete chunks: {incomplete[:10]}...")
+        if not force:
+            click.echo()
+            click.echo("Use --force to finalize anyway, or retry failed chunks:")
+            click.echo(f"  dh batch retry {job_id}")
+            raise SystemExit(1)
+        click.echo()
+        click.echo(click.style("Warning: Finalizing with incomplete chunks", fg="yellow"))
+
+    # Update status
+    manifest.status = JobStatus.FINALIZING
+    save_manifest(manifest, base_path)
+
+    # Finalize based on pipeline type
+    click.echo()
+    if manifest.pipeline in ("embed-t5", "embed"):
+        _finalize_embeddings(output_dir, output_path)
+    elif manifest.pipeline == "boltz":
+        _finalize_boltz(output_dir, output_path)
+    else:
+        _finalize_generic(output_dir, output_path)
+
+    # Update manifest
+    manifest.status = JobStatus.FINALIZED
+    if manifest.output:
+        manifest.output.destination = str(output_path)
+        manifest.output.finalized = True
+    save_manifest(manifest, base_path)
+
+    click.echo()
+    click.echo(click.style(f"✓ Results saved to: {output_path}", fg="green"))
+
+    # Clean up
+    if not keep_intermediates:
+        click.echo(f"Cleaning up job directory: {job_dir}")
+        delete_job_directory(job_id, base_path)
+        click.echo(click.style("✓ Job directory deleted", fg="green"))
+    else:
+        click.echo(f"Job directory preserved: {job_dir}")
+
+
+def _check_completion(job_id: str, base_path: str) -> list[int]:
+    """Check which chunks are incomplete (no .done marker)."""
+    job_dir = get_job_dir(job_id, base_path)
+    input_dir = job_dir / "input"
+    output_dir = job_dir / "output"
+
+    if not input_dir.exists():
+        return []
+
+    incomplete = []
+    for chunk_path in sorted(input_dir.glob("chunk_*.fasta")):
+        idx_str = chunk_path.stem.split("_")[1]
+        idx = int(idx_str)
+        done_marker = output_dir / f"embed_{idx:03d}.done"
+        if not done_marker.exists():
+            incomplete.append(idx)
+
+    return incomplete
+
+
+def _finalize_embeddings(output_dir: Path, output_path: Path):
+    """Combine H5 embedding files into a single output."""
+    h5_files = sorted(output_dir.glob("embed_*.h5"))
+
+    if not h5_files:
+        click.echo("No H5 files found in output directory.", err=True)
+        raise SystemExit(1)
+
+    click.echo(f"Found {len(h5_files)} H5 files to combine")
+
+    # Check if output already exists
+    if output_path.exists():
+        click.echo(f"Output file already exists: {output_path}", err=True)
+        raise SystemExit(1)
+
+    # Ensure output directory exists
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        from dayhoff_tools.h5 import combine_h5_files, deduplicate_h5_file, optimize_protein_embedding_chunks
+
+        # Combine H5 files
+        click.echo("Combining H5 files...")
+        combine_h5_files(
+            input_folder=str(output_dir),
+            output_file=str(output_path),
+            glob_pattern="embed_*.h5",
+        )
+
+        # Deduplicate
+        click.echo("Deduplicating...")
+        deduplicate_h5_file(str(output_path))
+
+        # Optimize chunks
+        click.echo("Optimizing chunks...")
+        optimize_protein_embedding_chunks(str(output_path))
+
+        click.echo(click.style("✓ H5 files combined successfully", fg="green"))
+
+    except ImportError:
+        # Fall back to simple concatenation
+        click.echo("h5 module not available, using simple copy...")
+        if len(h5_files) == 1:
+            shutil.copy2(h5_files[0], output_path)
+        else:
+            # For multiple files without h5 module, just copy first file
+            # This is a fallback - the h5 module should be available
+            click.echo(
+                click.style(
+                    "Warning: Cannot combine multiple H5 files without dayhoff_tools.h5 module. "
+                    "Only copying first file.",
+                    fg="yellow",
+                )
+            )
+            shutil.copy2(h5_files[0], output_path)
+
+
+def _finalize_boltz(output_dir: Path, output_path: Path):
+    """Move Boltz output directories to destination."""
+    # Find all output directories (one per complex)
+    complex_dirs = [d for d in output_dir.iterdir() if d.is_dir()]
+
+    if not complex_dirs:
+        click.echo("No output directories found.", err=True)
+        raise SystemExit(1)
+
+    click.echo(f"Found {len(complex_dirs)} structure predictions to move")
+
+    # Ensure output directory exists
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    for complex_dir in complex_dirs:
+        dest = output_path / complex_dir.name
+        if dest.exists():
+            click.echo(f"  Skipping {complex_dir.name} (already exists)")
+            continue
+        shutil.move(str(complex_dir), str(dest))
+        click.echo(f"  Moved {complex_dir.name}")
+
+    click.echo(click.style("✓ Structures moved successfully", fg="green"))
+
+
+def _finalize_generic(output_dir: Path, output_path: Path):
+    """Generic finalization - copy output directory."""
+    if output_path.exists():
+        click.echo(f"Output path already exists: {output_path}", err=True)
+        raise SystemExit(1)
+
+    click.echo(f"Copying output directory to {output_path}...")
+    shutil.copytree(output_dir, output_path)
+    click.echo(click.style("✓ Output copied successfully", fg="green"))

dayhoff_tools/cli/batch/commands/list_jobs.py

@@ -0,0 +1,78 @@
+"""List command for showing recent jobs."""
+
+import click
+
+from ..manifest import BATCH_JOBS_BASE, JobStatus, list_jobs as list_manifests
+from .status import format_status, format_time_ago
+
+
+@click.command("list")
+@click.option("--user", help="Filter by username")
+@click.option(
+    "--status",
+    "status_filter",
+    type=click.Choice([s.value for s in JobStatus]),
+    help="Filter by status",
+)
+@click.option("--pipeline", help="Filter by pipeline type")
+@click.option("--limit", default=20, type=int, help="Maximum number of jobs to show [default: 20]")
+@click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
+def list_jobs(user, status_filter, pipeline, limit, base_path):
+    """List recent batch jobs.
+
+    Shows a table of recent jobs with their status, pipeline type, and creation time.
+
+    \b
+    Examples:
+      dh batch list                      # All recent jobs
+      dh batch list --user dma           # Filter by user
+      dh batch list --status running     # Filter by status
+      dh batch list --pipeline embed-t5  # Filter by pipeline type
+      dh batch list --limit 50           # Show more jobs
+    """
+    status_enum = JobStatus(status_filter) if status_filter else None
+
+    manifests = list_manifests(
+        base_path=base_path,
+        user=user,
+        status=status_enum,
+        pipeline=pipeline,
+        limit=limit,
+    )
+
+    if not manifests:
+        click.echo("No jobs found.")
+        if user or status_filter or pipeline:
+            click.echo("Try removing filters to see all jobs.")
+        return
+
+    # Print header
+    click.echo()
+    click.echo(
+        f"{'JOB ID':<35} {'STATUS':<12} {'PIPELINE':<12} {'USER':<10} {'CREATED':<12}"
+    )
+    click.echo("-" * 85)
+
+    for manifest in manifests:
+        click.echo(
+            f"{manifest.job_id:<35} "
+            f"{format_status(manifest.status):<21} "  # Extra space for ANSI color codes
+            f"{manifest.pipeline:<12} "
+            f"{manifest.user:<10} "
+            f"{format_time_ago(manifest.created):<12}"
+        )
+
+    click.echo()
+    click.echo(f"Showing {len(manifests)} jobs.")
+
+    # Show filter hints
+    hints = []
+    if not user:
+        hints.append("--user <name>")
+    if not status_filter:
+        hints.append("--status <status>")
+    if not pipeline:
+        hints.append("--pipeline <type>")
+
+    if hints:
+        click.echo(f"Filter with: {' '.join(hints)}")

dayhoff_tools/cli/batch/commands/local.py

@@ -0,0 +1,95 @@
+"""Local command for debugging job chunks locally."""
+
+import subprocess
+
+import click
+
+from ..manifest import BATCH_JOBS_BASE, get_job_dir, load_manifest
+
+
+@click.command()
+@click.argument("job_id")
+@click.option("--index", required=True, type=int, help="Array index to run")
+@click.option("--shell", "run_shell", is_flag=True, help="Drop into shell instead of running command")
+@click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
+def local(job_id, index, run_shell, base_path):
+    """Run a job chunk locally for debugging.
+
+    Runs a specific array index of a job in a local Docker container,
+    allowing you to debug failed chunks or test changes.
+
+    \b
+    Examples:
+      dh batch local dma-embed-20260109-a3f2 --index 27
+      dh batch local dma-embed-20260109-a3f2 --index 27 --shell
+    """
+    # Load manifest
+    try:
+        manifest = load_manifest(job_id, base_path)
+    except FileNotFoundError:
+        click.echo(f"Job not found: {job_id}", err=True)
+        raise SystemExit(1)
+
+    # Get job directory and image
+    job_dir = get_job_dir(job_id, base_path)
+    image_uri = manifest.image_uri
+
+    if not image_uri:
+        click.echo("Job has no image URI, cannot run locally.", err=True)
+        raise SystemExit(1)
+
+    # Validate index
+    if manifest.input and manifest.input.num_chunks:
+        if index >= manifest.input.num_chunks:
+            click.echo(
+                f"Index {index} out of range. Job has {manifest.input.num_chunks} chunks (0-{manifest.input.num_chunks - 1}).",
+                err=True,
+            )
+            raise SystemExit(1)
+
+    click.echo(f"Running job {job_id} index {index} locally")
+    click.echo(f"Image: {image_uri}")
+    click.echo(f"Job directory: {job_dir}")
+    click.echo()
+
+    # Build Docker command
+    cmd = [
+        "docker", "run", "--rm",
+        "--gpus", "all",
+        "-v", "/primordial:/primordial",
+        "-v", f"{job_dir}:{job_dir}",
+        "-e", f"AWS_BATCH_JOB_ARRAY_INDEX={index}",
+        "-e", f"JOB_DIR={job_dir}",
+        "-e", f"JOB_ID={job_id}",
+    ]
+
+    if run_shell:
+        cmd.extend(["-it", "--entrypoint", "/bin/bash"])
+        click.echo("Dropping into container shell...")
+        click.echo(f"  JOB_DIR={job_dir}")
+        click.echo(f"  AWS_BATCH_JOB_ARRAY_INDEX={index}")
+    else:
+        click.echo("Running worker command...")
+
+    cmd.append(image_uri)
+
+    click.echo()
+    click.echo(f"Command: {' '.join(cmd)}")
+    click.echo()
+
+    try:
+        result = subprocess.run(cmd)
+        if result.returncode != 0:
+            click.echo(
+                click.style(f"Container exited with code {result.returncode}", fg="red"),
+                err=True,
+            )
+            raise SystemExit(result.returncode)
+        else:
+            click.echo(click.style("✓ Container completed successfully", fg="green"))
+    except FileNotFoundError:
+        click.echo(
+            click.style("Error: Docker not found. Is Docker installed and running?", fg="red"),
+            err=True,
+        )
+        raise SystemExit(1)