diff --git a/backend/scripts/run_text2sql_bird_parallel.py b/backend/scripts/run_text2sql_bird_parallel.py index 619b2e1..217a0bf 100644 --- a/backend/scripts/run_text2sql_bird_parallel.py +++ b/backend/scripts/run_text2sql_bird_parallel.py @@ -105,6 +105,17 @@ def assert_frozen_mappings(manifest: dict[str, Any]) -> None: class RecoverableShardResultError(ValueError): """A shard wrote a result, but one or more agents did not finish.""" + + +class TransientShardResultError(OSError): + """A shard's result.json is not yet ready (missing, mid-write, or unreadable). + + ``wait_for_shard_result`` treats this as a transient filesystem race and + retries it during the publication grace period. Semantic validation + failures (wrong IDs, missing completion timestamp, malformed payload) + raise plain ``ValueError`` so the scheduler can fail fast instead of + waiting for a grace period that can never resolve the error. + """ STATIC_CONFIG_KEYS = ( "provider", "mode", @@ -464,7 +475,13 @@ def shard_command( def validate_shard_result(shard: dict[str, Any]) -> None: result_path = Path(shard["result"]) if not result_path.exists(): - raise ValueError(f"shard {shard['index']} did not create result.json") + # The result file is not yet visible. This is the only "transient" + # condition that ``wait_for_shard_result`` should retry; semantic + # validation errors below raise ``ValueError`` directly so the + # scheduler can fail fast. + raise TransientShardResultError( + f"shard {shard['index']} did not create result.json" + ) payload = load_json(result_path) actual_ids = [int(item["question_id"]) for item in payload.get("results", [])] expected_ids = [int(value) for value in shard["question_ids"]] @@ -499,17 +516,23 @@ def wait_for_shard_result( *, grace_seconds: float = 5.0, ) -> None: - """Allow a just-exited Windows worker a brief result-publication grace period.""" + """Allow a just-exited Windows worker a brief result-publication grace period. + + Only transient filesystem conditions (missing file, mid-write parse error) + are retried during the grace period. Semantic validation failures (wrong + IDs, missing completion timestamp) are raised immediately so a corrupt + result cannot stall a worker slot for the full grace window. + """ deadline = time.monotonic() + grace_seconds while True: try: validate_shard_result(shard) return - except RecoverableShardResultError: - # A fully published unfinished result is not a filesystem race. Let - # the scheduler archive and requeue it immediately. + except (RecoverableShardResultError, ValueError): + # A fully published invalid/unfinished result is not a filesystem + # race. Let the scheduler archive and requeue it immediately. raise - except (OSError, ValueError, json.JSONDecodeError): + except (TransientShardResultError, json.JSONDecodeError): if time.monotonic() >= deadline: raise time.sleep(0.05) @@ -627,7 +650,11 @@ def run_shards( f"requeued from scratch ({exc}); archived={archive_path}", flush=True, ) - except ValueError: + except (TransientShardResultError, ValueError): + # Either the file disappeared between the exists() check + # and validate_shard_result, or the cached result is + # semantically invalid (wrong IDs / missing timestamp). + # Either way, fall through to relaunch the shard. pass else: completed.add(index) @@ -713,7 +740,11 @@ def run_shards( if return_code != 0: raise ValueError(f"process exited with code {return_code}") wait_for_shard_result(shard) - except ValueError as exc: + except (TransientShardResultError, ValueError) as exc: + # TransientShardResultError (missing file after grace) and + # plain ValueError (process non-zero exit or permanent + # validation failure) are both requeueable; the next loop + # pass decides whether to archive or relaunch. if retry_attempts and attempts[index] >= retry_attempts: failures.append(f"shard {index}: {exc}; log={shard['log']}") print(f"[parallel] shard {index:02d} failed permanently", flush=True) @@ -753,7 +784,9 @@ def merge_results(manifest_path: Path, output: Path) -> Path: if result_path.exists(): try: validate_shard_result(shard) - except ValueError: + except (TransientShardResultError, ValueError): + # Either the file vanished or the cached result is + # semantically invalid; the timeout check below will handle it. pass else: result_valid = True diff --git a/backend/tests/test_text2sql_bird_parallel.py b/backend/tests/test_text2sql_bird_parallel.py index 1e093c7..84408a0 100644 --- a/backend/tests/test_text2sql_bird_parallel.py +++ b/backend/tests/test_text2sql_bird_parallel.py @@ -677,7 +677,11 @@ def test_result_publication_grace_retries_transient_missing_artifact( def delayed_validate(_shard): calls.append(1) if len(calls) < 3: - raise ValueError("result.json is not visible yet") + # ``validate_shard_result`` signals "not yet visible" with a + # TransientShardResultError; the grace period must retry it. + raise parallel.TransientShardResultError( + "result.json is not visible yet" + ) monkeypatch.setattr(parallel, "validate_shard_result", delayed_validate) monkeypatch.setattr(parallel.time, "sleep", lambda _seconds: None) diff --git a/backend/tests/test_wait_for_shard_permanent_failure.py b/backend/tests/test_wait_for_shard_permanent_failure.py new file mode 100644 index 0000000..66887a7 --- /dev/null +++ b/backend/tests/test_wait_for_shard_permanent_failure.py @@ -0,0 +1,226 @@ +"""Regression tests for the parallel shard result publication logic. + +Background +---------- +``wait_for_shard_result`` gives a just-exited worker a short grace period +to publish its ``result.json``. It must distinguish two failure modes: + +* Transient: the file is missing or mid-write (e.g., on Windows where the + writing process can exit before the OS has fully flushed the file). The + grace period must retry these so a transient filesystem race does not + crash the run. +* Permanent: the file is present but semantically invalid (wrong IDs, + missing completion timestamp). The worker has already exited; the + result can never become valid during the grace window. Retrying wastes + up to ``grace_seconds`` per shard and stalls a worker slot. + +The previous implementation caught ``ValueError`` for both modes, which +caused permanent validation failures to spin for the full grace period +before raising. The fix introduces ``TransientShardResultError`` for the +transient case and lets permanent validation errors propagate immediately. +""" +from __future__ import annotations + +import importlib.util +import json +import sys +import time +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +SCRIPT = ROOT / "scripts" / "run_text2sql_bird_parallel.py" +SPEC = importlib.util.spec_from_file_location("run_text2sql_bird_parallel", SCRIPT) +parallel = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = parallel +SPEC.loader.exec_module(parallel) + + +def _make_shard(result_path: Path, question_ids: list[int], index: int = 1) -> dict: + return { + "index": index, + "question_ids": question_ids, + "result": str(result_path), + "output": str(result_path.parent), + } + + +def test_wait_for_shard_result_raises_immediately_on_wrong_ids( + tmp_path: Path, +) -> None: + """Wrong number of IDs is permanent; no grace-period spinning.""" + result_path = tmp_path / "result.json" + result_path.write_text(json.dumps({ + "config": {"finished_at": "2026-01-01T00:00:00+00:00"}, + "results": [], # 0 results instead of 2 expected + })) + shard = _make_shard(result_path, [10, 20]) + + start = time.monotonic() + raised = None + try: + parallel.wait_for_shard_result(shard, grace_seconds=2.0) + except ValueError as exc: + raised = exc + elapsed = time.monotonic() - start + + assert raised is not None, "expected a ValueError for permanent validation failure" + assert "incomplete or out of order" in str(raised) + # Permanent errors must propagate well under the 2.0s grace period. + assert elapsed < 0.5, ( + f"wait_for_shard_result retried a permanent validation error for " + f"{elapsed:.2f}s; expected immediate raise" + ) + + +def test_wait_for_shard_result_raises_immediately_on_missing_timestamp( + tmp_path: Path, +) -> None: + """Missing ``finished_at`` is permanent; no grace-period spinning.""" + result_path = tmp_path / "result.json" + result_path.write_text(json.dumps({ + "config": {}, # no finished_at + "results": [{"question_id": 1}], + })) + shard = _make_shard(result_path, [1]) + + start = time.monotonic() + raised = None + try: + parallel.wait_for_shard_result(shard, grace_seconds=2.0) + except ValueError as exc: + raised = exc + elapsed = time.monotonic() - start + + assert raised is not None + assert "no completion timestamp" in str(raised) + assert elapsed < 0.5 + + +def test_wait_for_shard_result_retries_transient_missing_file( + tmp_path: Path, + monkeypatch, +) -> None: + """Missing result.json (transient) must be retried within the grace period.""" + result_path = tmp_path / "result.json" + if result_path.exists(): + result_path.unlink() + + call_count = {"value": 0} + + def eventually_publish(_shard: dict) -> None: + call_count["value"] += 1 + if call_count["value"] == 2: + result_path.write_text(json.dumps({ + "config": {"finished_at": "2026-01-01T00:00:00+00:00"}, + "results": [{"question_id": 1}], + })) + # Let the real validator drive the missing-vs-present check. + parallel.validate_shard_result.__wrapped__(_shard) if hasattr( + parallel.validate_shard_result, "__wrapped__" + ) else None + + # The published-file check goes through Path.exists(); use a fresh + # validator that fakes "file appeared" on the second call. + original = parallel.validate_shard_result + + def wrapped(shard: dict) -> None: + call_count["value"] += 1 + if call_count["value"] >= 2 and not result_path.exists(): + result_path.write_text(json.dumps({ + "config": {"finished_at": "2026-01-01T00:00:00+00:00"}, + "results": [{"question_id": 1}], + })) + original(shard) + + monkeypatch.setattr(parallel, "validate_shard_result", wrapped) + monkeypatch.setattr(parallel.time, "sleep", lambda _seconds: None) + + shard = _make_shard(result_path, [1]) + parallel.wait_for_shard_result(shard, grace_seconds=2.0) + assert call_count["value"] >= 2 + + +def test_validate_shard_result_raises_transient_when_result_missing( + tmp_path: Path, +) -> None: + """``validate_shard_result`` must signal a missing file with the + transient exception so ``wait_for_shard_result`` knows to retry it. + """ + missing = tmp_path / "absent.json" + assert not missing.exists() + + raised = None + try: + parallel.validate_shard_result(_make_shard(missing, [1])) + except parallel.TransientShardResultError as exc: + raised = exc + assert raised is not None + assert "did not create result.json" in str(raised) + # Must also be catchable as OSError so existing call-sites that already + # catch OSError continue to work. + assert isinstance(raised, OSError) + + +def test_run_shards_treats_transient_and_permanent_during_relaunch( + tmp_path: Path, + monkeypatch, +) -> None: + """When ``validate_shard_result`` raises ``TransientShardResultError`` + on a cached result, ``run_shards`` must relaunch the shard (not + crash the whole run). + """ + manifest_path = tmp_path / "manifest.json" + shard = { + "index": 1, + "question_ids": [7], + "questions": str(tmp_path / "questions.json"), + "output": str(tmp_path / "shard-1" / "run"), + "result": str(tmp_path / "shard-1" / "run" / "result.json"), + "log": str(tmp_path / "shard-1" / "run.log"), + } + (tmp_path / "shard-1" / "run").mkdir(parents=True) + Path(shard["questions"]).write_text("[]") + Path(shard["result"]).write_text("{}") # exists but invalid JSON + manifest_path.write_text(json.dumps({ + "format": parallel.FORMAT, + "database_root": str(tmp_path / "databases"), + "mapping_root": str(tmp_path / "mappings"), + "shard_count": 1, + "run": { + "mode": "physical", + "model": "test-model", + "provider": "ollama", + "api_base": "http://localhost", + "api_key_env": "", + "max_output_tokens": 1024, + "seed": 17, + "max_steps": 3, + "require_plan": False, + "timeout": 30, + "shard_wall_timeout": 0, + "agent_wall_timeout": 0, + }, + "shards": [shard], + })) + + class FakeProcess: + pid = 1 + def poll(self): return 0 + + # First poll returns a valid result on the second invocation. + call_state = {"count": 0} + + def fake_popen(_command, **_kwargs): + call_state["count"] += 1 + if call_state["count"] >= 2: + Path(shard["result"]).write_text(json.dumps({ + "config": {"finished_at": "2026-01-01T00:00:00+00:00"}, + "results": [{"question_id": 7}], + })) + return FakeProcess() + + monkeypatch.setattr(parallel.subprocess, "Popen", fake_popen) + monkeypatch.setattr(parallel.time, "sleep", lambda _seconds: None) + + parallel.run_shards(manifest_path, worker_count=1, retry_attempts=0, retry_delay=0) + assert call_state["count"] >= 2