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

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

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

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

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

dh_cli-0.8.2/src/dh_cli/_identity.py

@@ -0,0 +1,88 @@
+"""Identity resolution for `dh` commands that read per-developer secrets.
+
+The `github_commands` and (future) `bedrock` commands both key per-user
+Secrets Manager entries on the caller's Dayhoff handle. This module
+resolves that handle from the current SSO session; the server-side
+resource policy makes the matching decision using `aws:userid` suffix-
+matching on the same handle (see
+blueprints/terraform/environments/dev/github_pat_secrets.tf header for
+the full story of why aws:userid, not aws:PrincipalTag/Email or
+aws:username or aws:PrincipalArn).
+
+Design note — where the handle comes from:
+
+    There is no SDK API that lets a session read its own `aws:userid`
+    directly in a structured way. The closest observable is the
+    assumed-role ARN's RoleSessionName, which IAM Identity Center sets
+    to the login username (the handle) for every DeveloperAccess
+    session — exactly the same string that IAM populates as the suffix
+    of `aws:userid` during policy evaluation. So "what handle am I?"
+    (this function) and "which secret can I read?" (the server-side
+    policy) are answered by the same identity fact by construction.
+
+    For historical reasons the RoleSessionName in this org is a bare
+    handle like `dma`, not `dma@dayhofflabs.com`. If Identity Center is
+    ever reconfigured to use emails as session names — or if ABAC is
+    turned on and the server-side policy flips to
+    aws:PrincipalTag/Email — the optional `domain` argument still
+    handles the email-style case (strips the suffix) without code
+    changes here.
+"""
+
+from __future__ import annotations
+
+import re
+
+# Kept for the email-style RoleSessionName case; unused in the current
+# org (handles are bare). If you ever need to reintroduce domain-
+# stripping, pass domain="dayhofflabs.com" explicitly.
+DEFAULT_DOMAIN = "dayhofflabs.com"
+
+_SSO_ASSUMED_ROLE_RE = re.compile(r"^arn:aws:sts::\d+:assumed-role/AWSReservedSSO_[^/]+/(?P<session>.+)$")
+
+
+class HandleResolutionError(RuntimeError):
+    """Raised when the current session's handle can't be determined.
+
+    The caller is expected to turn this into a user-facing error
+    pointing at `awslogin dev-devaccess`.
+    """
+
+
+def resolve_handle_from_session(session, *, domain: str = DEFAULT_DOMAIN) -> str:
+    """Return the dev handle (session-name portion of aws:userid).
+
+    Args:
+        session: a boto3 Session configured with the caller's SSO
+            credentials. The function calls `sts:GetCallerIdentity` on
+            it; if that call would fall through to the engine instance
+            role instead of the dev's SSO creds, the caller must have
+            already detected and errored on that before getting here
+            (see `github_commands._sso_session`).
+        domain: email domain to strip from the session name, for orgs
+            where IAM Identity Center uses emails as RoleSessionNames.
+            No-op for bare-handle RoleSessionNames (the current org).
+
+    Returns:
+        The handle (e.g. `"dma"`). This is the same string that appears
+        as the suffix of `aws:userid` during IAM policy evaluation for
+        the caller's session — i.e. the value the server-side policy
+        matches against in the `StringLike "aws:userid": "*:<handle>"`
+        condition.
+
+    Raises:
+        HandleResolutionError: the caller's ARN doesn't look like an
+            Identity Center DeveloperAccess session.
+    """
+    arn = session.client("sts").get_caller_identity()["Arn"]
+    match = _SSO_ASSUMED_ROLE_RE.match(arn)
+    if not match:
+        raise HandleResolutionError(
+            f"Caller ARN does not look like an AWS SSO session: {arn}. "
+            f"Run `awslogin dev-devaccess` (or pass --handle explicitly)."
+        )
+    session_name = match.group("session")
+    suffix = f"@{domain}"
+    if session_name.endswith(suffix):
+        return session_name[: -len(suffix)]
+    return session_name

{dh_cli-0.8.0 → dh_cli-0.8.2}/src/dh_cli/batch/commands/submit.py

@@ -2,6 +2,7 @@
 
 import click
 import yaml
+from click.core import ParameterSource
 
 from ..aws_batch import BatchClient, BatchError, resolve_dependency
 from ..job_id import generate_job_id, get_aws_username
@@ -34,7 +35,9 @@ DEFAULT_QUEUE = "t4-1x-spot"
 @click.option("--dry-run", is_flag=True, help="Show plan without submitting")
 @click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
 @click.option("--after", "after", multiple=True, help="Job ID(s) to wait for before starting")
