dayhoff-tools 1.14.10__py3-none-any.whl → 1.14.11__py3-none-any.whl

dayhoff_tools/batch/workers/base.py

@@ -30,12 +30,18 @@ def get_array_index() -> int:
 
     For array jobs, reads AWS_BATCH_JOB_ARRAY_INDEX.
     For retry jobs, maps from BATCH_RETRY_INDICES.
+    For resliced retry jobs, uses the raw array index (chunks are renumbered).
     For single jobs (array_size=1), defaults to 0.
 
     Returns:
         The array index this worker should process
     """
-    # Check for retry mode first
+    # For resliced retries, use raw array index (chunks are renumbered 0..N-1)
+    if os.environ.get("RESLICE_PREFIX"):
+        array_idx = os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", "0")
+        return int(array_idx)
+
+    # Check for retry mode (non-resliced)
     retry_indices = os.environ.get("BATCH_RETRY_INDICES")
     if retry_indices:
         # In retry mode, we have a list of indices and use array index to pick
@@ -69,6 +75,20 @@ def get_job_dir() -> Path:
     return Path(job_dir)
 
 
+def get_reslice_prefix() -> str | None:
+    """Get the reslice prefix from environment if set.
+
+    When RESLICE_PREFIX is set (e.g., 'r1'), files are named like:
+    - Input: chunk_r1_000.fasta
+    - Output: embed_r1_000.h5
+    - Done marker: embed_r1_000.done
+
+    Returns:
+        The reslice prefix or None if not in reslice mode
+    """
+    return os.environ.get("RESLICE_PREFIX")
+
+
 def get_input_file(index: int, job_dir: Path, prefix: str = "chunk") -> Path:
     """Get the input file path for a given index.
 
@@ -80,6 +100,9 @@ def get_input_file(index: int, job_dir: Path, prefix: str = "chunk") -> Path:
     Returns:
         Path to input file
     """
+    reslice = get_reslice_prefix()
+    if reslice:
+        return job_dir / "input" / f"{prefix}_{reslice}_{index:03d}.fasta"
     return job_dir / "input" / f"{prefix}_{index:03d}.fasta"
 
 
@@ -97,6 +120,9 @@ def get_output_file(
     Returns:
         Path to output file
     """
+    reslice = get_reslice_prefix()
+    if reslice:
+        return job_dir / "output" / f"{prefix}_{reslice}_{index:03d}{suffix}"
     return job_dir / "output" / f"{prefix}_{index:03d}{suffix}"
 
 
