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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dayhoff-tools
-Version: 1.14.8
+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.8 → dayhoff_tools-1.14.9}/dayhoff_tools/cli/batch/commands/clean.py

@@ -80,9 +80,9 @@ def clean(user, older_than, dry_run, force, base_path):
             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}")
-            click.echo("Cannot safely clean jobs without knowing their status.")
-            return
+            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 = []

{dayhoff_tools-1.14.8 → 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.8 → dayhoff_tools-1.14.9}/pyproject.toml

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