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

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:
2026-09-11 06:30:48 +00:00
parent 5a6b1acb49
commit 5ee98f9e4d
3 changed files with 295 additions and 36 deletions

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: