convergent-cli 0.1.0__py3-none-any.whl

convergent_cli/calibration_finish.py

@@ -0,0 +1,310 @@
+"""Validate recorded treatment identity, declare it, then finish each async send."""
+
+from __future__ import annotations
+
+import argparse
+import time
+from collections.abc import Mapping
+from concurrent.futures import FIRST_COMPLETED, CancelledError, Future, wait
+from dataclasses import dataclass
+from threading import Event, Thread
+
+from convergent_cli import calibration_runner as runner
+from convergent_cli.calibration_send import SendDeadline, SendTimeout
+from convergent_cli.commands import calibration as cal
+from convergent_cli.confirm import confirm_write
+from convergent_cli.customer_client import ControlPlaneRefusedError
+from convergent_cli.errors import ExitCode, UsageError
+from convergent_cli.output import emit_json
+
+
+@dataclass
+class _Run:
+    tenant: cal._ManifestTenant
+    source: str
+    run_id: str
+    recorded: Mapping[str, object]
+    gate_run_id: str | None = None
+
+
+def concurrency(value: str) -> int:
+    try:
+        bound = int(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError("--concurrency must be a positive integer") from exc
+    if bound < 1:
+        raise argparse.ArgumentTypeError("--concurrency must be a positive integer")
+    return bound
+
+
+def finish(args: argparse.Namespace) -> int:
+    treatment = runner.require_treatment(args.treatment)
+    k = runner.require_k(args.k)
+    pins = cal._parsed_pins(args, None)
+    label = cal._input_label(args)
+    named = cal._optional_digest(args.preprocessing_digest)
+    gate, gate_pins = cal._gate_choice(args, treatment)
+    manifest = cal._load_manifest(runner.require_manifest_id(args.corpus_manifest_id))
+    cal._require_scoped_preprocessing(manifest, named, "--preprocessing-digest")
+    cal._refuse_treatment_bound_elsewhere(treatment, manifest, k, pins, label, named)
+    runs: list[_Run] = []
+    for tenant in manifest.included:
+        source, run_id, recorded = cal._resolve_treatment_run(
+            treatment, manifest, tenant, k, pins, label, named
+        )
+        if recorded is None:
+            raise UsageError(f"treatment_run_missing: {tenant.tenant_tag}")
+        runs.append(_Run(tenant, source, run_id, recorded))
+    if not runs:
+        raise UsageError("treatment_runs_missing: no included tenants")
+    _settle(runs, wait=args.wait)
+    runtime, digest, suite = _identity(runs, k, pins)
+    if runner.runtime_has_decider(runtime) and gate is None and not args.ungated:
+        raise UsageError(
+            "gate_choice_required: pass --gate-baseline TREATMENT or --ungated. "
+            "The first send is final."
+        )
+    for run in runs:
+        tenant = run.tenant
+        cal._refuse_mismatch(
+            run.recorded,
+            tenant,
+            dataset_id=tenant.dataset_id,
+            version_id=tenant.version_id,
+            samples=k,
+            preprocessing_digest=cal._active_digest(tenant, named),
+            criterion_versions=pins,
+            treatment=treatment,
+            tenant_tag=tenant.tenant_tag,
+            manifest_id=manifest.manifest_id,
+            label_release=manifest.label_release,
+            source_run_id=run.source,
+        )
+        if gate is not None:
+            run.gate_run_id = cal._treatment_id(
+                gate,
+                manifest,
+                tenant,
+                k,
+                run.source,
+                gate_pins,
+                labeled=label is not None,
+                input_base=cal._input_base(manifest, tenant, label, named),
+                preprocessing_digest=named,
+            )
+            baseline = cal._get_run(tenant, run.gate_run_id)
+            if baseline is None or baseline.get("status") != "succeeded":
+                raise UsageError(f"gate_baseline_unavailable: {tenant.tenant_tag}")
+    confirm_write("This will declare the treatment and send each tenant's contribution.")
+    # The declaration endpoint compares the complete identity and is a no-op
+    # for an identical declaration. Revalidate even when every send is stored.
+    declaration_digest = cal._declaration_preprocessing_digest(manifest, named)
+    runtime = cal._resolved_declaration_runtime(
+        runs[0].recorded, runtime, runs[0].tenant.tenant_tag
+    )
+    cal._write_declaration(
+        treatment, manifest.manifest_id, [suite], runtime, k, 1000, declaration_digest
+    )
+    rows: list[dict[str, object]] = [
+        {"tenant_tag": run.tenant.tenant_tag, "declared": True, "sent": "not_sent", "reason": None}
+        for run in runs
+    ]
+    pending: dict[Future[dict[str, object]], int] = {}
+    remaining = iter(enumerate(runs))
+    stopped = Event()
+    try:
+        for index, run in remaining:
+            pending[_start_send(run, treatment, args, stopped)] = index
+            rows[index]["sent"] = "sending"
+            if len(pending) == args.concurrency:
+                break
+        while pending:
+            done, _ = wait(pending, return_when=FIRST_COMPLETED)
+            # Read every completed result before starting another tenant.
+            for future in done:
+                rows[pending[future]] = future.result()
+            for future in done:
+                del pending[future]
+            for index, run in remaining:
+                pending[_start_send(run, treatment, args, stopped)] = index
+                rows[index]["sent"] = "sending"
+                if len(pending) == args.concurrency:
+                    break
+        for row in rows:
+            if args.format != "json":
+                state, reason = row["sent"], row["reason"]
+                detail = f"{state} {reason}" if reason is not None else state
+                unchanged = (
+                    " (already declared and stored; unchanged)" if state == "duplicate" else ""
+                )
+                print(f"{row['tenant_tag']}: declared, sent: {detail}{unchanged}")
+                if "hint" in row:
+                    print(row["hint"])
+        shown = cal._staff_get(f"{cal._TREATMENTS_PATH}/{cal._quote(treatment)}")
+        summary = {
+            "complete": shown.get("complete"),
+            "tenants": len(cal._mappings(shown.get("tenants"))),
+            "grade_once_usd": shown.get("grade_once_usd"),
+        }
+    except (Exception, KeyboardInterrupt):
+        stopped.set()
+        for future in pending:
+            future.cancel()
+        for future, index in pending.items():
+            if future.cancelled():
+                rows[index]["sent"] = "not_sent"
+            elif future.done():
+                if future.exception() is None:
+                    rows[index] = future.result()
+                else:
+                    rows[index].update(sent="error", reason="interrupted")
+        if args.format == "json":
+            emit_json({"tenants": rows, "summary": None})
+        raise
+    if args.format == "json":
+        emit_json({"tenants": rows, "summary": summary})
+    else:
+        print(
+            f"complete={str(summary['complete']).lower()} tenants={summary['tenants']} "
+            f"grade_once_usd={summary['grade_once_usd']}"
+        )
+    refused = any(row["sent"] in {"refused", "error", "timeout"} for row in rows)
+    return int(ExitCode.REMOTE) if refused or summary["complete"] is not True else 0
+
+
+def _start_send(
+    run: _Run, treatment: str, args: argparse.Namespace, stopped: Event
+) -> Future[dict[str, object]]:
+    future: Future[dict[str, object]] = Future()
+
+    def send() -> None:
+        if not future.set_running_or_notify_cancel():
+            return
+        try:
+            future.set_result(_send(run, treatment, args, stopped))
+        except BaseException as exc:
+            future.set_exception(exc)
+
+    # A blocked HTTP request must not keep the CLI alive during interpreter exit.
+    Thread(target=send, daemon=True).start()
+    return future
+
+
+def _send(run: _Run, treatment: str, args: argparse.Namespace, stopped: Event) -> dict[str, object]:
+    row: dict[str, object] = {
+        "tenant_tag": run.tenant.tenant_tag,
+        "declared": True,
+        "sent": "sending",
+        "reason": None,
+    }
+    if stopped.is_set():
+        row["sent"] = "not_sent"
+        return row
+    deadline = SendDeadline(args.send_timeout)
+    try:
+        deadline.remaining()
+        state, reason = cal._send_contribution(
+            run.tenant,
+            run.run_id,
+            treatment,
+            run.gate_run_id,
+            ungated=args.ungated,
+            deadline=deadline,
+            cancelled=stopped.is_set,
+        )
+        if state == "stored":
+            state = "duplicate"
+        row.update(sent=state, reason=reason)
+        while state in {"sending", "counting"} and not stopped.is_set():
+            deadline.pause()
+            if stopped.is_set():
+                break
+            state, reason = cal._poll_contribution(run.tenant, run.run_id, deadline=deadline)
+            row.update(sent=state, reason=reason)
+    except CancelledError:
+        row.update(sent="not_sent", reason=None)
+    except SendTimeout:
+        row.update(
+            sent="timeout",
+            reason="send_timeout",
+            hint=f"Check cv calibration show {treatment}.",
+        )
+    except ControlPlaneRefusedError as exc:
+        row.update(sent="refused", reason=str(exc))
+    except Exception as exc:
+        # Exception text can contain customer response bodies or credentials.
+        row.update(sent="error", reason=type(exc).__name__)
+    return row
+
+
+def _settle(runs: list[_Run], *, wait: bool) -> None:
+    while True:
+        pending = [
+            run
+            for run in runs
+            if run.recorded.get("status")
+            not in {
+                "succeeded",
+                "failed",
+                "cancelled",
+            }
+        ]
+        if not pending or not wait:
+            break
+        time.sleep(30)
+        for run in pending:
+            recorded = cal._get_run(run.tenant, run.run_id)
+            if recorded is None:
+                raise UsageError(f"treatment_run_missing: {run.tenant.tenant_tag}")
+            run.recorded = recorded
+    for run in runs:
+        status = run.recorded.get("status")
+        if status != "succeeded":
+            raise UsageError(f"run_not_succeeded: {run.tenant.tenant_tag} ({status})")
+
+
+def _identity(
+    runs: list[_Run],
+    k: int,
+    pins: Mapping[str, str],
+) -> tuple[dict[str, object], str, str]:
+    runtimes: list[dict[str, object]] = []
+    digests: list[str] = []
+    suites: list[str] = []
+    for run in runs:
+        tag = run.tenant.tenant_tag
+        metadata = run.recorded.get("metadata")
+        if not isinstance(metadata, Mapping):
+            raise UsageError(f"run_metadata_missing: {tag}")
+        # Admission omits samples_per_task when it is 1.
+        samples = metadata.get("samples_per_task", 1)
+        if isinstance(samples, bool) or not isinstance(samples, int) or samples != k:
+            raise UsageError(f"k_mismatch: {tag}")
+        if metadata.get("platform_grader_criterion_versions", {}) != pins:
+            raise UsageError(f"criterion_pins_mismatch: {tag}")
+        runtime = metadata.get("platform_grader_runtime")
+        if not isinstance(runtime, dict) or not runtime:
+            raise UsageError(f"runtime_pin_missing: {tag}")
+        preprocessing = run.recorded.get("grader_preprocessing")
+        digest = (
+            preprocessing.get("grader_preprocessing_digest")
+            if isinstance(preprocessing, Mapping)
+            else None
+        )
+        if not isinstance(digest, str):
+            raise UsageError(f"preprocessing_digest_missing: {tag}")
+        suite = runner.suite_digest_from_run(run.recorded)
+        if suite is None:
+            raise UsageError(f"grader_suite_digest_missing: {tag}")
+        runtimes.append(runtime)
+        digests.append(runner.require_digest(digest, "recorded preprocessing digest"))
+        suites.append(suite)
+    tags = ", ".join(run.tenant.tenant_tag for run in runs)
+    for values, reason in (
+        (runtimes, "runtime_pin_mismatch"),
+        (digests, "preprocessing_digest_mismatch"),
+        (suites, "grader_suite_digest_mismatch"),
+    ):
+        if any(value != values[0] for value in values[1:]):
+            raise UsageError(f"{reason}: tenants disagree ({tags})")
+    return runtimes[0], digests[0], suites[0]