diff --git a/nemo_retriever/src/nemo_retriever/harness/artifact_writer.py b/nemo_retriever/src/nemo_retriever/harness/artifact_writer.py index c1b73ce31..b83101d35 100644 --- a/nemo_retriever/src/nemo_retriever/harness/artifact_writer.py +++ b/nemo_retriever/src/nemo_retriever/harness/artifact_writer.py @@ -10,6 +10,7 @@ import json import os from pathlib import Path +import re import shutil import sys import tempfile @@ -50,15 +51,17 @@ def append_jsonl(path: Path, payload: Mapping[str, Any]) -> None: def append_text(path: Path, text: str) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("a", encoding="utf-8") as handle: - handle.write(text) + try: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(text) + except OSError as exc: + raise artifact_write_error(exc) from exc @contextlib.contextmanager def capture_output_to_log(path: Path, *, label: str): """Capture noisy stdout/stderr to a persistent run log.""" - path.parent.mkdir(parents=True, exist_ok=True) append_text(path, f"\n## {utc_now()} {label} start\n") try: stdout_fd, stderr_fd = sys.stdout.fileno(), sys.stderr.fileno() @@ -99,16 +102,19 @@ def capture_output_to_log(path: Path, *, label: str): if buf is not None: buf.seek(0) captured = buf.read() - with path.open("ab") as handle: - if captured: - handle.write(captured) - if not captured.endswith(b"\n"): - handle.write(b"\n") - if failure_traceback: - handle.write(failure_traceback.encode("utf-8", errors="replace")) - if not failure_traceback.endswith("\n"): - handle.write(b"\n") - handle.write(f"## {utc_now()} {label} {'failed' if failed else 'complete'}\n".encode("utf-8")) + try: + with path.open("ab") as handle: + if captured: + handle.write(captured) + if not captured.endswith(b"\n"): + handle.write(b"\n") + if failure_traceback: + handle.write(failure_traceback.encode("utf-8", errors="replace")) + if not failure_traceback.endswith("\n"): + handle.write(b"\n") + handle.write(f"## {utc_now()} {label} {'failed' if failed else 'complete'}\n".encode("utf-8")) + except OSError as exc: + raise artifact_write_error(exc) from exc if failed and captured: sys.stderr.buffer.write(captured) sys.stderr.flush() @@ -121,12 +127,21 @@ def capture_output_to_log(path: Path, *, label: str): os.close(saved_stdout) -_SENSITIVE_KEY_PARTS = ("api_key", "password", "secret", "credential", "token", "webhook") +_SENSITIVE_KEY_PARTS = { + "credential", + "credentials", + "password", + "secret", + "token", + "webhook", +} def _is_sensitive_key(value: Any) -> bool: - normalized = str(value).lower() - return any(token in normalized for token in _SENSITIVE_KEY_PARTS) + snake_case = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", str(value)) + normalized = re.sub(r"[^a-z0-9]+", "_", snake_case.lower()).strip("_") + parts = set(normalized.split("_")) + return "api_key" in normalized or bool(parts & _SENSITIVE_KEY_PARTS) def redact(value: Any) -> Any: diff --git a/nemo_retriever/src/nemo_retriever/harness/beir_runner.py b/nemo_retriever/src/nemo_retriever/harness/beir_runner.py index 3150879cd..093d46f4b 100644 --- a/nemo_retriever/src/nemo_retriever/harness/beir_runner.py +++ b/nemo_retriever/src/nemo_retriever/harness/beir_runner.py @@ -17,7 +17,7 @@ FailurePayload, HarnessRunError, ) -from nemo_retriever.harness.json_io import write_json +from nemo_retriever.harness.json_io import artifact_write_error, write_json from nemo_retriever.query.options import QueryRequest from nemo_retriever.query.workflow import ResolvedQueryPlan from nemo_retriever.tools.recall.beir import ( @@ -30,12 +30,15 @@ def _write_trec_run(path: Path, run: Mapping[str, Mapping[str, float]]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("w", encoding="utf-8") as handle: - for query_id, docs in run.items(): - ordered = sorted(docs.items(), key=lambda item: (-item[1], item[0])) - for rank, (doc_id, score) in enumerate(ordered, start=1): - handle.write(f"{query_id} Q0 {doc_id} {rank} {float(score):.6f} retriever-harness\n") + try: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + for query_id, docs in run.items(): + ordered = sorted(docs.items(), key=lambda item: (-item[1], item[0])) + for rank, (doc_id, score) in enumerate(ordered, start=1): + handle.write(f"{query_id} Q0 {doc_id} {rank} {float(score):.6f} retriever-harness\n") + except OSError as exc: + raise artifact_write_error(exc) from exc def _write_query_result( @@ -165,8 +168,6 @@ def _dense_retrieve( raw_hits: list[list[dict[str, Any]]] = [] latencies_ms: list[float] = [] query_results_path = writer.path("query_results.jsonl") - if query_results_path.exists(): - query_results_path.unlink() for query_id, query_text in zip(dataset.query_ids, dataset.queries): start = time.perf_counter() try: @@ -251,8 +252,6 @@ def _agentic_retrieve( query_count = len(dataset.queries) per_query_ms = total_ms / query_count if query_count else 0.0 query_results_path = writer.path("query_results.jsonl") - if query_results_path.exists(): - query_results_path.unlink() for query_id, query_text, doc_ids in zip(dataset.query_ids, dataset.queries, ranked_doc_ids): _write_query_result( query_results_path, diff --git a/nemo_retriever/src/nemo_retriever/harness/execution.py b/nemo_retriever/src/nemo_retriever/harness/execution.py index fb3aa9438..66ae7c5ea 100644 --- a/nemo_retriever/src/nemo_retriever/harness/execution.py +++ b/nemo_retriever/src/nemo_retriever/harness/execution.py @@ -373,6 +373,8 @@ def run_prepared_benchmark( with capture_output_to_log(writer.path("run.log"), label="ingest"): ingest_summary = run_ingest_workflow(ingest_plan, dry_run=False) except Exception as exc: + if isinstance(exc, HarnessRunError) and exc.exit_code == EXIT_ARTIFACT_WRITE_FAILURE: + raise raise HarnessRunError( EXIT_INGEST_FAILURE, FailurePayload( diff --git a/nemo_retriever/tests/test_harness_artifacts.py b/nemo_retriever/tests/test_harness_artifacts.py index b5c791218..58f5e7364 100644 --- a/nemo_retriever/tests/test_harness_artifacts.py +++ b/nemo_retriever/tests/test_harness_artifacts.py @@ -2,11 +2,19 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from contextlib import contextmanager import json import pytest -from nemo_retriever.harness.artifact_writer import artifact_paths, ArtifactWriter, capture_output_to_log, redact +from nemo_retriever.harness.artifact_writer import ( + append_text, + artifact_paths, + ArtifactWriter, + capture_output_to_log, + redact, +) +from nemo_retriever.harness.beir_runner import _write_trec_run from nemo_retriever.harness.contracts import ( EXIT_ARTIFACT_WRITE_FAILURE, EXIT_INGEST_FAILURE, @@ -91,6 +99,34 @@ def test_redact_recurses_into_structured_override_values(): assert redact(override) == 'query={"reranker_api_key": "", "top_k": 10}' +def test_redact_preserves_token_limits(): + payload = { + "reranker_api_key": "secret-value", + "webhookUrl": "https://hooks.example.invalid/secret", + "caption_max_tokens": 77, + "text_chunk_overlap_tokens": 12, + "tokenizer_model": "model-name", + } + + assert redact(payload) == { + "reranker_api_key": "", + "webhookUrl": "", + "caption_max_tokens": 77, + "text_chunk_overlap_tokens": 12, + "tokenizer_model": "model-name", + } + + +def test_text_artifact_writes_are_classified(tmp_path): + with pytest.raises(HarnessRunError) as log_error: + append_text(tmp_path, "log line") + with pytest.raises(HarnessRunError) as trec_error: + _write_trec_run(tmp_path, {"query": {"document": 1.0}}) + + assert log_error.value.exit_code == EXIT_ARTIFACT_WRITE_FAILURE + assert trec_error.value.exit_code == EXIT_ARTIFACT_WRITE_FAILURE + + def test_invalid_run_config_preserves_existing_artifacts(tmp_path): (tmp_path / "results.json").write_text('{"old": true}', encoding="utf-8") (tmp_path / "lancedb").mkdir() @@ -222,6 +258,35 @@ def fail_writer(**kwargs): assert error.value.failure.message == "OSError: artifact directory unavailable" +def test_run_benchmark_preserves_run_log_write_failure(monkeypatch, tmp_path): + import nemo_retriever.harness.execution as execution + from nemo_retriever.harness.json_io import artifact_write_error + + class FakeIngestPlan: + documents = () + + @contextmanager + def fail_log_capture(*args, **kwargs): + raise artifact_write_error(OSError("run log unavailable")) + yield + + monkeypatch.setattr(execution, "resolve_ingest_plan", lambda request: FakeIngestPlan()) + monkeypatch.setattr(execution, "run_ingest_workflow", lambda plan, dry_run: {}) + monkeypatch.setattr(execution, "resolve_query_plan", lambda request: object()) + monkeypatch.setattr(execution, "query_plan_payload", lambda plan: {}) + monkeypatch.setattr(execution, "capture_output_to_log", fail_log_capture) + + outcome = run_benchmark( + "jp20_smoke", + output_dir=str(tmp_path / "run"), + overrides=(f'dataset.path="{tmp_path}"',), + ) + + assert outcome.exit_code == EXIT_ARTIFACT_WRITE_FAILURE + assert outcome.results["failure"]["failure_reason"] == "artifact_write_failed" + assert outcome.results["failure"]["message"] == "OSError: run log unavailable" + + def test_run_benchmark_keeps_non_write_failures_internal(monkeypatch, tmp_path): import nemo_retriever.harness.execution as execution