tapbench 0.1.0__py3-none-any.whl

tapbench/__init__.py

@@ -0,0 +1,5 @@
+"""Fixed-rate TAP async benchmarking."""
+
+from importlib.metadata import version
+
+__version__ = version("tapbench")

tapbench/__main__.py

@@ -0,0 +1,6 @@
+"""Module entry point for ``python -m tapbench``."""
+
+from tapbench.cli import main
+
+if __name__ == "__main__":
+    main()

tapbench/cli.py

@@ -0,0 +1,564 @@
+"""Command-line interface and output rendering for tapbench."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import hashlib
+import json
+import math
+import os
+import random
+import signal
+import sys
+import tempfile
+from collections.abc import Sequence
+from dataclasses import asdict
+from datetime import UTC, datetime
+from enum import Enum
+from pathlib import Path
+from typing import Any, cast
+from urllib.parse import urlsplit, urlunsplit
+
+from tapbench.models import JobRecord, RunConfig, RunResult, SchedulerResult
+
+JSON_SCHEMA = "tapbench.run/v1"
+
+
+class ConfigError(ValueError):
+    """Raised when command configuration is invalid."""
+
+
+def _positive_finite(value: str) -> float:
+    try:
+        parsed = float(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError("must be a number") from exc
+    if not math.isfinite(parsed) or parsed <= 0:
+        raise argparse.ArgumentTypeError("must be finite and greater than zero")
+    return parsed
+
+
+def _positive_int(value: str) -> int:
+    try:
+        parsed = int(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError("must be an integer") from exc
+    if parsed < 1:
+        raise argparse.ArgumentTypeError("must be at least one")
+    return parsed
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the command parser."""
+
+    parser = argparse.ArgumentParser(prog="tapbench", description=__doc__)
+    subparsers = parser.add_subparsers(dest="command", required=True)
+    run = subparsers.add_parser("run", help="run a fixed-rate TAP async benchmark")
+    run.add_argument("--url", required=True, help="TAP base URL")
+    run.add_argument(
+        "--query-file", required=True, type=Path, help="file containing one ADQL query"
+    )
+    run.add_argument("--rate", required=True, type=_positive_finite, help="submissions per second")
+    run.add_argument("--duration", required=True, type=_positive_finite, help="scheduling seconds")
+    run.add_argument("--max-inflight", type=_positive_int, default=200)
+    run.add_argument("--poll-interval", type=_positive_finite, default=0.5)
+    run.add_argument("--request-timeout", type=_positive_finite, default=30.0)
+    run.add_argument("--job-timeout", type=_positive_finite, default=300.0)
+    run.add_argument(
+        "--download-results",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="download completed results (default: enabled)",
+    )
+    run.add_argument("--keep-jobs", action="store_true", help="do not delete remote UWS jobs")
+    run.add_argument("--json-output", type=Path, help="write detailed JSON results")
+    run.add_argument("--seed", type=int, default=0, help="polling jitter seed")
+    run.add_argument(
+        "--progress",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="show live progress (default: enabled)",
+    )
+    return parser
+
+
+def _normalize_url(raw_url: str) -> str:
+    parsed = urlsplit(raw_url)
+    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
+        raise ConfigError("--url must be an absolute HTTP or HTTPS URL")
+    if parsed.username is not None or parsed.password is not None:
+        raise ConfigError("--url must not contain credentials")
+    if parsed.query or parsed.fragment:
+        raise ConfigError("--url must not contain a query string or fragment")
+    path = parsed.path.rstrip("/")
+    return urlunsplit((parsed.scheme.lower(), parsed.netloc, path, "", ""))
+
+
+def config_from_namespace(args: argparse.Namespace) -> tuple[RunConfig, str]:
+    """Validate parsed arguments and read the query file without network access."""
+
+    query_file = args.query_file.expanduser()
+    if not query_file.is_file():
+        raise ConfigError(f"query file is not a readable regular file: {query_file}")
+    try:
+        query = query_file.read_text(encoding="utf-8")
+    except OSError as exc:
+        raise ConfigError(f"cannot read query file: {exc.strerror or type(exc).__name__}") from None
+    if not query.strip():
+        raise ConfigError("query file is empty")
+
+    json_output = args.json_output.expanduser() if args.json_output else None
+    if json_output is not None:
+        if json_output.resolve() == query_file.resolve():
+            raise ConfigError("--json-output must not overwrite --query-file")
+        if not json_output.parent.is_dir():
+            raise ConfigError("--json-output parent directory does not exist")
+
+    config = RunConfig(
+        url=_normalize_url(args.url),
+        query_file=query_file,
+        rate=args.rate,
+        duration=args.duration,
+        max_inflight=args.max_inflight,
+        poll_interval=args.poll_interval,
+        request_timeout=args.request_timeout,
+        job_timeout=args.job_timeout,
+        download_results=args.download_results,
+        keep_jobs=args.keep_jobs,
+        json_output=json_output,
+        seed=args.seed,
+        progress=args.progress,
+    )
+    return config, query
+
+
+def load_credentials() -> tuple[str, str]:
+    """Read required Basic-auth credentials from the inherited environment."""
+
+    username = os.environ.get("TAPBENCH_USER")
+    password = os.environ.get("TAPBENCH_PASSWORD")
+    if not username or not password:
+        raise ConfigError("TAPBENCH_USER and TAPBENCH_PASSWORD must both be set")
+    return username, password
+
+
+def human_bytes(value: int) -> str:
+    """Render a byte count using binary units."""
+
+    amount = float(value)
+    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
+        if abs(amount) < 1024.0 or unit == "TiB":
+            return f"{amount:.0f} {unit}" if unit == "B" else f"{amount:.2f} {unit}"
+        amount /= 1024.0
+    raise AssertionError("unreachable")
+
+
+def progress_line(
+    records: Sequence[JobRecord],
+    *,
+    active: int,
+    elapsed_s: float,
+    duration_s: float,
+    download_results: bool,
+    scheduling_complete: bool = False,
+) -> str:
+    """Render one stable live-progress snapshot."""
+
+    scheduling = not scheduling_complete and elapsed_s < duration_s
+    label = (
+        f"schedule {min(elapsed_s, duration_s):5.1f}/{duration_s:.1f}s"
+        if scheduling
+        else f"drain {elapsed_s:5.1f}s"
+    )
+    completed = sum(record.final_uws_phase == "COMPLETED" for record in records)
+    errors = sum(record.final_uws_phase == "ERROR" for record in records)
+    aborted = sum(record.final_uws_phase == "ABORTED" for record in records)
+    values = [
+        f"[{label}]",
+        f"offered={len(records)}",
+        f"started={sum(record.actual_submission_start_offset_s is not None for record in records)}",
+        f"dropped={sum(record.dropped_locally for record in records)}",
+        f"active={active}",
+        f"created={sum(record.job_created for record in records)}",
+        f"completed={completed}",
+        f"error={errors}",
+        f"aborted={aborted}",
+        f"success={sum(record.success for record in records)}",
+        f"failed={sum(record.failure_category is not None for record in records)}",
+        f"polls={sum(record.polling_request_count for record in records)}",
+    ]
+    submission_rate_limits = sum(record.submission_rate_limit_count for record in records)
+    if submission_rate_limits:
+        values.append(f"submission_rate_limited={submission_rate_limits}")
+    start_rate_limits = sum(record.start_rate_limit_count for record in records)
+    if start_rate_limits:
+        values.append(f"start_rate_limited={start_rate_limits}")
+    rate_limits = sum(record.polling_rate_limit_count for record in records)
+    if rate_limits:
+        values.append(f"rate_limited={rate_limits}")
+    cleanup_rate_limits = sum(record.cleanup_rate_limit_count for record in records)
+    if cleanup_rate_limits:
+        values.append(f"cleanup_rate_limited={cleanup_rate_limits}")
+    if download_results:
+        values.append(
+            f"downloaded={human_bytes(sum(record.result_bytes or 0 for record in records))}"
+        )
+    return " ".join(values)
+
+
+def redact_output(value: Any, secrets: Sequence[str]) -> Any:
+    """Recursively redact secret values from output-bound data."""
+
+    if isinstance(value, str):
+        redacted = value
+        for secret in secrets:
+            if secret:
+                redacted = redacted.replace(secret, "<redacted>")
+        return redacted
+    if isinstance(value, dict):
+        return {str(key): redact_output(item, secrets) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [redact_output(item, secrets) for item in value]
+    return value
+
+
+def _json_safe(value: Any) -> Any:
+    if isinstance(value, Enum):
+        return value.value
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, dict):
+        return {str(key): _json_safe(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_json_safe(item) for item in value]
+    return value
+
+
+def atomic_write_json(path: Path, payload: dict[str, Any]) -> None:
+    """Atomically replace *path* with a complete JSON document."""
+
+    temporary: Path | None = None
+    try:
+        with tempfile.NamedTemporaryFile(
+            mode="w",
+            encoding="utf-8",
+            dir=path.parent,
+            prefix=f".{path.name}.",
+            suffix=".tmp",
+            delete=False,
+        ) as stream:
+            temporary = Path(stream.name)
+            json.dump(_json_safe(payload), stream, indent=2, sort_keys=True, allow_nan=False)
+            stream.write("\n")
+            stream.flush()
+            os.fsync(stream.fileno())
+        os.replace(temporary, path)
+        temporary = None
+    finally:
+        if temporary is not None:
+            temporary.unlink(missing_ok=True)
+
+
+def result_payload(
+    result: RunResult, summary: dict[str, Any], *, secrets: Sequence[str] = ()
+) -> dict[str, Any]:
+    """Build the versioned machine-readable result document."""
+
+    config = asdict(result.config)
+    config["authentication"] = "basic"
+    payload = {
+        "schema": JSON_SCHEMA,
+        "configuration": config,
+        "query_file": str(result.config.query_file),
+        "query_sha256": result.query_sha256,
+        "started_at": result.started_at,
+        "finished_at": result.finished_at,
+        "actual_scheduling_duration_s": result.scheduling_duration_s,
+        "elapsed_s": result.elapsed_s,
+        "interrupted": result.interrupted,
+        "summary": summary,
+        "records": [asdict(record) for record in result.records],
+    }
+    return cast(dict[str, Any], redact_output(payload, secrets))
+
+
+def _format_metric(value: float | int | None, *, decimals: int = 3) -> str:
+    return "unavailable" if value is None else f"{value:.{decimals}f}"
+
+
+def render_summary(result: RunResult, summary: dict[str, Any]) -> str:
+    """Render the final human-readable summary."""
+
+    counts = summary["counts"]
+    throughput = summary["throughput"]
+    mode = "download results" if result.config.download_results else "no download"
+    lines = [
+        "tapbench partial summary" if result.interrupted else "tapbench summary",
+        f"  mode:                         {mode}",
+        "  requested rate:              "
+        f"{result.config.rate:.3f} jobs/s for {result.config.duration:.3f} s",
+        f"  opportunities:               {counts['opportunities']}",
+        f"  started / dropped:           {counts['started']} / {counts['dropped']}",
+        f"  jobs created:                {counts['jobs_created']}",
+        "  completed / error / aborted: "
+        f"{counts['completed']} / {counts['error']} / {counts['aborted']}",
+        "  UWS timing samples:           "
+        f"{counts['server_timing_samples']} / {counts['completed']} completed",
+        f"  successful / failed:         {counts['successful']} / {counts['failed']}",
+        f"  peak in-flight:              {result.peak_inflight}",
+        f"  HTTP requests / polls:       {counts['http_requests']} / {counts['polls']}",
+        "  average polls / created job: "
+        f"{_format_metric(summary['average_polls_per_created_job'])}",
+    ]
+    if counts["submission_rate_limits"]:
+        lines.append(
+            "  submission rate limits / wait:"
+            f" {counts['submission_rate_limits']} / "
+            f"{_format_metric(counts['submission_rate_limit_wait_s'])} s"
+        )
+    if counts["start_rate_limits"]:
+        lines.append(
+            "  start rate limits / wait:      "
+            f"{counts['start_rate_limits']} / "
+            f"{_format_metric(counts['start_rate_limit_wait_s'])} s"
+        )
+    if counts["polling_rate_limits"]:
+        lines.append(
+            "  polling rate limits / wait:  "
+            f"{counts['polling_rate_limits']} / "
+            f"{_format_metric(counts['polling_rate_limit_wait_s'])} s"
+        )
+    if counts["cleanup_rate_limits"]:
+        lines.append(
+            "  cleanup rate limits / wait:  "
+            f"{counts['cleanup_rate_limits']} / "
+            f"{_format_metric(counts['cleanup_rate_limit_wait_s'])} s"
+        )
+    if result.config.download_results:
+        lines.append(f"  downloaded:                  {human_bytes(counts['downloaded_bytes'])}")
+    lines.extend(
+        [
+            "  achieved submission rate:    "
+            f"{_format_metric(throughput['achieved_submission_start_rate_jobs_s'])} jobs/s",
+            "  job creation throughput:      "
+            f"{_format_metric(throughput['job_creation_throughput_jobs_s'])} jobs/s",
+            "  server execution throughput:  "
+            f"{_format_metric(throughput['server_execution_throughput_jobs_s'])} jobs/s",
+        ]
+    )
+    if result.config.download_results:
+        lines.append(
+            "  result completion throughput: "
+            f"{_format_metric(throughput['result_completion_throughput_jobs_s'])} jobs/s"
+        )
+        lines.append(
+            "  download throughput:          "
+            f"{_format_metric(throughput['result_download_throughput_bytes_s'])} B/s"
+        )
+
+    percentiles = summary["latency"]
+    lines.extend(["", "latency (seconds)                 p50      p90      p95      p99      max"])
+    labels = {
+        "scheduling_delay_s": "scheduling delay",
+        "job_creation_s": "job creation",
+        "query_queue_s": "query queue (UWS)",
+        "query_execution_s": "query execution (UWS)",
+        "submission_to_terminal_s": "submission to terminal",
+        "result_response_headers_s": "result response headers",
+        "result_first_byte_s": "result first byte",
+        "result_download_s": "result download",
+        "submission_to_download_s": "submission to download",
+    }
+    for key, label in labels.items():
+        if not result.config.download_results and key.startswith(
+            ("result_", "submission_to_download")
+        ):
+            continue
+        row = percentiles[key]
+        rendered = " ".join(
+            f"{_format_metric(row[name]):>8}" for name in ("p50", "p90", "p95", "p99", "max")
+        )
+        lines.append(f"  {label:<28}{rendered}")
+
+    failures = summary["failure_breakdown"]
+    if failures or counts["cleanup_failures"]:
+        lines.extend(["", "failures"])
+        for category, count in sorted(failures.items()):
+            if count:
+                lines.append(f"  {category + ':':<30}{count}")
+        if counts["cleanup_failures"]:
+            lines.append(f"  {'cleanup failures:':<30}{counts['cleanup_failures']} (separate)")
+    return "\n".join(lines)
+
+
+async def _progress_reporter(
+    state: Any, config: RunConfig, started: float, done_event: asyncio.Event
+) -> None:
+    """Render scheduler state without participating in job execution."""
+
+    interactive = sys.stderr.isatty()
+    interval = 1.0 if interactive else 5.0
+    last_length = 0
+    loop = asyncio.get_running_loop()
+    next_emit = loop.time()
+    saw_scheduling_complete = state.scheduling_complete
+
+    def emit() -> None:
+        nonlocal last_length
+        line = progress_line(
+            state.records,
+            active=state.active_count,
+            elapsed_s=loop.time() - started,
+            duration_s=config.duration,
+            download_results=config.download_results,
+            scheduling_complete=state.scheduling_complete,
+        )
+        if interactive:
+            padded = line.ljust(last_length)
+            print(f"\r{padded}", file=sys.stderr, end="", flush=True)
+            last_length = len(line)
+        else:
+            print(f"tapbench progress: {line}", file=sys.stderr, flush=True)
+
+    emit()
+    next_emit = loop.time() + interval
+    while not done_event.is_set():
+        try:
+            await asyncio.wait_for(done_event.wait(), timeout=0.1)
+        except TimeoutError:
+            now = loop.time()
+            transitioned = state.scheduling_complete and not saw_scheduling_complete
+            if transitioned or now >= next_emit:
+                emit()
+                next_emit = now + interval
+            saw_scheduling_complete = state.scheduling_complete
+    if state.scheduling_complete and not saw_scheduling_complete:
+        emit()
+    if interactive and last_length:
+        print(file=sys.stderr)
+
+
+async def execute_run(config: RunConfig, query: str, username: str, password: str) -> RunResult:
+    """Execute one benchmark run and return all measurements."""
+
+    from tapbench.client import TapClient, make_client_session
+    from tapbench.scheduler import SchedulerState, run_fixed_rate
+
+    loop = asyncio.get_running_loop()
+    started_mono = loop.time()
+    started_at = datetime.now(UTC).isoformat()
+    stop_event = asyncio.Event()
+    state = SchedulerState()
+    jitter_rng = random.Random(config.seed)
+    reporter_done = asyncio.Event()
+
+    installed_signals: list[signal.Signals] = []
+    interrupt_count = 0
+
+    def interrupt(signum: signal.Signals) -> None:
+        nonlocal interrupt_count
+        interrupt_count += 1
+        if interrupt_count == 1:
+            stop_event.set()
+            return
+        loop.remove_signal_handler(signum)
+        signal.signal(signum, signal.SIG_DFL)
+        os.kill(os.getpid(), signum)
+
+    for signum in (signal.SIGINT, signal.SIGTERM):
+        try:
+            loop.add_signal_handler(signum, interrupt, signum)
+            installed_signals.append(signum)
+        except (NotImplementedError, RuntimeError):
+            pass
+
+    scheduler_result = None
+    progress_task: asyncio.Task[None] | None = None
+    try:
+        async with make_client_session(config) as session:
+            client = TapClient(config, session, started_mono, username, password)
+
+            async def run_one(record: JobRecord) -> None:
+                await client.run_job(query, record, jitter_rng)
+
+            if config.progress:
+                progress_task = asyncio.create_task(
+                    _progress_reporter(state, config, started_mono, reporter_done)
+                )
+            scheduler_result = await run_fixed_rate(
+                config,
+                run_one,
+                stop_event=stop_event,
+                state=state,
+                clock=loop.time,
+                sleep=asyncio.sleep,
+            )
+    except asyncio.CancelledError:
+        stop_event.set()
+        if scheduler_result is None:
+            scheduler_result = SchedulerResult(
+                records=state.records,
+                peak_inflight=state.peak_inflight,
+                scheduling_duration_s=loop.time() - started_mono,
+                interrupted=True,
+            )
+    finally:
+        reporter_done.set()
+        if progress_task is not None:
+            await progress_task
+        for signum in installed_signals:
+            loop.remove_signal_handler(signum)
+
+    assert scheduler_result is not None
+    finished_mono = loop.time()
+    return RunResult(
+        config=config,
+        query_sha256=hashlib.sha256(query.encode()).hexdigest(),
+        started_at=started_at,
+        finished_at=datetime.now(UTC).isoformat(),
+        scheduling_duration_s=scheduler_result.scheduling_duration_s,
+        elapsed_s=finished_mono - started_mono,
+        peak_inflight=scheduler_result.peak_inflight,
+        records=scheduler_result.records,
+        interrupted=scheduler_result.interrupted,
+    )
+
+
+def run_command(config: RunConfig, query: str) -> int:
+    """Run the async implementation and write terminal/JSON output."""
+
+    from tapbench.stats import aggregate_run
+
+    username, password = load_credentials()
+    result = asyncio.run(execute_run(config, query, username, password))
+    summary = aggregate_run(result)
+    print(render_summary(result, summary))
+    if config.json_output is not None:
+        atomic_write_json(
+            config.json_output,
+            result_payload(result, summary, secrets=(username, password)),
+        )
+    return 3 if result.interrupted else 0
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    """CLI entry point."""
+
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    try:
+        config, query = config_from_namespace(args)
+        status = run_command(config, query)
+    except ConfigError as exc:
+        parser.error(str(exc))
+    except OSError as exc:
+        print(f"tapbench: fatal output error: {type(exc).__name__}", file=sys.stderr)
+        status = 4
+    except KeyboardInterrupt:
+        print("tapbench: interrupted", file=sys.stderr)
+        status = 3
+    except Exception as exc:
+        print(f"tapbench: fatal run error: {type(exc).__name__}", file=sys.stderr)
+        status = 4
+    raise SystemExit(status)