dayhoff-tools 1.14.12__tar.gz → 1.14.14__tar.gz

{dayhoff_tools-1.14.12 → dayhoff_tools-1.14.14}/PKG-INFO

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

{dayhoff_tools-1.14.12 → dayhoff_tools-1.14.14}/dayhoff_tools/cli/batch/commands/cancel.py

@@ -1,5 +1,7 @@
 """Cancel command for stopping running jobs."""
 
+import re
+
 import click
 
 from ..aws_batch import BatchClient, BatchError
@@ -11,6 +13,14 @@ from ..manifest import (
 )
 
 
+def _parse_retry_job_id(job_id: str) -> tuple[str, str | None]:
+    """Parse a job ID to extract parent job ID and retry suffix."""
+    match = re.match(r"^(.+)(-r\d+)$", job_id)
+    if match:
+        return match.group(1), job_id
+    return job_id, None
+
+
 @click.command()
 @click.argument("job_id")
 @click.option("--force", is_flag=True, help="Force termination of running containers")
@@ -19,19 +29,29 @@ def cancel(job_id, force, base_path):
     """Cancel a running batch job.
 
     Cancels the job in AWS Batch and updates the manifest status.
+    Also supports cancelling retry jobs by their ID (e.g., job-id-r1).
 
     \b
     Examples:
       dh batch cancel dma-embed-20260109-a3f2
+      dh batch cancel dma-embed-20260109-a3f2-r1  # Cancel specific retry
       dh batch cancel dma-embed-20260109-a3f2 --force
     """
+    # Check if this is a retry job ID
+    parent_job_id, retry_id = _parse_retry_job_id(job_id)
+
     # Load manifest
     try:
-        manifest = load_manifest(job_id, base_path)
+        manifest = load_manifest(parent_job_id, base_path)
     except FileNotFoundError:
-        click.echo(f"Job not found: {job_id}", err=True)
+        click.echo(f"Job not found: {parent_job_id}", err=True)
         raise SystemExit(1)
 
