dayhoff-tools 1.14.7__tar.gz → 1.14.9__tar.gz

{dayhoff_tools-1.14.7 → dayhoff_tools-1.14.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dayhoff-tools
-Version: 1.14.7
+Version: 1.14.9
 Summary: Common tools for all the repos at Dayhoff Labs
 Author: Daniel Martin-Alarcon
 Author-email: dma@dayhofflabs.com

{dayhoff_tools-1.14.7 → dayhoff_tools-1.14.9}/dayhoff_tools/cli/batch/__init__.py

@@ -12,6 +12,7 @@ import click
 
 from .commands.boltz import boltz
 from .commands.cancel import cancel
+from .commands.clean import clean
 from .commands.embed_t5 import embed_t5
 from .commands.finalize import finalize
 from .commands.list_jobs import list_jobs
@@ -36,6 +37,7 @@ def batch_cli():
       finalize   Combine results and clean up
       local      Run a chunk locally for debugging
       list       List recent jobs
+      clean      Remove old completed job directories
 
     \b
     Embedding Pipelines:
@@ -77,6 +79,7 @@ batch_cli.add_command(retry)
 batch_cli.add_command(finalize)
 batch_cli.add_command(local)
 batch_cli.add_command(list_jobs, name="list")
+batch_cli.add_command(clean)
 
 # Register pipeline commands
 batch_cli.add_command(embed_t5, name="embed-t5")

{dayhoff_tools-1.14.7 → dayhoff_tools-1.14.9}/dayhoff_tools/cli/batch/aws_batch.py

@@ -242,6 +242,64 @@ class BatchClient:
             failed=status_summary.get("FAILED", 0),
         )
 
