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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dh-cli
-Version: 0.8.1
+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.1 → dh_cli-0.8.3}/pyproject.toml

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

{dh_cli-0.8.1 → 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.1 → dh_cli-0.8.3}/src/dh_cli/bedrock/commands.py

@@ -304,8 +304,11 @@ def bedrock_cost(
         sys.exit(1)
 
     import boto3
+    from botocore.config import Config
 
-    s3 = boto3.client("s3")
+    # Match the thread pool used by walk_logs so urllib3 doesn't block
+    # or warn when many parallel GETs are in flight.
+    s3 = boto3.client("s3", config=Config(max_pool_connections=32))
 
     my_handle: Optional[str] = None
     if me:

{dh_cli-0.8.1 → dh_cli-0.8.3}/src/dh_cli/bedrock/cost_report.py

@@ -27,6 +27,7 @@ from __future__ import annotations
 import datetime as dt
 import gzip
 import json
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Iterable, Iterator
@@ -443,14 +444,46 @@ def walk_logs(
     region: str,
     start: dt.date,
     end: dt.date,
+    max_workers: int = 32,
 ) -> Iterator[dict]:
+    """Yield every invocation record in `[start, end]` (inclusive, UTC days).
+
+    Object GETs are parallelised with a thread pool because each day's
+    prefix holds hundreds of tiny (~400-byte) gzipped objects and
+    per-request latency dominates wall time. Records within a single
+    object are yielded in their original NDJSON order; records *across*
+    objects may be reordered — downstream aggregation (`build_report`)
+    is order-insensitive.
+
+    `max_workers` caps in-flight S3 GETs per day. The caller's
+    `s3_client` should be configured with `max_pool_connections` >=
+    `max_workers` (see `botocore.config.Config`) to avoid urllib3
+    connection-pool contention.
+    """
     paginator = s3_client.get_paginator("list_objects_v2")
     seen_keys: set[str] = set()
+
+    def _fetch_and_parse(key: str) -> list[dict]:
+        body = s3_client.get_object(Bucket=bucket, Key=key)["Body"].read()
+        decompressed = gzip.decompress(body)
+        out: list[dict] = []
+        # Each object is one or more JSON records separated by
+        # newlines (NDJSON). Older Bedrock traffic produced
+        # one-record objects; multi-record objects appeared in
+        # our bucket on 2026-04-20. Parse line-by-line so both
+        # shapes work, and tolerate a trailing newline.
+        for line in decompressed.splitlines():
+            if not line.strip():
+                continue
+            out.append(json.loads(line))
+        return out
+
     for day in _iter_days(start, end):
         prefix = (
             f"invocation-logs/AWSLogs/{account}/BedrockModelInvocationLogs/"
             f"{region}/{day.year:04d}/{day.month:02d}/{day.day:02d}/"
         )
+        keys: list[str] = []
         for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
             for obj in page.get("Contents", []) or []:
                 key = obj["Key"]
@@ -461,17 +494,18 @@ def walk_logs(
                 if key in seen_keys:
                     continue
                 seen_keys.add(key)
-                body = s3_client.get_object(Bucket=bucket, Key=key)["Body"].read()
-                decompressed = gzip.decompress(body)
-                # Each object is one or more JSON records separated by
-                # newlines (NDJSON). Older Bedrock traffic produced
-                # one-record objects; multi-record objects appeared in
-                # our bucket on 2026-04-20. Parse line-by-line so both
-                # shapes work, and tolerate a trailing newline.
-                for line in decompressed.splitlines():
-                    if not line.strip():
-                        continue
-                    yield json.loads(line)
+                keys.append(key)
+        if not keys:
+            continue
+        # One pool per day bounds concurrent in-flight GETs and caps
+        # peak memory (at most ~max_workers decompressed objects held
+        # at once). ex.map preserves submission order, so the day's
+        # records stream out in a stable — though not chronological —
+        # order.
+        with ThreadPoolExecutor(max_workers=max_workers) as ex:
+            for records in ex.map(_fetch_and_parse, keys):
+                for rec in records:
+                    yield rec
 
 
 def reconcile_with_cost_explorer(

{dh_cli-0.8.1 → dh_cli-0.8.3}/tests/batch/test_submit_merge.py

@@ -91,9 +91,7 @@ class TestCliWinsWhenExplicitlyPassed:
         "cli_flag,kwarg,cli_default,yaml_alt,_",
         MERGE_FIELDS,
     )
-    def test_cli_at_default_beats_yaml(
-        self, cli_runner, tmp_path, cli_flag, kwarg, cli_default, yaml_alt, _
-    ):
+    def test_cli_at_default_beats_yaml(self, cli_runner, tmp_path, cli_flag, kwarg, cli_default, yaml_alt, _):
         """CLI flag set to the Click default still wins over a different YAML value."""
         yaml_key = cli_flag.lstrip("-")
         config_path = tmp_path / "job.yaml"
@@ -112,9 +110,7 @@ class TestCliWinsWhenExplicitlyPassed:
         "cli_flag,kwarg,_,yaml_alt,cli_alt",
         MERGE_FIELDS,
     )
-    def test_cli_at_non_default_beats_yaml(
-        self, cli_runner, tmp_path, cli_flag, kwarg, _, yaml_alt, cli_alt
-    ):
+    def test_cli_at_non_default_beats_yaml(self, cli_runner, tmp_path, cli_flag, kwarg, _, yaml_alt, cli_alt):
         """CLI flag at a non-default value also wins over YAML (regression check)."""
         yaml_key = cli_flag.lstrip("-")
         config_path = tmp_path / "job.yaml"
@@ -137,9 +133,7 @@ class TestYamlWinsWhenCliOmitted:
         "cli_flag,kwarg,_,yaml_alt,__",
         MERGE_FIELDS,
     )
-    def test_yaml_wins_when_cli_not_passed(
-        self, cli_runner, tmp_path, cli_flag, kwarg, _, yaml_alt, __
-    ):
+    def test_yaml_wins_when_cli_not_passed(self, cli_runner, tmp_path, cli_flag, kwarg, _, yaml_alt, __):
         yaml_key = cli_flag.lstrip("-")
         config_path = tmp_path / "job.yaml"
         config_path.write_text(yaml.dump({"command": "echo hi", yaml_key: yaml_alt}))
@@ -157,9 +151,7 @@ class TestDefaultWhenNeitherSet:
         "_,kwarg,cli_default,__,___",
         MERGE_FIELDS,
     )
-    def test_click_default_applies(
-        self, cli_runner, tmp_path, _, kwarg, cli_default, __, ___
-    ):
+    def test_click_default_applies(self, cli_runner, tmp_path, _, kwarg, cli_default, __, ___):
         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]
@@ -218,9 +210,7 @@ class TestGpuRegressionFromReport:
 
     def test_explicit_gpus_one_beats_yaml_gpus_zero(self, cli_runner, tmp_path):
         config_path = tmp_path / "job.yaml"
-        config_path.write_text(
-            yaml.dump({"command": "echo hi", "queue": "t4-1x-spot", "gpus": 0})
-        )
+        config_path.write_text(yaml.dump({"command": "echo hi", "queue": "t4-1x-spot", "gpus": 0}))
         result, mock_client = _invoke(
             cli_runner,
             ["-f", str(config_path), "--gpus", "1"],

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