dh-cli 0.8.2__tar.gz → 0.8.3__tar.gz

{dh_cli-0.8.2 → dh_cli-0.8.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dh-cli
-Version: 0.8.2
+Version: 0.8.3
 Summary: Dayhoff Labs developer CLI
 Author-email: Dayhoff Labs <dev@dayhofflabs.com>
 License: # PolyForm Noncommercial License 1.0.0

{dh_cli-0.8.2 → dh_cli-0.8.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dh-cli"
-version = "0.8.2"
+version = "0.8.3"
 description = "Dayhoff Labs developer CLI"
 requires-python = ">=3.11"
 readme = "README.md"

{dh_cli-0.8.2 → dh_cli-0.8.3}/src/dh_cli/batch/commands/finalize.py

@@ -1,6 +1,7 @@
 """Finalize command for combining results and cleaning up."""
 
 import shutil
+import tarfile
 import tempfile
 from pathlib import Path
 
@@ -94,7 +95,10 @@ def finalize(job_id, output, force, keep_intermediates, full_output, skip_dedup,
         s3_temp_dir = Path(tempfile.mkdtemp())
         s3_output_prefix = f"{manifest.s3_prefix}output/"
         click.echo("Downloading outputs from S3...")
-        download_directory(s3_output_prefix, s3_temp_dir)
+        if manifest.pipeline == "boltz":
+            _download_boltz_s3_output(s3_output_prefix, s3_temp_dir)
+        else:
+            download_directory(s3_output_prefix, s3_temp_dir)
         output_dir = s3_temp_dir
     else:
         output_dir = job_dir / "output"
@@ -175,6 +179,52 @@ def finalize(job_id, output, force, keep_intermediates, full_output, skip_dedup,
             click.echo(f"Job directory preserved: {job_dir}")
 
 
+def _download_boltz_s3_output(s3_output_prefix: str, local_dir: Path) -> None:
+    """Download Boltz tar outputs from S3 and extract into local_dir.
+
+    Workers produce one `boltz_results_<name>.tar` per prediction plus
+    `boltz_*.done` marker objects. This helper downloads only those keys
+    (ignoring anything else under the prefix), extracts each tar in place,
+    and removes the tar file afterwards so the resulting layout matches
+    what `_finalize_boltz` expects on Primordial.
+
+    Args:
+        s3_output_prefix: S3 URI prefix like s3://bucket/jobs/<id>/output/
+        local_dir: Local directory to extract into
+    """
+    from ..s3_transport import _get_client, parse_s3_uri
+
+    bucket, prefix_key = parse_s3_uri(s3_output_prefix)
+    client = _get_client()
+    local_dir.mkdir(parents=True, exist_ok=True)
+
+    paginator = client.get_paginator("list_objects_v2")
+    tar_count = 0
+    done_count = 0
+    for page in paginator.paginate(Bucket=bucket, Prefix=prefix_key):
+        for obj in page.get("Contents", []):
+            key = obj["Key"]
+            relative = key[len(prefix_key) :]
+            if not relative:
+                continue
+            basename = Path(relative).name
+            if basename.startswith("boltz_results_") and basename.endswith(".tar"):
+                tar_path = local_dir / basename
+                client.download_file(bucket, key, str(tar_path))
+                try:
+                    with tarfile.open(tar_path, mode="r") as tf:
+                        tf.extractall(local_dir)
+                finally:
+                    tar_path.unlink(missing_ok=True)
+                tar_count += 1
+            elif basename.startswith("boltz_") and basename.endswith(".done"):
+                done_path = local_dir / basename
+                client.download_file(bucket, key, str(done_path))
+                done_count += 1
+
+    click.echo(f"  Downloaded {tar_count} prediction tars, {done_count} done markers")
+
+
 def _check_completion(job_id: str, base_path: str, output_dir: Path | None = None) -> list[int]:
     """Check which chunks are incomplete (no .done marker).
 

dh_cli-0.8.3/tests/test_finalize_boltz_tar.py

@@ -0,0 +1,257 @@
+"""Tests for the Boltz S3 tar-aware finalize download path.
+
+When a Boltz job was run in S3 mode, workers upload one
+`boltz_results_<name>.tar` per prediction (plus per-worker done markers).
+Finalize must download those tars, extract them into `boltz_results_*/`
+directories matching the legacy on-disk layout, and leave the existing
+`_finalize_boltz` logic untouched.
+
+See plan: nutshell/plans/dma/05_2026/0512_boltz_s3_fanout_and_cross_az_dig.md.
+"""
+
+import tarfile
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from dh_cli.batch.commands.finalize import _download_boltz_s3_output
+
+
+@pytest.fixture
+def temp_dir():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        yield Path(tmpdir)
+
+
+def _build_essential_tar(tar_path: Path, complex_name: str) -> None:
+    """Build a boltz_results_<name>.tar like the worker produces."""
+    with tempfile.TemporaryDirectory() as src:
+        src_p = Path(src)
+        pred_subdir = src_p / f"boltz_results_{complex_name}" / "predictions" / complex_name
+        pred_subdir.mkdir(parents=True)
+        (pred_subdir / f"{complex_name}_model_0.cif").write_text(f"CIF {complex_name}\n")
+        (pred_subdir / f"confidence_{complex_name}_model_0.json").write_text(
+            f'{{"cx":"{complex_name}"}}'
+        )
+        with tarfile.open(tar_path, mode="w") as tf:
+            root = src_p / f"boltz_results_{complex_name}"
+            for f in sorted(root.rglob("*")):
+                if f.is_file():
+                    tf.add(f, arcname=f"boltz_results_{complex_name}/{f.relative_to(root)}")
+
+
+def _make_mock_s3_client(objects: dict[str, bytes]) -> MagicMock:
+    """Build a boto3 S3 client mock backed by an in-memory bucket.
+
+    `objects` is a mapping of key -> bytes. The mock implements:
+      - list_objects_v2 (via a paginator)
+      - download_file (writes object bytes to the local path)
+    """
+    client = MagicMock()
+
+    def _paginate(Bucket, Prefix, **kwargs):
+        matching = [k for k in objects if k.startswith(Prefix)]
+        yield {"Contents": [{"Key": k} for k in sorted(matching)]}
+
+    paginator = MagicMock()
+    paginator.paginate.side_effect = _paginate
+    client.get_paginator.return_value = paginator
+
+    def _download_file(bucket, key, local_path, *args, **kwargs):
+        Path(local_path).parent.mkdir(parents=True, exist_ok=True)
+        Path(local_path).write_bytes(objects[key])
+
+    client.download_file.side_effect = _download_file
+    return client
+
+
+class TestDownloadsAndExtracts:
+    def test_downloads_tars_and_extracts(self, temp_dir):
+        tar_a = temp_dir / "A.tar"
+        tar_b = temp_dir / "B.tar"
+        _build_essential_tar(tar_a, "A")
+        _build_essential_tar(tar_b, "B")
+
+        objects = {
+            "jobs/j/output/boltz_results_A.tar": tar_a.read_bytes(),
+            "jobs/j/output/boltz_results_B.tar": tar_b.read_bytes(),
+            "jobs/j/output/boltz_0.done": b"",
+        }
+        client = _make_mock_s3_client(objects)
+
+        dest = temp_dir / "extracted"
+        dest.mkdir()
+
+        with patch("dh_cli.batch.s3_transport._get_client", return_value=client):
+            _download_boltz_s3_output("s3://bucket/jobs/j/output/", dest)
+
+        assert (dest / "boltz_results_A" / "predictions" / "A" / "A_model_0.cif").read_text() == "CIF A\n"
+        assert (dest / "boltz_results_B" / "predictions" / "B" / "B_model_0.cif").read_text() == "CIF B\n"
+        assert (dest / "boltz_0.done").exists()
+
+    def test_local_tar_cleaned_up_after_extract(self, temp_dir):
+        tar_a = temp_dir / "A.tar"
+        _build_essential_tar(tar_a, "A")
+        objects = {"jobs/j/output/boltz_results_A.tar": tar_a.read_bytes()}
+        client = _make_mock_s3_client(objects)
+
+        dest = temp_dir / "extracted"
+        dest.mkdir()
+
+        with patch("dh_cli.batch.s3_transport._get_client", return_value=client):
+            _download_boltz_s3_output("s3://bucket/jobs/j/output/", dest)
+
+        leftover = list(dest.rglob("*.tar"))
+        assert leftover == []
+
+    def test_ignores_non_tar_non_done_keys(self, temp_dir):
+        """Extra objects in the output prefix that aren't tars or done markers are ignored."""
+        tar_a = temp_dir / "A.tar"
+        _build_essential_tar(tar_a, "A")
+        objects = {
+            "jobs/j/output/boltz_results_A.tar": tar_a.read_bytes(),
+            "jobs/j/output/stray.txt": b"ignore me",
+            "jobs/j/output/notes/readme.md": b"also ignored",
+        }
+        client = _make_mock_s3_client(objects)
+
+        dest = temp_dir / "extracted"
+        dest.mkdir()
+
+        with patch("dh_cli.batch.s3_transport._get_client", return_value=client):
+            _download_boltz_s3_output("s3://bucket/jobs/j/output/", dest)
+
+        assert (dest / "boltz_results_A" / "predictions" / "A" / "A_model_0.cif").exists()
+        assert not (dest / "stray.txt").exists()
+        assert not (dest / "notes").exists()
+
+
+class TestRoundTrip:
+    def test_worker_to_finalize_bit_identical(self, temp_dir):
+        """Worker's sync_boltz_essential_to_s3 output -> finalize download -> same bytes."""
+        from dh_batch.s3_sync import sync_boltz_essential_to_s3
+
+        job_dir = temp_dir / "job"
+        output_dir = job_dir / "output"
+        pred_a = output_dir / "boltz_results_A" / "predictions" / "A"
+        pred_a.mkdir(parents=True)
+        (pred_a / "A_model_0.cif").write_text("CIF ROUNDTRIP\n")
+        (pred_a / "confidence_A_model_0.json").write_text('{"r":1}')
+        (output_dir / "boltz_0.done").write_text("ok")
+
+        uploaded: dict[str, bytes] = {}
+        worker_client = MagicMock()
+
+        def _upload_file(local_path, bucket, key, *args, **kwargs):
+            uploaded[key] = Path(local_path).read_bytes()
+
+        worker_client.upload_file.side_effect = _upload_file
+
+        with patch("dh_batch.s3_transport._get_client", return_value=worker_client):
+            sync_boltz_essential_to_s3(job_dir, "s3://bucket/jobs/j/")
+
+        finalize_client = _make_mock_s3_client(uploaded)
+        dest = temp_dir / "extracted"
+        dest.mkdir()
+
+        with patch("dh_cli.batch.s3_transport._get_client", return_value=finalize_client):
+            _download_boltz_s3_output("s3://bucket/jobs/j/output/", dest)
+
+        cif = dest / "boltz_results_A" / "predictions" / "A" / "A_model_0.cif"
+        conf = dest / "boltz_results_A" / "predictions" / "A" / "confidence_A_model_0.json"
+        done = dest / "boltz_0.done"
+
+        assert cif.read_text() == "CIF ROUNDTRIP\n"
+        assert conf.read_text() == '{"r":1}'
+        assert done.read_text() == "ok"
+
+
+class TestFinalizeDispatches:
+    def test_boltz_s3_uses_tar_path(self, temp_dir):
+        """`dh batch finalize` for Boltz in S3 mode calls the tar-aware helper, not download_directory."""
+        from dh_cli.batch.commands.finalize import finalize as finalize_cmd
+        from dh_cli.batch.manifest import JobManifest, JobStatus
+
+        manifest = JobManifest(
+            job_id="test-boltz",
+            user="tester",
+            pipeline="boltz",
+            storage_mode="s3",
+            status=JobStatus.SUCCEEDED,
+            s3_prefix="s3://bucket/jobs/test-boltz/",
+        )
+
+        from click.testing import CliRunner
+
+        runner = CliRunner()
+
+        with patch("dh_cli.batch.commands.finalize.load_manifest", return_value=manifest):
+            with patch(
+                "dh_cli.batch.commands.finalize._download_boltz_s3_output"
+            ) as mock_tar_download:
+                with patch(
+                    "dh_cli.batch.commands.finalize._check_completion", return_value=[]
+                ):
+                    with patch("dh_cli.batch.commands.finalize._finalize_boltz") as mock_fb:
+                        with patch("dh_cli.batch.commands.finalize.save_manifest_s3"):
+                            result = runner.invoke(
+                                finalize_cmd,
+                                [
+                                    "test-boltz",
+                                    "--output",
+                                    str(temp_dir / "final"),
+                                    "--keep-intermediates",
+                                ],
+                                input="y\n",
+                            )
+
+        assert result.exit_code == 0, result.output
+        assert mock_tar_download.called, "tar-aware download helper must be called for Boltz S3"
+        assert mock_fb.called
+
+    def test_non_boltz_s3_uses_download_directory(self, temp_dir):
+        """Non-Boltz pipelines (e.g. embed-t5) in S3 mode keep using download_directory."""
+        from dh_cli.batch.commands.finalize import finalize as finalize_cmd
+        from dh_cli.batch.manifest import JobManifest, JobStatus
+
+        manifest = JobManifest(
+            job_id="test-embed",
+            user="tester",
+            pipeline="embed-t5",
+            storage_mode="s3",
+            status=JobStatus.SUCCEEDED,
+            s3_prefix="s3://bucket/jobs/test-embed/",
+        )
+
+        from click.testing import CliRunner
+
+        runner = CliRunner()
+
+        with patch("dh_cli.batch.commands.finalize.load_manifest", return_value=manifest):
+            with patch(
+                "dh_cli.batch.commands.finalize._download_boltz_s3_output"
+            ) as mock_tar_download:
+                with patch("dh_cli.batch.s3_transport.download_directory") as mock_dd:
+                    with patch(
+                        "dh_cli.batch.commands.finalize._check_completion", return_value=[]
+                    ):
+                        with patch(
+                            "dh_cli.batch.commands.finalize._finalize_embeddings"
+                        ) as mock_fe:
+                            with patch("dh_cli.batch.commands.finalize.save_manifest_s3"):
+                                result = runner.invoke(
+                                    finalize_cmd,
+                                    [
+                                        "test-embed",
+                                        "--output",
+                                        str(temp_dir / "out.h5"),
+                                        "--keep-intermediates",
+                                    ],
+                                )
+
+        assert result.exit_code == 0, result.output
+        assert mock_dd.called, "non-Boltz pipelines should keep using download_directory"
+        assert not mock_tar_download.called, "tar-aware helper must not fire for embed-t5"
+        assert mock_fe.called