548 lines
20 KiB
Python
548 lines
20 KiB
Python
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from tqdm.auto import tqdm
|
|
|
|
from .figures import render_figures
|
|
from .features import compute_features
|
|
from .layout import details_outputs_exist
|
|
from .parser import default_output_dir, infer_analysis_dataset_name, load_records, path_looks_like_release_trace
|
|
from .preparation import stream_prepare
|
|
from .report import write_features, write_normalized, write_report
|
|
from .reporting import write_reports
|
|
from .resume_advanced import collect_existing_detail_paths, run_advanced_from_existing
|
|
from .study import parse_input_length_bucket_thresholds, run_study
|
|
|
|
ANALYSIS_PROVENANCE_FILENAME = "analysis_provenance.json"
|
|
ANALYSIS_PROVENANCE_SCHEMA_VERSION = 1
|
|
|
|
|
|
def build_parser():
|
|
parser = argparse.ArgumentParser(description="Analyze coding-agent trace patterns.")
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
analyze_parser = subparsers.add_parser(
|
|
"analyze",
|
|
help="Run the full analysis workflow from one formatter-generated *-raw.jsonl trace.",
|
|
)
|
|
analyze_parser.add_argument("input", help="Path to the formatter-generated *-raw.jsonl trace.")
|
|
analyze_parser.add_argument(
|
|
"--release-input",
|
|
default=None,
|
|
help="Path to the formatter-generated release .jsonl with hash_ids. Defaults to the sibling file without the `-raw` suffix.",
|
|
)
|
|
analyze_parser.add_argument(
|
|
"--dataset-name",
|
|
default=None,
|
|
help="Dataset name used for output paths and figure titles. Defaults to the formatted trace stem.",
|
|
)
|
|
analyze_parser.add_argument(
|
|
"--output-dir",
|
|
default=None,
|
|
help="Explicit analysis output directory. Defaults to outputs/analysis/<dataset>/",
|
|
)
|
|
analyze_parser.add_argument("--output-root", default="outputs/analysis")
|
|
analyze_parser.add_argument(
|
|
"--figure-dir",
|
|
default=None,
|
|
help="Explicit figure directory. Defaults to <output-dir>/figures/.",
|
|
)
|
|
analyze_parser.add_argument(
|
|
"--block-size",
|
|
type=int,
|
|
default=256,
|
|
help="Block size for theoretical cache analysis.",
|
|
)
|
|
analyze_parser.add_argument(
|
|
"--segment-mode",
|
|
default="tokenizer",
|
|
choices=["bytes", "tokenizer"],
|
|
help="How to segment prompts for theoretical cache analysis.",
|
|
)
|
|
analyze_parser.add_argument(
|
|
"--tokenizer-path",
|
|
default=None,
|
|
help="Local path or model id for tokenizer mode. Defaults to the local resolved tokenizer path.",
|
|
)
|
|
analyze_parser.add_argument(
|
|
"--tokenizer-batch-size",
|
|
type=int,
|
|
default=64,
|
|
help="Batch size used by tokenizer-based theoretical cache analysis.",
|
|
)
|
|
analyze_parser.add_argument(
|
|
"--model-family",
|
|
default="auto",
|
|
help="Model family for tokenizer/chat-template metadata. Defaults to auto-detect.",
|
|
)
|
|
analyze_parser.add_argument(
|
|
"--model-meta-dir",
|
|
default=None,
|
|
help="Override the base directory that contains model_meta/<provider>/<model>/.",
|
|
)
|
|
analyze_parser.add_argument(
|
|
"--input-length-buckets",
|
|
default=None,
|
|
help="Semicolon-separated input-length bucket thresholds in tokens, such as `32768;87040;131072` or `32Ki;85Ki;128Ki`.",
|
|
)
|
|
|
|
parse_parser = subparsers.add_parser("parse", help="Normalize a formatter-generated *-raw.jsonl trace.")
|
|
_add_common_args(parse_parser)
|
|
parse_parser.add_argument(
|
|
"--format",
|
|
default="jsonl",
|
|
choices=["jsonl", "csv", "parquet"],
|
|
help="Normalized output format.",
|
|
)
|
|
|
|
features_parser = subparsers.add_parser("features", help="Extract request-level features.")
|
|
_add_common_args(features_parser)
|
|
|
|
report_parser = subparsers.add_parser("report", help="Generate markdown and json summary reports.")
|
|
_add_common_args(report_parser)
|
|
report_parser.add_argument(
|
|
"--normalized-format",
|
|
default="jsonl",
|
|
choices=["jsonl", "csv", "parquet"],
|
|
help="Also emit normalized records in this format.",
|
|
)
|
|
|
|
study_parser = subparsers.add_parser(
|
|
"study",
|
|
help="Generate data tables and CDF plots for lengths, cache reuse, and tool timing.",
|
|
)
|
|
_add_common_args(study_parser)
|
|
study_parser.add_argument(
|
|
"--normalized-format",
|
|
default="jsonl",
|
|
choices=["jsonl", "csv", "parquet"],
|
|
help="Normalized output format.",
|
|
)
|
|
study_parser.add_argument(
|
|
"--block-size",
|
|
type=int,
|
|
default=256,
|
|
help="Block size for theoretical cache analysis.",
|
|
)
|
|
study_parser.add_argument(
|
|
"--segment-mode",
|
|
default="tokenizer",
|
|
choices=["bytes", "tokenizer"],
|
|
help="How to segment prompts for theoretical cache analysis.",
|
|
)
|
|
study_parser.add_argument(
|
|
"--tokenizer-path",
|
|
default=None,
|
|
help="Local path or model id for tokenizer mode. Defaults to the local resolved tokenizer path.",
|
|
)
|
|
study_parser.add_argument(
|
|
"--tokenizer-batch-size",
|
|
type=int,
|
|
default=64,
|
|
help="Batch size used by tokenizer-based theoretical cache analysis.",
|
|
)
|
|
study_parser.add_argument(
|
|
"--model-family",
|
|
default="auto",
|
|
help="Model family for tokenizer/chat-template metadata. Defaults to auto-detect.",
|
|
)
|
|
study_parser.add_argument(
|
|
"--model-meta-dir",
|
|
default=None,
|
|
help="Override the base directory that contains model_meta/<provider>/<model>/.",
|
|
)
|
|
study_parser.add_argument(
|
|
"--input-length-buckets",
|
|
default=None,
|
|
help="Semicolon-separated input-length bucket thresholds in tokens, such as `32768;87040;131072` or `32Ki;85Ki;128Ki`.",
|
|
)
|
|
|
|
resume_parser = subparsers.add_parser(
|
|
"resume-details",
|
|
aliases=["resume-advanced"],
|
|
help="Reuse existing source trace (*-raw.jsonl or legacy normalized.jsonl) + features.csv and compute only detailed analysis outputs.",
|
|
)
|
|
resume_parser.add_argument("input", help="Path to formatter-generated *-raw.jsonl")
|
|
resume_parser.add_argument("features", help="Path to existing features.csv")
|
|
resume_parser.add_argument(
|
|
"--release-input",
|
|
default=None,
|
|
help="Path to the formatter-generated release .jsonl with hash_ids. Defaults to the sibling file without the `-raw` suffix.",
|
|
)
|
|
resume_parser.add_argument(
|
|
"--output-dir",
|
|
required=True,
|
|
help="Existing output directory to receive detailed analysis outputs.",
|
|
)
|
|
resume_parser.add_argument(
|
|
"--block-size",
|
|
type=int,
|
|
default=256,
|
|
help="Block size for theoretical cache analysis.",
|
|
)
|
|
resume_parser.add_argument(
|
|
"--segment-mode",
|
|
default="tokenizer",
|
|
choices=["bytes", "tokenizer"],
|
|
help="How to segment prompts for theoretical cache analysis.",
|
|
)
|
|
resume_parser.add_argument(
|
|
"--tokenizer-path",
|
|
default=None,
|
|
help="Local path or model id for tokenizer mode. Defaults to the local resolved tokenizer path.",
|
|
)
|
|
resume_parser.add_argument(
|
|
"--tokenizer-batch-size",
|
|
type=int,
|
|
default=64,
|
|
help="Batch size used by tokenizer-based theoretical cache analysis.",
|
|
)
|
|
resume_parser.add_argument(
|
|
"--model-family",
|
|
default="auto",
|
|
help="Model family for tokenizer/chat-template metadata. Defaults to auto-detect.",
|
|
)
|
|
resume_parser.add_argument(
|
|
"--model-meta-dir",
|
|
default=None,
|
|
help="Override the base directory that contains model_meta/<provider>/<model>/.",
|
|
)
|
|
resume_parser.add_argument(
|
|
"--limit",
|
|
type=int,
|
|
default=None,
|
|
help="Only process the first N source/features rows. Useful for throughput benchmarking.",
|
|
)
|
|
resume_parser.add_argument(
|
|
"--input-length-buckets",
|
|
default=None,
|
|
help="Semicolon-separated input-length bucket thresholds in tokens, such as `32768;87040;131072` or `32Ki;85Ki;128Ki`.",
|
|
)
|
|
|
|
return parser
|
|
|
|
|
|
def _add_common_args(parser):
|
|
parser.add_argument("input", help="Path to the formatter-generated *-raw.jsonl trace.")
|
|
parser.add_argument("--limit", type=int, default=None, help="Limit number of input lines.")
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
default=None,
|
|
help="Output directory. Defaults to outputs/analysis/<input_stem>/",
|
|
)
|
|
|
|
|
|
def resolve_output_dir(input_path, output_dir):
|
|
return Path(output_dir) if output_dir else default_output_dir(input_path)
|
|
|
|
|
|
def _normalize_dataset_name(name: str) -> str:
|
|
text = str(name)
|
|
return text[:-4] if text.endswith("-raw") else text
|
|
|
|
|
|
def _resolve_analysis_output_dir(args):
|
|
dataset_name = args.dataset_name or _normalize_dataset_name(infer_analysis_dataset_name(args.input))
|
|
output_dir = Path(args.output_dir) if args.output_dir else Path(args.output_root) / dataset_name
|
|
figure_dir = Path(args.figure_dir) if args.figure_dir else output_dir / "figures"
|
|
return dataset_name, output_dir, figure_dir
|
|
|
|
|
|
def _resolve_release_input_path(raw_input: str, release_input: str | None) -> Path:
|
|
if release_input:
|
|
return Path(release_input)
|
|
raw_path = Path(raw_input)
|
|
name = raw_path.name
|
|
if name.endswith("-raw.jsonl"):
|
|
candidate = raw_path.with_name(name[:-len("-raw.jsonl")] + ".jsonl")
|
|
else:
|
|
raise ValueError("Expected a formatter-generated *-raw.jsonl input, or pass --release-input explicitly.")
|
|
return candidate
|
|
|
|
|
|
def _resolve_existing_release_input_path(raw_input: str, release_input: str | None) -> Path | None:
|
|
candidate = _resolve_release_input_path(raw_input, release_input)
|
|
if path_looks_like_release_trace(candidate):
|
|
return candidate
|
|
return None
|
|
|
|
|
|
def _existing_base_outputs(output_dir):
|
|
features = output_dir / "features.csv"
|
|
report = output_dir / "report.md"
|
|
if features.exists():
|
|
return {
|
|
"features": features,
|
|
"report": report if report.exists() else None,
|
|
}
|
|
return None
|
|
|
|
|
|
def _existing_detail_outputs(output_dir):
|
|
if not details_outputs_exist(output_dir):
|
|
return None
|
|
return collect_existing_detail_paths(output_dir)
|
|
|
|
|
|
def _file_fingerprint(path: str | Path) -> dict:
|
|
resolved = Path(path).resolve()
|
|
stat = resolved.stat()
|
|
return {
|
|
"path": str(resolved),
|
|
"size": stat.st_size,
|
|
"mtime_ns": stat.st_mtime_ns,
|
|
}
|
|
|
|
|
|
def _path_option(value: str | None) -> str | None:
|
|
if value is None:
|
|
return None
|
|
path = Path(value)
|
|
return str(path.resolve()) if path.exists() else str(value)
|
|
|
|
|
|
def _expected_analysis_provenance(args, release_input_path: Path, input_length_bucket_thresholds: list[int]) -> dict:
|
|
return {
|
|
"schema_version": ANALYSIS_PROVENANCE_SCHEMA_VERSION,
|
|
"raw_input": _file_fingerprint(args.input),
|
|
"release_input": _file_fingerprint(release_input_path),
|
|
"options": {
|
|
"block_size": args.block_size,
|
|
"segment_mode": args.segment_mode,
|
|
"tokenizer_path": _path_option(args.tokenizer_path),
|
|
"tokenizer_batch_size": args.tokenizer_batch_size,
|
|
"model_family": args.model_family,
|
|
"model_meta_dir": _path_option(args.model_meta_dir),
|
|
"input_length_buckets": input_length_bucket_thresholds,
|
|
},
|
|
}
|
|
|
|
|
|
def _analysis_provenance_path(output_dir: Path) -> Path:
|
|
return output_dir / ANALYSIS_PROVENANCE_FILENAME
|
|
|
|
|
|
def _load_analysis_provenance(output_dir: Path) -> dict | None:
|
|
path = _analysis_provenance_path(output_dir)
|
|
if not path.exists():
|
|
return None
|
|
try:
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
data = json.load(handle)
|
|
except (OSError, json.JSONDecodeError):
|
|
return None
|
|
return data if isinstance(data, dict) else None
|
|
|
|
|
|
def _analysis_provenance_matches(output_dir: Path, expected: dict) -> bool:
|
|
return _load_analysis_provenance(output_dir) == expected
|
|
|
|
|
|
def _write_analysis_provenance(output_dir: Path, provenance: dict) -> None:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
destination = _analysis_provenance_path(output_dir)
|
|
temp_path = destination.with_suffix(destination.suffix + ".tmp")
|
|
with temp_path.open("w", encoding="utf-8") as handle:
|
|
json.dump(provenance, handle, ensure_ascii=False, indent=2, sort_keys=True)
|
|
handle.write("\n")
|
|
temp_path.replace(destination)
|
|
|
|
|
|
def _stage_message(progress, step: int, total_steps: int, message: str) -> None:
|
|
tqdm.write(f"Stage {step}/{total_steps}: {message}")
|
|
progress.update(1)
|
|
progress.set_postfix(current=message)
|
|
|
|
|
|
def main(argv=None):
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
|
|
if args.command == "analyze":
|
|
dataset_name, output_dir, figure_dir = _resolve_analysis_output_dir(args)
|
|
input_length_bucket_thresholds = parse_input_length_bucket_thresholds(args.input_length_buckets)
|
|
release_input_path = _resolve_existing_release_input_path(args.input, args.release_input)
|
|
if release_input_path is None:
|
|
raise FileNotFoundError(
|
|
f"Release trace not found for raw trace {args.input}. "
|
|
"Run `python -m trace_formatter build-release <raw-trace>` first, or pass --release-input."
|
|
)
|
|
expected_provenance = _expected_analysis_provenance(args, release_input_path, input_length_bucket_thresholds)
|
|
reuse_allowed = _analysis_provenance_matches(output_dir, expected_provenance)
|
|
total_steps = 4
|
|
progress = tqdm(
|
|
total=total_steps,
|
|
desc="Analyze trace",
|
|
unit="stage",
|
|
dynamic_ncols=True,
|
|
)
|
|
try:
|
|
prepare_result = None
|
|
reusable_base = _existing_base_outputs(output_dir) if reuse_allowed else None
|
|
if reusable_base:
|
|
_stage_message(progress, 1, total_steps, "reuse existing features.csv")
|
|
prepare_result = {
|
|
"features_path": str(reusable_base["features"]),
|
|
"reused": True,
|
|
}
|
|
else:
|
|
_stage_message(progress, 1, total_steps, "prepare features.csv")
|
|
prepare_result = stream_prepare(args.input, output_dir, show_progress=True)
|
|
reusable_details = _existing_detail_outputs(output_dir) if reuse_allowed else None
|
|
if reusable_details:
|
|
_stage_message(progress, 2, total_steps, "reuse existing details/")
|
|
advanced_paths = reusable_details
|
|
else:
|
|
_stage_message(
|
|
progress,
|
|
2,
|
|
total_steps,
|
|
"detailed analysis: request metrics, tool/session stats, kvcache stats",
|
|
)
|
|
advanced_paths = run_advanced_from_existing(
|
|
args.input,
|
|
release_input_path,
|
|
prepare_result["features_path"],
|
|
output_dir,
|
|
input_length_bucket_thresholds=input_length_bucket_thresholds,
|
|
show_progress=True,
|
|
)
|
|
_stage_message(progress, 3, total_steps, "reporting: summary.json, report.md, analysis_snapshot.json")
|
|
report_result = write_reports(
|
|
features_path=prepare_result["features_path"],
|
|
output_dir=output_dir,
|
|
pipeline_summary={
|
|
"dataset_name": dataset_name,
|
|
"formatted_path": str(Path(args.input)),
|
|
"release_path": str(release_input_path),
|
|
**{key: str(value) for key, value in advanced_paths.items()},
|
|
},
|
|
)
|
|
_stage_message(
|
|
progress,
|
|
4,
|
|
total_steps,
|
|
"figures: 13 approved request/session/tool/kvcache plots",
|
|
)
|
|
figure_result = render_figures(
|
|
analysis_dir=output_dir,
|
|
fig_dir=figure_dir,
|
|
dataset_title=dataset_name,
|
|
show_progress=True,
|
|
)
|
|
_write_analysis_provenance(output_dir, expected_provenance)
|
|
finally:
|
|
progress.close()
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"dataset_name": dataset_name,
|
|
"formatted_path": str(Path(args.input)),
|
|
"output_dir": str(output_dir),
|
|
"prepare": prepare_result,
|
|
"details": {key: str(value) for key, value in advanced_paths.items()},
|
|
"report": report_result,
|
|
"figures": figure_result,
|
|
"release_path": str(release_input_path),
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
)
|
|
return 0
|
|
|
|
if args.command in {"resume-details", "resume-advanced"}:
|
|
input_length_bucket_thresholds = parse_input_length_bucket_thresholds(args.input_length_buckets)
|
|
release_input_path = _resolve_existing_release_input_path(args.input, args.release_input)
|
|
if release_input_path is None:
|
|
raise FileNotFoundError(
|
|
f"Release trace not found for raw trace {args.input}. "
|
|
"Run `python -m trace_formatter build-release <raw-trace>` first, or pass --release-input."
|
|
)
|
|
paths = run_advanced_from_existing(
|
|
args.input,
|
|
release_input_path,
|
|
args.features,
|
|
args.output_dir,
|
|
input_length_bucket_thresholds=input_length_bucket_thresholds,
|
|
show_progress=True,
|
|
limit=args.limit,
|
|
)
|
|
for path in paths.values():
|
|
print(path)
|
|
return 0
|
|
|
|
output_dir = resolve_output_dir(args.input, args.output_dir)
|
|
if args.command == "study" and args.limit is None:
|
|
input_length_bucket_thresholds = parse_input_length_bucket_thresholds(args.input_length_buckets)
|
|
reusable = _existing_base_outputs(output_dir)
|
|
if reusable:
|
|
release_input_path = _resolve_existing_release_input_path(args.input, None)
|
|
if release_input_path is not None:
|
|
paths = _existing_detail_outputs(output_dir)
|
|
if paths is None:
|
|
paths = run_advanced_from_existing(
|
|
args.input,
|
|
release_input_path,
|
|
reusable["features"],
|
|
output_dir,
|
|
input_length_bucket_thresholds=input_length_bucket_thresholds,
|
|
show_progress=True,
|
|
)
|
|
for path in paths.values():
|
|
print(path)
|
|
return 0
|
|
|
|
show_progress = args.command == "study"
|
|
records = load_records(
|
|
args.input,
|
|
limit=args.limit,
|
|
show_progress=show_progress,
|
|
progress_desc="Load trace",
|
|
)
|
|
|
|
if args.command == "parse":
|
|
path = write_normalized(records, output_dir, output_format=args.format)
|
|
print(path)
|
|
return 0
|
|
|
|
features = compute_features(records)
|
|
if args.command == "features":
|
|
path = write_features(features, output_dir)
|
|
print(path)
|
|
return 0
|
|
|
|
if args.command == "study":
|
|
input_length_bucket_thresholds = parse_input_length_bucket_thresholds(args.input_length_buckets)
|
|
paths = run_study(
|
|
records,
|
|
output_dir,
|
|
normalized_format=args.normalized_format,
|
|
source_path=args.input,
|
|
block_size=args.block_size,
|
|
segment_mode=args.segment_mode,
|
|
tokenizer_path=args.tokenizer_path,
|
|
model_family=args.model_family,
|
|
model_meta_dir=args.model_meta_dir,
|
|
input_length_bucket_thresholds=input_length_bucket_thresholds,
|
|
show_progress=show_progress,
|
|
tokenizer_batch_size=args.tokenizer_batch_size,
|
|
)
|
|
for path in paths.values():
|
|
print(path)
|
|
return 0
|
|
|
|
normalized_path = write_normalized(records, output_dir, output_format=args.normalized_format)
|
|
features_path = write_features(features, output_dir)
|
|
summary_path, report_path = write_report(records, features, output_dir)
|
|
print(normalized_path)
|
|
print(features_path)
|
|
print(summary_path)
|
|
print(report_path)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|