+    def get_job_statuses_batch(self, job_ids: list[str]) -> dict[str, str]:
+        """Get status for multiple jobs in a single API call.
+
+        AWS Batch allows up to 100 job IDs per describe_jobs call.
+        This method handles batching for larger lists.
+
+        Args:
+            job_ids: List of AWS Batch job IDs
+
+        Returns:
+            Dictionary mapping job_id -> status string
+            Status will be one of: SUBMITTED, PENDING, RUNNABLE, STARTING,
+            RUNNING, SUCCEEDED, FAILED, or "UNKNOWN" if not found.
+            For array jobs, derives overall status from child statuses.
+        """
+        if not job_ids:
+            return {}
+
+        results = {}
+        batch_size = 100  # AWS Batch limit
+
+        for i in range(0, len(job_ids), batch_size):
+            batch = job_ids[i : i + batch_size]
+            try:
+                response = self.batch.describe_jobs(jobs=batch)
+                for job in response.get("jobs", []):
+                    job_id = job.get("jobId")
+                    status = job.get("status", "UNKNOWN")
+
+                    # For array jobs, derive overall status from children
+                    if "arrayProperties" in job:
+                        summary = job["arrayProperties"].get("statusSummary", {})
+                        total = job["arrayProperties"].get("size", 0)
+                        succeeded = summary.get("SUCCEEDED", 0)
+                        failed = summary.get("FAILED", 0)
+
+                        if succeeded + failed == total:
+                            # All children complete
+                            status = "SUCCEEDED" if failed == 0 else "FAILED"
+                        elif summary.get("RUNNING", 0) > 0:
+                            status = "RUNNING"
+                        elif summary.get("STARTING", 0) > 0:
+                            status = "STARTING"
+                        elif summary.get("RUNNABLE", 0) > 0:
+                            status = "RUNNABLE"
+                        elif summary.get("PENDING", 0) > 0:
+                            status = "PENDING"
+
+                    results[job_id] = status
+            except ClientError as e:
+                logger.warning(f"Failed to describe batch of jobs: {e}")
+                # Mark these as unknown
+                for job_id in batch:
+                    if job_id not in results:
+                        results[job_id] = "UNKNOWN"
+
+        return results
+
     def get_failed_indices(self, job_id: str) -> list[int]:
         """Get the array indices that failed for an array job.
 

dayhoff_tools-1.14.9/dayhoff_tools/cli/batch/commands/clean.py

@@ -0,0 +1,139 @@
+"""Clean command for removing old job directories."""
+
+import click
+
+from ..aws_batch import BatchClient, BatchError
+from ..manifest import (
+    BATCH_JOBS_BASE,
+    JobStatus,
+    delete_job_directory,
+    list_jobs as list_manifests,
+)
+from .status import format_time_ago, _aws_status_to_job_status
+
+
+@click.command("clean")
+@click.option("--user", help="Only clean jobs for this user")
+@click.option(
+    "--older-than",
+    type=int,
+    default=7,
+    help="Only clean jobs older than N days [default: 7]",
+)
+@click.option("--dry-run", is_flag=True, help="Show what would be cleaned without deleting")
+@click.option("--force", is_flag=True, help="Delete without confirmation")
+@click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
+def clean(user, older_than, dry_run, force, base_path):
+    """Remove completed job directories to free up space.
+
+    Only removes jobs that have SUCCEEDED or FAILED in AWS Batch.
+    Jobs that are still running or pending are never removed.
+
+    \b
+    Examples:
+      dh batch clean                     # Clean jobs older than 7 days
+      dh batch clean --older-than 1      # Clean jobs older than 1 day
+      dh batch clean --dry-run           # Show what would be cleaned
+      dh batch clean --user dma          # Only clean dma's jobs
+    """
+    from datetime import datetime, timedelta, timezone
+
+    cutoff = datetime.now(timezone.utc) - timedelta(days=older_than)
+
+    # Get all manifests
+    manifests = list_manifests(
+        base_path=base_path,
+        user=user,
+        status=None,
+        pipeline=None,
+        limit=500,
+    )
+
+    if not manifests:
+        click.echo("No jobs found.")
+        return
+
+    # Filter to old jobs
+    old_manifests = []
+    for m in manifests:
+        created = m.created
+        if created.tzinfo is None:
+            created = created.replace(tzinfo=timezone.utc)
+        if created < cutoff:
+            old_manifests.append(m)
+
+    if not old_manifests:
+        click.echo(f"No jobs older than {older_than} days found.")
+        return
+
+    # Get live statuses for old jobs
+    batch_job_ids = []
+    manifest_to_batch_id = {}
+    for m in old_manifests:
+        if m.batch and m.batch.job_id:
+            batch_job_ids.append(m.batch.job_id)
+            manifest_to_batch_id[m.job_id] = m.batch.job_id
+
+    live_statuses = {}
+    if batch_job_ids:
+        try:
+            client = BatchClient()
+            live_statuses = client.get_job_statuses_batch(batch_job_ids)
+        except BatchError as e:
+            click.echo(f"Error: Could not fetch status from AWS Batch: {e}", err=True)
+            click.echo("Cannot safely clean jobs without knowing their status.", err=True)
+            raise SystemExit(1)
+
+    # Find jobs that are safe to clean (SUCCEEDED or FAILED)
+    safe_to_clean = []
+    for manifest in old_manifests:
+        if manifest.job_id in manifest_to_batch_id:
+            batch_id = manifest_to_batch_id[manifest.job_id]
+            aws_status = live_statuses.get(batch_id, "UNKNOWN")
+            if aws_status in ("SUCCEEDED", "FAILED"):
+                safe_to_clean.append((manifest, aws_status))
+        elif manifest.status in (JobStatus.FINALIZED, JobStatus.CANCELLED):
+            # Already finalized or cancelled - safe to clean
+            safe_to_clean.append((manifest, manifest.status.value.upper()))
+
+    if not safe_to_clean:
+        click.echo(f"No completed jobs older than {older_than} days to clean.")
+        return
+
+    # Show what will be cleaned
+    click.echo()
+    click.echo(f"{'JOB ID':<35} {'STATUS':<12} {'CREATED':<12}")
+    click.echo("-" * 65)
+
+    for manifest, status in safe_to_clean:
+        click.echo(
+            f"{manifest.job_id:<35} "
+            f"{status:<12} "
+            f"{format_time_ago(manifest.created):<12}"
+        )
+
+    click.echo()
+    click.echo(f"Found {len(safe_to_clean)} completed jobs to clean.")
+
+    if dry_run:
+        click.echo("(dry-run: no changes made)")
+        return
+
+    # Confirm before deleting
+    if not force:
+        if not click.confirm("Delete these job directories?"):
+            click.echo("Cancelled.")
+            return
+
+    # Delete job directories
+    deleted = 0
+    for manifest, _ in safe_to_clean:
+        try:
+            delete_job_directory(manifest.job_id, base_path)
+            deleted += 1
+            click.echo(f"  Deleted: {manifest.job_id}")
+        except Exception as e:
+            click.echo(f"  Failed to delete {manifest.job_id}: {e}")
+
+    click.echo()
+    click.echo(f"Cleaned {deleted} job directories.")

{dayhoff_tools-1.14.7 → dayhoff_tools-1.14.9}/dayhoff_tools/cli/batch/commands/finalize.py

@@ -29,17 +29,30 @@ from ..manifest import (
     is_flag=True,
     help="Don't delete job directory after finalizing",
 )
+@click.option(
+    "--full-output",
+    is_flag=True,
+    help="For Boltz: copy entire output directory (default: only essential files)",
+)
 @click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
-def finalize(job_id, output, force, keep_intermediates, base_path):
+def finalize(job_id, output, force, keep_intermediates, full_output, 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.
+    For Boltz jobs, extracts essential files (CIF structures and confidence JSON).
 
     \b
     Examples:
+      # Embedding job - combine H5 files
       dh batch finalize dma-embed-20260109-a3f2 --output /primordial/embeddings.h5
-      dh batch finalize dma-embed-20260109-a3f2 --output /primordial/embeddings.h5 --force
+
+      # Boltz job - extract essential files only (default)
+      dh batch finalize dma-boltz-20260113-190a --output /primordial/structures/
+
+      # Boltz job - copy all output files
+      dh batch finalize dma-boltz-20260113-190a --output /primordial/structures/ --full-output
+
+      # Keep job directory after finalizing
       dh batch finalize dma-embed-20260109-a3f2 --output /primordial/out.h5 --keep-intermediates
     """
     # Load manifest
