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

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

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

{dayhoff_tools-1.14.9 → dayhoff_tools-1.14.10}/dayhoff_tools/batch/workers/boltz.py

@@ -341,10 +341,31 @@ class BoltzProcessor:
         return expected_output_dir
 
 
+def _get_done_marker_for_file(job_dir: Path, file_stem: str) -> Path:
+    """Get the done marker path for a specific input file."""
+    return job_dir / "output" / f"boltz_{file_stem}.done"
+
+
+def _check_file_complete(job_dir: Path, file_stem: str) -> bool:
+    """Check if a specific file has been processed."""
+    return _get_done_marker_for_file(job_dir, file_stem).exists()
+
+
+def _mark_file_complete(job_dir: Path, file_stem: str):
+    """Mark a specific file as complete."""
+    done_marker = _get_done_marker_for_file(job_dir, file_stem)
+    done_marker.parent.mkdir(parents=True, exist_ok=True)
+    done_marker.touch()
+    logger.info(f"File {file_stem} marked complete: {done_marker}")
+
+
 def main():
-    """Boltz worker main entrypoint for AWS Batch array jobs."""
+    """Boltz worker main entrypoint for AWS Batch array jobs.
+
+    Each worker processes multiple files based on array index and total workers.
+    With N files and M workers, worker i processes files where file_index % M == i.
+    """
     from .base import (
-        check_already_complete,
         configure_worker_logging,
         get_array_index,
         get_job_dir,
@@ -358,34 +379,36 @@ def main():
         # Get configuration from environment
         index = get_array_index()
         job_dir = get_job_dir()
+        array_size = int(os.environ.get("BATCH_ARRAY_SIZE", "1"))
+        num_files = int(os.environ.get("BATCH_NUM_FILES", "0"))
 
         logger.info(f"Worker configuration:")
         logger.info(f"  Array index: {index}")
+        logger.info(f"  Array size: {array_size}")
+        logger.info(f"  Total files: {num_files}")
         logger.info(f"  Job directory: {job_dir}")
 
-        # Check idempotency
-        if check_already_complete(index, job_dir, prefix="boltz"):
-            logger.info("Exiting - complex already processed")
-            return
-
-        # Find input file by index
+        # Find all input files
         input_dir = job_dir / "input"
         input_files = sorted(input_dir.glob("*.yaml"))
+        total_files = len(input_files)
 
-        if index >= len(input_files):
-            logger.error(
-                f"Index {index} out of range. Found {len(input_files)} input files."
-            )
-            raise RuntimeError(f"Index {index} out of range")
+        if total_files == 0:
+            logger.error("No input files found")
+            raise RuntimeError("No input files found")
 
-        input_file = input_files[index]
-        logger.info(f"  Input file: {input_file}")
+        # Calculate which files this worker should process
+        # Worker i processes files where file_index % array_size == index
+        my_files = [f for i, f in enumerate(input_files) if i % array_size == index]
 
-        # Determine output directory
-        output_base = input_file.stem
-        output_dir = job_dir / "output" / output_base
+        logger.info(f"  Files assigned to this worker: {len(my_files)}")
 
-        # Get MSA directories
+        if not my_files:
+            logger.info("No files assigned to this worker, exiting successfully")
+            mark_complete(index, job_dir, prefix="boltz")
+            return
+
+        # Get MSA directories (shared across all files)
         job_msa_dir = job_dir / "msas"
         global_msa_dir = Path(os.environ.get("MSA_DIR", "/primordial/.cache/msas"))
 
@@ -408,7 +431,7 @@ def main():
         if boltz_options:
             logger.info(f"  Boltz options: {boltz_options}")
 
-        # Create processor and run
+        # Create processor (reused for all files)
         processor = BoltzProcessor(
             num_workers=None,  # Auto-detect
             boltz_options=boltz_options,
@@ -416,16 +439,60 @@ def main():
             cache_dir=cache_dir,
         )
 
-        # Ensure output directory exists
-        output_dir.parent.mkdir(parents=True, exist_ok=True)
+        # Process each assigned file
+        completed = 0
+        failed = 0
+
+        for file_idx, input_file in enumerate(my_files):
+            file_stem = input_file.stem
+
+            # Check if this file is already complete (idempotency)
+            if _check_file_complete(job_dir, file_stem):
+                logger.info(
+                    f"[{file_idx + 1}/{len(my_files)}] {file_stem}: "
+                    "already complete, skipping"
+                )
+                completed += 1
+                continue
+
+            logger.info(
+                f"[{file_idx + 1}/{len(my_files)}] Processing {file_stem}..."
+            )
+
+            try:
+                # Determine output directory
+                output_dir = job_dir / "output" / file_stem
+                output_dir.parent.mkdir(parents=True, exist_ok=True)
+
+                result_dir = processor.run(str(input_file), str(output_dir))
+
+                # Mark this file as complete
+                _mark_file_complete(job_dir, file_stem)
+
+                logger.info(
+                    f"[{file_idx + 1}/{len(my_files)}] {file_stem}: "
+                    f"completed successfully -> {result_dir}"
+                )
+                completed += 1
+
+            except Exception as e:
+                logger.error(
+                    f"[{file_idx + 1}/{len(my_files)}] {file_stem}: "
+                    f"failed with error: {e}"
+                )
+                failed += 1
+                # Continue processing other files even if one fails
 
-        result_dir = processor.run(str(input_file), str(output_dir))
+        # Summary
+        logger.info(f"Worker {index} finished: {completed} completed, {failed} failed")
 
-        # Mark as complete
+        # Mark worker as complete
         mark_complete(index, job_dir, prefix="boltz")
 
-        logger.info(f"Complex {input_file.stem} completed successfully")
-        logger.info(f"Output: {result_dir}")
+        if failed > 0:
+            logger.warning(f"{failed} file(s) failed to process")
+            # Don't exit with error - some files succeeded and are marked complete
+            # The failed files can be retried later
 
     except Exception as e:
         logger.exception(f"Worker failed with error: {e}")

{dayhoff_tools-1.14.9 → dayhoff_tools-1.14.10}/dayhoff_tools/cli/batch/commands/boltz.py

@@ -232,6 +232,8 @@ def _submit_batch_job(
             "JOB_ID": job_id,
             "BOLTZ_CACHE": "/primordial/.cache/boltz",
             "MSA_DIR": "/primordial/.cache/msas",
+            "BATCH_ARRAY_SIZE": str(array_size),
+            "BATCH_NUM_FILES": str(num_files),
         }
 
         batch_job_id = client.submit_job(

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

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