diff --git a/pyrit/score/scorer_evaluation/scorer_metrics_io.py b/pyrit/score/scorer_evaluation/scorer_metrics_io.py index 2cc9596424..021735dc22 100644 --- a/pyrit/score/scorer_evaluation/scorer_metrics_io.py +++ b/pyrit/score/scorer_evaluation/scorer_metrics_io.py @@ -8,6 +8,9 @@ import json import logging +import os +import secrets +import stat import threading from dataclasses import asdict from pathlib import Path @@ -330,6 +333,130 @@ def _append_jsonl_entry(file_path: Path, lock: threading.Lock, entry: dict[str, raise +def _read_registry_lines(file_path: Path) -> list[tuple[str, dict[str, Any] | None]]: + """ + Load a registry as (raw line, parsed entry) pairs, keeping unparseable lines. + + Unlike ``_load_jsonl``, which is a lookup helper that may ignore what it cannot + read, this is the source of truth for a rewrite: every line in the file has to be + accounted for, so a line that is not valid JSON is returned with ``None`` and the + caller decides what to do with it. Read errors propagate instead of yielding a + short list, because a rewrite built from a partial read would delete the rest of + the registry. + + Args: + file_path (Path): Path to the JSONL file. + + Returns: + list[tuple[str, dict[str, Any] | None]]: One pair per line, with the raw line + (including its original whitespace and line ending) and the parsed entry + (or ``None`` when the line is not a JSON object). Empty when the file + does not exist. + + Raises: + OSError: If the file exists but cannot be read. + UnicodeDecodeError: If the file is not valid UTF-8. + """ + if not file_path.exists(): + logger.debug(f"Registry file not found: {file_path}") + return [] + + lines: list[tuple[str, dict[str, Any] | None]] = [] + with open(file_path, encoding="utf-8", newline="") as f: + for line_num, raw_line in enumerate(f, 1): + stripped = raw_line.strip() + if not stripped: + lines.append((raw_line, None)) + continue + try: + parsed = json.loads(stripped) + except json.JSONDecodeError as e: + logger.warning(f"Invalid JSON at line {line_num} in {file_path}: {e}") + lines.append((raw_line, None)) + continue + lines.append((raw_line, parsed if isinstance(parsed, dict) else None)) + return lines + + +def _create_staging_file(file_path: Path, mode: int = 0o666) -> tuple[Path, int]: + """ + Create an exclusive staging file whose requested mode is filtered by the process umask. + + Returns: + tuple[Path, int]: The staging path and its open file descriptor. + + Raises: + FileExistsError: If no unique staging name could be created. + """ + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_BINARY", 0) + for _ in range(100): + temp_path = file_path.with_name(f"{file_path.name}.tmp-{secrets.token_hex(8)}") + try: + return temp_path, os.open(temp_path, flags, mode) + except FileExistsError: + continue + raise FileExistsError(f"Could not create a unique staging file next to {file_path}") + + +def _cleanup_staging_file(temp_path: Path) -> None: + """Remove a staging file without masking an earlier write/replace error.""" + try: + temp_path.unlink(missing_ok=True) + except PermissionError: + try: + # Windows refuses to unlink a read-only file. Make only this disposable + # staging file writable; the registry's permissions remain untouched. + temp_path.chmod(stat.S_IREAD | stat.S_IWRITE) + temp_path.unlink(missing_ok=True) + except OSError as cleanup_error: + logger.warning("Failed to clean up staging file %s: %s", temp_path, cleanup_error) + except OSError as cleanup_error: + logger.warning("Failed to clean up staging file %s: %s", temp_path, cleanup_error) + + +def _rewrite_jsonl_atomically(file_path: Path, lines: list[str]) -> None: + """ + Replace a registry file's contents in one step that readers either see whole. + + Args: + file_path (Path): Path to the JSONL file to rewrite. + lines (list[str]): Raw lines to write, including their original line endings. + + Raises: + OSError: If the file cannot be written or moved into place. + """ + file_path.parent.mkdir(parents=True, exist_ok=True) + existing_mode = stat.S_IMODE(file_path.stat().st_mode) if file_path.exists() else None + + # O_EXCL prevents independent writers from sharing a staging path. Create + # staging files for existing registries privately from the start; changing + # permissions after creation cannot protect readers that already opened it. + # New registries use 0666 so the OS applies the current umask as before. + creation_mode = 0o600 if existing_mode is not None else 0o666 + temp_path, temp_fd = _create_staging_file(file_path, mode=creation_mode) + try: + if existing_mode is not None and hasattr(os, "fchmod"): + # Apply the registry's permissions before copying its contents so a + # private registry never has a more readable staging copy. + os.fchmod(temp_fd, existing_mode) + + with os.fdopen(temp_fd, "w", encoding="utf-8", newline="") as temp_file: + temp_fd = -1 + for line in lines: + temp_file.write(line) + + if existing_mode is not None: + # Writing can clear special permission bits, so restore the final mode. + os.chmod(temp_path, existing_mode) + # The staging file is closed before replace for platforms that cannot replace + # an open file. + os.replace(temp_path, file_path) + finally: + if temp_fd >= 0: + os.close(temp_fd) + _cleanup_staging_file(temp_path) + + def replace_evaluation_results( *, file_path: Path, @@ -344,11 +471,18 @@ def replace_evaluation_results( and adds the new entry. Only one entry per eval_hash is maintained in the registry, ensuring we always track the highest-fidelity evaluation. + Lines that could not be parsed are kept verbatim rather than dropped: a corrupt line + is not an entry this call is allowed to delete, and pre-computed metrics for other + scorers cost hours of model calls to regenerate. + Args: file_path (Path): The full path to the JSONL file. scorer_identifier (ComponentIdentifier): The scorer's configuration identifier. eval_hash (str): The pre-computed evaluation hash for grouping. metrics (ScorerMetrics): The computed metrics (ObjectiveScorerMetrics or HarmScorerMetrics). + + Raises: + OSError: If the registry exists but cannot be read, written, or moved into place. """ # Get or create lock for this file path file_path_str = str(file_path) @@ -361,29 +495,31 @@ def replace_evaluation_results( new_entry["metrics"] = _metrics_to_registry_dict(metrics) with _file_write_locks[file_path_str]: - try: - # Load existing entries - existing_entries = _load_jsonl(file_path) - - # Filter out entries with the same hash - filtered_entries = [e for e in existing_entries if e.get("eval_hash") != eval_hash] - - # Add the new entry - filtered_entries.append(new_entry) - - # Rewrite the file atomically - file_path.parent.mkdir(parents=True, exist_ok=True) - with open(file_path, "w", encoding="utf-8") as f: - for entry in filtered_entries: - f.write(json.dumps(entry) + "\n") - - replaced = len(existing_entries) != len(filtered_entries) - action = "Replaced" if replaced else "Added" - logger.info( - f"{action} metrics for {scorer_identifier.class_name}" - f" (eval_hash={eval_hash[:8]}...) in {file_path.name}" - ) - - except Exception as e: - logger.error(f"Failed to replace entry in registry {file_path}: {e}") - raise + # Load existing entries, keeping track of which lines could not be parsed + existing_lines = _read_registry_lines(file_path) + + # Keep every line that is not the entry being replaced, including unparseable ones + preserved = [raw for raw, parsed in existing_lines if parsed is None or parsed.get("eval_hash") != eval_hash] + + # Keep the registry's existing line-ending style and only add a separator when + # the final preserved line did not have one. + line_ending = next( + (ending for raw in reversed(preserved) for ending in ("\r\n", "\n", "\r") if raw.endswith(ending)), + "\n", + ) + output_lines = [*preserved] + if output_lines and not output_lines[-1].endswith(("\n", "\r")): + output_lines[-1] += line_ending + output_lines.append(json.dumps(new_entry) + line_ending) + + # Rewrite the file with the surviving lines plus the new entry + _rewrite_jsonl_atomically(file_path, output_lines) + + replaced = len(preserved) != len(existing_lines) + action = "Replaced" if replaced else "Added" + dropped = sum(1 for _, parsed in existing_lines if parsed is None) + if dropped: + action += f" (kept {dropped} unparseable line(s) verbatim)" + logger.info( + f"{action} metrics for {scorer_identifier.class_name} (eval_hash={eval_hash[:8]}...) in {file_path.name}" + ) diff --git a/tests/unit/score/test_scorer_metrics_io.py b/tests/unit/score/test_scorer_metrics_io.py index 8e5069c953..4d68eeaac0 100644 --- a/tests/unit/score/test_scorer_metrics_io.py +++ b/tests/unit/score/test_scorer_metrics_io.py @@ -1,7 +1,10 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import errno import json +import os +import stat from pathlib import Path from unittest.mock import patch @@ -17,6 +20,7 @@ _append_jsonl_entry, _load_jsonl, _metrics_to_registry_dict, + _rewrite_jsonl_atomically, add_evaluation_results, find_harm_metrics_by_eval_hash, find_objective_metrics_by_eval_hash, @@ -467,3 +471,333 @@ def test_replace_evaluation_results_preserves_other_entries(tmp_path): assert replaced["metrics"]["accuracy"] == 0.99 finally: sio._file_write_locks = original_locks + + +def test_replace_evaluation_results_keeps_unparseable_lines(tmp_path): + import pyrit.score.scorer_evaluation.scorer_metrics_io as sio + + original_locks = sio._file_write_locks.copy() + try: + path = tmp_path / "test_metrics.jsonl" + add_evaluation_results( + file_path=path, + scorer_identifier=_make_identifier(class_name="A"), + eval_hash="keep_me", + metrics=_make_objective_metrics(accuracy=0.70), + ) + torn = '{"hash_b": "b", "metrics": {"acc' + with open(path, "a", encoding="utf-8") as f: + f.write(torn + "\n") + + replace_evaluation_results( + file_path=path, + scorer_identifier=_make_identifier(class_name="B_new"), + eval_hash="new_hash", + metrics=_make_objective_metrics(accuracy=0.99), + ) + + assert torn in path.read_text(encoding="utf-8"), "the rewrite deleted a line it could not read" + hashes = {entry["eval_hash"] for entry in _load_jsonl(path)} + assert hashes == {"keep_me", "new_hash"} + finally: + sio._file_write_locks = original_locks + + +def test_replace_evaluation_results_leaves_registry_intact_after_a_failed_read(tmp_path): + import pyrit.score.scorer_evaluation.scorer_metrics_io as sio + + original_locks = sio._file_write_locks.copy() + try: + path = tmp_path / "test_metrics.jsonl" + add_evaluation_results( + file_path=path, + scorer_identifier=_make_identifier(class_name="A"), + eval_hash="hash_a", + metrics=_make_objective_metrics(accuracy=0.70), + ) + add_evaluation_results( + file_path=path, + scorer_identifier=_make_identifier(class_name="C"), + eval_hash="hash_c", + metrics=_make_objective_metrics(accuracy=0.80), + ) + undecodable = path.read_bytes() + b'{"eval_hash": "bad", "metrics": \xff\xfe}\n' + path.write_bytes(undecodable) + + with pytest.raises(UnicodeDecodeError): + replace_evaluation_results( + file_path=path, + scorer_identifier=_make_identifier(class_name="New"), + eval_hash="new_hash", + metrics=_make_objective_metrics(accuracy=0.99), + ) + + assert path.read_bytes() == undecodable, "a partial read must not be rewritten over the registry" + finally: + sio._file_write_locks = original_locks + + +def test_replace_evaluation_results_preserves_raw_lines_and_endings(tmp_path): + import pyrit.score.scorer_evaluation.scorer_metrics_io as sio + + original_locks = sio._file_write_locks.copy() + try: + path = tmp_path / "test_metrics.jsonl" + original = ( + b' {"eval_hash": "keep", "metrics": {"value": 1}}\r\n' + b"\t \r\n" + b"\n" + b"\tnot json \n" + b' {"eval_hash": "other", "metrics": {"value": 2}}' + ) + path.write_bytes(original) + + replace_evaluation_results( + file_path=path, + scorer_identifier=_make_identifier(), + eval_hash="new_hash", + metrics=_make_objective_metrics(accuracy=0.99), + ) + + rewritten = path.read_bytes() + assert rewritten.startswith(original) + assert rewritten.endswith(b"\n") + assert [entry["eval_hash"] for entry in _load_jsonl(path)] == ["keep", "other", "new_hash"] + finally: + sio._file_write_locks = original_locks + + +@pytest.mark.skipif(os.name == "nt", reason="POSIX permission semantics") +def test_replace_evaluation_results_preserves_existing_permissions(tmp_path): + import pyrit.score.scorer_evaluation.scorer_metrics_io as sio + + original_locks = sio._file_write_locks.copy() + try: + path = tmp_path / "test_metrics.jsonl" + path.write_text('{"eval_hash": "keep", "metrics": {}}\n', encoding="utf-8") + path.chmod(0o664) + + replace_evaluation_results( + file_path=path, + scorer_identifier=_make_identifier(), + eval_hash="new_hash", + metrics=_make_objective_metrics(accuracy=0.99), + ) + + assert stat.S_IMODE(path.stat().st_mode) == 0o664 + finally: + sio._file_write_locks = original_locks + + +@pytest.mark.skipif(os.name == "nt", reason="POSIX permission semantics") +def test_existing_registry_permissions_are_applied_before_staging_write(tmp_path): + import pyrit.score.scorer_evaluation.scorer_metrics_io as sio + + path = tmp_path / "test_metrics.jsonl" + path.write_text('{"value": 1}\n', encoding="utf-8") + path.chmod(0o600) + + real_create_staging_file = sio._create_staging_file + real_fdopen = os.fdopen + staging_paths = [] + + def record_staging_file(file_path, mode=0o666): + staging_path, fd = real_create_staging_file(file_path, mode=mode) + staging_paths.append(staging_path) + return staging_path, fd + + class PermissionCheckingWriter: + def __init__(self, stream): + self.stream = stream + + def __enter__(self): + self.stream.__enter__() + return self + + def __exit__(self, *args): + return self.stream.__exit__(*args) + + def write(self, content): + assert stat.S_IMODE(staging_paths[0].stat().st_mode) == 0o600 + return self.stream.write(content) + + with ( + patch.object(sio, "_create_staging_file", side_effect=record_staging_file), + patch.object( + os, + "fdopen", + side_effect=lambda *args, **kwargs: PermissionCheckingWriter(real_fdopen(*args, **kwargs)), + ), + ): + _rewrite_jsonl_atomically(path, ['{"value": 2}\n']) + + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + + +@pytest.mark.skipif(os.name == "nt", reason="POSIX permission semantics") +def test_existing_registry_staging_file_is_private_at_creation(tmp_path): + path = tmp_path / "test_metrics.jsonl" + path.write_text('{"value": 1}\n', encoding="utf-8") + path.chmod(0o600) + + original_umask = os.umask(0o022) + real_open = os.open + creation_modes = [] + + def record_creation_mode(file_path, flags, mode=0o777, **kwargs): + fd = real_open(file_path, flags, mode, **kwargs) + if Path(file_path).name.startswith(f"{path.name}.tmp-"): + creation_modes.append(stat.S_IMODE(Path(file_path).stat().st_mode)) + return fd + + try: + with patch.object(os, "open", side_effect=record_creation_mode): + _rewrite_jsonl_atomically(path, ['{"value": 2}\n']) + finally: + os.umask(original_umask) + + assert creation_modes == [0o600] + + +@pytest.mark.skipif(os.name == "nt", reason="POSIX umask semantics") +def test_new_registry_uses_umask_permissions(tmp_path): + path = tmp_path / "test_metrics.jsonl" + original_umask = os.umask(0o022) + try: + _rewrite_jsonl_atomically(path, [json.dumps({"value": 1}) + "\n"]) + assert stat.S_IMODE(path.stat().st_mode) == 0o644 + finally: + os.umask(original_umask) + + +def test_cleanup_preserves_replace_error_and_removes_read_only_staging_file(tmp_path): + path = tmp_path / "test_metrics.jsonl" + original = b'{"value": 1}\n' + path.write_bytes(original) + real_unlink = Path.unlink + cleanup_attempts: list[str] = [] + + def fail_once_for_staging_file(self, missing_ok=False): + if self.name.startswith(f"{path.name}.tmp-") and not cleanup_attempts: + cleanup_attempts.append(self.name) + raise PermissionError("read-only staging file") + return real_unlink(self, missing_ok=missing_ok) + + with ( + patch( + "pyrit.score.scorer_evaluation.scorer_metrics_io.os.replace", + side_effect=PermissionError("replace failed"), + ), + patch.object(Path, "unlink", new=fail_once_for_staging_file), + ): + with pytest.raises(PermissionError, match="replace failed"): + _rewrite_jsonl_atomically(path, [json.dumps({"value": 2}) + "\n"]) + + assert cleanup_attempts + assert path.read_bytes() == original + assert not list(tmp_path.glob(f"{path.name}.tmp-*")) + + +def test_rewrite_jsonl_atomically_uses_distinct_staging_files(tmp_path): + path = tmp_path / "test_metrics.jsonl" + names: list[str] = [] + real_replace = os.replace + + def record_replace(source, destination): + names.append(Path(source).name) + real_replace(source, destination) + + with patch("pyrit.score.scorer_evaluation.scorer_metrics_io.os.replace", side_effect=record_replace): + _rewrite_jsonl_atomically(path, [json.dumps({"value": 1}) + "\n"]) + _rewrite_jsonl_atomically(path, [json.dumps({"value": 2}) + "\n"]) + + assert len(names) == 2 + assert len(set(names)) == 2 + + +def test_rewrite_jsonl_atomically_retries_staging_collisions(tmp_path: Path) -> None: + path = tmp_path / "test_metrics.jsonl" + path.write_bytes(b'{"value": 1}\n') + collision = tmp_path / "test_metrics.jsonl.tmp-occupied" + other_writer = b'{"value": "other writer"}\n' + collision.write_bytes(other_writer) + + with patch( + "pyrit.score.scorer_evaluation.scorer_metrics_io.secrets.token_hex", + side_effect=["occupied", "available"], + ) as token_hex: + _rewrite_jsonl_atomically(path, ['{"value": 2}\n']) + + assert token_hex.call_count == 2 + assert path.read_bytes() == b'{"value": 2}\n' + assert collision.read_bytes() == other_writer + assert set(tmp_path.iterdir()) == {path, collision} + + +def test_rewrite_jsonl_atomically_preserves_files_when_staging_collisions_exhaust_retries(tmp_path: Path) -> None: + path = tmp_path / "test_metrics.jsonl" + original = b'{"value": 1}\n' + path.write_bytes(original) + collision = tmp_path / "test_metrics.jsonl.tmp-occupied" + other_writer = b'{"value": "other writer"}\n' + collision.write_bytes(other_writer) + + with patch( + "pyrit.score.scorer_evaluation.scorer_metrics_io.secrets.token_hex", return_value="occupied" + ) as token_hex: + with pytest.raises(FileExistsError, match="Could not create a unique staging file"): + _rewrite_jsonl_atomically(path, ['{"value": 2}\n']) + + assert token_hex.call_count == 100 + assert path.read_bytes() == original + assert collision.read_bytes() == other_writer + assert set(tmp_path.iterdir()) == {path, collision} + + +def test_rewrite_jsonl_atomically_closes_descriptor_after_fdopen_failure(tmp_path: Path) -> None: + path = tmp_path / "test_metrics.jsonl" + original = b'{"value": 1}\n' + path.write_bytes(original) + with patch.object(os, "fdopen", side_effect=OSError("fdopen failed")) as fdopen: + with pytest.raises(OSError, match="fdopen failed"): + _rewrite_jsonl_atomically(path, ['{"value": 2}\n']) + + with pytest.raises(OSError) as error: + os.fstat(fdopen.call_args.args[0]) + assert error.value.errno == errno.EBADF + assert path.read_bytes() == original + assert not list(tmp_path.glob(f"{path.name}.tmp-*")) + + +@pytest.mark.parametrize( + "unlink_errors", + [ + pytest.param([OSError("cleanup failed")], id="unlink"), + pytest.param([PermissionError("read-only staging file"), OSError("cleanup failed")], id="read-only-retry"), + ], +) +def test_cleanup_failure_is_logged_without_masking_replace_error( + *, tmp_path: Path, caplog: pytest.LogCaptureFixture, unlink_errors: list[OSError] +) -> None: + path = tmp_path / "test_metrics.jsonl" + original = b'{"value": 1}\n' + path.write_bytes(original) + original_mode = stat.S_IMODE(path.stat().st_mode) + replace_error = PermissionError("replace failed") + + with ( + patch.object(os, "replace", side_effect=replace_error), + patch.object(Path, "unlink", autospec=True, side_effect=unlink_errors) as unlink, + ): + with pytest.raises(PermissionError, match="replace failed") as error: + _rewrite_jsonl_atomically(path, ['{"value": 2}\n']) + + assert error.value is replace_error + assert path.read_bytes() == original + assert stat.S_IMODE(path.stat().st_mode) == original_mode + staging_paths = list(tmp_path.glob(f"{path.name}.tmp-*")) + assert len(staging_paths) == 1 + staging_path = staging_paths[0] + assert [call.args[0] for call in unlink.call_args_list] == [staging_path] * len(unlink_errors) + assert caplog.messages == [f"Failed to clean up staging file {staging_path}: cleanup failed"] + staging_path.unlink()