@@ -81,7 +94,7 @@ def finalize(job_id, output, force, keep_intermediates, base_path):
     if manifest.pipeline in ("embed-t5", "embed"):
         _finalize_embeddings(output_dir, output_path)
     elif manifest.pipeline == "boltz":
-        _finalize_boltz(output_dir, output_path)
+        _finalize_boltz(output_dir, output_path, full_output=full_output)
     else:
         _finalize_generic(output_dir, output_path)
 
@@ -196,29 +209,100 @@ def _finalize_embeddings(output_dir: Path, output_path: Path):
             shutil.copy2(h5_files[0], output_path)
 
 
-def _finalize_boltz(output_dir: Path, output_path: Path):
-    """Move Boltz output directories to destination."""
+def _finalize_boltz(output_dir: Path, output_path: Path, full_output: bool = False):
+    """Move Boltz output to destination.
+    
+    Args:
+        output_dir: Source directory containing boltz_results_* folders
+        output_path: Destination directory for outputs
+        full_output: If True, copy entire output directories. If False (default),
+                    extract only essential files (CIF structures and confidence JSON).
+    """
     # Find all output directories (one per complex)
-    complex_dirs = [d for d in output_dir.iterdir() if d.is_dir()]
+    complex_dirs = [d for d in output_dir.iterdir() if d.is_dir() and d.name.startswith("boltz_results_")]
 
     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")
+    click.echo(f"Found {len(complex_dirs)} structure predictions")
+    
+    if full_output:
+        click.echo("Mode: Copying full output (all files)")
+    else:
+        click.echo("Mode: Extracting essential files only (CIF + confidence JSON)")
+        click.echo("       Use --full-output to copy all files")
+    
+    # Confirm before proceeding
+    click.echo()
+    if not click.confirm(f"Copy results to {output_path}?"):
+        click.echo("Cancelled.")
+        raise SystemExit(0)
 
     # Ensure output directory exists
     output_path.mkdir(parents=True, exist_ok=True)
 
+    copied_count = 0
+    skipped_count = 0
+    
     for complex_dir in complex_dirs:
-        dest = output_path / complex_dir.name
+        complex_name = complex_dir.name.replace("boltz_results_", "")
+        dest = output_path / complex_name
+        
         if dest.exists():
-            click.echo(f"  Skipping {complex_dir.name} (already exists)")
+            click.echo(f"  Skipping {complex_name} (already exists)")
+            skipped_count += 1
             continue
-        shutil.move(str(complex_dir), str(dest))
-        click.echo(f"  Moved {complex_dir.name}")
+        
+        if full_output:
+            # Copy entire directory
+            shutil.copytree(complex_dir, dest)
+            click.echo(f"  Copied {complex_name} (full output)")
+        else:
+            # Extract only essential files
+            _extract_essential_boltz_files(complex_dir, dest, complex_name)
+            click.echo(f"  Extracted {complex_name} (essential files)")
+        
+        copied_count += 1
 
