dayhoff-tools 1.14.4__tar.gz → 1.14.5__tar.gz

{dayhoff_tools-1.14.4 → dayhoff_tools-1.14.5}/PKG-INFO

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

{dayhoff_tools-1.14.4 → dayhoff_tools-1.14.5}/dayhoff_tools/batch/workers/base.py

@@ -30,12 +30,10 @@ def get_array_index() -> int:
 
     For array jobs, reads AWS_BATCH_JOB_ARRAY_INDEX.
     For retry jobs, maps from BATCH_RETRY_INDICES.
+    For single jobs (array_size=1), defaults to 0.
 
     Returns:
         The array index this worker should process
-
-    Raises:
-        RuntimeError: If no array index can be determined
     """
     # Check for retry mode first
     retry_indices = os.environ.get("BATCH_RETRY_INDICES")
@@ -49,15 +47,11 @@ def get_array_index() -> int:
             )
         return indices[array_idx]
 
-    # Standard array job mode
-    array_idx = os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX")
-    if array_idx is not None:
-        return int(array_idx)
-
-    raise RuntimeError(
-        "Could not determine array index. "
-        "Set AWS_BATCH_JOB_ARRAY_INDEX or BATCH_RETRY_INDICES environment variable."
-    )
+    # Standard array job mode - default to 0 for single jobs
+    # Note: When array_size=1, AWS Batch runs a single job (not an array),
+    # so AWS_BATCH_JOB_ARRAY_INDEX is not set. Default to 0.
+    array_idx = os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", "0")
+    return int(array_idx)
 
 
 def get_job_dir() -> Path:

{dayhoff_tools-1.14.4 → dayhoff_tools-1.14.5}/dayhoff_tools/cli/batch/commands/boltz.py

@@ -273,46 +273,84 @@ def _submit_batch_job(
 
 
 def _run_local_mode(input_path: Path):
-    """Run Boltz locally for a single complex."""
-    click.echo("Running Boltz locally...")
+    """Run Boltz locally in a Docker container.
+
+    This runs the boltz container with the normal entrypoint, processing
+    the first YAML file (index 0) for testing purposes.
+    """
+    import subprocess
+
+    click.echo("Running Boltz locally in container...")
     click.echo(f"Input directory: {input_path}")
 
-    # Find first YAML file
+    # Find YAML files
     yaml_files = list(input_path.glob("*.yaml"))
     if not yaml_files:
         click.echo(click.style("Error: No YAML files found", fg="red"), err=True)
         raise SystemExit(1)
 
-    input_file = yaml_files[0]
-    click.echo(f"Processing: {input_file.name}")
-    click.echo()
+    click.echo(f"Found {len(yaml_files)} YAML files, will process: {yaml_files[0].name}")
 
-    try:
-        from dayhoff_tools.batch.workers.boltz import BoltzProcessor
+    # Create a temporary job directory structure
+    temp_job_dir = input_path / ".local_boltz_job"
+    temp_input_dir = temp_job_dir / "input"
+    temp_output_dir = temp_job_dir / "output"
 
-        processor = BoltzProcessor(
-            num_workers=None,  # Auto-detect
-            msa_folder=None,
-            cache_dir=(
-                "/primordial/.cache/boltz" if os.path.exists("/primordial") else None
-            ),
-        )
+    # Clean up any previous run
+    if temp_job_dir.exists():
+        shutil.rmtree(temp_job_dir)
 
-        result_dir = processor.run(str(input_file))
+    temp_input_dir.mkdir(parents=True)
+    temp_output_dir.mkdir(parents=True)
 
-        click.echo()
-        click.echo(click.style("✓ Prediction complete!", fg="green"))
-        click.echo(f"Output: {result_dir}")
+    # Copy YAML files to input directory
+    for yaml_file in yaml_files:
+        shutil.copy2(yaml_file, temp_input_dir / yaml_file.name)
+
+    click.echo(f"Output will be at: {temp_output_dir}/")
+    click.echo()
 
-    except ImportError as e:
+    cmd = [
+        "docker", "run", "--rm",
+        "--gpus", "all",
+        "-v", "/primordial:/primordial",
+        "-v", f"{temp_job_dir}:{temp_job_dir}",
+        "-e", f"JOB_DIR={temp_job_dir}",
+        "-e", "AWS_BATCH_JOB_ARRAY_INDEX=0",
+        "-e", "BOLTZ_CACHE=/primordial/.cache/boltz",
+        "-e", "MSA_DIR=/primordial/.cache/msas",
+        DEFAULT_IMAGE_URI,
+    ]
+
+    click.echo(f"Running: {' '.join(cmd)}")
+    click.echo()
+
+    try:
+        result = subprocess.run(cmd)
+        if result.returncode != 0:
+            click.echo(
+                click.style(f"Container exited with code {result.returncode}", fg="red"),
+                err=True,
+            )
+            raise SystemExit(result.returncode)
+
+        # Check for output
+        output_dirs = list(temp_output_dir.iterdir()) if temp_output_dir.exists() else []
+        if output_dirs:
+            click.echo()
+            click.echo(click.style("✓ Prediction complete!", fg="green"))
+            click.echo(f"Output directory: {temp_output_dir}")
+            for d in output_dirs:
+                click.echo(f"  - {d.name}")
+        else:
+            click.echo(click.style("Warning: No output found", fg="yellow"))
+
+    except FileNotFoundError:
         click.echo(
-            click.style(f"Error: Missing dependency: {e}", fg="red"),
+            click.style("Error: Docker not found. Is Docker installed and running?", fg="red"),
             err=True,
         )
         raise SystemExit(1)
-    except Exception as e:
-        click.echo(click.style(f"Error: {e}", fg="red"), err=True)
-        raise SystemExit(1)
 
 
 def _run_shell_mode(input_path: Path):

{dayhoff_tools-1.14.4 → dayhoff_tools-1.14.5}/pyproject.toml

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