diff --git a/src/teich/anonymize.py b/src/teich/anonymize.py index 4df3627..beac800 100644 --- a/src/teich/anonymize.py +++ b/src/teich/anonymize.py @@ -2,13 +2,16 @@ from __future__ import annotations +from concurrent.futures import ProcessPoolExecutor from dataclasses import dataclass, field from pathlib import Path import hashlib +import os import json import re import shutil import string +import sys from tempfile import NamedTemporaryFile from typing import Any @@ -24,6 +27,18 @@ ".log", } +# ponytail: process startup isn't free — only fan out when there are enough +# files for the parallelism to pay for itself. +_MIN_FILES_FOR_PARALLEL = 8 +_WINDOWS_MAX_PROCESS_WORKERS = 61 + + +def _process_worker_count(file_count: int) -> int: + workers = min(os.cpu_count() or 1, file_count) + if sys.platform == "win32": + workers = min(workers, _WINDOWS_MAX_PROCESS_WORKERS) + return workers + @dataclass class AnonymizeFileReport: @@ -73,11 +88,17 @@ def anonymize_path(input_path: Path, output_path: Path, *, in_place: bool = Fals report.files.append(file_report) return report - for source_file in sorted(path for path in input_path.rglob("*") if path.is_file()): - relative_path = source_file.relative_to(input_path) - destination = source_file if in_place else output_path / relative_path - file_report = _anonymize_file(source_file, destination) - report.files.append(file_report) + source_files = sorted(path for path in input_path.rglob("*") if path.is_file()) + destinations = [source if in_place else output_path / source.relative_to(input_path) for source in source_files] + workers = _process_worker_count(len(source_files)) + if workers > 1 and len(source_files) >= _MIN_FILES_FOR_PARALLEL: + # Each file is anonymized independently (fresh TraceAnonymizer per + # file), so files can be processed in parallel safely. + with ProcessPoolExecutor(max_workers=workers) as executor: + report.files.extend(executor.map(_anonymize_file, source_files, destinations)) + else: + for source_file, destination in zip(source_files, destinations): + report.files.append(_anonymize_file(source_file, destination)) return report diff --git a/tests/test_extract_anonymize_cli.py b/tests/test_extract_anonymize_cli.py index 01de554..c765458 100644 --- a/tests/test_extract_anonymize_cli.py +++ b/tests/test_extract_anonymize_cli.py @@ -7,6 +7,7 @@ from typer.testing import CliRunner +from teich import anonymize as anonymize_module from teich.converter import convert_traces_to_training_data from teich import extract as extract_module from teich.extract import CURSOR_EXTRACTION_NOTICE, default_session_sources @@ -1642,6 +1643,13 @@ def test_anonymize_replaces_emails_keys_and_home_usernames_consistently(tmp_path assert "api_key=1" in result.output +def test_anonymize_parallel_worker_count_caps_windows(monkeypatch): + monkeypatch.setattr(anonymize_module.os, "cpu_count", lambda: 128) + monkeypatch.setattr(anonymize_module.sys, "platform", "win32") + + assert anonymize_module._process_worker_count(100) == 61 + + def test_anonymize_jsonl_preserves_valid_json_after_escaped_path_replacements(tmp_path: Path): input_dir = tmp_path / "input" input_dir.mkdir()