@@ -111,6 +137,9 @@ def get_done_marker(index: int, job_dir: Path, prefix: str = "embed") -> Path:
     Returns:
         Path to done marker file
     """
+    reslice = get_reslice_prefix()
+    if reslice:
+        return job_dir / "output" / f"{prefix}_{reslice}_{index:03d}.done"
     return job_dir / "output" / f"{prefix}_{index:03d}.done"
 
 

dayhoff_tools/cli/batch/commands/finalize.py

@@ -118,7 +118,12 @@ def finalize(job_id, output, force, keep_intermediates, full_output, base_path):
 
 
 def _check_completion(job_id: str, base_path: str) -> list[int]:
-    """Check which chunks are incomplete (no .done marker)."""
+    """Check which chunks are incomplete (no .done marker).
+
+    Handles both original chunks (chunk_000.fasta) and resliced chunks
+    (chunk_r1_000.fasta). For original chunks that were resliced in a retry,
+    checks if all resliced chunks completed.
+    """
     job_dir = get_job_dir(job_id, base_path)
     input_dir = job_dir / "input"
     output_dir = job_dir / "output"
@@ -126,13 +131,53 @@ def _check_completion(job_id: str, base_path: str) -> list[int]:
     if not input_dir.exists():
         return []
 
+    # Load manifest to check for resliced retries
+    try:
+        manifest = load_manifest(job_id, base_path)
+        resliced_indices: set[int] = set()
+        reslice_info: dict[str, int] = {}  # prefix -> expected count
+
+        for retry in manifest.retries:
+            if retry.reslice_prefix and retry.reslice_count:
+                resliced_indices.update(retry.indices)
+                reslice_info[retry.reslice_prefix] = retry.reslice_count
+    except FileNotFoundError:
+        resliced_indices = set()
+        reslice_info = {}
+
     incomplete = []
-    for chunk_path in sorted(input_dir.glob("chunk_*.fasta")):
+
+    # Check original chunks (chunk_000.fasta pattern)
+    for chunk_path in sorted(input_dir.glob("chunk_[0-9][0-9][0-9].fasta")):
         idx_str = chunk_path.stem.split("_")[1]
         idx = int(idx_str)
+
+        # Check for original done marker
         done_marker = output_dir / f"embed_{idx:03d}.done"
-        if not done_marker.exists():
-            incomplete.append(idx)
+        if done_marker.exists():
+            continue
+
+        # Check if this chunk was resliced
+        if idx in resliced_indices:
+            # Find which retry covered this index and check if complete
+            is_covered = False
+            for retry in manifest.retries:
+                if (
+                    retry.reslice_prefix
+                    and retry.reslice_count
+                    and idx in retry.indices
+                ):
+                    # Check if all resliced chunks for this retry completed
+                    done_count = len(
+                        list(output_dir.glob(f"embed_{retry.reslice_prefix}_*.done"))
+                    )
+                    if done_count >= retry.reslice_count:
+                        is_covered = True
+                        break
+            if is_covered:
+                continue
+
+        incomplete.append(idx)
 
     return incomplete
 

dayhoff_tools/cli/batch/commands/retry.py

@@ -1,6 +1,7 @@
 """Retry command for re-running failed chunks."""
 
 from datetime import datetime
+from pathlib import Path
 
 import click
 
@@ -19,20 +20,37 @@ from ..manifest import (
 @click.command()
 @click.argument("job_id")
 @click.option("--indices", help="Specific indices to retry (comma-separated)")
+@click.option(
+    "--queue",
+    help="Override job queue (e.g., 't4-1x' for on-demand instead of spot)",
+)
+@click.option(
+    "--reslice",
+    type=int,
+    help="Reslice failed chunks into N thinner chunks (reduces interruption risk)",
+)
 @click.option(
     "--dry-run", is_flag=True, help="Show what would be retried without submitting"
 )
 @click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
-def retry(job_id, indices, dry_run, base_path):
+def retry(job_id, indices, queue, reslice, dry_run, base_path):
     """Retry failed chunks of a batch job.
 
     Identifies failed array indices and submits a new job to retry only
-    those specific indices.
+    those specific indices. Outputs go to the same job directory, so
+    finalization works normally after retries complete.
+
+    The --reslice option concatenates failed chunks and re-splits them into
+    thinner slices, reducing the time per worker and thus the risk of spot
+    interruptions. Resliced outputs are named with a prefix (e.g., embed_r1_000.h5)
+    and are automatically included in finalization.
 
     \b
     Examples:
       dh batch retry dma-embed-20260109-a3f2              # Retry all failed
       dh batch retry dma-embed-20260109-a3f2 --indices 5,12,27  # Retry specific indices
