diff --git a/README.md b/README.md index 7780129..eda1450 100644 --- a/README.md +++ b/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. diff --git a/changelog.md b/changelog.md index c6dee0b..cd2f581 100644 --- a/changelog.md +++ b/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 diff --git a/src/swe_data_processing/cli.py b/src/swe_data_processing/cli.py index 2d15b95..7f56dc0 100644 --- a/src/swe_data_processing/cli.py +++ b/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, diff --git a/src/swe_data_processing/io.py b/src/swe_data_processing/io.py index 8f4bd50..40281dd 100644 --- a/src/swe_data_processing/io.py +++ b/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: diff --git a/tests/test_cli.py b/tests/test_cli.py index 93f184a..dd3c165 100644 --- a/tests/test_cli.py +++ b/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() == ""