thirst 0.3.0__py3-none-any.whl

thirst/__init__.py

@@ -0,0 +1,3 @@
+"""Thirst CLI: plan, run and measure GPU jobs."""
+
+__version__ = "0.3.0"

thirst/__main__.py

@@ -0,0 +1,153 @@
+"""Thirst CLI entry point: argument parsing, dispatch and plain-words error handling (plan 04.0).
+
+Exit codes: 0 ok, 2 bad arguments, 3 config, 4 API, 5 GPU, 130 Ctrl-C. An expected failure prints
+one line ("thirst: ...") and never a stack trace; a surprise prints a generic line and, with
+``--verbose``, the log (which never contains the site key or a signature).
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+from collections.abc import Sequence
+
+from . import __version__, commands
+from .context import Context, default_context, stderr_writer, stdout_writer
+from .errors import EXIT_BAD_ARGS, EXIT_INTERRUPTED, EXIT_OK, CliError
+from .planning import MAX_WORK_HOURS
+
+logger = logging.getLogger("thirst")
+logger.addHandler(logging.NullHandler())
+
+HANDLERS = {
+    "register": "cmd_register",
+    "plan": "cmd_plan",
+    "run": "cmd_run",
+    "sweep": "cmd_sweep",
+    "doctor": "cmd_doctor",
+    "sync": "cmd_sync",
+    "receipt": "cmd_receipt",
+}
+
+
+def _job_options(sub: argparse.ArgumentParser) -> None:
+    sub.add_argument(
+        "--hours", type=float, required=True, help=f"work at full power (max {MAX_WORK_HOURS:g})"
+    )
+    sub.add_argument("--deadline", required=True, help='e.g. "2026-10-10 09:00" (India time)')
+    sub.add_argument("--pausable", action="store_true", help="allow pauses (single-process jobs)")
+    sub.add_argument("--alpha", type=float, default=0.5, help="0 = least CO2, 1 = least water")
+    sub.add_argument("--gpu-model", default="auto", help="GPU model id; default: detect")
+    sub.add_argument("--gpus", type=int, default=None, help="number of GPUs; default: detect")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="thirst", description="Plan, run and measure GPU jobs.")
+    parser.add_argument("--version", action="version", version=f"thirst {__version__}")
+    parser.add_argument("--verbose", action="store_true", help="log requests (never secrets)")
+    commands_ = parser.add_subparsers(dest="command_name")
+
+    reg = commands_.add_parser("register", help="create this lab's site")
+    reg.add_argument("--api-url", default=None, help="API address (or THIRST_API_URL)")
+    reg.add_argument("--name", required=True)
+    reg.add_argument("--city", required=True)
+    reg.add_argument("--cooling", required=True, choices=["room_ac", "air_chiller", "evaporative"])
+    reg.add_argument("--gpus", type=int, required=True, help="GPUs in the lab")
+    reg.add_argument("--org", default=None)
+    reg.add_argument("--coverage", type=float, default=0.0, help="share of work covered, 0 to 1")
+    reg.add_argument("--private", action="store_true", help="keep the site off the leaderboard")
+    reg.add_argument("--public", action="store_true", help="list the site (the default)")
+    reg.add_argument("--force", action="store_true", help="replace an existing registration")
+
+    plan = commands_.add_parser("plan", help="show the plan for a job without running it")
+    _job_options(plan)
+    plan.add_argument("--json", action="store_true", help="print the raw server response")
+
+    run = commands_.add_parser("run", help="plan, then run your command: thirst run ... -- cmd")
+    _job_options(run)
+    run.add_argument("--now", action="store_true", help="start right away (no waiting)")
+    run.add_argument("--modelled", action="store_true", help="no GPU: estimate the energy")
+    run.add_argument("--no-cap", action="store_true", help="never change the power limit")
+    run.add_argument("--dry-run", action="store_true", help="show the plan, start nothing")
+    run.add_argument("--fallback-clocks", action="store_true", help="lock clocks instead of power")
+    run.add_argument("--slot-seconds", type=float, default=None, help=argparse.SUPPRESS)
+
+    swp = commands_.add_parser("sweep", help="measure the power-cap sweep (needs a GPU and torch)")
+    swp.add_argument("--levels", default="100,85,70")
+    swp.add_argument("--steps", type=int, default=300)
+    swp.add_argument("--warmup", type=int, default=30)
+    swp.add_argument("--idle-seconds", type=float, default=60.0)
+    swp.add_argument("--gpu-model", default="auto")
+    swp.add_argument("--out", default=None, help="powercap_sweep.json to update")
+
+    doc = commands_.add_parser("doctor", help="check the setup")
+    doc.add_argument("--restore", action="store_true", help="undo a power cap left behind")
+
+    commands_.add_parser("sync", help="post receipts waiting in the outbox")
+
+    rec = commands_.add_parser("receipt", help="saved receipts")
+    rec_sub = rec.add_subparsers(dest="receipt_action", required=True)
+    show = rec_sub.add_parser("show", help="print a saved receipt")
+    show.add_argument("job_id")
+    return parser
+
+
+def _split_command(argv: list[str]) -> tuple[list[str], list[str] | None]:
+    """Split ``thirst run ... -- cmd args`` at the first ``--`` (only for `run`)."""
+    if "--" in argv and "run" in argv[: argv.index("--")]:
+        cut = argv.index("--")
+        return argv[:cut], argv[cut + 1 :]
+    return argv, None
+
+
+def _parse(argv: list[str], ctx: Context) -> argparse.Namespace | int:
+    head, command = _split_command(argv)
+    parser = build_parser()
+    try:
+        args = parser.parse_args(head)
+    except SystemExit as exc:
+        return exc.code if isinstance(exc.code, int) else EXIT_BAD_ARGS
+    if args.command_name is None:
+        for line in parser.format_help().rstrip("\n").split("\n"):
+            ctx.out(line)
+        return EXIT_OK
+    args.command = command
+    return args
+
+
+def _setup_logging(verbose: bool) -> None:
+    if verbose:
+        handler = logging.StreamHandler(sys.stderr)
+        handler.setFormatter(logging.Formatter("%(asctime)s %(name)s %(levelname)s %(message)s"))
+        logger.addHandler(handler)
+        logger.setLevel(logging.INFO)
+
+
+def main(argv: Sequence[str] | None = None, *, context: Context | None = None) -> int:
+    """Run the CLI and return the process exit code."""
+    ctx = context or default_context(stdout_writer(), stderr_writer())
+    parsed = _parse(list(sys.argv[1:] if argv is None else argv), ctx)
+    if isinstance(parsed, int):
+        return parsed
+    _setup_logging(parsed.verbose)
+    handler = getattr(commands, HANDLERS[parsed.command_name])
+    try:
+        return int(handler(parsed, ctx))
+    except CliError as err:
+        ctx.err(f"thirst: {err.message}")
+        return err.exit_code
+    except KeyboardInterrupt:
+        ctx.err("thirst: cancelled.")
+        return EXIT_INTERRUPTED
+    except Exception:
+        logger.exception("unexpected error")
+        ctx.err(
+            "thirst: something unexpected went wrong. "
+            "Run again with --verbose and share the output with your team."
+        )
+        return 1
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(main())

thirst/client.py

@@ -0,0 +1,196 @@
+"""HTTP client for the Thirst API (plan 04.2, contract 08 §3, TRD 7.1).
+
+Receipts are signed: ``signature = hex(HMAC_SHA256(token_ascii, ts + "." + raw_body))``. The body is
+serialised once to bytes and exactly those bytes are signed and sent. Retries (1, 2, 4 s) only
+happen on network errors and 5xx, never on 4xx. The token and signature are never logged.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import json
+import logging
+import re
+import time
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from datetime import datetime
+from email.utils import parsedate_to_datetime
+from typing import Any
+
+import httpx
+
+from . import __version__
+from .errors import ApiError
+
+logger = logging.getLogger(__name__)
+
+CONNECT_TIMEOUT_S = 5.0
+READ_TIMEOUT_S = 15.0
+BACKOFF_SECONDS = (1.0, 2.0, 4.0)
+_SITE_ID = re.compile(r"s_[a-z2-7]{10}")
+_GENERIC = "The Thirst server had a problem. Please try again in a few minutes."
+
+
+@dataclass(frozen=True)
+class ReceiptPost:
+    """Result of posting a receipt. ``duplicate`` means the server already had this job."""
+
+    response: dict[str, Any] | None
+    duplicate: bool
+
+
+def sign(token: str, timestamp: str, body: bytes) -> str:
+    """Hex HMAC-SHA256 of ``timestamp + "." + body`` keyed by the token's ASCII bytes."""
+    message = timestamp.encode("ascii") + b"." + body
+    return hmac.new(token.encode("ascii"), message, hashlib.sha256).hexdigest()
+
+
+def encode_body(payload: Mapping[str, Any]) -> bytes:
+    """The single serialisation of a request body. Refuses NaN and Infinity."""
+    try:
+        text = json.dumps(payload, separators=(",", ":"), sort_keys=True, allow_nan=False)
+    except (TypeError, ValueError):
+        raise ApiError(
+            "Could not prepare the request: a value was missing or not a finite number.",
+            code="bad_request",
+        ) from None
+    return text.encode("ascii")  # ensure_ascii is on, so this cannot fail
+
+
+def _error_for(response: httpx.Response) -> ApiError:
+    """Turn an error response into a readable ``ApiError``, using the server's own words."""
+    status = response.status_code
+    try:
+        data = response.json()
+    except ValueError:
+        data = None
+    if isinstance(data, dict) and isinstance(data.get("error"), str):
+        message = data.get("message")
+        text = message if isinstance(message, str) and message else _GENERIC
+        return ApiError(text, code=data["error"], status=status)
+    if status >= 500:
+        return ApiError(_GENERIC, code="server_error", status=status)
+    return ApiError(f"The server refused the request (HTTP {status}).", "http_error", status)
+
+
+class Client:
+    def __init__(
+        self,
+        api_url: str,
+        token: str = "",
+        site_id: str = "",
+        *,
+        transport: httpx.BaseTransport | None = None,
+        sleep: Callable[[float], None] = time.sleep,
+        clock: Callable[[], float] = time.time,
+    ) -> None:
+        self._base = api_url.rstrip("/")
+        self._token = token
+        self._site_id = site_id
+        self._sleep = sleep
+        self._clock = clock
+        self._http = httpx.Client(
+            transport=transport,
+            timeout=httpx.Timeout(READ_TIMEOUT_S, connect=CONNECT_TIMEOUT_S),
+            follow_redirects=False,
+        )
+
+    def __repr__(self) -> str:
+        return f"Client(api_url={self._base!r}, site_id={self._site_id!r})"
+
+    def close(self) -> None:
+        self._http.close()
+
+    # ------------------------------------------------------------ public calls
+
+    def post_plan(self, payload: Mapping[str, Any]) -> dict[str, Any]:
+        return self._json_call("POST", "/plan", encode_body(payload), signed=False)
+
+    def post_site(self, payload: Mapping[str, Any]) -> dict[str, Any]:
+        return self._json_call("POST", "/sites", encode_body(payload), signed=False)
+
+    def get_site(self, site_id: str) -> dict[str, Any]:
+        if not _SITE_ID.fullmatch(site_id):
+            raise ApiError("That is not a valid site id.", code="bad_request")
+        return self._json_call("GET", f"/sites/{site_id}", None, signed=False)
+
+    def post_receipt(self, payload: Mapping[str, Any]) -> ReceiptPost:
+        """Post a signed receipt. A 409 ``duplicate_job`` is not an error: it was recorded."""
+        try:
+            data = self._json_call("POST", "/receipts", encode_body(payload), signed=True)
+        except ApiError as err:
+            if err.code == "duplicate_job":
+                return ReceiptPost(response=None, duplicate=True)
+            raise
+        return ReceiptPost(response=data, duplicate=False)
+
+    def get_server_time(self) -> datetime:
+        """The server's clock, from the ``Date`` header of ``GET /assumptions``."""
+        response = self._send("GET", "/assumptions", None, signed=False)
+        header = response.headers.get("date")
+        try:
+            if not header:
+                raise ValueError("no date")
+            return parsedate_to_datetime(header)
+        except (TypeError, ValueError):
+            raise ApiError(
+                "The server did not report its time.", code="bad_response", status=200
+            ) from None
+
+    # ------------------------------------------------------------ plumbing
+
+    def _json_call(
+        self, method: str, path: str, body: bytes | None, *, signed: bool
+    ) -> dict[str, Any]:
+        response = self._send(method, path, body, signed=signed)
+        try:
+            data = response.json()
+        except ValueError:
+            data = None
+        if not isinstance(data, dict):
+            raise ApiError(
+                "The server sent a reply the CLI could not read.",
+                code="bad_response",
+                status=response.status_code,
+            )
+        return data
+
+    def _headers(self, body: bytes | None, signed: bool) -> dict[str, str]:
+        headers = {"User-Agent": f"thirst-cli/{__version__}"}
+        if body is not None:
+            headers["Content-Type"] = "application/json"
+        if signed:
+            stamp = str(int(self._clock()))
+            headers["X-Thirst-Timestamp"] = stamp
+            headers["X-Thirst-Signature"] = sign(self._token, stamp, body or b"")
+            headers["X-Thirst-Site"] = self._site_id
+        return headers
+
+    def _send(self, method: str, path: str, body: bytes | None, *, signed: bool) -> httpx.Response:
+        delays = (*BACKOFF_SECONDS, None)
+        for delay in delays:
+            started = time.monotonic()
+            try:
+                response = self._http.request(
+                    method, self._base + path, content=body, headers=self._headers(body, signed)
+                )
+            except httpx.HTTPError as exc:
+                logger.info("%s %s -> network error (%s)", method, path, type(exc).__name__)
+                failure = ApiError(
+                    "Could not reach the Thirst server. Check your internet connection.",
+                    code="network_error",
+                )
+            else:
+                took_ms = int((time.monotonic() - started) * 1000)
+                logger.info("%s %s -> %d in %d ms", method, path, response.status_code, took_ms)
+                if response.status_code < 400:
+                    return response
+                failure = _error_for(response)
+                if response.status_code < 500:
+                    raise failure
+            if delay is None:
+                raise failure
+            self._sleep(delay)
+        raise AssertionError("unreachable")  # pragma: no cover