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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dayhoff-tools
-Version: 1.14.3
+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.3 → 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.3 → 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.3 → dayhoff_tools-1.14.5}/dayhoff_tools/cli/batch/commands/embed_t5.py

@@ -1,6 +1,8 @@
 """T5 embedding pipeline command."""
 
 import os
+import shutil
+import subprocess
 from pathlib import Path
 
 import click
@@ -232,38 +234,76 @@ def _submit_batch_job(input_path: Path, workers: int, queue: str, seqs_per_chunk
 
 
 def _run_local_mode(input_path: Path):
-    """Run embedding locally for a single chunk."""
-    import subprocess
+    """Run embedding locally in a Docker container.
 
-    click.echo("Running T5 embedding locally...")
+    This runs the embed container with the normal entrypoint, processing
+    a single chunk (index 0) for testing purposes.
+    """
+    click.echo("Running T5 embedding locally in container...")
     click.echo(f"Input: {input_path}")
 
-    # Check if we have the embedder available
-    try:
-        from dayhoff_tools.embedders import T5Embedder
+    input_dir = input_path.parent
 
-        output_file = input_path.with_suffix(".h5")
-        click.echo(f"Output: {output_file}")
-        click.echo()
+    # Create a temporary job directory structure in the input directory
+    # The worker expects JOB_DIR/input/chunk_000.fasta format
+    temp_job_dir = input_dir / ".local_embed_job"
+    temp_input_dir = temp_job_dir / "input"
+    temp_output_dir = temp_job_dir / "output"
 
-        embedder = T5Embedder(
-            max_seq_length=4500,
-            large_protein_threshold=2500,
-            batch_residue_limit=4500,
-        )
-        embedder.run(str(input_path), str(output_file))
+    # Clean up any previous run
+    if temp_job_dir.exists():
+        shutil.rmtree(temp_job_dir)
 
-        click.echo()
-        click.echo(click.style("✓ Embedding complete!", fg="green"))
-        click.echo(f"Output: {output_file}")
+    temp_input_dir.mkdir(parents=True)
+    temp_output_dir.mkdir(parents=True)
+
+    # Symlink or copy the input file as chunk_000.fasta
+    chunk_path = temp_input_dir / "chunk_000.fasta"
+    chunk_path.symlink_to(input_path.resolve())
+
+    click.echo(f"Output will be at: {temp_output_dir}/embed_000.h5")
+    click.echo()
+
+    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",
+        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_file = temp_output_dir / "embed_000.h5"
+        if output_file.exists():
+            # Move output to final location
+            final_output = input_path.with_suffix(".h5")
+            shutil.move(str(output_file), str(final_output))
+            click.echo()
+            click.echo(click.style("✓ Embedding complete!", fg="green"))
+            click.echo(f"Output: {final_output}")
+        else:
+            click.echo(click.style("Warning: No output file found", fg="yellow"))
+
+        # Clean up temp directory
+        shutil.rmtree(temp_job_dir, ignore_errors=True)
 
-    except ImportError:
+    except FileNotFoundError:
         click.echo(
-            click.style(
-                "Error: T5Embedder requires 'embedders' extra. "
-                "Install with: pip install 'dayhoff-tools[embedders]'",
-                fg="red",
-            ),
+            click.style("Error: Docker not found. Is Docker installed and running?", fg="red"),
             err=True,
         )
         raise SystemExit(1)
@@ -271,8 +311,6 @@ def _run_local_mode(input_path: Path):
 
 def _run_shell_mode(input_path: Path):
     """Drop into container shell for debugging."""
-    import subprocess
-
     click.echo("Dropping into container shell...")
     click.echo(f"Input will be available at: /input/{input_path.name}")
     click.echo()

{dayhoff_tools-1.14.3 → dayhoff_tools-1.14.5}/dayhoff_tools/cli/engines_studios/studio_commands.py

@@ -434,7 +434,10 @@ def list_studios(env: Optional[str]):
     help="Environment (dev, sand, prod) - auto-detected if not specified",
 )
 def attach_studio(
-    engine_name_or_id: str, skip_ssh_config: bool, user: Optional[str], env: Optional[str]
+    engine_name_or_id: str,
+    skip_ssh_config: bool,
+    user: Optional[str],
+    env: Optional[str],
 ):
     """Attach your studio to an engine with progress tracking."""
 

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

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