Skip to content

[harness] Preserve harness artifact values and write exit codes #2286

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jioffe502 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[harness] Preserve harness artifact values and write exit codes #2286

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 33 additions & 17 deletions nemo_retriever/src/nemo_retriever/harness/artifact_writer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import json
import os
from pathlib import Path
import re
import shutil
import sys
import tempfile
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand All @@ -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)
)
Comment on lines +142 to +145

@greptile-apps greptile-apps Bot Jul 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Token-plural qualifier checks only the immediately preceding word

The check parts[index - 1] in _TOKEN_CREDENTIAL_QUALIFIERS only examines the word at index - 1. A key like service_account_tokens would not be redacted because "account" is the immediate predecessor of "tokens", even though "service" (a qualifier) appears at index - 2. Similarly oauth_app_tokens, api_service_tokens, etc. would slip through. The fix is to check any part between position 0 and index - 1 for a qualifier, rather than only index - 1.

Prompt To Fix With AI 
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/harness/artifact_writer.py
Line: 142-145

Comment:
**Token-plural qualifier checks only the immediately preceding word**

The check `parts[index - 1] in _TOKEN_CREDENTIAL_QUALIFIERS` only examines the word at `index - 1`. A key like `service_account_tokens` would not be redacted because "account" is the immediate predecessor of "tokens", even though "service" (a qualifier) appears at `index - 2`. Similarly `oauth_app_tokens`, `api_service_tokens`, etc. would slip through. The fix is to check any part between position 0 and `index - 1` for a qualifier, rather than only `index - 1`.

How can I resolve this? If you propose a fix, please make it concise.



def redact(value: Any) -> Any:
Expand Down
21 changes: 10 additions & 11 deletions nemo_retriever/src/nemo_retriever/harness/beir_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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(
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions nemo_retriever/src/nemo_retriever/harness/execution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
75 changes: 74 additions & 1 deletion nemo_retriever/tests/test_harness_artifacts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Comment on lines +276 to +279

@greptile-apps greptile-apps Bot Jul 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 The yield after raise is unreachable at runtime but is semantically required: Python classifies a function as a generator function only when it contains at least one yield expression, regardless of reachability. Without the yield, @contextmanager would receive a plain function and would fail when __enter__ calls next(). A brief comment prevents future confusion or spurious linter suppression.

Suggested change
@contextmanager
def fail_log_capture(*args, **kwargs):
raise artifact_write_error(OSError("run log unavailable"))
yield
@contextmanager
def fail_log_capture(*args, **kwargs):
raise artifact_write_error(OSError("run log unavailable"))
yield # unreachable; required to make this a generator function for @contextmanager
Prompt To Fix With AI 
This is a comment left during a code review.
Path: nemo_retriever/tests/test_harness_artifacts.py
Line: 276-279

Comment:
The `yield` after `raise` is unreachable at runtime but is semantically required: Python classifies a function as a generator function only when it contains at least one `yield` expression, regardless of reachability. Without the `yield`, `@contextmanager` would receive a plain function and would fail when `__enter__` calls `next()`. A brief comment prevents future confusion or spurious linter suppression.

```suggestion
    @contextmanager
    def fail_log_capture(*args, **kwargs):
        raise artifact_write_error(OSError("run log unavailable"))
        yield  # unreachable; required to make this a generator function for @contextmanager
```

How can I resolve this? If you propose a fix, please make it concise.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!


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

Expand Down
Loading