Skip to content

Simplify standalone SourceHunt resume around work-item memoization #143

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
auroter wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Simplify standalone SourceHunt resume around work-item memoization #143

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 120 additions & 1 deletion clearwing/llm/budget.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ def __init__(
default_max_output_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
manifest_filename: str = "manifest.json",
endpoint: LLMEndpoint | None = None,
resume: bool = False,
) -> None:
if not math.isfinite(limit_usd) or limit_usd < 0:
raise BudgetConfigurationError("LLM budget must be a finite value >= 0")
Expand Down Expand Up @@ -188,16 +189,133 @@ def __init__(
self.ledger_path = session_dir / "spend-ledger.jsonl"
self.manifest_path = session_dir / manifest_filename
with self._lock:
if resume:
self._restore_history_locked()
self._persist_event_locked(
{
"event": "run_started",
"event": "run_resumed" if resume else "run_started",
"session_id": self.session_id,
"budget_usd": self.limit_usd,
"carried_forward_usd": self._spent_usd,
"timestamp": self._timestamp(),
}
)
self._persist_snapshot_locked()

def _restore_history_locked(self) -> None:
"""Restore settled calls and conservatively close interrupted reservations."""

if not self.ledger_path.is_file():
raise BudgetConfigurationError(
f"Cannot resume session {self.session_id!r}: spend ledger is missing"
)
try:
lines = self.ledger_path.read_text(encoding="utf-8").splitlines()
except (OSError, UnicodeDecodeError) as exc:
raise BudgetConfigurationError(
f"Cannot resume session {self.session_id!r}: spend ledger is unreadable"
) from exc

reservations: dict[str, dict[str, Any]] = {}
settlements: dict[str, dict[str, Any]] = {}
has_run_header = False
for index, line in enumerate(lines):
try:
event = json.loads(line)
except json.JSONDecodeError as exc:
if index == len(lines) - 1:
break
raise BudgetConfigurationError(
f"Cannot resume session {self.session_id!r}: spend ledger is corrupt"
) from exc
if not isinstance(event, dict):
continue
if event.get("event") in {"run_started", "run_resumed"}:
has_run_header = has_run_header or event.get("session_id") == self.session_id
call_id = str(event.get("call_id") or "")
if not call_id:
continue
if event.get("event") == "call_reserved":
reservations.setdefault(call_id, event)
elif event.get("event") == "call_settled":
settlements.setdefault(call_id, event)

if not has_run_header:
raise BudgetConfigurationError(
f"Cannot resume session {self.session_id!r}: spend ledger is corrupt"
)

for call_id in reservations.keys() | settlements.keys():
if call_id in settlements:
self._restore_settlement_locked(settlements[call_id])
else:
self._recover_reservation_locked(call_id, reservations[call_id])
if self.enforcing and self._spent_usd >= self.limit_usd - self._EPSILON:
self._exhausted = True
self._status = "budget_exhausted"

def _restore_settlement_locked(self, event: dict[str, Any]) -> None:
try:
cost = float(event["cost_usd"])
input_tokens = int(event.get("input_tokens", 0))
output_tokens = int(event.get("output_tokens", 0))
cached_tokens = int(event.get("cached_input_tokens", 0))
except (KeyError, TypeError, ValueError) as exc:
raise BudgetConfigurationError(
f"Cannot resume session {self.session_id!r}: spend ledger is corrupt"
) from exc
if (
not math.isfinite(cost)
or min(cost, input_tokens, output_tokens, cached_tokens) < 0
):
raise BudgetConfigurationError(
f"Cannot resume session {self.session_id!r}: spend ledger is corrupt"
)
if not isinstance(event.get("metadata"), dict):
event = {**event, "metadata": {}}
self._records.append(event)
self._spent_usd += cost
self._input_tokens += input_tokens
self._output_tokens += output_tokens
self._cached_input_tokens += cached_tokens

def _recover_reservation_locked(
self, call_id: str, reservation: dict[str, Any]
) -> None:
try:
reserved_usd = float(reservation["reserved_usd"])
except (KeyError, TypeError, ValueError) as exc:
raise BudgetConfigurationError(
f"Cannot resume session {self.session_id!r}: spend ledger is corrupt"
) from exc
if not math.isfinite(reserved_usd) or reserved_usd < 0:
raise BudgetConfigurationError(
f"Cannot resume session {self.session_id!r}: spend ledger is corrupt"
)
charged = reserved_usd if reservation.get("budget_enforcing", reserved_usd > 0) else 0.0
event = {
"event": "call_settled",
"call_id": call_id,
"timestamp": self._timestamp(),
"stage": reservation.get("stage"),
"model": reservation.get("model"),
"provider": reservation.get("provider"),
"status": "recovered_ambiguous_failure",
"reserved_usd": reserved_usd,
"cost_usd": charged,
"cost_source": "reservation" if charged else "none",
"input_tokens": 0,
"cached_input_tokens": 0,
"output_tokens": 0,
"metadata": (
reservation["metadata"] if isinstance(reservation.get("metadata"), dict) else {}
),
"error": "Process exited before spend settlement; charged on resume",
}
self._records.append(event)
self._spent_usd += charged
self._persist_event_locked(event)

@property
def enforcing(self) -> bool:
"""Whether this run has a finite non-zero dollar cap."""
Expand Down Expand Up @@ -349,6 +467,7 @@ def reserve_call(
"stage": stage,
"model": model,
"provider": provider,
"budget_enforcing": self.enforcing,
"reserved_usd": reserved_usd,
"input_token_upper_bound": input_token_upper_bound,
"max_output_tokens": effective_max_tokens,
Expand Down
3 changes: 3 additions & 0 deletions clearwing/sourcehunt/findings_pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,9 @@ def count(self) -> int:
async def add(self, finding: Finding) -> Finding:
"""Add a finding, classify primitive, run dedup, return updated finding."""
async with self._lock:
existing = self._findings.get(finding.get("id", ""))
if existing is not None:
return existing
if not finding.get("primitive_type"):
finding.primitive_type = await self._classify_primitive(finding)

Expand Down
91 changes: 71 additions & 20 deletions clearwing/sourcehunt/pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,19 +72,33 @@ class WorkItem:
seed_transcript: str | None = None
entry_point: Any = None # EntryPoint | None — spec 004
seed_context: str | None = None # spec 004 seed corpus
context_id: str = "" # hash of per-file seeded-crash/static prompt inputs

def stable_identifier(self, run_id: str) -> str:
def stable_identifier(self, run_id: str, tier: str = "") -> str:
entry_point = self.entry_point
entry_point_id = (
{
"file": getattr(entry_point, "file_path", ""),
"function": getattr(entry_point, "function_name", ""),
"start": getattr(entry_point, "start_line", 0),
"end": getattr(entry_point, "end_line", 0),
"type": getattr(entry_point, "entry_type", ""),
}
if entry_point is not None
else None
)
return stable_run_id(
"work",
{
"run_id": run_id,
"file": self.file_target.get("path", ""),
"tier": tier,
"band": self.band,
"attempt": self.attempt,
"entry_point": (
getattr(entry_point, "function_name", "") if entry_point is not None else ""
),
"entry_point": entry_point_id,
"seed_context": self.seed_context,
"seed_transcript": self.seed_transcript,
"context_id": self.context_id,
},
)

Expand Down Expand Up @@ -212,6 +226,7 @@ class HuntPoolConfig:
findings_pool: Any = None # FindingsPool | None — spec 005
trajectory_root: str | Path | None = None
instrumentation: Any = None # SourceHuntInstrumentation | None
checkpoint: Any = None # SourceHuntCheckpoint | None


def _format_seed_context(entries: list) -> str | None:
Expand Down Expand Up @@ -297,6 +312,11 @@ def _expand_to_work_items(self, files: list[FileTarget], band: str) -> list[Work
else []
)
seed_entries = self.config.seed_corpus_by_file.get(file_path, [])
context = {
"seeded_crash": self.config.seeded_crashes_by_file.get(file_path),
"static_hints": self.config.semgrep_hints_by_file.get(file_path, []),
}
context_id = stable_run_id("context", context) if any(context.values()) else ""

if entry_points:
for ep in entry_points:
Expand All @@ -314,6 +334,7 @@ def _expand_to_work_items(self, files: list[FileTarget], band: str) -> list[Work
attempt=attempt,
entry_point=ep,
seed_context=seed_ctx,
context_id=context_id,
)
)
else:
Expand All @@ -325,6 +346,7 @@ def _expand_to_work_items(self, files: list[FileTarget], band: str) -> list[Work
band=band,
attempt=attempt,
seed_context=seed_ctx,
context_id=context_id,
)
)
return items
Expand Down Expand Up @@ -448,7 +470,7 @@ async def _run_tier_phase(
else None
)
spent = 0.0
in_flight: dict[asyncio.Task[TargetResult], WorkItem] = {}
in_flight: dict[asyncio.Future[TargetResult], tuple[WorkItem, str, bool]] = {}
item_iter = iter(work_items)
promotion_queue: list[WorkItem] = []