+      dh batch retry dma-embed-20260109-a3f2 --queue t4-1x     # Use on-demand (no spot interruptions)
+      dh batch retry dma-embed-20260109-a3f2 --reslice 40      # Reslice into 40 thinner chunks
       dh batch retry dma-embed-20260109-a3f2 --dry-run   # Show what would be retried
     """
     # Load manifest
@@ -56,48 +74,90 @@ def retry(job_id, indices, dry_run, base_path):
 
     click.echo(f"Found {len(retry_indices)} chunks to retry: {retry_indices}")
 
-    if dry_run:
-        click.echo()
-        click.echo(click.style("Dry run - job not submitted", fg="yellow"))
-        return
-
     # Check if we have the required info
     if not manifest.batch:
         click.echo("Job has no batch configuration.", err=True)
         raise SystemExit(1)
 
-    # Generate retry job ID
-    retry_id = f"{job_id}-r{len(manifest.retries) + 1}"
+    # Generate retry job ID and reslice prefix
+    retry_num = len(manifest.retries) + 1
+    retry_id = f"{job_id}-r{retry_num}"
+    reslice_prefix = f"r{retry_num}" if reslice else None
+
+    job_dir = get_job_dir(job_id, base_path)
+
+    if reslice:
+        # Count sequences in failed chunks to estimate split
+        total_seqs = _count_sequences_in_chunks(job_dir, retry_indices)
+        seqs_per_chunk = max(1, total_seqs // reslice)
+        click.echo(f"Total sequences in failed chunks: {total_seqs:,}")
+        click.echo(f"Reslicing into {reslice} chunks (~{seqs_per_chunk:,} seqs each)")
+
+    if dry_run:
+        click.echo()
+        click.echo(click.style("Dry run - job not submitted", fg="yellow"))
+        return
 
     click.echo()
     click.echo(f"Retry job ID: {retry_id}")
 
+    # Handle reslicing if requested
+    if reslice:
+        click.echo(f"Reslice prefix: {reslice_prefix}")
+        actual_chunks = _reslice_failed_chunks(
+            job_dir, retry_indices, reslice_prefix, reslice
+        )
+        click.echo(f"Created {actual_chunks} resliced chunks")
+        array_size = actual_chunks
+    else:
+        array_size = len(retry_indices)
+
     # Submit retry job
     try:
         client = BatchClient()
-        job_dir = get_job_dir(job_id, base_path)
 
         environment = {
             "JOB_DIR": str(job_dir),
             "JOB_ID": job_id,
-            "BATCH_RETRY_INDICES": ",".join(str(i) for i in retry_indices),
         }
 
-        batch_job_id = client.submit_array_job_with_indices(
-            job_name=retry_id,
-            job_definition=manifest.batch.job_definition or "dayhoff-embed-t5",
-            job_queue=manifest.batch.queue,
-            indices=retry_indices,
-            environment=environment,
-            timeout_seconds=6 * 3600,
-            retry_attempts=3,
-        )
+        # Use provided queue or fall back to original
+        job_queue = queue or manifest.batch.queue
+        if queue and queue != manifest.batch.queue:
+            click.echo(f"Using queue: {job_queue} (original: {manifest.batch.queue})")
+
+        if reslice:
+            # Resliced retry: use RESLICE_PREFIX, sequential indices 0..N-1
+            environment["RESLICE_PREFIX"] = reslice_prefix
+            batch_job_id = client.submit_job(
+                job_name=retry_id,
+                job_definition=manifest.batch.job_definition or "dayhoff-embed-t5",
+                job_queue=job_queue,
+                array_size=array_size,
+                environment=environment,
+                timeout_seconds=6 * 3600,
+                retry_attempts=5,
+            )
+        else:
+            # Standard retry: use BATCH_RETRY_INDICES mapping
+            environment["BATCH_RETRY_INDICES"] = ",".join(str(i) for i in retry_indices)
+            batch_job_id = client.submit_array_job_with_indices(
+                job_name=retry_id,
+                job_definition=manifest.batch.job_definition or "dayhoff-embed-t5",
+                job_queue=job_queue,
+                indices=retry_indices,
+                environment=environment,
+                timeout_seconds=6 * 3600,
+                retry_attempts=5,
+            )
 
         # Update manifest with retry info
         retry_info = RetryInfo(
             retry_id=retry_id,
             indices=retry_indices,
             batch_job_id=batch_job_id,
+            reslice_prefix=reslice_prefix,
+            reslice_count=array_size if reslice else None,
             created=datetime.utcnow(),
         )
         manifest.retries.append(retry_info)
@@ -129,8 +189,8 @@ def _find_incomplete_chunks(job_id: str, base_path: str) -> list[int]:
     if not input_dir.exists():
         return []
 
-    # Find all input chunks
-    input_chunks = sorted(input_dir.glob("chunk_*.fasta"))
+    # Find all original input chunks (not resliced ones)
+    input_chunks = sorted(input_dir.glob("chunk_[0-9][0-9][0-9].fasta"))
     incomplete = []
 
     for chunk_path in input_chunks:
@@ -144,3 +204,85 @@ def _find_incomplete_chunks(job_id: str, base_path: str) -> list[int]:
             incomplete.append(idx)
 
     return incomplete
+
+
+def _count_sequences_in_chunks(job_dir: Path, indices: list[int]) -> int:
+    """Count total sequences in the specified chunk files."""
+    input_dir = job_dir / "input"
+    total = 0
+
+    for idx in indices:
+        chunk_path = input_dir / f"chunk_{idx:03d}.fasta"
+        if chunk_path.exists():
+            with open(chunk_path) as f:
+                for line in f:
+                    if line.startswith(">"):
+                        total += 1
+
+    return total
+
+
+def _reslice_failed_chunks(
+    job_dir: Path, indices: list[int], reslice_prefix: str, num_chunks: int
+) -> int:
+    """Concatenate failed chunks and re-split into thinner slices.
+
+    Creates new chunk files named chunk_{prefix}_000.fasta, etc.
+
+    Args:
+        job_dir: Job directory path
+        indices: List of failed chunk indices
+        reslice_prefix: Prefix for new chunk files (e.g., 'r1')
+        num_chunks: Target number of new chunks
+
+    Returns:
+        Actual number of chunks created
+    """
+    from dayhoff_tools.fasta import split_fasta
+    import tempfile
+
+    input_dir = job_dir / "input"
+
+    # Concatenate all failed chunks into a temp file
+    with tempfile.NamedTemporaryFile(
+        mode="w", suffix=".fasta", delete=False
+    ) as tmp_file:
+        tmp_path = tmp_file.name
+        total_seqs = 0
+
+        for idx in indices:
+            chunk_path = input_dir / f"chunk_{idx:03d}.fasta"
+            if chunk_path.exists():
+                with open(chunk_path) as f:
+                    for line in f:
+                        tmp_file.write(line)
+                        if line.startswith(">"):
+                            total_seqs += 1
+
+    try:
+        # Calculate sequences per chunk
+        seqs_per_chunk = max(1, (total_seqs + num_chunks - 1) // num_chunks)
+
+        # Split into new chunks with reslice prefix
+        # split_fasta creates files like: chunk_r1_1.fasta, chunk_r1_2.fasta, etc.
+        actual_chunks = split_fasta(
+            fasta_file=tmp_path,
+            target_folder=str(input_dir),
+            base_name=f"chunk_{reslice_prefix}",
+            sequences_per_file=seqs_per_chunk,
+            max_files=num_chunks,
+            show_progress=True,
+        )
+
+        # Rename to zero-padded indices (chunk_r1_000.fasta, etc.)
+        for i in range(1, actual_chunks + 1):
+            old_name = input_dir / f"chunk_{reslice_prefix}_{i}.fasta"
+            new_name = input_dir / f"chunk_{reslice_prefix}_{i-1:03d}.fasta"
+            if old_name.exists():
+                old_name.rename(new_name)
+
+        return actual_chunks
+
+    finally:
+        # Clean up temp file
+        Path(tmp_path).unlink(missing_ok=True)

dayhoff_tools/cli/batch/manifest.py

@@ -62,6 +62,12 @@ class RetryInfo(BaseModel):
     retry_id: str = Field(..., description="Retry job ID")
     indices: list[int] = Field(..., description="Array indices being retried")
     batch_job_id: str | None = Field(None, description="AWS Batch job ID for retry")
+    reslice_prefix: str | None = Field(
+        None, description="Reslice prefix if chunks were resliced (e.g., 'r1')"
+    )
+    reslice_count: int | None = Field(
+        None, description="Number of resliced chunks created"
+    )
     created: datetime = Field(default_factory=datetime.utcnow)
 
 

{dayhoff_tools-1.14.10.dist-info → dayhoff_tools-1.14.11.dist-info}/METADATA

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

{dayhoff_tools-1.14.10.dist-info → dayhoff_tools-1.14.11.dist-info}/RECORD

@@ -1,7 +1,7 @@
 dayhoff_tools/__init__.py,sha256=M5zThPyEBRYa5CfwlzKhcqTevWn3OKu62cjV6Zqie2A,469
 dayhoff_tools/batch/__init__.py,sha256=N7v1pUg3rp68W6J7cX2Gssxdfp57Z7G0WeFJh7gKeiM,163
 dayhoff_tools/batch/workers/__init__.py,sha256=DJTtiBSE5k9J6qwG_MsXDn8hgz_lvEdaMwqKa6YyqXo,369
-dayhoff_tools/batch/workers/base.py,sha256=Jie5ScJrPSRkmLQdcM2hPJTmk3pIeLL0oA8Y9pkP0HA,4269
+dayhoff_tools/batch/workers/base.py,sha256=ZveUds6GW-aInj9LL0_VEYkXxOyZWRcOeJdacnEb8xo,5346
 dayhoff_tools/batch/workers/boltz.py,sha256=breCGITkX-_Xj9I_C1R-mKkG9RUCzSw3Kcdbbx_S5LA,18405
 dayhoff_tools/batch/workers/embed_t5.py,sha256=A5WqsQa1WZ7_la5X5wt0XUP-VwOglH04WyEINhwipeY,2750
 dayhoff_tools/chemistry/standardizer.py,sha256=uMn7VwHnx02nc404eO6fRuS4rsl4dvSPf2ElfZDXEpY,11188
@@ -14,15 +14,15 @@ dayhoff_tools/cli/batch/commands/boltz.py,sha256=N0LksmtOpkvnEsR0SAUHxtksKPAsQjR
 dayhoff_tools/cli/batch/commands/cancel.py,sha256=ZnHAJVzMGC0_1EQGpMSdYUlzm9yi-E9NxRJKBsetYW8,3111
 dayhoff_tools/cli/batch/commands/clean.py,sha256=nWOKbVM2nDuLMpyC038Q9aylOQxk2bq4N0JF65qJg-s,4570
 dayhoff_tools/cli/batch/commands/embed_t5.py,sha256=QXFydAw0wndevdzXF1cxikxMmvn1BuQ5p9lwutQFajU,11453
-dayhoff_tools/cli/batch/commands/finalize.py,sha256=Ir8XeA62lk_nUcWCLjjsYeqS58BBM1LTaEorEmQxOhA,11118
+dayhoff_tools/cli/batch/commands/finalize.py,sha256=xr3GFcMbvtU6UYiJI3UXhQqeaACSFzKIZOxz4GK-Dmo,12785
 dayhoff_tools/cli/batch/commands/list_jobs.py,sha256=COfxZddDVUAHeTayNAB3ruYNhgrE3osgFxY2qzf33cg,4284
 dayhoff_tools/cli/batch/commands/local.py,sha256=dZeKhNakaM1jS-EoByAwg1nWspRRoOmYzcwzjEKBaIA,3226
 dayhoff_tools/cli/batch/commands/logs.py,sha256=ctgJksdzFmqBdD18ePPsZe2BpuJYtHz2xAaMPnUplmQ,5293
-dayhoff_tools/cli/batch/commands/retry.py,sha256=PbstKr7iP7Uu4csisp3cvM8al9YOFw27XVlc591jLbE,4509
+dayhoff_tools/cli/batch/commands/retry.py,sha256=JgOMiwuESNG4Wp_fQ_vxy1RwyOZPt_kmVbLYdxYTVBY,9897
 dayhoff_tools/cli/batch/commands/status.py,sha256=k8ltke62vKUX_4MmnFDEd6H2s8KmVFA71UUXLbleA7o,9790
 dayhoff_tools/cli/batch/commands/submit.py,sha256=AXbvSReN8fLlzR5swE81pH7yvbCC1GUMCbsDrfoHAws,6786
 dayhoff_tools/cli/batch/job_id.py,sha256=mqr8SwcPlUWIYLaR_C4kACmL2ZFK8jaddd7B-45-XaQ,4246
-dayhoff_tools/cli/batch/manifest.py,sha256=blUw2ASxFTmdLV5ikz4VD91n5CN0V37rlyX11M0H2VE,8624
+dayhoff_tools/cli/batch/manifest.py,sha256=eLyiOFXonAUh5rfHcXyR2CzmRKXQz9-tTbJcWnVbbmE,8857
 dayhoff_tools/cli/cloud_commands.py,sha256=xrWQZm48e09GTh8RmYh3JqXUzW8mMy07dHk_77LqZx8,41199
 dayhoff_tools/cli/engine1/__init__.py,sha256=RE45X2IPCNUvQp6_OHI2lbjwoqexkH3roUPccIHkeJo,9540
 dayhoff_tools/cli/engine1/engine_core.py,sha256=TECO6766GCZTtbxOY7QRggwmjpMH9JZSuxjbwF46unU,27061
@@ -71,7 +71,7 @@ dayhoff_tools/intake/uniprot.py,sha256=BZYJQF63OtPcBBnQ7_P9gulxzJtqyorgyuDiPeOJq
 dayhoff_tools/logs.py,sha256=DKdeP0k0kliRcilwvX0mUB2eipO5BdWUeHwh-VnsICs,838
 dayhoff_tools/sqlite.py,sha256=jV55ikF8VpTfeQqqlHSbY8OgfyfHj8zgHNpZjBLos_E,18672
 dayhoff_tools/warehouse.py,sha256=UETBtZD3r7WgvURqfGbyHlT7cxoiVq8isjzMuerKw8I,24475
-dayhoff_tools-1.14.10.dist-info/METADATA,sha256=fhnkPfMm1G4qVxsu1zrGjLweThYGZpOh8XU61le3zVQ,3185
-dayhoff_tools-1.14.10.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
-dayhoff_tools-1.14.10.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
-dayhoff_tools-1.14.10.dist-info/RECORD,,
+dayhoff_tools-1.14.11.dist-info/METADATA,sha256=-3MjX0CiSeDoQ0MiF0mQlw700sWzZgeyVowtbljZ5qc,3185
+dayhoff_tools-1.14.11.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
+dayhoff_tools-1.14.11.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
+dayhoff_tools-1.14.11.dist-info/RECORD,,