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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dh-cli
-Version: 0.8.2
+Version: 0.8.4
 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.4}/pyproject.toml

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

{dh_cli-0.8.2 → dh_cli-0.8.4}/src/dh_cli/batch/aws_batch.py

@@ -17,6 +17,11 @@ from dh_cli.batch.manifest import (
 
 logger = logging.getLogger(__name__)
 
+_ECR_IMAGE_RE = re.compile(
+    r"^\d{12}\.dkr\.ecr\.[a-z0-9-]+\.amazonaws\.com/[^:@\s]+"
+    r"(?::[^\s]+|@sha256:[a-f0-9]{64})?$"
+)
+
 
 class BatchError(Exception):
     """Error interacting with AWS Batch."""
@@ -165,14 +170,27 @@ class BatchClient:
         create a new revision of the base definition with only the image
         changed, preserving all other containerProperties (roles, volumes, etc.).
 
-        Returns the ARN of the new revision.
+        Returns the ARN of the new revision, or the ARN of the latest existing
+        revision if its image already matches.
         """
+        if not _ECR_IMAGE_RE.match(image):
+            raise BatchError(
+                f"Image override must be a fully-qualified ECR URL "
+                f"(<account>.dkr.ecr.<region>.amazonaws.com/<repo>[:tag|@digest]); "
+                f"got: {image!r}"
+            )
+
         resp = self.batch.describe_job_definitions(jobDefinitionName=base_definition, status="ACTIVE")
         definitions = resp.get("jobDefinitions", [])
         if not definitions:
             raise BatchError(f"Job definition not found: {base_definition}")
 
         latest = sorted(definitions, key=lambda d: d["revision"])[-1]
+        if latest["containerProperties"].get("image") == image:
+            arn = latest["jobDefinitionArn"]
+            logger.info(f"Reusing job definition {arn}; image already matches {image}")
+            return arn
+
         container_props = latest["containerProperties"].copy()
         container_props["image"] = image
 

{dh_cli-0.8.2 → dh_cli-0.8.4}/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.2 → dh_cli-0.8.4}/src/dh_cli/batch/commands/submit.py

@@ -4,7 +4,7 @@ import click
 import yaml
 from click.core import ParameterSource
 
-from ..aws_batch import BatchClient, BatchError, resolve_dependency
+from ..aws_batch import _ECR_IMAGE_RE, BatchClient, BatchError, resolve_dependency
 from ..job_id import generate_job_id, get_aws_username
 from ..manifest import (
     BATCH_JOBS_BASE,
@@ -114,6 +114,13 @@ def submit(
     job_retry = _pick("retry", retry)
     job_timeout = _pick("timeout", timeout)
     job_image = image or config.get("image")
+    if job_image and not _ECR_IMAGE_RE.match(job_image):
+        raise click.BadParameter(
+            f"--image must be a fully-qualified ECR URL "
+            f"(<account>.dkr.ecr.<region>.amazonaws.com/<repo>[:tag|@digest]); "
+            f"got: {job_image!r}",
+            param_hint="--image",
+        )
 
     # Parse environment variables
     job_env = dict(config.get("env", {}))

dh_cli-0.8.4/tests/batch/test_image_override.py

@@ -0,0 +1,141 @@
+"""Tests for BatchClient._register_image_override.
+
+Pins two contracts:
+
+1. Bare names (anything not a fully-qualified ECR URL) are rejected, so
+   typos like ``dh batch submit --image dayhoff-generic`` cannot drift the
+   floating ref by registering a revision pointing at Docker Hub.
+2. When the latest active revision already carries the requested image,
+   no new revision is registered — the existing ARN is returned. Without
+   this, every override invocation bumps the JD revision counter.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+def _make_client():
+    with patch("dh_cli.batch.aws_batch.boto3") as mock_boto3:
+        mock_batch = MagicMock()
+        mock_logs = MagicMock()
+        mock_boto3.client.side_effect = [mock_batch, mock_logs]
+
+        from dh_cli.batch.aws_batch import BatchClient
+
+        return BatchClient(), mock_batch
+
+
+_ECR_URL = "123456789012.dkr.ecr.us-east-1.amazonaws.com/dayhoff-generic:abc123"
+
+
+class TestImageValidation:
+    """Bare image names must be rejected before any AWS call."""
+
+    @pytest.mark.parametrize(
+        "bad_image",
+        [
+            "dayhoff-generic",
+            "dayhoff-generic:latest",
+            "docker.io/library/python:3.12",
+            "ghcr.io/foo/bar:1.0",
+            "123456789012.dkr.ecr.us-east-1.amazonaws.com/dayhoff-generic with-space",
+        ],
+    )
+    def test_rejects_non_ecr_image(self, bad_image):
+        from dh_cli.batch.aws_batch import BatchError
+
+        client, mock_batch = _make_client()
+        with pytest.raises(BatchError, match="fully-qualified ECR URL"):
+            client._register_image_override("dayhoff-generic", bad_image)
+
+        mock_batch.describe_job_definitions.assert_not_called()
+        mock_batch.register_job_definition.assert_not_called()
+
+    @pytest.mark.parametrize(
+        "good_image",
+        [
+            "123456789012.dkr.ecr.us-east-1.amazonaws.com/dayhoff-generic:abc123",
+            "123456789012.dkr.ecr.us-west-2.amazonaws.com/foo/bar:latest",
+            "123456789012.dkr.ecr.us-east-1.amazonaws.com/x@sha256:"
+            + "a" * 64,
+        ],
+    )
+    def test_accepts_ecr_url(self, good_image):
+        client, mock_batch = _make_client()
+        mock_batch.describe_job_definitions.return_value = {
+            "jobDefinitions": [
+                {
+                    "jobDefinitionArn": "arn:aws:batch:...:job-definition/dayhoff-generic:1",
+                    "revision": 1,
+                    "type": "container",
+                    "containerProperties": {"image": "old-image"},
+                }
+            ]
+        }
+        mock_batch.register_job_definition.return_value = {
+            "jobDefinitionArn": "arn:aws:batch:...:job-definition/dayhoff-generic:2"
+        }
+
+        client._register_image_override("dayhoff-generic", good_image)
+        mock_batch.register_job_definition.assert_called_once()
+
+
+class TestNoOpReuse:
+    """If latest revision already matches the requested image, reuse it."""
+
+    def test_reuses_latest_when_image_matches(self):
+        client, mock_batch = _make_client()
+        mock_batch.describe_job_definitions.return_value = {
+            "jobDefinitions": [
+                {
+                    "jobDefinitionArn": "arn:...:job-definition/dayhoff-generic:42",
+                    "revision": 42,
+                    "type": "container",
+                    "containerProperties": {"image": _ECR_URL},
+                },
+                {
+                    "jobDefinitionArn": "arn:...:job-definition/dayhoff-generic:41",
+                    "revision": 41,
+                    "type": "container",
+                    "containerProperties": {"image": "older"},
+                },
+            ]
+        }
+
+        arn = client._register_image_override("dayhoff-generic", _ECR_URL)
+
+        assert arn == "arn:...:job-definition/dayhoff-generic:42"
+        mock_batch.register_job_definition.assert_not_called()
+
+    def test_registers_new_revision_when_image_differs(self):
+        client, mock_batch = _make_client()
+        mock_batch.describe_job_definitions.return_value = {
+            "jobDefinitions": [
+                {
+                    "jobDefinitionArn": "arn:...:job-definition/dayhoff-generic:42",
+                    "revision": 42,
+                    "type": "container",
+                    "containerProperties": {"image": "different-image"},
+                }
+            ]
+        }
+        mock_batch.register_job_definition.return_value = {
+            "jobDefinitionArn": "arn:...:job-definition/dayhoff-generic:43"
+        }
+
+        arn = client._register_image_override("dayhoff-generic", _ECR_URL)
+
+        assert arn == "arn:...:job-definition/dayhoff-generic:43"
+        mock_batch.register_job_definition.assert_called_once()
+        kwargs = mock_batch.register_job_definition.call_args[1]
+        assert kwargs["containerProperties"]["image"] == _ECR_URL
+
+    def test_raises_when_no_active_definition(self):
+        from dh_cli.batch.aws_batch import BatchError
+
+        client, mock_batch = _make_client()
+        mock_batch.describe_job_definitions.return_value = {"jobDefinitions": []}
+
+        with pytest.raises(BatchError, match="Job definition not found"):
+            client._register_image_override("dayhoff-generic", _ECR_URL)

dh_cli-0.8.4/tests/batch/test_submit_image_validation.py

@@ -0,0 +1,90 @@
+"""CLI-layer image validation tests for `dh batch submit`.
+
+Pins the contract: a non-ECR --image (or YAML image: ...) is rejected at
+parse time, before any AWS call. Without this, bare names like
+``--image dayhoff-generic`` reach BatchClient and silently mint a JD
+revision pointing at Docker Hub.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+import yaml
+from click.testing import CliRunner
+
+
+@pytest.fixture
+def cli_runner():
+    return CliRunner()
+
+
+def _invoke(cli_runner, args, tmp_path):
+    base = tmp_path / "jobs"
+    with (
+        patch("dh_cli.batch.commands.submit.get_aws_username", return_value="jason"),
+        patch("dh_cli.batch.commands.submit.BatchClient") as mock_batch_cls,
+        patch(
+            "dh_cli.batch.commands.submit.generate_job_id",
+            return_value="jason-batch-20260519-img00001",
+        ),
+    ):
+        mock_client = MagicMock()
+        mock_client.submit_job.return_value = "aws-uuid-img"
+        mock_batch_cls.return_value = mock_client
+
+        from dh_cli.batch.commands.submit import submit
+
+        result = cli_runner.invoke(submit, args + ["--base-path", str(base)])
+        return result, mock_client, mock_batch_cls
+
+
+class TestCliImageValidation:
+    @pytest.mark.parametrize(
+        "bad_image",
+        ["dayhoff-generic", "dayhoff-generic:latest", "ghcr.io/foo/bar:1.0"],
+    )
+    def test_bare_image_rejected_before_submit(self, cli_runner, tmp_path, bad_image):
+        result, _, mock_batch_cls = _invoke(
+            cli_runner,
+            ["--command", "echo hi", "--image", bad_image],
+            tmp_path,
+        )
+        assert result.exit_code != 0
+        assert "fully-qualified ECR URL" in result.output
+        mock_batch_cls.assert_not_called()
+
+    def test_yaml_bare_image_rejected_before_submit(self, cli_runner, tmp_path):
+        config_path = tmp_path / "job.yaml"
+        config_path.write_text(
+            yaml.dump(
+                {
+                    "command": "echo hi",
+                    "image": "dayhoff-generic",
+                }
+            )
+        )
+        result, _, mock_batch_cls = _invoke(
+            cli_runner, ["-f", str(config_path)], tmp_path
+        )
+        assert result.exit_code != 0
+        assert "fully-qualified ECR URL" in result.output
+        mock_batch_cls.assert_not_called()
+
+    def test_valid_ecr_url_accepted(self, cli_runner, tmp_path):
+        good = "123456789012.dkr.ecr.us-east-1.amazonaws.com/dayhoff-generic:abc"
+        result, mock_client, _ = _invoke(
+            cli_runner,
+            ["--command", "echo hi", "--image", good],
+            tmp_path,
+        )
+        assert result.exit_code == 0, result.output
+        call_kwargs = mock_client.submit_job.call_args[1]
+        assert call_kwargs["image_override"] == good
+
+    def test_no_image_skips_validation(self, cli_runner, tmp_path):
+        result, mock_client, _ = _invoke(
+            cli_runner, ["--command", "echo hi"], tmp_path
+        )
+        assert result.exit_code == 0, result.output
+        call_kwargs = mock_client.submit_job.call_args[1]
+        assert call_kwargs.get("image_override") is None

dh_cli-0.8.4/tests/test_finalize_boltz_tar.py

@@ -0,0 +1,245 @@
+"""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