+@click.pass_context
 def submit(
+    ctx,
     config_file,
     command,
     queue,
@@ -52,7 +55,10 @@ def submit(
 ):
     """Submit a custom batch job.
 
-    Jobs can be defined via a config file (-f) or inline options.
+    Jobs can be defined via a config file (-f) or inline options. When
+    both are provided, a CLI flag takes precedence over the
+    corresponding YAML field only if the user actually passes the
+    flag; otherwise the YAML value wins.
 
     \b
     Examples:
@@ -90,13 +96,23 @@ def submit(
     if not job_command:
         raise click.UsageError("Must specify --command or provide config file with 'command' field")
 
-    job_queue = queue if queue != DEFAULT_QUEUE else config.get("queue", queue)
-    job_memory = memory if memory != "30G" else config.get("memory", memory)
-    job_vcpus = vcpus if vcpus != 8 else config.get("vcpus", vcpus)
-    job_gpus = gpus if gpus != 1 else config.get("gpus", gpus)
-    job_array = array if array != 1 else config.get("array", array)
-    job_retry = retry if retry != 3 else config.get("retry", retry)
-    job_timeout = timeout if timeout != "6h" else config.get("timeout", timeout)
+    # Merge CLI flags with YAML. CLI wins iff the user actually passed the
+    # flag; otherwise YAML if set; otherwise the Click default. Uses
+    # ParameterSource to tell "user typed --gpus 1" from "Click filled in
+    # the default 1", which a bare value comparison cannot do.
+    def _pick(param_name, cli_value, yaml_key=None):
+        yaml_key = yaml_key or param_name
+        if ctx.get_parameter_source(param_name) == ParameterSource.COMMANDLINE:
+            return cli_value
+        return config.get(yaml_key, cli_value)
+
+    job_queue = _pick("queue", queue)
+    job_memory = _pick("memory", memory)
+    job_vcpus = _pick("vcpus", vcpus)
+    job_gpus = _pick("gpus", gpus)
+    job_array = _pick("array", array)
+    job_retry = _pick("retry", retry)
+    job_timeout = _pick("timeout", timeout)
     job_image = image or config.get("image")
 
     # Parse environment variables

{dh_cli-0.8.0 → dh_cli-0.8.2}/src/dh_cli/bedrock/__init__.py

@@ -1,4 +1,5 @@
 """`dh bedrock` command group — key delivery + per-user cost reporting."""
+
 from .commands import bedrock
 
 __all__ = ["bedrock"]

{dh_cli-0.8.0 → dh_cli-0.8.2}/src/dh_cli/bedrock/commands.py

@@ -12,6 +12,7 @@ Two user-facing commands:
 Both commands default to reading the caller's identity via STS to
 resolve their own handle, so the common case is parameter-free.
 """
+
 from __future__ import annotations
 
 import datetime as dt
@@ -82,8 +83,7 @@ def _resolve_handle_from_sts() -> str:
     if principal.principal_type in ("claude-code", "cursor"):
         return principal.principal_name
     raise click.ClickException(
-        f"Couldn't infer a developer handle from your identity ({arn}). "
-        "Pass --handle explicitly."
+        f"Couldn't infer a developer handle from your identity ({arn}). Pass --handle explicitly."
     )
 
 
@@ -153,9 +153,7 @@ def bedrock_key(handle: Optional[str], region: str, mode: str):
         ak = payload.get("access_key_id", "")
         sk = payload.get("secret_access_key", "")
         if not ak or not sk:
-            raise click.ClickException(
-                f"Secret `{secret_id}` is missing access_key_id/secret_access_key fields."
-            )
+            raise click.ClickException(f"Secret `{secret_id}` is missing access_key_id/secret_access_key fields.")
         click.echo(f"export AWS_ACCESS_KEY_ID='{ak}'")
         click.echo(f"export AWS_SECRET_ACCESS_KEY='{sk}'")
         click.echo(f"export AWS_DEFAULT_REGION='{payload.get('region', region)}'")
@@ -296,9 +294,7 @@ def bedrock_cost(
     if start is None:
         start = end - dt.timedelta(days=days - 1)
     if start > end:
-        raise click.BadParameter(
-            f"--start ({start}) must be on or before --end ({end})."
-        )
+        raise click.BadParameter(f"--start ({start}) must be on or before --end ({end}).")
 
     pricing_file = pricing_path or cr.default_pricing_path()
     try:
@@ -308,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:
@@ -329,10 +328,9 @@ def bedrock_cost(
             # mode, including 'model' and 'principal_type' which collapse
             # principal_name to "" in the output rows.
             records = (
-                rec for rec in records
-                if cr.classify_arn(
-                    rec.get("identity", {}).get("arn", "")
-                ).principal_name == my_handle
+                rec
+                for rec in records
+                if cr.classify_arn(rec.get("identity", {}).get("arn", "")).principal_name == my_handle
             )
         report = cr.build_report(records, pricing, group_by=group_by)
     except cr.UnknownModel as exc:
@@ -393,16 +391,10 @@ def bedrock_cost(
     # Keep reconcile output minimal in csv/markdown modes so the body
     # of the output stays pipe-friendly — stderr, not stdout.
     stream_err = output_format in ("csv", "markdown")
-    delta_pct = (
-        f"{result.delta_fraction * 100:.1f}%"
-        if result.delta_fraction != float("inf")
-        else "n/a"
-    )
+    delta_pct = f"{result.delta_fraction * 100:.1f}%" if result.delta_fraction != float("inf") else "n/a"
     status = "OK" if result.ok else "DRIFT"
     reconcile_line = (
-        f"\nReconcile: estimate ${estimate_total:,.2f}  "
-        f"Cost Explorer ${ce_total:,.2f}  "
-        f"delta {delta_pct}  [{status}]"
+        f"\nReconcile: estimate ${estimate_total:,.2f}  Cost Explorer ${ce_total:,.2f}  delta {delta_pct}  [{status}]"
     )
     click.echo(reconcile_line, err=stream_err)
     # Absolute-dollar floor on the drift exit code: below $1 of discrepancy,

{dh_cli-0.8.0 → dh_cli-0.8.2}/src/dh_cli/bedrock/cost_report.py

@@ -21,11 +21,13 @@ Exported API:
     fetch_cost_explorer_total(start, end) -> float
     default_pricing_path() -> Path
 """
+
 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
@@ -108,7 +110,7 @@ def resolve_base_model(model_id: str) -> str:
     stripped = model_id
     for prefix in ("us.", "global.", "eu.", "apac."):
         if stripped.startswith(prefix):
-            stripped = stripped[len(prefix):]
+            stripped = stripped[len(prefix) :]
             break
     for base in _BASE_MODELS:
         if base in stripped:
@@ -165,9 +167,7 @@ def build_report(
     group_by: str = "user+model",
 ) -> Report:
     if group_by not in _VALID_GROUP_BY:
-        raise ValueError(
-            f"group_by={group_by!r} is not one of {sorted(_VALID_GROUP_BY)}"
-        )
+        raise ValueError(f"group_by={group_by!r} is not one of {sorted(_VALID_GROUP_BY)}")
     agg: dict[tuple, dict[str, Any]] = {}
     for rec in records:
         # Bedrock emits records for failed validations / throttles with
@@ -264,17 +264,45 @@ def render_markdown(report: Report) -> str:
 # Columns that are constant/empty for the chosen grouping are dropped so
 # the output table stays narrow and scannable in a terminal.
 _PRETTY_COLUMNS_BY_GROUP = {
-    "user": ("principal_type", "principal_name", "invocations",
-             "input_tokens", "output_tokens", "cache_read",
-             "cache_write", "estimated_cost_usd"),
-    "user+model": ("principal_type", "principal_name", "model", "invocations",
-                   "input_tokens", "output_tokens", "cache_read",
-                   "cache_write", "estimated_cost_usd"),
-    "model": ("model", "invocations", "input_tokens", "output_tokens",
-              "cache_read", "cache_write", "estimated_cost_usd"),
-    "principal_type": ("principal_type", "invocations", "input_tokens",
-                       "output_tokens", "cache_read", "cache_write",
-                       "estimated_cost_usd"),
+    "user": (
+        "principal_type",
+        "principal_name",
+        "invocations",
+        "input_tokens",
+        "output_tokens",
+        "cache_read",
+        "cache_write",
+        "estimated_cost_usd",
+    ),
+    "user+model": (
+        "principal_type",
+        "principal_name",
+        "model",
+        "invocations",
+        "input_tokens",
+        "output_tokens",
+        "cache_read",
+        "cache_write",
+        "estimated_cost_usd",
+    ),
+    "model": (
+        "model",
+        "invocations",
+        "input_tokens",
+        "output_tokens",
+        "cache_read",
+        "cache_write",
+        "estimated_cost_usd",
+    ),
+    "principal_type": (
+        "principal_type",
+        "invocations",
+        "input_tokens",
+        "output_tokens",
+        "cache_read",
+        "cache_write",
+        "estimated_cost_usd",
+    ),
 }
 
 # Nicer column headers for the pretty renderer.
@@ -290,10 +318,16 @@ _PRETTY_HEADERS = {
     "estimated_cost_usd": "cost",
 }
 
-_NUMERIC_COLUMNS = frozenset({
-    "invocations", "input_tokens", "output_tokens",
-    "cache_read", "cache_write", "estimated_cost_usd",
-})
+_NUMERIC_COLUMNS = frozenset(
+    {
+        "invocations",
+        "input_tokens",
+        "output_tokens",
+        "cache_read",
+        "cache_write",
+        "estimated_cost_usd",
+    }
+)
 
 
 def _format_cell(column: str, row: ReportRow) -> str:
@@ -323,9 +357,7 @@ def render_pretty(report: Report, *, group_by: str = "user+model") -> str:
     # Totals footer — numeric columns sum, non-numeric columns are
     # blank except the first one which gets "TOTAL".
     if report.rows:
-        totals: dict[str, int] = {
-            c: 0 for c in _NUMERIC_COLUMNS if c in columns and c != "estimated_cost_usd"
-        }
+        totals: dict[str, int] = {c: 0 for c in _NUMERIC_COLUMNS if c in columns and c != "estimated_cost_usd"}
         cost_total = 0.0
         for row in report.rows:
             for c in totals:
@@ -378,17 +410,19 @@ def render_csv(report: Report) -> str:
     writer = csv.writer(buf)
     writer.writerow(_COLUMNS)
     for row in report.rows:
-        writer.writerow([
-            row.principal_type,
-            row.principal_name,
-            row.model,
-            row.invocations,
-            row.input_tokens,
-            row.output_tokens,
-            row.cache_read,
-            row.cache_write,
-            f"{row.estimated_cost_usd:.6f}",
-        ])
+        writer.writerow(
+            [
+                row.principal_type,
+                row.principal_name,
+                row.model,
+                row.invocations,
+                row.input_tokens,
+                row.output_tokens,
+                row.cache_read,
+                row.cache_write,
+                f"{row.estimated_cost_usd:.6f}",
+            ]
+        )
     return buf.getvalue()
 
 
@@ -410,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"]
@@ -428,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(
@@ -451,8 +518,9 @@ def reconcile_with_cost_explorer(
     if estimate_total == 0:
         return ReconcileResult(False, float("inf"), estimate_total, ce_total, threshold)
     if ce_total == 0:
-        return ReconcileResult(False, abs(ce_total - estimate_total) / estimate_total,
-                               estimate_total, ce_total, threshold)
+        return ReconcileResult(
+            False, abs(ce_total - estimate_total) / estimate_total, estimate_total, ce_total, threshold
+        )
     delta = abs(ce_total - estimate_total) / estimate_total
     return ReconcileResult(delta <= threshold, delta, estimate_total, ce_total, threshold)
 
@@ -479,8 +547,7 @@ def fetch_cost_explorer_total(start: dt.date, end: dt.date) -> float:
         Dimension="SERVICE",
     )
     bedrock_services = [
-        v["Value"] for v in dim.get("DimensionValues", [])
-        if v["Value"].endswith("(Amazon Bedrock Edition)")
+        v["Value"] for v in dim.get("DimensionValues", []) if v["Value"].endswith("(Amazon Bedrock Edition)")
     ]
     if not bedrock_services:
         # No Bedrock-family spend at all in the window — CE honestly

{dh_cli-0.8.0 → dh_cli-0.8.2}/src/dh_cli/engines_studios/engine_commands.py

@@ -606,11 +606,7 @@ def list_engines(env: Optional[str]):
             return left + mid.join("─" * (w + 1) for w in cols) + right
 
         click.echo(border("╭", "┬", "╮"))
-        click.echo(
-            "│"
-            + "│".join(f" {h:{a}{w}}" for h, a, w in zip(headers, aligns, cols))
-            + "│"
-        )
+        click.echo("│" + "│".join(f" {h:{a}{w}}" for h, a, w in zip(headers, aligns, cols)) + "│")
         click.echo(border("├", "┼", "┤"))
 
         for i, engine in enumerate(engines):
@@ -649,9 +645,7 @@ def list_engines(env: Optional[str]):
             else:
                 disk_d = f"\033[32m{disk_text:>{dw}}\033[0m"
 
-            click.echo(
-                f"│ {name_d}│ {state_d}│ {user:<{uw}}│ {etype:<{tw}}│ {uptime_d}│ {disk_d}│"
-            )
+            click.echo(f"│ {name_d}│ {state_d}│ {user:<{uw}}│ {etype:<{tw}}│ {uptime_d}│ {disk_d}│")
 
         click.echo(border("╰", "┴", "╯"))
         click.echo(f"Total: {len(engines)}\n")