gahow/ali-trace-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(formatter): resumable chunk merge with bounded fan-in

Browse Source 
Collapse chunk_*.jsonl in rounds with a bounded merge fan-in so the number of
simultaneously open chunk iterators stays flat regardless of how many chunks a
trace produces, and write a resume manifest so --resume-tmp-dir can continue
from an existing formatter temp directory without rescanning the raw trace.
This commit is contained in:
Gahow Wang
2026-09-11 06:30:48 +00:00
parent 5a6b1acb49
commit 5ee98f9e4d
3 changed files with 295 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

94
tests/test_ali_trace_pipeline.py
View File

@@ -8,6 +8,7 @@ from datetime import datetime, timezone
from pathlib import Path
from unittest import mock
import trace_formatter.formatting as formatting
from trace_analyzer.cli import main as analyzer_main
from trace_analyzer.parser import load_records
from trace_formatter.cli import main as formatter_main
@@ -684,6 +685,99 @@ class AliTracePipelineTest(unittest.TestCase):
self.assertEqual(stats["row_count"], 1)
formatted_rows = [json.loads(line) for line in output_path.read_text(encoding="utf-8").splitlines()]
self.assertEqual([row["meta"]["request_id"] for row in formatted_rows], ["req-zst"])
def test_format_and_sort_trace_limits_simultaneous_chunk_iterators_during_merge(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
input_dir = root / "raw"
input_dir.mkdir()
output_path = root / "formatted.jsonl"
with (input_dir / "0417-1500-1530.jsonl").open("w", encoding="utf-8") as handle:
for index in range(5):
handle.write(
json.dumps(
make_raw_row(
f"req-{index}",
utc_ms(f"2026-04-17 15:00:0{index + 1}.000"),
)
)
+ "\n"
)
original_iter_chunk_rows = formatting._iter_chunk_rows
active_iterators = 0
max_active_iterators = 0
def counting_iter_chunk_rows(path):
nonlocal active_iterators, max_active_iterators
active_iterators += 1
max_active_iterators = max(max_active_iterators, active_iterators)
try:
yield from original_iter_chunk_rows(path)
finally:
active_iterators -= 1
with mock.patch("trace_formatter.formatting._resolve_merge_fan_in", return_value=2):
with mock.patch("trace_formatter.formatting._iter_chunk_rows", side_effect=counting_iter_chunk_rows):
stats = format_and_sort_trace(input_dir=input_dir, output_path=output_path, chunk_bytes=256)
self.assertEqual(stats["chunk_count"], 5)
self.assertEqual(max_active_iterators, 2)
self.assertEqual(
[record.meta.request_id for record in load_records(output_path)],
["req-0", "req-1", "req-2", "req-3", "req-4"],
)
def test_format_and_sort_trace_can_resume_from_existing_chunk_dir_without_rescanning(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
input_dir = root / "raw"
input_dir.mkdir()
output_path = root / "formatted.jsonl"
resume_dir = root / "resume-tmp"
resume_dir.mkdir()
raw_rows = [
make_raw_row("req-1", utc_ms("2026-04-17 15:00:01.000")),
make_raw_row("req-2", utc_ms("2026-04-17 15:00:02.000"), user_id="user-2"),
]
with (input_dir / "0417-1500-1530.jsonl").open("w", encoding="utf-8") as handle:
for row in raw_rows:
handle.write(json.dumps(row) + "\n")
chunk_rows = []
for seq, raw in enumerate(reversed(raw_rows)):
request_params, response_params, request_payload, attributes, messages = formatting._extract_request_components(raw)
sort_time_ms, user_id, request_id, sequence_hashes, roles, row_json = formatting._build_unified_row_payload_from_components(
raw,
request_params=request_params,
response_params=response_params,
request_payload=request_payload,
request_attributes=attributes,
messages=messages,
source_file="0417-1500-1530.jsonl",
source_line=seq + 1,
)
chunk_rows.append((sort_time_ms, seq, user_id, request_id, sequence_hashes, roles, row_json))
formatting._write_chunk([chunk_rows[0]], resume_dir, 0)
formatting._write_chunk([chunk_rows[1]], resume_dir, 1)
with mock.patch("trace_formatter.formatting.open_trace_text", side_effect=AssertionError("unexpected rescan")):
stats = format_and_sort_trace(
input_dir=input_dir,
output_path=output_path,
resume_tmp_dir=resume_dir,
)
self.assertEqual(stats["row_count"], 2)
self.assertEqual(stats["chunk_count"], 2)
self.assertIsNone(stats["filtered_row_count"])
self.assertEqual(
[record.meta.request_id for record in load_records(output_path)],
["req-1", "req-2"],
)
self.assertFalse(resume_dir.exists())
def test_export_release_ready_trace_defaults_temp_dir_to_output_parent(self):
with tempfile.TemporaryDirectory() as temp_dir:

6
trace_formatter/cli.py
View File

@@ -64,6 +64,11 @@ def build_parser() -> argparse.ArgumentParser:
help="Base directory used when --output is omitted. Defaults to a sibling trace-*-formatted directory.",
)
format_parser.add_argument("--tmp-dir", default=None)
format_parser.add_argument(
"--resume-tmp-dir",
default=None,
help="Resume from an existing formatter temp directory containing chunk_*.jsonl and skip raw trace scanning.",
)
format_parser.add_argument("--chunk-bytes", type=int, default=128 * 1024 * 1024)
format_parser.add_argument(
"--log-file",
@@ -146,6 +151,7 @@ def main(argv: list[str] | None = None) -> int:
input_dir=args.input,
output_path=raw_output_path,
tmp_dir=args.tmp_dir,
resume_tmp_dir=args.resume_tmp_dir,
chunk_bytes=args.chunk_bytes,
start_time=args.start_time,
end_time=args.end_time,

231
trace_formatter/formatting.py
View File

@@ -12,7 +12,7 @@ from concurrent.futures import ProcessPoolExecutor, as_completed
from contextlib import contextmanager, nullcontext
from dataclasses import asdict
from pathlib import Path
from typing import Iterator
from typing import Iterable, Iterator
from trace_analyzer.helpers import normalize_unicode_text, parse_jsonish, safe_int
from tokenizers import Tokenizer
@@ -308,6 +308,10 @@ def _build_unified_row_payload_from_components(
def _write_chunk(rows: list[tuple[int, int, str, str, str, str, str]], tmp_dir: Path, chunk_index: int) -> Path:
rows.sort(key=lambda item: (item[0], item[1]))
path = tmp_dir / f"chunk_{chunk_index:05d}.jsonl"
return _write_chunk_rows(rows, path)
def _write_chunk_rows(rows: Iterable[tuple[int, int, str, str, str, str, str]], path: Path) -> Path:
with path.open("w", encoding="utf-8") as handle:
for sort_time_ms, seq, user_id, request_id, sequence_hashes, roles, row_json in rows:
handle.write(f"{sort_time_ms}\t{seq}\t{user_id}\t{request_id}\t{sequence_hashes}\t{roles}\t{row_json}\n")
@@ -321,6 +325,136 @@ def _iter_chunk_rows(path: Path) -> Iterator[tuple[int, int, str, str, str, str,
yield int(sort_text), int(seq_text), user_id, request_id, sequence_hashes, roles, row_json
def _merge_chunk_group(chunk_paths: list[Path], output_path: Path) -> Path:
iterators = [_iter_chunk_rows(path) for path in chunk_paths]
return _write_chunk_rows(
heapq.merge(*iterators, key=lambda item: (item[0], item[1])),
output_path,
)
def _resume_manifest_path(chunk_root: Path) -> Path:
return chunk_root / "resume-manifest.json"
def _write_resume_manifest(chunk_root: Path, payload: dict) -> None:
_resume_manifest_path(chunk_root).write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
def _load_resume_manifest(chunk_root: Path) -> dict | None:
path = _resume_manifest_path(chunk_root)
if not path.exists():
return None
return json.loads(path.read_text(encoding="utf-8"))
def _resolve_merge_fan_in() -> int:
fallback_fan_in = 128
max_fan_in = 512
reserved_fds = 32
try:
import resource
soft_limit, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
except (ImportError, OSError, ValueError):
return fallback_fan_in
if soft_limit < 0:
return max_fan_in
return max(2, min(max_fan_in, soft_limit - reserved_fds))
def _collapse_chunk_paths(chunk_paths: list[Path], tmp_dir: Path, *, merge_fan_in: int | None = None) -> list[Path]:
current_paths = list(chunk_paths)
fan_in = max(2, merge_fan_in if merge_fan_in is not None else _resolve_merge_fan_in())
round_index = 0
while len(current_paths) > fan_in:
next_round: list[Path] = []
for group_index, start in enumerate(range(0, len(current_paths), fan_in)):
group = current_paths[start : start + fan_in]
if len(group) == 1:
next_round.append(group[0])
continue
merged_path = tmp_dir / f"merge_{round_index:02d}_{group_index:05d}.jsonl"
_merge_chunk_group(group, merged_path)
for path in group:
path.unlink(missing_ok=True)
next_round.append(merged_path)
current_paths = next_round
round_index += 1
return current_paths
def _discover_resume_chunk_paths(chunk_root: Path) -> list[Path]:
chunk_paths = sorted(chunk_root.glob("chunk_*.jsonl"))
if not chunk_paths:
raise FileNotFoundError(f"No resumable chunk_*.jsonl files found under {chunk_root}")
return chunk_paths
def _merge_chunks_into_output(
*,
chunk_paths: list[Path],
chunk_root: Path,
destination: Path,
merge_fan_in: int | None,
progress_stream,
show_progress: bool,
row_count_hint: int | None = None,
progress_chunk_count: int | None = None,
) -> dict:
raw_chunk_count = len(chunk_paths)
visible_chunk_count = progress_chunk_count if progress_chunk_count is not None else raw_chunk_count
temp_raw_destination = chunk_root / "formatted-raw.tmp.jsonl"
reduced_chunk_paths = _collapse_chunk_paths(chunk_paths, chunk_root, merge_fan_in=merge_fan_in)
iterators = [_iter_chunk_rows(path) for path in reduced_chunk_paths]
sessionizer = LogicalSessionizer()
row_count = 0
min_sort_time_ms: int | None = None
max_sort_time_ms: int | None = None
user_scoped_rows = 0
merge_progress = tqdm(
total=row_count_hint,
desc="Merge formatted trace",
unit="row",
dynamic_ncols=True,
file=progress_stream,
disable=not show_progress,
)
try:
with temp_raw_destination.open("w", encoding="utf-8") as output_handle:
for sort_time_ms, _, user_id, request_id, sequence_hashes, roles, row_json in heapq.merge(
*iterators, key=lambda item: (item[0], item[1])
):
assignment = sessionizer.assign_precomputed(
user_id=user_id,
request_id=request_id,
sequence_hashes=decode_prefix_hashes(sequence_hashes),
roles=decode_roles(roles),
)
output_handle.write(_apply_session_assignment_to_row_json(row_json, assignment))
output_handle.write("\n")
row_count += 1
min_sort_time_ms = sort_time_ms if min_sort_time_ms is None else min(min_sort_time_ms, sort_time_ms)
max_sort_time_ms = sort_time_ms if max_sort_time_ms is None else max(max_sort_time_ms, sort_time_ms)
if user_id:
user_scoped_rows += 1
if show_progress:
merge_progress.update(1)
merge_progress.set_postfix(rows=row_count if row_count_hint is None else row_count_hint, chunks=visible_chunk_count)
finally:
if show_progress:
merge_progress.close()
shutil.move(str(temp_raw_destination), str(destination))
return {
"row_count": row_count,
"chunk_count": raw_chunk_count,
"min_sort_time_ms": min_sort_time_ms or 0,
"max_sort_time_ms": max_sort_time_ms or 0,
"rows_with_user_id": user_scoped_rows,
}
def _replace_json_field_once(row_json: str, *, key: str, value_text: str) -> str:
target = f'"{key}":""'
if target in row_json:
@@ -654,7 +788,9 @@ def format_and_sort_trace(
input_dir: str | Path,
output_path: str | Path,
tmp_dir: str | Path | None = None,
resume_tmp_dir: str | Path | None = None,
chunk_bytes: int = 128 * 1024 * 1024,
merge_fan_in: int | None = None,
start_time: str | None = None,
end_time: str | None = None,
truncate_to_window: bool = True,
@@ -670,6 +806,38 @@ def format_and_sort_trace(
total_input_bytes = sum(path.stat().st_size for path in source_files if path.suffix != ".zst")
has_zst = any(path.suffix == ".zst" for path in source_files)
if resume_tmp_dir is not None:
chunk_root = Path(resume_tmp_dir)
manifest = _load_resume_manifest(chunk_root) or {}
with _open_progress_stream(log_file) as progress_stream:
merge_stats = _merge_chunks_into_output(
chunk_paths=_discover_resume_chunk_paths(chunk_root),
chunk_root=chunk_root,
destination=destination,
merge_fan_in=merge_fan_in,
progress_stream=progress_stream,
show_progress=show_progress,
row_count_hint=manifest.get("row_count"),
progress_chunk_count=manifest.get("chunk_count"),
)
shutil.rmtree(chunk_root, ignore_errors=True)
return {
"output_path": str(destination),
"row_count": merge_stats["row_count"],
"source_file_count": len(source_files),
"chunk_count": manifest.get("chunk_count", merge_stats["chunk_count"]),
"min_sort_time_ms": merge_stats["min_sort_time_ms"],
"max_sort_time_ms": merge_stats["max_sort_time_ms"],
"rows_with_user_id": merge_stats["rows_with_user_id"],
"truncated_row_count": manifest.get("truncated_row_count"),
"filtered_row_count": manifest.get("filtered_row_count"),
"filtered_empty_messages_row_count": manifest.get("filtered_empty_messages_row_count"),
"filtered_empty_response_params_row_count": manifest.get("filtered_empty_response_params_row_count"),
"window_start_ms": time_window.start_ms if time_window is not None else None,
"window_end_ms": time_window.end_ms if time_window is not None else None,
"resumed_from_tmp_dir": str(chunk_root),
}
with _open_progress_stream(log_file) as progress_stream, tempfile.TemporaryDirectory(dir=temp_root_dir) as temp_root:
temp_raw_destination = Path(temp_root) / "formatted-raw.tmp.jsonl"
chunk_root = Path(temp_root)
@@ -783,46 +951,37 @@ def format_and_sort_trace(
if chunk_rows:
chunk_paths.append(_write_chunk(chunk_rows, chunk_root, len(chunk_paths)))
iterators = [_iter_chunk_rows(path) for path in chunk_paths]
sessionizer = LogicalSessionizer()
merge_progress = tqdm(
total=total_rows,
desc="Merge formatted trace",
unit="row",
dynamic_ncols=True,
file=progress_stream,
disable=not show_progress,
raw_chunk_count = len(chunk_paths)
_write_resume_manifest(
chunk_root,
{
"row_count": total_rows,
"chunk_count": raw_chunk_count,
"truncated_row_count": truncated_rows,
"filtered_row_count": filtered_rows,
"filtered_empty_messages_row_count": filtered_empty_messages_rows,
"filtered_empty_response_params_row_count": filtered_empty_response_rows,
},
)
merge_stats = _merge_chunks_into_output(
chunk_paths=chunk_paths,
chunk_root=chunk_root,
destination=destination,
merge_fan_in=merge_fan_in,
progress_stream=progress_stream,
show_progress=show_progress,
row_count_hint=total_rows,
progress_chunk_count=raw_chunk_count,
)
try:
with temp_raw_destination.open("w", encoding="utf-8") as output_handle:
for _, _, user_id, request_id, sequence_hashes, roles, row_json in heapq.merge(
*iterators, key=lambda item: (item[0], item[1])
):
assignment = sessionizer.assign_precomputed(
user_id=user_id,
request_id=request_id,
sequence_hashes=decode_prefix_hashes(sequence_hashes),
roles=decode_roles(roles),
)
output_handle.write(_apply_session_assignment_to_row_json(row_json, assignment))
output_handle.write("\n")
if show_progress:
merge_progress.update(1)
merge_progress.set_postfix(rows=total_rows, chunks=len(chunk_paths))
finally:
if show_progress:
merge_progress.close()
shutil.move(str(temp_raw_destination), str(destination))
return {
"output_path": str(destination),
"row_count": total_rows,
"row_count": merge_stats["row_count"],
"source_file_count": len(source_files),
"chunk_count": len(chunk_paths),
"min_sort_time_ms": min_sort_time_ms or 0,
"max_sort_time_ms": max_sort_time_ms or 0,
"rows_with_user_id": user_scoped_rows,
"chunk_count": raw_chunk_count,
"min_sort_time_ms": merge_stats["min_sort_time_ms"],
"max_sort_time_ms": merge_stats["max_sort_time_ms"],
"rows_with_user_id": merge_stats["rows_with_user_id"],
"truncated_row_count": truncated_rows,
"filtered_row_count": filtered_rows,
"filtered_empty_messages_row_count": filtered_empty_messages_rows,