Parallelize profiling by Parquet shard

README.md

@@ -103,7 +103,7 @@ swe-qc profile \
   --output qc_outputs/deterministic_v3/metrics.jsonl \
   --errors qc_outputs/deterministic_v3/profile.errors.jsonl \
   --tokenizer-json "$TOKENIZER_JSON" \
-  --workers 8 \
+  --workers 24 \
   --resume
 
 swe-qc summarize \
@@ -119,12 +119,15 @@ swe-qc sample \
   --seed 20260818
 ```
 
-`profile --resume` skips sample IDs already present in the metrics file. Writes
-are append-only and flushed per record. `summarize` and `sample` atomically
-replace their outputs so a partial file is never mistaken for a complete run.
-Any dataset subdirectory containing Parquet shards is also a valid `--input`,
-so the four trajectory families can be profiled in separate CPU processes and
-their metrics concatenated before `summarize`.
+For a Parquet directory and `--workers > 1`, each worker is a separate process
+that reads one shard at a time and writes an atomic shard part. This bypasses
+the Python GIL, bounds live trajectory memory to roughly one small Parquet batch
+per process, and lets `profile --resume` reuse completed shard parts after an
+interruption. The final metrics and error manifests are concatenated atomically
+in stable shard order. For JSONL or a limited pilot, `--resume` instead skips
+sample IDs already present in the append-only output. `summarize` and `sample`
+also replace their outputs atomically.
+
 Profile rows and decisions record the absolute source Parquet shard. `sample`
 uses that provenance to read only shards containing selected IDs instead of
 scanning the complete dataset.

changelog.md

@@ -4,6 +4,19 @@ This file records strategy changes that materially affect dataset decisions or
 training-data semantics. Generated audit manifests are not treated as stable API
 contracts.
 
+## 3.0.1 - 2026-08-18
+
+### CPU profiling
+
+- Replaced directory-level Python thread concurrency with Parquet shard-level
+  processes so CPU-bound tokenization and metric extraction bypass the GIL.
+- Each process streams batches of 16 records to an atomic shard part; the parent
+  merges completed parts in source order without loading them into memory.
+- `profile --resume` now reuses completed shard parts, while interrupted
+  temporary parts are safely recomputed.
+- Review sampling uses stored Parquet provenance and reads only shards that
+  contain selected trajectories.
+
 ## 3.0.0 - 2026-08-18
 
 ### Strategy replacement

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,

src/swe_data_processing/io.py

@@ -36,20 +36,34 @@ def iter_jsonl(path: Path) -> Iterator[dict[str, Any]]:
             yield value
 
 
-def iter_parquet_dataset(dataset_dir: Path, *, batch_size: int = 16) -> Iterator[dict[str, Any]]:
-    """Stream records from every Open-SWE-Traces Parquet shard in stable order."""
+def parquet_files(dataset_dir: Path) -> list[Path]:
+    """Return all Parquet shards in stable source order."""
 
     files = sorted(dataset_dir.rglob("*.parquet"))
     if not files:
         raise FileNotFoundError(f"No Parquet shards found under {dataset_dir}")
-    for path in files:
-        parquet_file = pq.ParquetFile(path)
-        for batch in parquet_file.iter_batches(batch_size=batch_size):
-            for record in batch.to_pylist():
-                if isinstance(record, dict):
-                    value = dict(record)
-                    value["_source_parquet"] = str(path.resolve())
-                    yield value
+    return files
+
+
+def iter_parquet_file(path: Path, *, batch_size: int = 16) -> Iterator[dict[str, Any]]:
+    """Stream one Parquet shard without materializing it in memory."""
+
+    parquet_file = pq.ParquetFile(path)
+    for batch in parquet_file.iter_batches(batch_size=batch_size):
+        for record in batch.to_pylist():
+            if isinstance(record, dict):
+                value = dict(record)
+                value["_source_parquet"] = str(path.resolve())
+                yield value
+
+
+def iter_parquet_dataset(
+    dataset_dir: Path, *, batch_size: int = 16
+) -> Iterator[dict[str, Any]]:
+    """Stream records from every Open-SWE-Traces Parquet shard in stable order."""
+
+    for path in parquet_files(dataset_dir):
+        yield from iter_parquet_file(path, batch_size=batch_size)
 
 
 def iter_records(path: Path) -> Iterator[dict[str, Any]]:
@@ -57,6 +71,8 @@ def iter_records(path: Path) -> Iterator[dict[str, Any]]:
 
     if path.is_dir():
         yield from iter_parquet_dataset(path)
+    elif path.suffix.lower() == ".parquet":
+        yield from iter_parquet_file(path)
     elif path.suffix.lower() in {".jsonl", ".json"}:
         yield from iter_jsonl(path)
     else:

tests/test_cli.py

@@ -3,9 +3,17 @@
 from __future__ import annotations
 
 import json
+from argparse import Namespace
 from pathlib import Path
 
-from swe_data_processing.cli import _run_streaming_stage, build_parser
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from swe_data_processing.cli import (
+    _run_parallel_parquet_profile,
+    _run_streaming_stage,
+    build_parser,
+)
 
 
 def test_streaming_stage_writes_each_record(tmp_path: Path) -> None:
@@ -36,3 +44,34 @@ def test_cli_has_only_deterministic_pipeline_commands() -> None:
     help_text = parser.format_help()
     assert "classify" not in help_text
     assert "repair" not in help_text
+
+
+def test_parallel_parquet_profile_uses_atomic_shard_parts(tmp_path: Path) -> None:
+    dataset = tmp_path / "dataset"
+    dataset.mkdir()
+    for index in range(2):
+        pq.write_table(
+            pa.Table.from_pylist(
+                [{"trajectory_id": f"sample-{index}", "resolved": 1}]
+            ),
+            dataset / f"part-{index}.parquet",
+        )
+    output = tmp_path / "metrics.jsonl"
+    errors = tmp_path / "errors.jsonl"
+    status = _run_parallel_parquet_profile(
+        Namespace(
+            input=dataset,
+            output=output,
+            errors=errors,
+            tokenizer_json=None,
+            workers=2,
+            resume=False,
+        )
+    )
+
+    assert status == 0
+    assert [json.loads(line)["sample_id"] for line in output.read_text().splitlines()] == [
+        "sample-0",
+        "sample-1",
+    ]
+    assert errors.read_text() == ""