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

@@ -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:

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,

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]
+        raw_chunk_count = len(chunk_paths)
-        sessionizer = LogicalSessionizer()
+        _write_resume_manifest(
-        merge_progress = tqdm(
+            chunk_root,
-            total=total_rows,
+            {
-            desc="Merge formatted trace",
+                "row_count": total_rows,
-            unit="row",
+                "chunk_count": raw_chunk_count,
-            dynamic_ncols=True,
+                "truncated_row_count": truncated_rows,
-            file=progress_stream,
+                "filtered_row_count": filtered_rows,
-            disable=not show_progress,
+                "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),
+        "chunk_count": raw_chunk_count,
-        "min_sort_time_ms": min_sort_time_ms or 0,
+        "min_sort_time_ms": merge_stats["min_sort_time_ms"],
-        "max_sort_time_ms": max_sort_time_ms or 0,
+        "max_sort_time_ms": merge_stats["max_sort_time_ms"],
-        "rows_with_user_id": user_scoped_rows,
+        "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,