352 lines
12 KiB
Python
352 lines
12 KiB
Python
"""Command-line interface for deterministic Open-SWE-Traces profiling."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import shutil
|
|
import sys
|
|
from collections.abc import Callable, Iterator
|
|
from concurrent.futures import (
|
|
FIRST_COMPLETED,
|
|
Future,
|
|
ProcessPoolExecutor,
|
|
ThreadPoolExecutor,
|
|
as_completed,
|
|
wait,
|
|
)
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .heuristics import classify_metrics
|
|
from .io import (
|
|
append_jsonl,
|
|
get_sample_id,
|
|
iter_jsonl,
|
|
iter_parquet_file,
|
|
iter_records,
|
|
load_completed_ids,
|
|
parquet_files,
|
|
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
|
|
|
|
_PROFILE_COUNTER: TokenCounter | None = None
|
|
|
|
|
|
def _initialize_profile_process(tokenizer_json: str | None) -> None:
|
|
"""Create one tokenizer per process instead of per trajectory."""
|
|
|
|
global _PROFILE_COUNTER
|
|
_PROFILE_COUNTER = TokenCounter(Path(tokenizer_json)) if tokenizer_json else None
|
|
|
|
|
|
def _profile_parquet_shard(
|
|
shard: Path,
|
|
metrics_part: Path,
|
|
errors_part: Path | None,
|
|
) -> tuple[int, int]:
|
|
"""Profile one shard into atomic part files with bounded memory."""
|
|
|
|
processed = 0
|
|
failures = 0
|
|
metrics_temporary = metrics_part.with_suffix(metrics_part.suffix + ".tmp")
|
|
errors_temporary = (
|
|
errors_part.with_suffix(errors_part.suffix + ".tmp") if errors_part else None
|
|
)
|
|
metrics_part.parent.mkdir(parents=True, exist_ok=True)
|
|
with metrics_temporary.open("w", encoding="utf-8") as metrics_handle:
|
|
errors_handle = (
|
|
errors_temporary.open("w", encoding="utf-8")
|
|
if errors_temporary is not None
|
|
else None
|
|
)
|
|
try:
|
|
for record in iter_parquet_file(shard):
|
|
sample_id = get_sample_id(record)
|
|
try:
|
|
result = extract_metrics(record, _PROFILE_COUNTER)
|
|
metrics_handle.write(
|
|
json.dumps(result, ensure_ascii=False, default=str) + "\n"
|
|
)
|
|
except Exception as exc: # noqa: BLE001 - isolate bad records.
|
|
failures += 1
|
|
if errors_handle is not None:
|
|
error = {
|
|
"sample_id": sample_id,
|
|
"error_type": type(exc).__name__,
|
|
"error": str(exc)[:2000],
|
|
"source_parquet": str(shard.resolve()),
|
|
}
|
|
errors_handle.write(
|
|
json.dumps(error, ensure_ascii=False) + "\n"
|
|
)
|
|
processed += 1
|
|
finally:
|
|
if errors_handle is not None:
|
|
errors_handle.close()
|
|
os.replace(metrics_temporary, metrics_part)
|
|
if errors_temporary is not None and errors_part is not None:
|
|
os.replace(errors_temporary, errors_part)
|
|
return processed, failures
|
|
|
|
|
|
def _merge_parts(parts: list[Path], output: Path) -> None:
|
|
"""Concatenate completed shard outputs atomically in dataset order."""
|
|
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = output.with_suffix(output.suffix + ".tmp")
|
|
with temporary.open("wb") as destination:
|
|
for part in parts:
|
|
with part.open("rb") as source:
|
|
shutil.copyfileobj(source, destination, length=1024 * 1024)
|
|
os.replace(temporary, output)
|
|
|
|
|
|
def _run_parallel_parquet_profile(args: argparse.Namespace) -> int:
|
|
"""Profile independent Parquet shards in separate CPU processes."""
|
|
|
|
shards = parquet_files(args.input)
|
|
parts_dir = args.output.parent / f".{args.output.name}.parts"
|
|
parts_dir.mkdir(parents=True, exist_ok=True)
|
|
metrics_parts = [parts_dir / f"{index:05d}.metrics.jsonl" for index in range(len(shards))]
|
|
error_parts = [parts_dir / f"{index:05d}.errors.jsonl" for index in range(len(shards))]
|
|
|
|
if not args.resume:
|
|
for path in parts_dir.glob("*.jsonl"):
|
|
path.unlink()
|
|
for path in parts_dir.glob("*.tmp"):
|
|
path.unlink()
|
|
|
|
pending = []
|
|
processed = 0
|
|
failures = 0
|
|
tokenizer_json = str(args.tokenizer_json) if args.tokenizer_json else None
|
|
with ProcessPoolExecutor(
|
|
max_workers=args.workers,
|
|
initializer=_initialize_profile_process,
|
|
initargs=(tokenizer_json,),
|
|
) as executor:
|
|
for shard, metrics_part, errors_part in zip(
|
|
shards, metrics_parts, error_parts, strict=True
|
|
):
|
|
if args.resume and metrics_part.exists() and errors_part.exists():
|
|
continue
|
|
pending.append(
|
|
executor.submit(
|
|
_profile_parquet_shard,
|
|
shard,
|
|
metrics_part,
|
|
errors_part if args.errors else None,
|
|
)
|
|
)
|
|
for completed_shards, future in enumerate(as_completed(pending), 1):
|
|
shard_processed, shard_failures = future.result()
|
|
processed += shard_processed
|
|
failures += shard_failures
|
|
print(
|
|
f"profile: shards={completed_shards}/{len(pending)} "
|
|
f"processed={processed} failures={failures}",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
missing = [path for path in metrics_parts if not path.exists()]
|
|
if missing:
|
|
raise RuntimeError(f"Missing {len(missing)} completed metric parts")
|
|
_merge_parts(metrics_parts, args.output)
|
|
if args.errors:
|
|
missing_errors = [path for path in error_parts if not path.exists()]
|
|
if missing_errors:
|
|
raise RuntimeError(f"Missing {len(missing_errors)} completed error parts")
|
|
_merge_parts(error_parts, args.errors)
|
|
print(f"profile: completed={processed} failures={failures}", file=sys.stderr)
|
|
return 1 if failures else 0
|
|
|
|
|
|
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,
|
|
"error_type": type(error).__name__,
|
|
"error": str(error)[:2000],
|
|
}
|
|
if path is None:
|
|
print(json.dumps(value, ensure_ascii=False), file=sys.stderr)
|
|
else:
|
|
append_jsonl(path, value)
|
|
|
|
|
|
def _run_streaming_stage(
|
|
*,
|
|
input_path: Path,
|
|
output_path: Path,
|
|
errors_path: Path | None,
|
|
limit: int | None,
|
|
resume: bool,
|
|
workers: int,
|
|
processor: Callable[[dict[str, Any]], dict[str, Any]],
|
|
) -> int:
|
|
"""Profile records concurrently while serializing writes in one thread."""
|
|
|
|
if workers < 1:
|
|
raise ValueError("workers must be at least 1")
|
|
completed = load_completed_ids(output_path) if resume else set()
|
|
processed = 0
|
|
failures = 0
|
|
|
|
def selected_records() -> Iterator[dict[str, Any]]:
|
|
selected = 0
|
|
for record in iter_records(input_path):
|
|
if get_sample_id(record) in completed:
|
|
continue
|
|
if limit is not None and selected >= limit:
|
|
break
|
|
selected += 1
|
|
yield record
|
|
|
|
def process_one(
|
|
record: dict[str, Any],
|
|
) -> tuple[str, dict[str, Any] | None, Exception | None]:
|
|
sample_id = get_sample_id(record)
|
|
try:
|
|
return sample_id, processor(record), None
|
|
except Exception as exc: # noqa: BLE001 - isolate individual records.
|
|
return sample_id, None, exc
|
|
|
|
def record_outcome(
|
|
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, error)
|
|
elif result is not None:
|
|
append_jsonl(output_path, result)
|
|
processed += 1
|
|
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="profile") as executor:
|
|
pending: set[
|
|
Future[tuple[str, dict[str, Any] | None, Exception | None]]
|
|
] = set()
|
|
for _ in range(workers):
|
|
try:
|
|
pending.add(executor.submit(process_one, next(records)))
|
|
except StopIteration:
|
|
break
|
|
while pending:
|
|
finished, pending = wait(pending, return_when=FIRST_COMPLETED)
|
|
for future in finished:
|
|
record_outcome(future.result())
|
|
try:
|
|
pending.add(executor.submit(process_one, next(records)))
|
|
except StopIteration:
|
|
pass
|
|
|
|
print(f"profile: completed={processed} failures={failures}", file=sys.stderr)
|
|
return 1 if failures else 0
|
|
|
|
|
|
def command_profile(args: argparse.Namespace) -> int:
|
|
"""Extract per-trajectory metrics without modifying source records."""
|
|
|
|
if args.input.is_dir() and args.workers > 1 and args.limit is None:
|
|
return _run_parallel_parquet_profile(args)
|
|
counter = TokenCounter(args.tokenizer_json) if args.tokenizer_json else None
|
|
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,
|
|
processor=lambda record: extract_metrics(record, counter),
|
|
)
|
|
|
|
|
|
def command_summarize(args: argparse.Namespace) -> int:
|
|
"""Derive global thresholds, summary statistics, and decisions."""
|
|
|
|
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 command_sample(args: argparse.Namespace) -> int:
|
|
"""Materialize a deterministic stratified review sample."""
|
|
|
|
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:
|
|
"""Build the intentionally small three-command interface."""
|
|
|
|
parser = argparse.ArgumentParser(prog="swe-qc", description=__doc__)
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
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)
|
|
|
|
summarize = subparsers.add_parser(
|
|
"summarize", help="Build summary and deterministic decisions"
|
|
)
|
|
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)
|
|
|
|
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:
|
|
"""Run the selected command."""
|
|
|
|
args = build_parser().parse_args(argv)
|
|
return int(args.func(args))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|