-    click.echo(click.style("✓ Structures moved successfully", fg="green"))
+    click.echo()
+    if skipped_count > 0:
+        click.echo(f"Copied {copied_count} predictions, skipped {skipped_count} existing")
+    else:
+        click.echo(click.style(f"✓ Copied {copied_count} structure predictions successfully", fg="green"))
+
+
+def _extract_essential_boltz_files(source_dir: Path, dest_dir: Path, complex_name: str):
+    """Extract only essential files from Boltz output.
+    
+    Essential files are:
+    - predictions/*/*.cif (structure files)
+    - predictions/*/confidence_*.json (confidence metrics)
+    
+    Args:
+        source_dir: Source boltz_results_* directory
+        dest_dir: Destination directory to create
+        complex_name: Name of the complex (for better error messages)
+    """
+    dest_dir.mkdir(parents=True, exist_ok=True)
+    
+    predictions_dir = source_dir / "predictions"
+    if not predictions_dir.exists():
+        click.echo(f"    Warning: No predictions directory found for {complex_name}", err=True)
+        return
+    
+    # Find all subdirectories in predictions/ (usually just one named after the complex)
+    for pred_subdir in predictions_dir.iterdir():
+        if not pred_subdir.is_dir():
+            continue
+        
+        # Copy CIF files (structures)
+        for cif_file in pred_subdir.glob("*.cif"):
+            shutil.copy2(cif_file, dest_dir / cif_file.name)
+        
+        # Copy confidence JSON files
+        for json_file in pred_subdir.glob("confidence_*.json"):
+            shutil.copy2(json_file, dest_dir / json_file.name)
 
 
 def _finalize_generic(output_dir: Path, output_path: Path):

{dayhoff_tools-1.14.7 → dayhoff_tools-1.14.9}/dayhoff_tools/cli/batch/commands/list_jobs.py

@@ -2,8 +2,9 @@
 
 import click
 
+from ..aws_batch import BatchClient, BatchError
 from ..manifest import BATCH_JOBS_BASE, JobStatus, list_jobs as list_manifests
-from .status import format_status, format_time_ago
+from .status import format_status, format_time_ago, _aws_status_to_job_status
 
 
 @click.command("list")
@@ -23,6 +24,7 @@ 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.
+    Status is fetched live from AWS Batch.
 
     \b
     Examples:
@@ -34,12 +36,13 @@ def list_jobs(user, status_filter, pipeline, limit, base_path):
     """
     status_enum = JobStatus(status_filter) if status_filter else None
 
+    # Fetch more manifests than requested to allow filtering by live status
     manifests = list_manifests(
         base_path=base_path,
         user=user,
-        status=status_enum,
+        status=None,  # Don't filter by status yet - will filter after getting live status
         pipeline=pipeline,
-        limit=limit,
+        limit=limit * 3,  # Fetch extra to account for status filtering
     )
 
     if not manifests:
@@ -48,6 +51,51 @@ def list_jobs(user, status_filter, pipeline, limit, base_path):
             click.echo("Try removing filters to see all jobs.")
         return
 
+    # Collect AWS Batch job IDs for live status lookup
+    batch_job_ids = []
+    manifest_to_batch_id = {}
+    for m in manifests:
+        if m.batch and m.batch.job_id:
+            batch_job_ids.append(m.batch.job_id)
+            manifest_to_batch_id[m.job_id] = m.batch.job_id
+
+    # Fetch live statuses from AWS Batch
+    live_statuses = {}
+    if batch_job_ids:
+        try:
+            client = BatchClient()
+            live_statuses = client.get_job_statuses_batch(batch_job_ids)
+        except BatchError as e:
+            click.echo(f"Warning: Could not fetch live status from AWS Batch: {e}")
+
+    # Build display data with live status
+    display_data = []
+    for manifest in manifests:
+        # Use live status if available, otherwise fall back to manifest status
+        if manifest.job_id in manifest_to_batch_id:
+            batch_id = manifest_to_batch_id[manifest.job_id]
+            aws_status = live_statuses.get(batch_id)
+            if aws_status:
+                live_status = _aws_status_to_job_status(aws_status)
+            else:
+                live_status = manifest.status
+        else:
+            live_status = manifest.status
+
+        # Apply status filter if specified
+        if status_enum and live_status != status_enum:
+            continue
+
+        display_data.append((manifest, live_status))
+
+        # Stop once we have enough
+        if len(display_data) >= limit:
+            break
+
+    if not display_data:
+        click.echo("No jobs found matching filters.")
+        return
+
     # Print header
     click.echo()
     click.echo(
@@ -55,17 +103,17 @@ def list_jobs(user, status_filter, pipeline, limit, base_path):
     )
     click.echo("-" * 85)
 
-    for manifest in manifests:
+    for manifest, live_status in display_data:
         click.echo(
             f"{manifest.job_id:<35} "
-            f"{format_status(manifest.status):<21} "  # Extra space for ANSI color codes
+            f"{format_status(live_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.")
+    click.echo(f"Showing {len(display_data)} jobs.")
 
     # Show filter hints
     hints = []

{dayhoff_tools-1.14.7 → dayhoff_tools-1.14.9}/dayhoff_tools/cli/batch/commands/status.py

@@ -80,21 +80,79 @@ def status(job_id, user, status_filter, pipeline, base_path):
         _show_job_list(user, status_filter, pipeline, base_path)
 
 
+def _aws_status_to_job_status(aws_status: str) -> JobStatus:
+    """Convert AWS Batch status to JobStatus enum."""
+    mapping = {
+        "SUBMITTED": JobStatus.SUBMITTED,
+        "PENDING": JobStatus.PENDING,
+        "RUNNABLE": JobStatus.RUNNING,  # Runnable means waiting for compute
+        "STARTING": JobStatus.RUNNING,
+        "RUNNING": JobStatus.RUNNING,
+        "SUCCEEDED": JobStatus.SUCCEEDED,
+        "FAILED": JobStatus.FAILED,
+    }
+    return mapping.get(aws_status, JobStatus.SUBMITTED)
+
+
 def _show_job_list(user, status_filter, pipeline, base_path):
     """Show a list of recent jobs."""
     status_enum = JobStatus(status_filter) if status_filter else None
     manifests = list_manifests(
         base_path=base_path,
         user=user,
-        status=status_enum,
+        status=None,  # Don't filter yet - we'll filter after getting live status
         pipeline=pipeline,
-        limit=20,
+        limit=50,  # Fetch more, filter later
     )
 
     if not manifests:
         click.echo("No jobs found.")
         return
 
+    # Collect AWS Batch job IDs for live status lookup
+    batch_job_ids = []
+    manifest_to_batch_id = {}
+    for m in manifests:
+        if m.batch and m.batch.job_id:
+            batch_job_ids.append(m.batch.job_id)
+            manifest_to_batch_id[m.job_id] = m.batch.job_id
+
+    # Fetch live statuses from AWS Batch
+    live_statuses = {}
+    if batch_job_ids:
+        try:
+            client = BatchClient()
+            live_statuses = client.get_job_statuses_batch(batch_job_ids)
+        except BatchError as e:
+            click.echo(f"Warning: Could not fetch live status from AWS Batch: {e}")
+
+    # Build display data with live status
+    display_data = []
+    for manifest in manifests:
+        # Use live status if available, otherwise fall back to manifest status
+        if manifest.job_id in manifest_to_batch_id:
+            batch_id = manifest_to_batch_id[manifest.job_id]
+            aws_status = live_statuses.get(batch_id)
+            if aws_status:
+                live_status = _aws_status_to_job_status(aws_status)
+            else:
+                live_status = manifest.status
+        else:
+            live_status = manifest.status
+
+        # Apply status filter if specified
+        if status_enum and live_status != status_enum:
+            continue
+
+        display_data.append((manifest, live_status))
+
+    if not display_data:
+        click.echo("No jobs found matching filters.")
+        return
+
+    # Limit to 20 after filtering
+    display_data = display_data[:20]
+
     # Print header
     click.echo()
     click.echo(
@@ -102,17 +160,17 @@ def _show_job_list(user, status_filter, pipeline, base_path):
     )
     click.echo("-" * 85)
 
-    for manifest in manifests:
+    for manifest, live_status in display_data:
         click.echo(
             f"{manifest.job_id:<35} "
-            f"{format_status(manifest.status):<21} "  # Extra space for color codes
+            f"{format_status(live_status):<21} "  # Extra space for 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)} most recent jobs.")
+    click.echo(f"Showing {len(display_data)} most recent jobs.")
     click.echo("Use 'dh batch status <job-id>' for details.")
 
 

{dayhoff_tools-1.14.7 → dayhoff_tools-1.14.9}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "dayhoff-tools"
-version = "1.14.7"
+version = "1.14.9"
 description = "Common tools for all the repos at Dayhoff Labs"
 authors = [
     {name = "Daniel Martin-Alarcon", email = "dma@dayhofflabs.com"}