From 5ee98f9e4dc0cde1233802de666953a66f10fb5c Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Fri, 11 Sep 2026 06:30:48 +0000 Subject: [PATCH] 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. --- tests/test_ali_trace_pipeline.py | 94 +++++++++++++ trace_formatter/cli.py | 6 + trace_formatter/formatting.py | 231 ++++++++++++++++++++++++++----- 3 files changed, 295 insertions(+), 36 deletions(-) diff --git a/tests/test_ali_trace_pipeline.py b/tests/test_ali_trace_pipeline.py index e104514..a484af4 100644 --- a/tests/test_ali_trace_pipeline.py +++ b/tests/test_ali_trace_pipeline.py @@ -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: diff --git a/trace_formatter/cli.py b/trace_formatter/cli.py index 66d4102..2849247 100644 --- a/trace_formatter/cli.py +++ b/trace_formatter/cli.py @@ -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, diff --git a/trace_formatter/formatting.py b/trace_formatter/formatting.py index e15f46c..ca9e4dd 100644 --- a/trace_formatter/formatting.py +++ b/trace_formatter/formatting.py @@ -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,