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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dayhoff-tools
-Version: 1.14.3
+Version: 1.14.4
 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.4}/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.4}/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.4}/pyproject.toml

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