fix(harness): preserve artifact values and write exits

Author: jioffe502

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,22 @@ 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_MARKERS = ("api_key", "password", "secret", "credential", "webhook")
+_TOKEN_CREDENTIAL_QUALIFIERS = {"access", "api", "auth", "bearer", "oauth", "refresh", "service", "session"}
 
 
 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("_")
+    if any(marker in normalized for marker in _SENSITIVE_KEY_MARKERS):
+        return True
+    parts = normalized.split("_")
+    if "token" in parts:
+        return True
+    return any(
+        part == "tokens" and index > 0 and parts[index - 1] in _TOKEN_CREDENTIAL_QUALIFIERS
+        for index, part in enumerate(parts)
+    )
 
 
 def redact(value: Any) -> Any:

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,

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(

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,42 @@ def test_redact_recurses_into_structured_override_values():
     assert redact(override) == 'query={"reranker_api_key": "<redacted>", "top_k": 10}'
 
 
+def test_redact_preserves_token_limits():
+    payload = {
+        "reranker_api_key": "secret-value",
+        "webhookUrl": "https://hooks.example.invalid/secret",
+        "webhooks": ["https://hooks.example.invalid/one"],
+        "passwords": ["password-value"],
+        "client_secrets": ["client-secret"],
+        "access_tokens": ["access-token"],
+        "caption_max_tokens": 77,
+        "text_chunk_overlap_tokens": 12,
+        "tokenizer_model": "model-name",
+    }
+
+    assert redact(payload) == {
+        "reranker_api_key": "<redacted>",
+        "webhookUrl": "<redacted>",
+        "webhooks": "<redacted>",
+        "passwords": "<redacted>",
+        "client_secrets": "<redacted>",
+        "access_tokens": "<redacted>",
+        "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 +266,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