Parallelize profiling by Parquet shard

src/swe_data_processing/cli.py

@@ -4,9 +4,18 @@ from __future__ import annotations
 
 import argparse
 import json
+import os
+import shutil
 import sys
 from collections.abc import Callable, Iterator
-from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
+from concurrent.futures import (
+    FIRST_COMPLETED,
+    Future,
+    ProcessPoolExecutor,
+    ThreadPoolExecutor,
+    as_completed,
+    wait,
+)
 from pathlib import Path
 from typing import Any
 
@@ -15,8 +24,10 @@ from .io import (
     append_jsonl,
     get_sample_id,
     iter_jsonl,
+    iter_parquet_file,
     iter_records,
     load_completed_ids,
+    parquet_files,
     write_json,
     write_jsonl,
 )
@@ -25,6 +36,137 @@ from .sampling import iter_review_records, select_review_ids
 from .summary import build_summary
 from .tokenization import TokenCounter
 
+_PROFILE_COUNTER: TokenCounter | None = None
+
+
+def _initialize_profile_process(tokenizer_json: str | None) -> None:
+    """Create one tokenizer per process instead of per trajectory."""
+
+    global _PROFILE_COUNTER
+    _PROFILE_COUNTER = TokenCounter(Path(tokenizer_json)) if tokenizer_json else None
+
+
+def _profile_parquet_shard(
+    shard: Path,
+    metrics_part: Path,
+    errors_part: Path | None,
+) -> tuple[int, int]:
+    """Profile one shard into atomic part files with bounded memory."""
+
+    processed = 0
+    failures = 0
+    metrics_temporary = metrics_part.with_suffix(metrics_part.suffix + ".tmp")
+    errors_temporary = (
+        errors_part.with_suffix(errors_part.suffix + ".tmp") if errors_part else None
+    )
+    metrics_part.parent.mkdir(parents=True, exist_ok=True)
+    with metrics_temporary.open("w", encoding="utf-8") as metrics_handle:
+        errors_handle = (
+            errors_temporary.open("w", encoding="utf-8")
+            if errors_temporary is not None
+            else None
+        )
+        try:
+            for record in iter_parquet_file(shard):
+                sample_id = get_sample_id(record)
+                try:
+                    result = extract_metrics(record, _PROFILE_COUNTER)
+                    metrics_handle.write(
+                        json.dumps(result, ensure_ascii=False, default=str) + "\n"
+                    )
+                except Exception as exc:  # noqa: BLE001 - isolate bad records.
+                    failures += 1
+                    if errors_handle is not None:
+                        error = {
+                            "sample_id": sample_id,
+                            "error_type": type(exc).__name__,
+                            "error": str(exc)[:2000],
+                            "source_parquet": str(shard.resolve()),
+                        }
+                        errors_handle.write(
+                            json.dumps(error, ensure_ascii=False) + "\n"
+                        )
+                processed += 1
+        finally:
+            if errors_handle is not None:
+                errors_handle.close()
+    os.replace(metrics_temporary, metrics_part)
+    if errors_temporary is not None and errors_part is not None:
+        os.replace(errors_temporary, errors_part)
+    return processed, failures
+
+
+def _merge_parts(parts: list[Path], output: Path) -> None:
+    """Concatenate completed shard outputs atomically in dataset order."""
+
+    output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = output.with_suffix(output.suffix + ".tmp")
+    with temporary.open("wb") as destination:
+        for part in parts:
+            with part.open("rb") as source:
+                shutil.copyfileobj(source, destination, length=1024 * 1024)
+    os.replace(temporary, output)
+
+
+def _run_parallel_parquet_profile(args: argparse.Namespace) -> int:
+    """Profile independent Parquet shards in separate CPU processes."""
+
+    shards = parquet_files(args.input)
+    parts_dir = args.output.parent / f".{args.output.name}.parts"
+    parts_dir.mkdir(parents=True, exist_ok=True)
+    metrics_parts = [parts_dir / f"{index:05d}.metrics.jsonl" for index in range(len(shards))]
+    error_parts = [parts_dir / f"{index:05d}.errors.jsonl" for index in range(len(shards))]
+
+    if not args.resume:
+        for path in parts_dir.glob("*.jsonl"):
+            path.unlink()
+        for path in parts_dir.glob("*.tmp"):
+            path.unlink()
+
+    pending = []
+    processed = 0
+    failures = 0
+    tokenizer_json = str(args.tokenizer_json) if args.tokenizer_json else None
+    with ProcessPoolExecutor(
+        max_workers=args.workers,
+        initializer=_initialize_profile_process,
+        initargs=(tokenizer_json,),
+    ) as executor:
+        for shard, metrics_part, errors_part in zip(
+            shards, metrics_parts, error_parts, strict=True
+        ):
+            if args.resume and metrics_part.exists() and errors_part.exists():
+                continue
+            pending.append(
+                executor.submit(
+                    _profile_parquet_shard,
+                    shard,
+                    metrics_part,
+                    errors_part if args.errors else None,
+                )
+            )
+        for completed_shards, future in enumerate(as_completed(pending), 1):
+            shard_processed, shard_failures = future.result()
+            processed += shard_processed
+            failures += shard_failures
+            print(
+                f"profile: shards={completed_shards}/{len(pending)} "
+                f"processed={processed} failures={failures}",
+                file=sys.stderr,
+            )
+
+    missing = [path for path in metrics_parts if not path.exists()]
+    if missing:
+        raise RuntimeError(f"Missing {len(missing)} completed metric parts")
+    _merge_parts(metrics_parts, args.output)
+    if args.errors:
+        missing_errors = [path for path in error_parts if not path.exists()]
+        if missing_errors:
+            raise RuntimeError(f"Missing {len(missing_errors)} completed error parts")
+        _merge_parts(error_parts, args.errors)
+    print(f"profile: completed={processed} failures={failures}", file=sys.stderr)
+    return 1 if failures else 0
+
 
 def _write_error(path: Path | None, sample_id: str, error: Exception) -> None:
     """Write a compact error without interrupting a long profile run."""
@@ -121,6 +263,8 @@ def _run_streaming_stage(
 def command_profile(args: argparse.Namespace) -> int:
     """Extract per-trajectory metrics without modifying source records."""
 
+    if args.input.is_dir() and args.workers > 1 and args.limit is None:
+        return _run_parallel_parquet_profile(args)
     counter = TokenCounter(args.tokenizer_json) if args.tokenizer_json else None
     return _run_streaming_stage(
         input_path=args.input,