+    # If cancelling a specific retry job
+    if retry_id:
+        _cancel_retry_job(manifest, retry_id, force, base_path)
+        return
+
     # Check if job can be cancelled
     if manifest.status in (
         JobStatus.SUCCEEDED,
@@ -94,3 +114,47 @@ def cancel(job_id, force, base_path):
     except BatchError as e:
         click.echo(click.style(f"✗ Failed to cancel job: {e}", fg="red"), err=True)
         raise SystemExit(1)
+
+
+def _cancel_retry_job(manifest, retry_id: str, force: bool, base_path: str):
+    """Cancel a specific retry job."""
+    # Find the retry info
+    retry_info = None
+    for retry in manifest.retries:
+        if retry.retry_id == retry_id:
+            retry_info = retry
+            break
+
+    if not retry_info:
+        click.echo(f"Retry job not found: {retry_id}", err=True)
+        click.echo(f"Known retries: {[r.retry_id for r in manifest.retries]}", err=True)
+        raise SystemExit(1)
+
+    if not retry_info.batch_job_id:
+        click.echo(f"Retry job {retry_id} has no AWS Batch job ID.", err=True)
+        raise SystemExit(1)
+
+    # Cancel in AWS Batch
+    try:
+        client = BatchClient()
+
+        if force:
+            click.echo(f"Terminating retry job {retry_info.batch_job_id}...")
+            client.terminate_job(
+                retry_info.batch_job_id,
+                reason="Terminated by user via dh batch cancel --force",
+            )
+        else:
+            click.echo(f"Cancelling retry job {retry_info.batch_job_id}...")
+            client.cancel_job(
+                retry_info.batch_job_id,
+                reason="Cancelled by user via dh batch cancel",
+            )
+
+        click.echo()
+        click.echo(click.style(f"✓ Retry job {retry_id} cancelled successfully", fg="green"))
+        click.echo(f"Parent job: {manifest.job_id}")
+
+    except BatchError as e:
+        click.echo(click.style(f"✗ Failed to cancel retry job: {e}", fg="red"), err=True)
+        raise SystemExit(1)

{dayhoff_tools-1.14.12 → dayhoff_tools-1.14.14}/dayhoff_tools/cli/batch/commands/finalize.py

@@ -34,8 +34,13 @@ from ..manifest import (
     is_flag=True,
     help="For Boltz: copy entire output directory (default: only essential files)",
 )
+@click.option(
+    "--skip-dedup",
+    is_flag=True,
+    help="Skip deduplication step (use if input has no duplicates)",
+)
 @click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
-def finalize(job_id, output, force, keep_intermediates, full_output, base_path):
+def finalize(job_id, output, force, keep_intermediates, full_output, skip_dedup, base_path):
     """Combine results and clean up job intermediates.
 
     For embedding jobs, combines H5 files into a single output file.
@@ -46,6 +51,9 @@ def finalize(job_id, output, force, keep_intermediates, full_output, base_path):
       # Embedding job - combine H5 files
       dh batch finalize dma-embed-20260109-a3f2 --output /primordial/embeddings.h5
 
+      # Skip deduplication (faster if input has no duplicates)
+      dh batch finalize dma-embed-20260109-a3f2 --output /primordial/embeddings.h5 --skip-dedup
+
       # Boltz job - extract essential files only (default)
       dh batch finalize dma-boltz-20260113-190a --output /primordial/structures/
 
@@ -92,7 +100,7 @@ def finalize(job_id, output, force, keep_intermediates, full_output, base_path):
     # Finalize based on pipeline type
     click.echo()
     if manifest.pipeline in ("embed-t5", "embed"):
-        _finalize_embeddings(output_dir, output_path)
+        _finalize_embeddings(output_dir, output_path, skip_dedup=skip_dedup)
     elif manifest.pipeline == "boltz":
         _finalize_boltz(output_dir, output_path, full_output=full_output)
     else:
@@ -182,7 +190,7 @@ def _check_completion(job_id: str, base_path: str) -> list[int]:
     return incomplete
 
 
-def _finalize_embeddings(output_dir: Path, output_path: Path):
+def _finalize_embeddings(output_dir: Path, output_path: Path, skip_dedup: bool = False):
     """Combine H5 embedding files into a single output."""
     h5_files = sorted(output_dir.glob("embed_*.h5"))
 
@@ -191,6 +199,8 @@ def _finalize_embeddings(output_dir: Path, output_path: Path):
         raise SystemExit(1)
 
     click.echo(f"Found {len(h5_files)} H5 files to combine")
+    if skip_dedup:
+        click.echo("Skipping deduplication (--skip-dedup)")
 
     # Check if output already exists
     if output_path.exists():
@@ -213,10 +223,9 @@ def _finalize_embeddings(output_dir: Path, output_path: Path):
             click.echo("Single chunk - copying directly...")
             shutil.copy2(h5_files[0], output_path)
         else:
-            # Multiple files - combine, deduplicate, and optimize
+            # Multiple files - combine and optionally deduplicate
             with tempfile.TemporaryDirectory() as tmpdir:
                 combined_path = Path(tmpdir) / "combined.h5"
-                deduped_path = Path(tmpdir) / "deduped.h5"
 
                 # Combine H5 files
                 click.echo("Combining H5 files...")
@@ -226,13 +235,17 @@ def _finalize_embeddings(output_dir: Path, output_path: Path):
                     output_file=str(combined_path),
                 )
 
-                # Deduplicate
-                click.echo("Deduplicating...")
-                deduplicate_h5_file(str(combined_path), str(deduped_path))
-
-                # Optimize chunks
-                click.echo("Optimizing chunks...")
-                optimize_protein_embedding_chunks(str(deduped_path), str(output_path))
+                if skip_dedup:
+                    # Skip dedup - optimize directly from combined
+                    click.echo("Optimizing chunks...")
+                    optimize_protein_embedding_chunks(str(combined_path), str(output_path))
+                else:
+                    # Full pipeline: combine -> dedup -> optimize
+                    deduped_path = Path(tmpdir) / "deduped.h5"
+                    click.echo("Deduplicating...")
+                    deduplicate_h5_file(str(combined_path), str(deduped_path))
+                    click.echo("Optimizing chunks...")
+                    optimize_protein_embedding_chunks(str(deduped_path), str(output_path))
 
         click.echo(click.style("✓ H5 files combined successfully", fg="green"))
 

{dayhoff_tools-1.14.12 → dayhoff_tools-1.14.14}/pyproject.toml

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