Replace LLM cleanup with deterministic profiling
This commit is contained in:
jiachun
+80
-237
@@ -1,4 +1,4 @@
|
||||
"""Command-line interface for static Open-SWE-Traces quality control."""
|
||||
"""Command-line interface for deterministic Open-SWE-Traces profiling."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -10,20 +10,27 @@ from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .client import GLMClient
|
||||
from .config import Settings
|
||||
from .features import extract_static_signals
|
||||
from .io import append_jsonl, get_sample_id, iter_records, load_completed_ids, load_jsonl_index
|
||||
from .repair import apply_static_repair
|
||||
from .workflow import audit_trajectory, classify_record, plan_static_repair, review_repaired_record
|
||||
from .heuristics import classify_metrics
|
||||
from .io import (
|
||||
append_jsonl,
|
||||
get_sample_id,
|
||||
iter_jsonl,
|
||||
iter_records,
|
||||
load_completed_ids,
|
||||
write_json,
|
||||
write_jsonl,
|
||||
)
|
||||
from .metrics import extract_metrics
|
||||
from .sampling import iter_review_records, select_review_ids
|
||||
from .summary import build_summary
|
||||
from .tokenization import TokenCounter
|
||||
|
||||
|
||||
def _write_error(path: Path | None, sample_id: str, stage: str, error: Exception) -> None:
|
||||
"""Write a compact non-sensitive error record for later retry."""
|
||||
def _write_error(path: Path | None, sample_id: str, error: Exception) -> None:
|
||||
"""Write a compact error without interrupting a long profile run."""
|
||||
|
||||
value = {
|
||||
"sample_id": sample_id,
|
||||
"stage": stage,
|
||||
"error_type": type(error).__name__,
|
||||
"error": str(error)[:2000],
|
||||
}
|
||||
@@ -41,10 +48,9 @@ def _run_streaming_stage(
|
||||
limit: int | None,
|
||||
resume: bool,
|
||||
workers: int,
|
||||
stage_name: str,
|
||||
processor: Callable[[dict[str, Any]], dict[str, Any] | None],
|
||||
processor: Callable[[dict[str, Any]], dict[str, Any]],
|
||||
) -> int:
|
||||
"""Run one append-only stage with resume support and progress reporting."""
|
||||
"""Profile records concurrently while serializing writes in one thread."""
|
||||
|
||||
if workers < 1:
|
||||
raise ValueError("workers must be at least 1")
|
||||
@@ -55,8 +61,7 @@ def _run_streaming_stage(
|
||||
def selected_records() -> Iterator[dict[str, Any]]:
|
||||
selected = 0
|
||||
for record in iter_records(input_path):
|
||||
sample_id = get_sample_id(record)
|
||||
if sample_id in completed:
|
||||
if get_sample_id(record) in completed:
|
||||
continue
|
||||
if limit is not None and selected >= limit:
|
||||
break
|
||||
@@ -68,58 +73,55 @@ def _run_streaming_stage(
|
||||
) -> tuple[str, dict[str, Any] | None, Exception | None]:
|
||||
sample_id = get_sample_id(record)
|
||||
try:
|
||||
result = processor(record)
|
||||
except Exception as exc: # noqa: BLE001 - each sample must fail independently.
|
||||
return sample_id, processor(record), None
|
||||
except Exception as exc: # noqa: BLE001 - isolate individual records.
|
||||
return sample_id, None, exc
|
||||
return sample_id, result, None
|
||||
|
||||
def record_outcome(
|
||||
outcome: tuple[str, dict[str, Any] | None, Exception | None],
|
||||
outcome: tuple[str, dict[str, Any] | None, Exception | None]
|
||||
) -> None:
|
||||
nonlocal processed, failures
|
||||
sample_id, result, error = outcome
|
||||
if error is not None:
|
||||
failures += 1
|
||||
_write_error(errors_path, sample_id, stage_name, error)
|
||||
_write_error(errors_path, sample_id, error)
|
||||
elif result is not None:
|
||||
append_jsonl(output_path, result)
|
||||
processed += 1
|
||||
if processed % 10 == 0:
|
||||
print(f"{stage_name}: processed={processed} failures={failures}", file=sys.stderr)
|
||||
if processed % 1000 == 0:
|
||||
print(f"profile: processed={processed} failures={failures}", file=sys.stderr)
|
||||
|
||||
records = iter(selected_records())
|
||||
if workers == 1:
|
||||
for record in records:
|
||||
record_outcome(process_one(record))
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=workers, thread_name_prefix=stage_name) as executor:
|
||||
with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="profile") as executor:
|
||||
pending: set[
|
||||
Future[tuple[str, dict[str, Any] | None, Exception | None]]
|
||||
] = set()
|
||||
for _ in range(workers):
|
||||
try:
|
||||
record = next(records)
|
||||
pending.add(executor.submit(process_one, next(records)))
|
||||
except StopIteration:
|
||||
break
|
||||
pending.add(executor.submit(process_one, record))
|
||||
|
||||
while pending:
|
||||
finished, pending = wait(pending, return_when=FIRST_COMPLETED)
|
||||
for future in finished:
|
||||
record_outcome(future.result())
|
||||
try:
|
||||
record = next(records)
|
||||
pending.add(executor.submit(process_one, next(records)))
|
||||
except StopIteration:
|
||||
continue
|
||||
pending.add(executor.submit(process_one, record))
|
||||
pass
|
||||
|
||||
print(f"{stage_name}: completed={processed} failures={failures}", file=sys.stderr)
|
||||
print(f"profile: completed={processed} failures={failures}", file=sys.stderr)
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
def command_features(args: argparse.Namespace) -> int:
|
||||
"""Extract deterministic features without requiring an API key."""
|
||||
def command_profile(args: argparse.Namespace) -> int:
|
||||
"""Extract per-trajectory metrics without modifying source records."""
|
||||
|
||||
counter = TokenCounter(args.tokenizer_json) if args.tokenizer_json else None
|
||||
return _run_streaming_stage(
|
||||
input_path=args.input,
|
||||
output_path=args.output,
|
||||
@@ -127,236 +129,77 @@ def command_features(args: argparse.Namespace) -> int:
|
||||
limit=args.limit,
|
||||
resume=args.resume,
|
||||
workers=args.workers,
|
||||
stage_name="features",
|
||||
processor=lambda record: {
|
||||
"sample_id": get_sample_id(record),
|
||||
"resolved": record.get("resolved"),
|
||||
"static_signals": extract_static_signals(record),
|
||||
},
|
||||
processor=lambda record: extract_metrics(record, counter),
|
||||
)
|
||||
|
||||
|
||||
def command_classify(args: argparse.Namespace) -> int:
|
||||
"""Classify records through GLM-5.2 and local policy enforcement."""
|
||||
def command_summarize(args: argparse.Namespace) -> int:
|
||||
"""Derive global thresholds, summary statistics, and decisions."""
|
||||
|
||||
settings = Settings.from_env()
|
||||
with GLMClient(settings) as client:
|
||||
return _run_streaming_stage(
|
||||
input_path=args.input,
|
||||
output_path=args.output,
|
||||
errors_path=args.errors,
|
||||
limit=args.limit,
|
||||
resume=args.resume,
|
||||
workers=args.workers,
|
||||
stage_name="classification",
|
||||
processor=lambda record: classify_record(record, client),
|
||||
)
|
||||
|
||||
|
||||
def command_audit(args: argparse.Namespace) -> int:
|
||||
"""Locate a safe boundary, materialize its prefix, then score the prefix."""
|
||||
|
||||
settings = Settings.from_env()
|
||||
with GLMClient(settings) as client:
|
||||
return _run_streaming_stage(
|
||||
input_path=args.input,
|
||||
output_path=args.output,
|
||||
errors_path=args.errors,
|
||||
limit=args.limit,
|
||||
resume=args.resume,
|
||||
workers=args.workers,
|
||||
stage_name="two_call_audit",
|
||||
processor=lambda record: audit_trajectory(record, client),
|
||||
)
|
||||
|
||||
|
||||
def command_repair_plan(args: argparse.Namespace) -> int:
|
||||
"""Ask GLM-5.2 for allowlisted static repair plans."""
|
||||
|
||||
classifications = load_jsonl_index(args.classifications)
|
||||
settings = Settings.from_env()
|
||||
with GLMClient(settings) as client:
|
||||
def processor(record: dict[str, Any]) -> dict[str, Any] | None:
|
||||
sample_id = get_sample_id(record)
|
||||
classification = classifications[sample_id]
|
||||
if classification["qc_decision"] != "STATIC_REPAIR" and not args.include_nonrepair:
|
||||
return None
|
||||
return plan_static_repair(record, classification, client)
|
||||
|
||||
return _run_streaming_stage(
|
||||
input_path=args.input,
|
||||
output_path=args.output,
|
||||
errors_path=args.errors,
|
||||
limit=args.limit,
|
||||
resume=args.resume,
|
||||
workers=args.workers,
|
||||
stage_name="repair_plan",
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
|
||||
def command_apply_repair(args: argparse.Namespace) -> int:
|
||||
"""Apply validated static repair plans deterministically without API calls."""
|
||||
|
||||
plans = load_jsonl_index(args.plans)
|
||||
completed = load_completed_ids(args.output) if args.resume else set()
|
||||
processed = 0
|
||||
failures = 0
|
||||
for record in iter_records(args.input):
|
||||
sample_id = get_sample_id(record)
|
||||
if sample_id in completed or sample_id not in plans:
|
||||
continue
|
||||
if args.limit is not None and processed >= args.limit:
|
||||
break
|
||||
try:
|
||||
repaired, diff = apply_static_repair(record, plans[sample_id])
|
||||
repaired["sample_id"] = sample_id
|
||||
repaired["static_repair_diff"] = diff
|
||||
append_jsonl(args.output, repaired)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
failures += 1
|
||||
_write_error(args.errors, sample_id, "apply_repair", exc)
|
||||
processed += 1
|
||||
print(f"apply_repair: completed={processed} failures={failures}", file=sys.stderr)
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
def command_review(args: argparse.Namespace) -> int:
|
||||
"""Run independent GLM review of applied static repairs."""
|
||||
|
||||
classifications = load_jsonl_index(args.classifications)
|
||||
plans = load_jsonl_index(args.plans)
|
||||
repaired_records = load_jsonl_index(args.repaired)
|
||||
settings = Settings.from_env()
|
||||
with GLMClient(settings) as client:
|
||||
def processor(record: dict[str, Any]) -> dict[str, Any] | None:
|
||||
sample_id = get_sample_id(record)
|
||||
if sample_id not in repaired_records:
|
||||
return None
|
||||
repaired = repaired_records[sample_id]
|
||||
return review_repaired_record(
|
||||
original_record=record,
|
||||
classification=classifications[sample_id],
|
||||
repair_plan=plans[sample_id],
|
||||
repaired_record=repaired,
|
||||
structured_diff=repaired.get("static_repair_diff") or {},
|
||||
client=client,
|
||||
)
|
||||
|
||||
return _run_streaming_stage(
|
||||
input_path=args.input,
|
||||
output_path=args.output,
|
||||
errors_path=args.errors,
|
||||
limit=args.limit,
|
||||
resume=args.resume,
|
||||
workers=args.workers,
|
||||
stage_name="repair_review",
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
|
||||
def command_smoke_test(_: argparse.Namespace) -> int:
|
||||
"""Perform one minimal authenticated structured-output request."""
|
||||
|
||||
settings = Settings.from_env()
|
||||
schema = {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"required": ["status", "model"],
|
||||
"properties": {
|
||||
"status": {"const": "ok"},
|
||||
"model": {"type": "string", "minLength": 1},
|
||||
},
|
||||
}
|
||||
with GLMClient(settings) as client:
|
||||
response = client.invoke_json(
|
||||
system_prompt=(
|
||||
"Return a JSON object with status='ok' and model set to the model name you were asked to use."
|
||||
),
|
||||
payload={"requested_model": settings.model},
|
||||
schema=schema,
|
||||
)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"endpoint": settings.endpoint,
|
||||
"model": settings.model,
|
||||
"response": response.data,
|
||||
"request_id": response.request_id,
|
||||
"usage": response.usage,
|
||||
"compatibility_fallback_used": response.compatibility_fallback_used,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
summary, threshold = build_summary(args.metrics)
|
||||
write_json(args.summary, summary)
|
||||
write_jsonl(
|
||||
args.decisions,
|
||||
(
|
||||
classify_metrics(metrics, threshold)
|
||||
for metrics in iter_jsonl(args.metrics)
|
||||
),
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def _add_stream_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
"""Add arguments shared by streaming pipeline stages."""
|
||||
def command_sample(args: argparse.Namespace) -> int:
|
||||
"""Materialize a deterministic stratified review sample."""
|
||||
|
||||
parser.add_argument("--input", type=Path, required=True, help="JSONL file or dataset directory")
|
||||
parser.add_argument("--output", type=Path, required=True, help="Append-only output JSONL")
|
||||
parser.add_argument("--errors", type=Path, help="Optional error JSONL")
|
||||
parser.add_argument("--limit", type=int, help="Maximum number of new records")
|
||||
parser.add_argument("--resume", action="store_true", help="Skip IDs already present in output")
|
||||
parser.add_argument(
|
||||
"--workers",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Concurrent record workers; API stages can use this for parallel requests",
|
||||
selected, decisions = select_review_ids(args.decisions, args.per_group, args.seed)
|
||||
write_jsonl(
|
||||
args.output,
|
||||
iter_review_records(args.input, selected, decisions),
|
||||
)
|
||||
print(f"sample: selected={len(selected)}", file=sys.stderr)
|
||||
return 0
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
"""Construct the complete command-line parser."""
|
||||
"""Build the intentionally small three-command interface."""
|
||||
|
||||
parser = argparse.ArgumentParser(prog="swe-qc", description=__doc__)
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
features = subparsers.add_parser("features", help="Extract deterministic static QC features")
|
||||
_add_stream_arguments(features)
|
||||
features.set_defaults(func=command_features)
|
||||
profile = subparsers.add_parser("profile", help="Extract deterministic metrics")
|
||||
profile.add_argument("--input", type=Path, required=True)
|
||||
profile.add_argument("--output", type=Path, required=True)
|
||||
profile.add_argument("--errors", type=Path)
|
||||
profile.add_argument("--tokenizer-json", type=Path)
|
||||
profile.add_argument("--limit", type=int)
|
||||
profile.add_argument("--resume", action="store_true")
|
||||
profile.add_argument("--workers", type=int, default=1)
|
||||
profile.set_defaults(func=command_profile)
|
||||
|
||||
classify = subparsers.add_parser("classify", help="Classify trajectories through GLM-5.2")
|
||||
_add_stream_arguments(classify)
|
||||
classify.set_defaults(func=command_classify)
|
||||
|
||||
audit = subparsers.add_parser(
|
||||
"audit", help="Locate a safe boundary, then score only the retained prefix"
|
||||
summarize = subparsers.add_parser(
|
||||
"summarize", help="Build summary and deterministic decisions"
|
||||
)
|
||||
_add_stream_arguments(audit)
|
||||
audit.set_defaults(func=command_audit)
|
||||
summarize.add_argument("--metrics", type=Path, required=True)
|
||||
summarize.add_argument("--summary", type=Path, required=True)
|
||||
summarize.add_argument("--decisions", type=Path, required=True)
|
||||
summarize.set_defaults(func=command_summarize)
|
||||
|
||||
repair_plan = subparsers.add_parser("repair-plan", help="Create allowlisted static repair plans")
|
||||
_add_stream_arguments(repair_plan)
|
||||
repair_plan.add_argument("--classifications", type=Path, required=True)
|
||||
repair_plan.add_argument("--include-nonrepair", action="store_true")
|
||||
repair_plan.set_defaults(func=command_repair_plan)
|
||||
|
||||
apply_repair = subparsers.add_parser("apply-repair", help="Apply static plans deterministically")
|
||||
_add_stream_arguments(apply_repair)
|
||||
apply_repair.add_argument("--plans", type=Path, required=True)
|
||||
apply_repair.set_defaults(func=command_apply_repair)
|
||||
|
||||
review = subparsers.add_parser("review", help="Independently review repaired records")
|
||||
_add_stream_arguments(review)
|
||||
review.add_argument("--classifications", type=Path, required=True)
|
||||
review.add_argument("--plans", type=Path, required=True)
|
||||
review.add_argument("--repaired", type=Path, required=True)
|
||||
review.set_defaults(func=command_review)
|
||||
|
||||
smoke_test = subparsers.add_parser("smoke-test", help="Test API authentication and JSON output")
|
||||
smoke_test.set_defaults(func=command_smoke_test)
|
||||
sample = subparsers.add_parser(
|
||||
"sample", help="Sample full trajectories by rule and length bucket"
|
||||
)
|
||||
sample.add_argument("--input", type=Path, required=True)
|
||||
sample.add_argument("--decisions", type=Path, required=True)
|
||||
sample.add_argument("--output", type=Path, required=True)
|
||||
sample.add_argument("--per-group", type=int, default=20)
|
||||
sample.add_argument("--seed", type=int, default=20260818)
|
||||
sample.set_defaults(func=command_sample)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
"""CLI entry point used by the ``swe-qc`` console script."""
|
||||
"""Run the selected command."""
|
||||
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
args = build_parser().parse_args(argv)
|
||||
return int(args.func(args))
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user