dayhoff-tools 1.14.4__py3-none-any.whl → 1.14.6__py3-none-any.whl

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:
@@ -89,7 +83,9 @@ def get_input_file(index: int, job_dir: Path, prefix: str = "chunk") -> Path:
     return job_dir / "input" / f"{prefix}_{index:03d}.fasta"
 
 
-def get_output_file(index: int, job_dir: Path, prefix: str = "embed", suffix: str = ".h5") -> Path:
+def get_output_file(
+    index: int, job_dir: Path, prefix: str = "embed", suffix: str = ".h5"
+) -> Path:
     """Get the output file path for a given index.
 
     Args:

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/cli/batch/commands/finalize.py

@@ -136,10 +136,11 @@ def _finalize_embeddings(output_dir: Path, output_path: Path):
 
         # Combine H5 files
         click.echo("Combining H5 files...")
+        # Get list of h5 file paths as strings
+        h5_file_paths = [str(f) for f in h5_files]
         combine_h5_files(
-            input_folder=str(output_dir),
+            input_files=h5_file_paths,
             output_file=str(output_path),
-            glob_pattern="embed_*.h5",
         )
 
         # Deduplicate

{dayhoff_tools-1.14.4.dist-info → dayhoff_tools-1.14.6.dist-info}/METADATA

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

{dayhoff_tools-1.14.4.dist-info → dayhoff_tools-1.14.6.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=eimShYucbCVs8MlLXdtPsxwQRPgXFXqABXIBE5kieJI,4305
+dayhoff_tools/batch/workers/base.py,sha256=Jie5ScJrPSRkmLQdcM2hPJTmk3pIeLL0oA8Y9pkP0HA,4269
 dayhoff_tools/batch/workers/boltz.py,sha256=QTIqMWuV_o3v1tYaZhvjyCl5F65ZlPcc0-TXtitMaFQ,14663
 dayhoff_tools/batch/workers/embed_t5.py,sha256=A5WqsQa1WZ7_la5X5wt0XUP-VwOglH04WyEINhwipeY,2750
 dayhoff_tools/chemistry/standardizer.py,sha256=uMn7VwHnx02nc404eO6fRuS4rsl4dvSPf2ElfZDXEpY,11188
@@ -10,10 +10,10 @@ dayhoff_tools/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU
 dayhoff_tools/cli/batch/__init__.py,sha256=gEuiDPWio8ihrIfF_Q6-5hKOnliPTxDKQrKN8-4Y3Ac,2320
 dayhoff_tools/cli/batch/aws_batch.py,sha256=L6lNpUaxyhMvLcM3dF8yEi2yCgTNdviaNMmG75bpBIc,12719
 dayhoff_tools/cli/batch/commands/__init__.py,sha256=1xRzzL_mc1hz1Pv0OWNr-g6fkL5XbEsOTGHzrqddLCA,458
-dayhoff_tools/cli/batch/commands/boltz.py,sha256=M7UNebaV7EAoHY_Fhml7JGnTZjM5o1nkYL22pSmqnWA,10422
+dayhoff_tools/cli/batch/commands/boltz.py,sha256=WARoklK5QTo71echUkm-anyj2AlEI9AEtWaiD_ds-EU,11759
 dayhoff_tools/cli/batch/commands/cancel.py,sha256=5FHLRoq8nvOr8HuKxXUZqAyAdjQP8seaNSaeAdhOoE8,2890
 dayhoff_tools/cli/batch/commands/embed_t5.py,sha256=UtF8ulJcz15sdd1QzvPVPeioN_AMks1DZ7c6vKtwhDE,11052
-dayhoff_tools/cli/batch/commands/finalize.py,sha256=2kQTtHxdW52s0lrLUlc7tx7Ab092sxs8uSc9-f2utUg,7198
+dayhoff_tools/cli/batch/commands/finalize.py,sha256=i9In09vrzc9jQlCA6F82ceuZB9mHUXeOaYYyRUrHQCc,7254
 dayhoff_tools/cli/batch/commands/list_jobs.py,sha256=y90_XqFI2qcvV2XUC01P3P2aIdtl5hXP3cBDEofsRtU,2432
 dayhoff_tools/cli/batch/commands/local.py,sha256=vOTojTAp4YZXIMLKNqc76oIE0oWMjB1mtTe6r19LrEc,3075
 dayhoff_tools/cli/batch/commands/logs.py,sha256=FOYpANQCsQTJKigq77k1xbhgOqtRPw9hyHkljILWlmo,5241
@@ -70,7 +70,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.4.dist-info/METADATA,sha256=DO3OUdkIlcX9vFR1LbE6BUB7uO-Qsmhi2EjhzA3LkXE,3184
-dayhoff_tools-1.14.4.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-dayhoff_tools-1.14.4.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
-dayhoff_tools-1.14.4.dist-info/RECORD,,
+dayhoff_tools-1.14.6.dist-info/METADATA,sha256=3rnOAUIQQN2JYIy6MtfSgAraglhmDY64Aq7fQWs-iRo,3184
+dayhoff_tools-1.14.6.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+dayhoff_tools-1.14.6.dist-info/entry_points.txt,sha256=iAf4jteNqW3cJm6CO6czLxjW3vxYKsyGLZ8WGmxamSc,49
+dayhoff_tools-1.14.6.dist-info/RECORD,,