diff --git a/rampart/pytest_plugin/_xdist.py b/rampart/pytest_plugin/_xdist.py index 3f85dec..8c4f486 100644 --- a/rampart/pytest_plugin/_xdist.py +++ b/rampart/pytest_plugin/_xdist.py @@ -21,6 +21,7 @@ import json import logging import math +from dataclasses import dataclass from datetime import datetime from typing import TYPE_CHECKING, Any, cast @@ -48,6 +49,7 @@ if TYPE_CHECKING: from collections.abc import Callable + from pathlib import Path import pytest @@ -63,6 +65,8 @@ _TRUNCATED_MARKER: str = "rampart_truncated" +SHARD_DIR_KEY: str = "rampart_shard_dir" + class WorkerOutputError(Exception): """Base error for xdist worker output processing failures.""" @@ -141,6 +145,43 @@ def get_worker_count(*, config: pytest.Config) -> int: return int(numprocesses) if numprocesses else 0 +def _is_local_popen_spec(*, spec: str) -> bool: + """Return whether one ``--tx`` gateway spec targets a local popen worker. + + Args: + spec (str): An xdist execnet gateway spec (e.g. ``"popen"`` or + ``"popen//python=python3.11"``). + + Returns: + bool: True when the spec launches a same-host popen worker with + no proxy routing (no ``ssh=``/``socket=``/``via=``). + """ + head = spec.split("//", 1)[0].strip() + return head == "popen" and "via=" not in spec + + +def shard_eligible(*, config: pytest.Config) -> bool: + """Return whether durable shard transport is usable for this run. + + Shard transport requires every worker to share the controller's + filesystem, which holds exactly when all ``--tx`` gateway specs are + local ``popen`` workers. Remote (``ssh=``/``socket=``) or proxied + (``via=``) gateways may not see the controller's shard directory, so + those runs fall back to the in-memory ``workeroutput`` transport. + + Args: + config (pytest.Config): The pytest configuration object. + + Returns: + bool: True when at least one gateway is configured and every + configured gateway is a local popen worker. + """ + tx = getattr(config.option, "tx", None) + if not tx: + return False + return all(_is_local_popen_spec(spec=str(spec)) for spec in tx) + + def _size_limit(*, config: pytest.Config) -> int: """Resolve the worker payload size cap from pytest config or default. @@ -832,6 +873,230 @@ def deserialize_worker_data(*, data: object) -> dict[str, list[Result]]: return out +def _shard_path(*, shard_dir: Path, worker_id: str) -> Path: + """Return the shard file path for a worker. + + Args: + shard_dir (Path): Directory holding per-worker shard files. + worker_id (str): The xdist worker id (e.g. ``"gw0"``). + + Returns: + Path: The ``worker-.jsonl`` path under ``shard_dir``. + """ + return shard_dir / f"worker-{worker_id}.jsonl" + + +def _shard_record_line(*, nodeid: str, index: int, result: Result) -> str: + """Serialize one result into a shard record JSON line. + + Args: + nodeid (str): The originating test nodeid. + index (int): The per-nodeid result index. + result (Result): The result to persist. + + Returns: + str: A JSON object encoding the record envelope and result body. + """ + return json.dumps( + { + "schema": SCHEMA_VERSION, + "nodeid": nodeid, + "index": index, + "result": _serialize_result(result=result, nodeid=nodeid), + }, + ) + + +def _shard_truncation_line(*, nodeid: str, index: int) -> str: + """Serialize an oversized-record marker into a shard JSON line. + + Args: + nodeid (str): The originating test nodeid. + index (int): The per-nodeid result index. + + Returns: + str: A JSON object flagging the dropped record's nodeid/index. + """ + return json.dumps( + { + "schema": SCHEMA_VERSION, + "nodeid": nodeid, + "index": index, + _TRUNCATED_MARKER: True, + }, + ) + + +@dataclass(frozen=True, kw_only=True) +class DroppedRecord: + """A shard record the controller could not recover. + + Attributes: + reason (str): Why the record was dropped (``"oversized"`` when a + worker wrote a truncation marker, ``"corrupt"`` when a line + failed to parse or deserialize). + nodeid (str | None): The originating test nodeid, when known. + index (int | None): The per-nodeid result index, when known. + """ + + reason: str + nodeid: str | None = None + index: int | None = None + + +class ShardWriter: + """Append-only writer for one worker's durable result shard. + + Each finished result is written as a single JSON line and flushed + immediately, so results a worker has already produced survive an + abrupt worker death (e.g. SIGKILL). An oversized result is written as + a marker line naming its nodeid/index rather than discarding the + whole shard. The writer holds the shard file open for the worker's + lifetime; callers must call :meth:`close` at session end. + """ + + def __init__(self, *, shard_dir: Path, worker_id: str, size_limit: int) -> None: + """Open the worker's shard file for appending. + + Args: + shard_dir (Path): Directory holding per-worker shard files. + worker_id (str): The xdist worker id (e.g. ``"gw0"``). + size_limit (int): Per-record byte cap; larger records are + written as truncation markers. + """ + self._size_limit = size_limit + self._path = _shard_path(shard_dir=shard_dir, worker_id=worker_id) + self._file = self._path.open("a", encoding="utf-8") + + def write(self, *, nodeid: str, index: int, result: Result) -> None: + """Append one result record and flush it to disk. + + Args: + nodeid (str): The originating test nodeid. + index (int): The per-nodeid result index. + result (Result): The finished result to persist. + """ + line = _shard_record_line(nodeid=nodeid, index=index, result=result) + if len(line.encode("utf-8")) > self._size_limit: + line = _shard_truncation_line(nodeid=nodeid, index=index) + self._file.write(line + "\n") + self._file.flush() + + def close(self) -> None: + """Close the underlying shard file.""" + self._file.close() + + +def _absorb_shard_record( + *, + record: dict[str, Any], + results: dict[str, list[Result]], + dropped: list[DroppedRecord], +) -> None: + """Route one validated shard record into ``results`` or ``dropped``. + + The record's nodeid and index are taken authoritatively from the + envelope (never from the serialized result body) so cross-worker + ordering cannot be influenced by untrusted payload content. + + Args: + record (dict[str, Any]): A parsed shard record with a matching + schema version. + results (dict[str, list[Result]]): Accumulator for recovered + results, grouped by nodeid. + dropped (list[DroppedRecord]): Accumulator for dropped records. + """ + nodeid = str(record.get("nodeid", "")) + raw_index = record.get("index") + index = raw_index if isinstance(raw_index, int) else None + if record.get(_TRUNCATED_MARKER): + dropped.append(DroppedRecord(reason="oversized", nodeid=nodeid, index=index)) + return + try: + result = _deserialize_result(data=record.get("result")) + except WorkerOutputError: + dropped.append(DroppedRecord(reason="corrupt", nodeid=nodeid, index=index)) + return + result.metadata["_pytest_nodeid"] = nodeid + if index is not None: + result.metadata["_rampart_result_index"] = index + results.setdefault(nodeid, []).append(result) + + +def _read_shard_line( + *, + line: str, + results: dict[str, list[Result]], + dropped: list[DroppedRecord], +) -> None: + """Parse one shard line into ``results`` or ``dropped`` in place. + + Blank lines are ignored. A line that is not valid JSON, not an + object, or carries the wrong schema version is recorded as a corrupt + drop — this tolerates a truncated final line from an abrupt worker + death. + + Args: + line (str): A single raw line from the shard file. + results (dict[str, list[Result]]): Accumulator for recovered + results, grouped by nodeid. + dropped (list[DroppedRecord]): Accumulator for dropped records. + """ + stripped = line.strip() + if not stripped: + return + try: + record = json.loads(stripped) + except json.JSONDecodeError: + dropped.append(DroppedRecord(reason="corrupt")) + return + if not isinstance(record, dict) or record.get("schema") != SCHEMA_VERSION: + dropped.append(DroppedRecord(reason="corrupt")) + return + _absorb_shard_record( + record=cast("dict[str, Any]", record), + results=results, + dropped=dropped, + ) + + +def read_worker_shard( + *, + shard_dir: Path, + worker_id: str, +) -> tuple[dict[str, list[Result]], list[DroppedRecord]]: + """Recover a worker's results from its durable shard file. + + Parses the worker's ``worker-.jsonl`` shard line by line, + tolerating a truncated or partially written final line. Corrupt and + oversized records are collected as :class:`DroppedRecord` entries so + the caller can mark the run incomplete; every cleanly parsed result + is returned. A missing shard file yields empty results and no drops. + + As with :func:`deserialize_worker_data`, each result's + ``_pytest_nodeid`` and ``_rampart_result_index`` metadata are set + authoritatively from the record envelope, never trusting values + embedded in the (attacker-influenced) serialized result body. + + Args: + shard_dir (Path): Directory holding per-worker shard files. + worker_id (str): The xdist worker id (e.g. ``"gw0"``). + + Returns: + tuple[dict[str, list[Result]], list[DroppedRecord]]: Recovered + results grouped by nodeid, and the list of dropped records. + """ + path = _shard_path(shard_dir=shard_dir, worker_id=worker_id) + if not path.exists(): + return {}, [] + results: dict[str, list[Result]] = {} + dropped: list[DroppedRecord] = [] + with path.open("r", encoding="utf-8") as handle: + for line in handle: + _read_shard_line(line=line, results=results, dropped=dropped) + return results, dropped + + def deserialize_trial_specs(*, data: object) -> dict[str, TrialSpec]: """Deserialize the ``trial_specs`` section of a worker payload. diff --git a/tests/unit/pytest_plugin/test_xdist.py b/tests/unit/pytest_plugin/test_xdist.py index a50b338..4b4cc0a 100644 --- a/tests/unit/pytest_plugin/test_xdist.py +++ b/tests/unit/pytest_plugin/test_xdist.py @@ -9,7 +9,7 @@ import logging import math from datetime import UTC, datetime -from typing import Any +from typing import TYPE_CHECKING, Any from unittest.mock import MagicMock import pytest @@ -36,9 +36,12 @@ DEFAULT_SIZE_LIMIT_BYTES, MAX_METADATA_DEPTH, SCHEMA_VERSION, + SHARD_DIR_KEY, SIZE_LIMIT_OPTION, WORKEROUTPUT_KEY, + DroppedRecord, SchemaVersionError, + ShardWriter, SizeLimitError, WorkerOutputError, _sanitize, @@ -52,10 +55,15 @@ handle_testnodedown, is_xdist_controller, is_xdist_worker, + read_worker_shard, serialize_worker_data, + shard_eligible, ) from rampart.reporting.sink import ReportSink, TestRunReport +if TYPE_CHECKING: + from pathlib import Path + def _make_result( *, @@ -1032,7 +1040,268 @@ def test_schema_version_is_v1(self) -> None: def test_workeroutput_key_namespaced(self) -> None: assert WORKEROUTPUT_KEY.startswith("rampart_xdist") + def test_shard_dir_key_namespaced(self) -> None: + assert SHARD_DIR_KEY == "rampart_shard_dir" + class TestTestRunReportTestable: def test_test_run_report_excluded_from_collection(self) -> None: assert TestRunReport.__test__ is False + + +class TestShardEligible: + def test_all_popen_specs_eligible(self) -> None: + config = _make_config(tx=["popen", "popen"]) + assert shard_eligible(config=config) is True + + def test_bare_popen_eligible(self) -> None: + config = _make_config(tx=["popen"]) + assert shard_eligible(config=config) is True + + def test_popen_with_python_option_eligible(self) -> None: + config = _make_config(tx=["popen//python=python3.11"]) + assert shard_eligible(config=config) is True + + def test_ssh_spec_not_eligible(self) -> None: + config = _make_config(tx=["ssh=host//python=python3"]) + assert shard_eligible(config=config) is False + + def test_socket_spec_not_eligible(self) -> None: + config = _make_config(tx=["socket=192.0.2.1:8888"]) + assert shard_eligible(config=config) is False + + def test_via_proxy_not_eligible(self) -> None: + config = _make_config(tx=["popen//via=proxy"]) + assert shard_eligible(config=config) is False + + def test_mixed_local_and_remote_not_eligible(self) -> None: + config = _make_config(tx=["popen", "ssh=host"]) + assert shard_eligible(config=config) is False + + def test_no_tx_not_eligible(self) -> None: + config = _make_config(tx=None) + assert shard_eligible(config=config) is False + + def test_empty_tx_not_eligible(self) -> None: + config = _make_config(tx=[]) + assert shard_eligible(config=config) is False + + +class TestShardWriter: + def test_write_appends_one_line_per_result(self, tmp_path: Path) -> None: + writer = ShardWriter( + shard_dir=tmp_path, + worker_id="gw0", + size_limit=DEFAULT_SIZE_LIMIT_BYTES, + ) + writer.write(nodeid="test::a", index=0, result=_make_result(summary="one")) + writer.write(nodeid="test::a", index=1, result=_make_result(summary="two")) + writer.close() + lines = ( + (tmp_path / "worker-gw0.jsonl") + .read_text( + encoding="utf-8", + ) + .splitlines() + ) + assert len(lines) == 2 + first = json.loads(lines[0]) + assert first["schema"] == SCHEMA_VERSION + assert first["nodeid"] == "test::a" + assert first["index"] == 0 + assert first["result"]["summary"] == "one" + + def test_write_flushes_before_close(self, tmp_path: Path) -> None: + writer = ShardWriter( + shard_dir=tmp_path, + worker_id="gw0", + size_limit=DEFAULT_SIZE_LIMIT_BYTES, + ) + writer.write(nodeid="n", index=0, result=_make_result()) + content = (tmp_path / "worker-gw0.jsonl").read_text(encoding="utf-8") + assert json.loads(content.strip())["nodeid"] == "n" + + def test_oversized_record_becomes_marker(self, tmp_path: Path) -> None: + writer = ShardWriter(shard_dir=tmp_path, worker_id="gw0", size_limit=10) + writer.write( + nodeid="test::big", + index=3, + result=_make_result(summary="x" * 500), + ) + writer.close() + record = json.loads( + (tmp_path / "worker-gw0.jsonl").read_text(encoding="utf-8").strip(), + ) + assert record["nodeid"] == "test::big" + assert record["index"] == 3 + assert "result" not in record + + def test_separate_files_per_worker(self, tmp_path: Path) -> None: + for worker in ("gw0", "gw1"): + writer = ShardWriter( + shard_dir=tmp_path, + worker_id=worker, + size_limit=DEFAULT_SIZE_LIMIT_BYTES, + ) + writer.write(nodeid="n", index=0, result=_make_result()) + writer.close() + assert (tmp_path / "worker-gw0.jsonl").exists() + assert (tmp_path / "worker-gw1.jsonl").exists() + + +class TestReadWorkerShard: + def test_round_trip_via_writer(self, tmp_path: Path) -> None: + writer = ShardWriter( + shard_dir=tmp_path, + worker_id="gw0", + size_limit=DEFAULT_SIZE_LIMIT_BYTES, + ) + writer.write(nodeid="test::a", index=0, result=_make_result(summary="one")) + writer.write(nodeid="test::b", index=0, result=_make_result(summary="two")) + writer.close() + results, dropped = read_worker_shard(shard_dir=tmp_path, worker_id="gw0") + assert dropped == [] + assert set(results) == {"test::a", "test::b"} + assert results["test::a"][0].summary == "one" + + def test_missing_file_returns_empty(self, tmp_path: Path) -> None: + results, dropped = read_worker_shard(shard_dir=tmp_path, worker_id="ghost") + assert results == {} + assert dropped == [] + + def test_metadata_set_authoritatively_from_envelope( + self, + tmp_path: Path, + ) -> None: + bogus = _make_result( + metadata={"_pytest_nodeid": "spoofed::node", "_rampart_result_index": 99}, + ) + writer = ShardWriter( + shard_dir=tmp_path, + worker_id="gw0", + size_limit=DEFAULT_SIZE_LIMIT_BYTES, + ) + writer.write(nodeid="real::node", index=4, result=bogus) + writer.close() + results, _ = read_worker_shard(shard_dir=tmp_path, worker_id="gw0") + recovered = results["real::node"][0] + assert recovered.metadata["_pytest_nodeid"] == "real::node" + assert recovered.metadata["_rampart_result_index"] == 4 + + def test_truncated_final_line_recovered_as_dropped( + self, + tmp_path: Path, + ) -> None: + writer = ShardWriter( + shard_dir=tmp_path, + worker_id="gw0", + size_limit=DEFAULT_SIZE_LIMIT_BYTES, + ) + writer.write(nodeid="test::a", index=0, result=_make_result(summary="done")) + writer.close() + with (tmp_path / "worker-gw0.jsonl").open("a", encoding="utf-8") as handle: + handle.write('{"schema": "rampart.xdist.v1", "nodeid": "test::b", "inde') + results, dropped = read_worker_shard(shard_dir=tmp_path, worker_id="gw0") + assert results["test::a"][0].summary == "done" + assert len(dropped) == 1 + assert dropped[0].reason == "corrupt" + + def test_oversized_marker_recorded_as_dropped(self, tmp_path: Path) -> None: + writer = ShardWriter(shard_dir=tmp_path, worker_id="gw0", size_limit=10) + writer.write( + nodeid="test::big", + index=2, + result=_make_result(summary="x" * 500), + ) + writer.close() + results, dropped = read_worker_shard(shard_dir=tmp_path, worker_id="gw0") + assert results == {} + assert dropped == [ + DroppedRecord(reason="oversized", nodeid="test::big", index=2), + ] + + def test_corrupt_line_recorded_as_dropped(self, tmp_path: Path) -> None: + (tmp_path / "worker-gw0.jsonl").write_text( + "not json at all\n", + encoding="utf-8", + ) + results, dropped = read_worker_shard(shard_dir=tmp_path, worker_id="gw0") + assert results == {} + assert dropped == [DroppedRecord(reason="corrupt")] + + def test_wrong_schema_recorded_as_dropped(self, tmp_path: Path) -> None: + line = json.dumps( + {"schema": "other.v9", "nodeid": "n", "index": 0, "result": {}}, + ) + (tmp_path / "worker-gw0.jsonl").write_text(line + "\n", encoding="utf-8") + results, dropped = read_worker_shard(shard_dir=tmp_path, worker_id="gw0") + assert results == {} + assert dropped[0].reason == "corrupt" + + def test_non_dict_result_body_dropped_with_ids(self, tmp_path: Path) -> None: + line = json.dumps( + { + "schema": SCHEMA_VERSION, + "nodeid": "test::x", + "index": 7, + "result": 123, + }, + ) + (tmp_path / "worker-gw0.jsonl").write_text(line + "\n", encoding="utf-8") + results, dropped = read_worker_shard(shard_dir=tmp_path, worker_id="gw0") + assert results == {} + assert dropped == [ + DroppedRecord(reason="corrupt", nodeid="test::x", index=7), + ] + + def test_blank_lines_ignored(self, tmp_path: Path) -> None: + writer = ShardWriter( + shard_dir=tmp_path, + worker_id="gw0", + size_limit=DEFAULT_SIZE_LIMIT_BYTES, + ) + writer.write(nodeid="test::a", index=0, result=_make_result()) + writer.close() + with (tmp_path / "worker-gw0.jsonl").open("a", encoding="utf-8") as handle: + handle.write("\n \n") + results, dropped = read_worker_shard(shard_dir=tmp_path, worker_id="gw0") + assert set(results) == {"test::a"} + assert dropped == [] + + def test_multiple_results_same_nodeid_grouped(self, tmp_path: Path) -> None: + writer = ShardWriter( + shard_dir=tmp_path, + worker_id="gw0", + size_limit=DEFAULT_SIZE_LIMIT_BYTES, + ) + writer.write(nodeid="test::a", index=0, result=_make_result(summary="first")) + writer.write(nodeid="test::a", index=1, result=_make_result(summary="second")) + writer.close() + results, dropped = read_worker_shard(shard_dir=tmp_path, worker_id="gw0") + assert dropped == [] + indexes = [r.metadata["_rampart_result_index"] for r in results["test::a"]] + assert indexes == [0, 1] + + def test_rich_result_survives_shard_round_trip(self, tmp_path: Path) -> None: + eval_result = _make_eval_result(outcome=EvalOutcome.DETECTED, confidence=0.8) + turn = _make_turn(eval_result=eval_result, turn_number=2) + result = _make_result( + summary="rich", + harm_category=HarmCategory.JAILBREAK, + turns=[turn], + observability_level=ObservabilityLevel.TOOL_AND_SIDE_EFFECTS, + ) + writer = ShardWriter( + shard_dir=tmp_path, + worker_id="gw0", + size_limit=DEFAULT_SIZE_LIMIT_BYTES, + ) + writer.write(nodeid="test::rich", index=0, result=result) + writer.close() + results, _ = read_worker_shard(shard_dir=tmp_path, worker_id="gw0") + recovered = results["test::rich"][0] + assert recovered.harm_category is HarmCategory.JAILBREAK + assert recovered.observability_level is ObservabilityLevel.TOOL_AND_SIDE_EFFECTS + assert len(recovered.turns) == 1 + assert recovered.turns[0].eval_result is not None + assert recovered.turns[0].eval_result.outcome is EvalOutcome.DETECTED