Expand All @@ -465,20 +487,35 @@ def _submit_next() -> bool:
if wi is None:
return False
band_cost = self.config.band_budget.for_band(wi.band)
work_item_id = wi.stable_identifier(self.config.session_id_prefix)
task = asyncio.create_task(
self._run_file_task(
wi.file_target,
cost_limit=band_cost,
tier=tier,
band=wi.band,
seed_transcript=wi.seed_transcript,
entry_point=wi.entry_point,
seed_context=wi.seed_context,
work_item_id=work_item_id,
work_item_id = wi.stable_identifier(self.config.session_id_prefix, tier)
cached = self.config.checkpoint.load(work_item_id) if self.config.checkpoint else None
if cached is not None and (
cached.target != wi.file_target.get("path", "")
or cached.tier != tier
or cached.band != wi.band
):
logger.warning("Ignoring mismatched cached hunter work %s", work_item_id)
cached = None
if cached is not None:
logger.info("Reusing completed hunter work for %s", wi.file_target.get("path", ""))
task = asyncio.get_running_loop().create_future()
task.set_result(cached)
from_cache = True
else:
task = asyncio.create_task(
self._run_file_task(
wi.file_target,
cost_limit=band_cost,
tier=tier,
band=wi.band,
seed_transcript=wi.seed_transcript,
entry_point=wi.entry_point,
seed_context=wi.seed_context,
work_item_id=work_item_id,
)
)
)
in_flight[task] = wi
from_cache = False
in_flight[task] = (wi, work_item_id, from_cache)
return True

for _ in range(max(1, self.config.max_parallel)):
Expand All @@ -498,7 +535,7 @@ def _submit_next() -> bool:
timeout,
len(in_flight),
)
for task, wi in list(in_flight.items()):
for task, (wi, _work_item_id, _from_cache) in list(in_flight.items()):
task.cancel()
key = wi.file_target.get("path", "")
async with self._state_lock:
Expand All @@ -513,7 +550,7 @@ def _submit_next() -> bool:
return spent

for task in done:
wi = in_flight.pop(task)
wi, work_item_id, from_cache = in_flight.pop(task)
key = wi.file_target.get("path", "")
try:
result = await task
Expand Down Expand Up @@ -544,6 +581,19 @@ def _submit_next() -> bool:
tier=tier,
band=wi.band,
)
if (
result.status == "completed"
and not from_cache
and self.config.checkpoint is not None
):
try:
self.config.checkpoint.save(work_item_id, result)
except Exception:
logger.warning(
"Could not checkpoint completed work %s",
work_item_id,
exc_info=True,
)
ep_suffix = f":{wi.entry_point.function_name}" if wi.entry_point else ""
async with self._state_lock:
self._results[f"{key}{ep_suffix}:{wi.band}:{wi.attempt}"] = result
Expand Down Expand Up @@ -624,6 +674,7 @@ def _submit_next() -> bool:
band=next_band,
attempt=wi.attempt,
seed_transcript=_extract_transcript(result),
context_id=wi.context_id,
)
)

Expand Down
Loading