mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-17 03:59:28 +00:00
learn: experiment fill -- fill a rule experiment's response cache with a hard ceiling, dry-run estimate
Implemented with Codex; review fix: the bounded client keeps the configured thinking setting so a fill answers as the run would. Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
Marius Ciepluch
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
212a391872
commit
8d14b0c13f
@@ -60,18 +60,24 @@ def add_commands(learn_sub: Any) -> None:
|
||||
("register", "record the manifest's and the units file's sha256 as the experiment row (both committed at HEAD)"),
|
||||
("units", "freeze the units of a program experiment from the memory (no model): the families' evaluation candidates"),
|
||||
("run", "run every arm on every unit, paired; arm A replays from the response cache when the manifest says so"),
|
||||
("fill", "fill a rule experiment's response cache under a separate budget, or estimate with --dry-run"),
|
||||
("analyse", "paired cluster bootstrap, exact McNemar, the two looks and the adoption verdict (counts and intervals only)"),
|
||||
):
|
||||
sub = experiment_sub.add_parser(name, help=text)
|
||||
sub.add_argument("manifest", type=Path, help="plane/experiments/<id>.yaml")
|
||||
sub.add_argument("--memory", help="the store (default OUSAST_MEMORY, then results_root()/memory)")
|
||||
if name in ("units", "run"):
|
||||
if name in ("units", "run", "fill"):
|
||||
sub.add_argument("--profile", choices=("static", "plane"), default="static")
|
||||
if name == "run":
|
||||
if name in ("run", "fill"):
|
||||
sub.add_argument(
|
||||
"--manifest-model", type=Path, default=Path("plane/models/deepseek-flash.yaml"), help="the chat Model (prices)"
|
||||
)
|
||||
if name == "run":
|
||||
sub.add_argument("--no-early-stop", action="store_true", help="take no first look during the run")
|
||||
if name == "fill":
|
||||
sub.add_argument("--budget-usd", type=float, help="required for paid fill; includes the smoke call")
|
||||
sub.add_argument("--dry-run", action="store_true", help="inspect and estimate only; no calls or writes, including --out")
|
||||
sub.add_argument("--out", type=Path, help="also write the fill report as JSON (ignored for --dry-run)")
|
||||
if name in ("run", "analyse"):
|
||||
sub.add_argument("--resamples", type=int, help="bootstrap resamples (default: the manifest's)")
|
||||
sub.add_argument("--out", type=Path, help="write the counts-only report here (analyse --record: benchmarks/experiments/<id>/)")
|
||||
@@ -154,6 +160,33 @@ def _experiment(args: argparse.Namespace) -> int:
|
||||
data = store.get_blob("excerpts", sha)
|
||||
return data.decode("utf-8") if data is not None else None
|
||||
|
||||
if args.experiment_command == "fill":
|
||||
from ..tool_hunter import ChatClient
|
||||
from .fill import fill_rule
|
||||
|
||||
def resolve_fill_client() -> ChatClient:
|
||||
from ..config import load_config
|
||||
from ..model.endpoint import resolve_chat_endpoint
|
||||
|
||||
resolved = resolve_chat_endpoint(load_config(None))
|
||||
if resolved is None:
|
||||
raise ex.ExperimentError("no chat endpoint configured; nothing was spent")
|
||||
return resolved[0]
|
||||
|
||||
report = fill_rule(
|
||||
store,
|
||||
manifest,
|
||||
memory,
|
||||
parameters,
|
||||
excerpt_text,
|
||||
vectors,
|
||||
client_factory=resolve_fill_client,
|
||||
budget_usd=args.budget_usd,
|
||||
dry_run=args.dry_run,
|
||||
)
|
||||
_write(report, None if args.dry_run else args.out)
|
||||
return {"done": 0, "dry_run": 0, "unfinished": 3, "failed": 1}[report["status"]]
|
||||
|
||||
if manifest.kind == "rule": # zero cost: cache only, behind a zero-ceiling meter; no endpoint is resolved
|
||||
caller, meter = ex.zero_cost_caller(_arm_model(manifest, manifest.arms["A"], store), parameters, store)
|
||||
summary = ex.run_rule(store, manifest, memory, caller, meter, excerpt_text, vectors)
|
||||
|
||||
@@ -0,0 +1,205 @@
|
||||
"""Fill rule-experiment responses without writing scores or changing the frozen population.
|
||||
|
||||
Estimates use Prompt.estimated_tokens and the family's mean cached completion usage,
|
||||
or 300 output tokens when no usage is available. Unknown answers may require one retry;
|
||||
dry-run counts those separately as conditional requests, never fabricating an answer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from ..model.endpoint import DeepSeekChatClient
|
||||
from ..plane.budget import BudgetExhausted, MeteredClient, prices_from
|
||||
from ..plane.memory import MemoryStore
|
||||
from ..tool_hunter import ChatClient, ChatResponse
|
||||
from . import experiments as ex
|
||||
from .examples import Example
|
||||
from .program import Caller, Program, Prompt, parse_answer, request_key
|
||||
|
||||
DEFAULT_OUTPUT_TOKENS = 300
|
||||
DEFAULT_MARGIN_USD = 0.01
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Request:
|
||||
family: str
|
||||
prompt: Prompt
|
||||
temperature: float
|
||||
sample: str
|
||||
|
||||
def key(self, caller: Caller) -> str:
|
||||
return request_key(caller.model, caller.params_digest, self.prompt.messages, self.temperature, self.sample, True)
|
||||
|
||||
def ask(self, caller: Caller) -> str:
|
||||
return caller.ask(self.prompt.messages, temperature=self.temperature, sample=self.sample)
|
||||
|
||||
def retry(self) -> Request:
|
||||
return Request(self.family, self.prompt, 0.0, f"{self.sample}:retry")
|
||||
|
||||
|
||||
class FillMeter(MeteredClient):
|
||||
"""Reserve a whole call before admission, including the first (smoke) call.
|
||||
|
||||
The fixed cent is a floor, never replaced by a smaller observed cost. For the
|
||||
production adapter, UTF-8 bytes conservatively bound input tokens (plus 256 framing
|
||||
tokens per message), and the provider enforces the output limit. Injected clients
|
||||
must honour the same token/usage contract. No hidden transport retries are allowed.
|
||||
"""
|
||||
|
||||
def __init__(self, client: ChatClient, parameters: Mapping[str, Any], budget: float, output_limit: int) -> None:
|
||||
super().__init__(client, prices=parameters, budget_usd=budget)
|
||||
self.limit = budget
|
||||
self.output_limit = output_limit
|
||||
self.largest_call_usd = 0.0
|
||||
|
||||
def complete(self, **kwargs: Any) -> ChatResponse:
|
||||
assert self.prices is not None
|
||||
messages = kwargs["messages"]
|
||||
input_bound = sum(len(str(m["content"]).encode("utf-8")) + 256 for m in messages)
|
||||
bound = (
|
||||
input_bound * max(self.prices.input_per_m, self.prices.cache_hit_per_m) + self.output_limit * self.prices.output_per_m
|
||||
) / 1_000_000
|
||||
reserve = max(DEFAULT_MARGIN_USD, self.largest_call_usd, bound)
|
||||
self.budget_usd = max(0.0, self.limit - reserve)
|
||||
before = self.usd or 0.0
|
||||
tokens_before = sum(self.usage.values())
|
||||
try:
|
||||
response = super().complete(**kwargs)
|
||||
except BudgetExhausted:
|
||||
raise
|
||||
except Exception:
|
||||
self.calls += 1 # attempted provider calls count even when no response was returned
|
||||
raise
|
||||
finally:
|
||||
self.largest_call_usd = max(self.largest_call_usd, (self.usd or 0.0) - before)
|
||||
if sum(self.usage.values()) <= tokens_before:
|
||||
raise RuntimeError("provider returned no metered token usage")
|
||||
return response
|
||||
|
||||
|
||||
def fill_rule(
|
||||
store: MemoryStore,
|
||||
manifest: ex.Manifest,
|
||||
memory: Sequence[Example],
|
||||
parameters: Mapping[str, Any],
|
||||
excerpt_text: Callable[[str], str | None],
|
||||
vectors: Mapping[str, Sequence[float]] | None = None,
|
||||
*,
|
||||
client: ChatClient | None = None,
|
||||
client_factory: Callable[[], ChatClient] | None = None,
|
||||
budget_usd: float | None = None,
|
||||
dry_run: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Inspect or fill every frozen unit. The first cache miss is the smoke call.
|
||||
|
||||
A successful smoke is retained in the cache; any provider error stops immediately
|
||||
and is reported without its message, which may contain credentials. A factory is
|
||||
resolved only after validation, and only when a paid request is actually missing.
|
||||
"""
|
||||
from .evaluate import candidate_from_example
|
||||
|
||||
if manifest.kind != "rule":
|
||||
raise ex.ExperimentError("experiment fill supports kind: rule only")
|
||||
if not dry_run and (budget_usd is None or not math.isfinite(budget_usd) or budget_usd < 0):
|
||||
raise ex.ExperimentError("--budget-usd must be a finite non-negative number (or use --dry-run)")
|
||||
prices = prices_from(parameters)
|
||||
if prices is None or any(not math.isfinite(p) or p < 0 for p in (prices.input_per_m, prices.output_per_m, prices.cache_hit_per_m)):
|
||||
raise ex.ExperimentError("fill requires finite, non-negative Model prices")
|
||||
ex.check_registered(store, manifest)
|
||||
units, folds, _ = ex._checked_units(manifest, memory)
|
||||
programs = {f: Program(ex.arm_spec(manifest, manifest.arms["A"], f, store), memory, excerpt_text, vectors) for f in manifest.families}
|
||||
models = {p.spec.model for p in programs.values()}
|
||||
if len(models) != 1:
|
||||
raise ex.ExperimentError("the rule programs must use the same model")
|
||||
caller = Caller(None, next(iter(models)), parameters, store=store)
|
||||
index = {e.id: e for e in memory}
|
||||
candidate_of = candidate_from_example(excerpt_text, vectors)
|
||||
requests: list[Request] = []
|
||||
for unit in sorted(units, key=lambda u: (u.group, u.family, u.unit)):
|
||||
candidate = candidate_of(index[unit.unit])
|
||||
if candidate is None:
|
||||
raise ex.ExperimentError(f"{manifest.id}: unit {unit.unit[:12]} has no excerpt in the store")
|
||||
for prompt, _, temperature, sample in programs[unit.family].prepare_samples(
|
||||
candidate, folds[unit.fold], seed=int(manifest.folds.get("seed", 0))
|
||||
):
|
||||
requests.append(Request(unit.family, prompt, temperature, sample))
|
||||
|
||||
# Request counts are distinct cache keys per family, even when multiple units share one.
|
||||
def inventory() -> dict[str, dict[str, Request]]:
|
||||
result: dict[str, dict[str, Request]] = {f: {} for f in manifest.families}
|
||||
for req in requests:
|
||||
result[req.family][req.key(caller)] = req
|
||||
entry = caller._cached(req.key(caller))
|
||||
if entry is not None and parse_answer(entry.get("content"), req.prompt.valid_lines) is None:
|
||||
retry = req.retry()
|
||||
result[req.family][retry.key(caller)] = retry
|
||||
return result
|
||||
|
||||
initial = inventory()
|
||||
initially_cached = {key for cells in initial.values() for key in cells if caller._cached(key) is not None}
|
||||
families: dict[str, Any] = {}
|
||||
for family, cells in initial.items():
|
||||
outputs = [
|
||||
float((entry.get("usage") or {})["completion_tokens"])
|
||||
for key in cells
|
||||
if (entry := caller._cached(key)) is not None and "completion_tokens" in (entry.get("usage") or {})
|
||||
]
|
||||
mean_output = sum(outputs) / len(outputs) if outputs else DEFAULT_OUTPUT_TOKENS
|
||||
missing = [req for key, req in cells.items() if key not in initially_cached]
|
||||
families[family] = {
|
||||
"units": sum(u.family == family for u in units),
|
||||
"requests": len(cells),
|
||||
"cached": len(cells) - len(missing),
|
||||
"missing": len(missing),
|
||||
"filled": 0,
|
||||
"still_missing": len(missing),
|
||||
"estimated_usd": sum(
|
||||
(r.prompt.estimated_tokens * prices.input_per_m + mean_output * prices.output_per_m) / 1_000_000 for r in missing
|
||||
),
|
||||
"estimated_output_tokens": mean_output,
|
||||
"conditional_retries": sum(not r.sample.endswith(":retry") for r in missing),
|
||||
}
|
||||
report: dict[str, Any] = {
|
||||
"experiment": manifest.id,
|
||||
"status": "dry_run" if dry_run else "done",
|
||||
"families": families,
|
||||
"client_calls": 0,
|
||||
"usd": 0.0,
|
||||
"reason": "",
|
||||
}
|
||||
if dry_run or not any(c["missing"] for c in families.values()):
|
||||
return report
|
||||
meter: FillMeter | None = None
|
||||
try:
|
||||
if client is None and client_factory is not None:
|
||||
client = client_factory()
|
||||
if client is None:
|
||||
raise ex.ExperimentError("no chat endpoint configured; nothing was spent")
|
||||
output_limit = max(p.spec.max_output_tokens for p in programs.values())
|
||||
if isinstance(client, DeepSeekChatClient):
|
||||
client = client.bounded(output_limit)
|
||||
assert budget_usd is not None
|
||||
meter = FillMeter(client, parameters, budget_usd, output_limit)
|
||||
caller.client = meter
|
||||
for req in requests:
|
||||
content = req.ask(caller)
|
||||
if parse_answer(content, req.prompt.valid_lines) is None:
|
||||
req.retry().ask(caller)
|
||||
except BudgetExhausted:
|
||||
report.update(status="unfinished", reason="budget exhausted; resume with a new fill budget")
|
||||
except ex.ExperimentError:
|
||||
raise
|
||||
except Exception:
|
||||
# Never echo provider exceptions: response bodies and URLs may contain secrets.
|
||||
report.update(status="failed", reason="provider or cache failure; fill stopped")
|
||||
final = inventory()
|
||||
for family, cells in final.items():
|
||||
cached = {key for key in cells if caller._cached(key) is not None}
|
||||
families[family].update(requests=len(cells), filled=len(cached - initially_cached), still_missing=len(cells) - len(cached))
|
||||
if meter is not None:
|
||||
report.update(client_calls=meter.calls, usd=meter.usd)
|
||||
return report
|
||||
@@ -506,20 +506,24 @@ class Program:
|
||||
out.append((prompt, got))
|
||||
return out
|
||||
|
||||
def decide(self, candidate: Candidate, fold: Fold, caller: Caller, *, seed: int = 0) -> Decision:
|
||||
def prepare_samples(self, candidate: Candidate, fold: Fold, *, seed: int = 0) -> list[tuple[Prompt, Retrieval, float, str]]:
|
||||
"""The exact prompts, temperatures and sample labels shared by execution and cache planning."""
|
||||
if self.spec.sampling == RETRIEVAL_ENSEMBLE:
|
||||
prepared = self.prepare_ensemble(candidate, fold, seed=seed)
|
||||
answers = [
|
||||
sample_answer(prompt, caller, temperature=TEMPERATURES[0], sample=str(j), families=self.families)
|
||||
for j, (prompt, _) in enumerate(prepared)
|
||||
]
|
||||
got = prepared[0][1]
|
||||
return replace(aggregate(answers), retrieval=got.mode, neighbours=got.neighbours)
|
||||
return [(prompt, got, TEMPERATURES[0], str(j)) for j, (prompt, got) in enumerate(prepared)]
|
||||
if self.spec.sampling != SAMPLING[0]:
|
||||
raise ValueError(f"unknown sampling {self.spec.sampling!r} (one of {', '.join(SAMPLING)})")
|
||||
prompt, got = self.prepare(candidate, fold, seed=seed)
|
||||
decision = classify(prompt, caller, self.spec.k, families=self.families)
|
||||
return replace(decision, retrieval=got.mode, neighbours=got.neighbours)
|
||||
return [(prompt, got, TEMPERATURES[0] if j == 0 else TEMPERATURES[1], str(j)) for j in range(self.spec.k)]
|
||||
|
||||
def decide(self, candidate: Candidate, fold: Fold, caller: Caller, *, seed: int = 0) -> Decision:
|
||||
prepared = self.prepare_samples(candidate, fold, seed=seed)
|
||||
answers = [
|
||||
sample_answer(prompt, caller, temperature=temperature, sample=sample, families=self.families)
|
||||
for prompt, _, temperature, sample in prepared
|
||||
]
|
||||
got = prepared[0][1]
|
||||
return replace(aggregate(answers), retrieval=got.mode, neighbours=got.neighbours)
|
||||
|
||||
|
||||
__all__ = [
|
||||
|
||||
@@ -96,6 +96,21 @@ class DeepSeekChatClient:
|
||||
self._disable_thinking = disable_thinking
|
||||
self.endpoint = endpoint
|
||||
self.usage: list[dict[str, object]] = [] # one entry per call, for the round's cost meter
|
||||
self._output_limit: int | None = None
|
||||
|
||||
def bounded(self, max_output_tokens: int) -> DeepSeekChatClient:
|
||||
"""A single-attempt, output-capped copy for callers reserving a hard per-call budget.
|
||||
|
||||
Empty JSON retries belong to the caller's meter, and transport failures must
|
||||
stop that caller before another potentially billable attempt.
|
||||
"""
|
||||
from dataclasses import replace
|
||||
|
||||
if max_output_tokens <= 0:
|
||||
raise ValueError("max_output_tokens must be positive")
|
||||
client = DeepSeekChatClient(replace(self._client, max_attempts=1), disable_thinking=self._disable_thinking, endpoint=self.endpoint)
|
||||
client._output_limit = max_output_tokens
|
||||
return client
|
||||
|
||||
def complete(
|
||||
self,
|
||||
@@ -117,7 +132,7 @@ class DeepSeekChatClient:
|
||||
logprobs=logprobs,
|
||||
temperature=temperature,
|
||||
)
|
||||
if json_object and not (response.content or "").strip():
|
||||
if self._output_limit is None and json_object and not (response.content or "").strip():
|
||||
response = self._call(
|
||||
model=model,
|
||||
messages=messages,
|
||||
@@ -147,6 +162,8 @@ class DeepSeekChatClient:
|
||||
temperature: float | None = None,
|
||||
) -> ChatResponse:
|
||||
extra: dict[str, object] = {}
|
||||
if self._output_limit is not None:
|
||||
extra["max_tokens"] = self._output_limit
|
||||
if self._disable_thinking:
|
||||
extra["thinking"] = {"type": "disabled"} # otherwise temperature is silently ignored
|
||||
if temperature is not None:
|
||||
|
||||
Reference in New Issue
Block a user