From a613c7bdb884a071017244f984d8aab790b3f73e Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 26 Jun 2026 15:49:08 +0800 Subject: [PATCH] feat(submit): configurable confirm retries and idempotent confirm Harden confirm_submit_job for unreliable test environments and transient spark-submit failures: - Add confirm_max_retries and confirm_retry_delay_seconds settings (env-configurable). - SUBMITTED pending returns cached result without re-running spark-submit. - FAILED pending is reset to PENDING and retried. - CANCELLED pending is still rejected. - On success, persist pending as SUBMITTED before creating the in-memory Job so a job_store failure cannot leave the record as PENDING while the YARN app is already running. - Persist tracking_url on PendingSubmission for idempotent returns. Tests cover retry-then-success, max-retries-exceeded, idempotency, FAILED reset, and pending-saved-before-job-store. --- common/config.py | 15 ++++ spark_executor/models.py | 2 + spark_executor/tools/submit.py | 106 ++++++++++++++++------ tests/unit/test_pending_store.py | 1 + tests/unit/test_submit_tool.py | 145 ++++++++++++++++++++++++++++++- 5 files changed, 240 insertions(+), 29 deletions(-) diff --git a/common/config.py b/common/config.py index 19afa64..75996c9 100644 --- a/common/config.py +++ b/common/config.py @@ -76,6 +76,13 @@ class Settings: # The binary is resolved against $PATH (set by docker-entrypoint.sh). spark_submit_bin: str = "spark-submit" + # --- Confirm retry policy --- + # confirm_submit_job retries spark-submit on SparkSubmitError to tolerate + # transient client/network issues in test environments or unstable clusters. + # Max retries (not counting the first attempt) and delay between attempts. + confirm_max_retries: int = 3 + confirm_retry_delay_seconds: float = 5.0 + # --- SSL/TLS defaults for YARN REST calls --- ssl_verify_default: bool = True ssl_ca_bundle_default: str | None = None @@ -106,6 +113,12 @@ class Settings: ssl_ca_bundle_default=( os.environ.get("SPARK_EXECUTOR_SSL_CA_BUNDLE_DEFAULT") or None ), + confirm_max_retries=int( + os.environ.get("SPARK_EXECUTOR_CONFIRM_MAX_RETRIES", "3") + ), + confirm_retry_delay_seconds=float( + os.environ.get("SPARK_EXECUTOR_CONFIRM_RETRY_DELAY_SECONDS", "5.0") + ), ) def reload(self) -> "Settings": @@ -123,6 +136,8 @@ class Settings: self.spark_submit_bin = fresh.spark_submit_bin self.ssl_verify_default = fresh.ssl_verify_default self.ssl_ca_bundle_default = fresh.ssl_ca_bundle_default + self.confirm_max_retries = fresh.confirm_max_retries + self.confirm_retry_delay_seconds = fresh.confirm_retry_delay_seconds return self diff --git a/spark_executor/models.py b/spark_executor/models.py index 8291609..b32645b 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -87,6 +87,7 @@ class Connection(BaseModel): class PendingSubmission(BaseModel): pending_id: str app_name: str + connection: str master: str deploy_mode: str @@ -103,3 +104,4 @@ class PendingSubmission(BaseModel): error: str | None = None job_id: str | None = None application_id: str | None = None + tracking_url: str | None = None diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py index be6585d..75aa74f 100644 --- a/spark_executor/tools/submit.py +++ b/spark_executor/tools/submit.py @@ -5,9 +5,11 @@ """ import os import secrets +import time import uuid from datetime import datetime +from common.config import settings from common.logging import logger from common.sql_guard import validate_pyspark_code from spark_executor.core.connection_store import store as conn_store @@ -148,10 +150,29 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: pending = pending_store.get(pending_id) if pending is None: raise KeyError(f"Unknown pending_id: {pending_id}") + + if pending.status == "SUBMITTED": + logger.info( + f"confirm_submit_job idempotent pending_id={pending_id} " + f"job_id={pending.job_id} application_id={pending.application_id}" + ) + return SubmitResult( + job_id=pending.job_id, + application_id=pending.application_id, + tracking_url=pending.tracking_url, + ) + if pending.status == "CANCELLED": + raise ValueError(f"pending_id {pending_id} is CANCELLED, cannot confirm") + if pending.status == "FAILED": + pending.status = "PENDING" + pending.error = None + pending_store.save(pending) + logger.info(f"confirm_submit_job reset pending_id={pending_id} from FAILED to PENDING") if pending.status != "PENDING": raise ValueError( f"pending_id {pending_id} is in status {pending.status!r}, not PENDING" ) + # Defense in depth: re-verify the script still exists. A user could # delete the file between prepare and confirm (or an external cleanup # job could remove it). 400 via the ValueError -> 400 handler. @@ -172,43 +193,72 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: f"confirm_submit_job start pending_id={pending_id} " f"application_target={pending.master} script_path={pending.script_path}" ) - try: - result = run_spark_submit(cmd) - except SparkSubmitError as exc: - pending.status = "FAILED" - pending.error = str(exc) + + last_error: SparkSubmitError | None = None + max_attempts = settings.confirm_max_retries + 1 + for attempt in range(1, max_attempts + 1): + try: + result = run_spark_submit(cmd) + except SparkSubmitError as exc: + last_error = exc + logger.warning( + f"confirm_submit_job attempt {attempt}/{max_attempts} failed " + f"pending_id={pending_id} err={exc}" + ) + if attempt < max_attempts: + time.sleep(settings.confirm_retry_delay_seconds) + continue + except Exception: + # Unexpected failure (not a spark-submit error): fail fast without retry. + logger.exception( + f"confirm_submit_job unexpected error pending_id={pending_id}" + ) + raise + + application_id, tracking_url = parse_spark_submit_output(result.stderr) + job_id = uuid.uuid4().hex[:12] + + # Persist pending as SUBMITTED before creating the in-memory Job so a + # job_store failure cannot leave the record as PENDING while the YARN + # app is already running. + pending.status = "SUBMITTED" + pending.job_id = job_id + pending.application_id = application_id + pending.tracking_url = tracking_url pending_store.save(pending) - logger.error(f"confirm_submit_job failed pending_id={pending_id} err={exc}") - raise - application_id, tracking_url = parse_spark_submit_output(result.stderr) - job_id = uuid.uuid4().hex[:12] + job_store.put( + Job( + job_id=job_id, + application_id=application_id, + script_path=pending.script_path, + queue=pending.queue, + submit_time=datetime.utcnow(), + connection=pending.connection, + yarn_rm_url=pending.yarn_rm_url, + ) + ) - job_store.put( - Job( + logger.info( + f"confirm_submit_job ok pending_id={pending_id} job_id={job_id} " + f"application_id={application_id}" + ) + return SubmitResult( job_id=job_id, application_id=application_id, - script_path=pending.script_path, - queue=pending.queue, - submit_time=datetime.utcnow(), - connection=pending.connection, - yarn_rm_url=pending.yarn_rm_url, + tracking_url=tracking_url, ) - ) - pending.status = "SUBMITTED" - pending.job_id = job_id - pending.application_id = application_id + # Exhausted all retries. + assert last_error is not None + pending.status = "FAILED" + pending.error = str(last_error) pending_store.save(pending) - logger.info( - f"confirm_submit_job ok pending_id={pending_id} job_id={job_id} " - f"application_id={application_id}" - ) - return SubmitResult( - job_id=job_id, - application_id=application_id, - tracking_url=tracking_url, + logger.error( + f"confirm_submit_job failed pending_id={pending_id} after {max_attempts} attempts " + f"err={last_error}" ) + raise last_error def list_pending_jobs() -> list[dict[str, object]]: diff --git a/tests/unit/test_pending_store.py b/tests/unit/test_pending_store.py index 73c3b0f..d7004aa 100644 --- a/tests/unit/test_pending_store.py +++ b/tests/unit/test_pending_store.py @@ -1,4 +1,5 @@ # coding=utf-8 +import json from datetime import datetime from pathlib import Path diff --git a/tests/unit/test_submit_tool.py b/tests/unit/test_submit_tool.py index 29a07bc..a761a10 100644 --- a/tests/unit/test_submit_tool.py +++ b/tests/unit/test_submit_tool.py @@ -1,13 +1,20 @@ # coding=utf-8 from pathlib import Path from unittest.mock import patch +import time import pytest +from common.config import settings from spark_executor.core import connection_store, pending_store from spark_executor.core.pending_store import PendingStore from spark_executor.models import Connection from spark_executor.tools import submit +from spark_executor.core.spark_submit import SparkSubmitError + + +def _spark_proc(stderr: str = "tracking URL: http://rm:8088/proxy/application_17400000001/\n"): + return type("P", (), {"returncode": 0, "stderr": stderr})() @pytest.fixture(autouse=True) @@ -331,7 +338,6 @@ def test_confirm_refuses_non_pending_status(monkeypatch, real_script): def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script): - from spark_executor.core.spark_submit import SparkSubmitError submit.prepare_submit_job( connection="prod", script_path=str(real_script), @@ -345,6 +351,7 @@ def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script): def _raise(_cmd): raise SparkSubmitError("boom") monkeypatch.setattr(submit, "run_spark_submit", _raise) + settings.confirm_max_retries = 0 with pytest.raises(SparkSubmitError): submit.confirm_submit_job(pending_id=pid) p = submit.pending_store.get(pid) @@ -352,6 +359,142 @@ def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script): assert "boom" in (p.error or "") +# --- confirm retry / idempotency --- + +def test_confirm_retries_spark_submit_failure_then_succeeds(monkeypatch, real_script): + monkeypatch.setattr(time, "sleep", lambda _s: None) + settings.confirm_max_retries = 2 + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) + pid = _last_pending_id() + calls = [] + def _maybe_fail(cmd): + calls.append(cmd) + if len(calls) < 3: + raise SparkSubmitError("transient") + return _spark_proc() + monkeypatch.setattr(submit, "run_spark_submit", _maybe_fail) + result = submit.confirm_submit_job(pending_id=pid) + assert result.application_id == "application_17400000001" + assert len(calls) == 3 + p = submit.pending_store.get(pid) + assert p.status == "SUBMITTED" + + +def test_confirm_exhausts_retries_and_sets_failed(monkeypatch, real_script): + monkeypatch.setattr(time, "sleep", lambda _s: None) + settings.confirm_max_retries = 2 + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) + pid = _last_pending_id() + calls = [] + def _always_fail(cmd): + calls.append(cmd) + raise SparkSubmitError("persistent") + monkeypatch.setattr(submit, "run_spark_submit", _always_fail) + with pytest.raises(SparkSubmitError, match="persistent"): + submit.confirm_submit_job(pending_id=pid) + assert len(calls) == 3 + p = submit.pending_store.get(pid) + assert p.status == "FAILED" + assert "persistent" in (p.error or "") + + +def test_confirm_idempotent_when_submitted(monkeypatch, real_script): + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) + pid = _last_pending_id() + monkeypatch.setattr(submit, "run_spark_submit", lambda _cmd: _spark_proc()) + first = submit.confirm_submit_job(pending_id=pid) + + call_count = [0] + def _counting_run(cmd): + call_count[0] += 1 + return _spark_proc() + monkeypatch.setattr(submit, "run_spark_submit", _counting_run) + second = submit.confirm_submit_job(pending_id=pid) + assert call_count[0] == 0 + assert second.job_id == first.job_id + assert second.application_id == first.application_id + assert second.tracking_url == first.tracking_url + + +def test_confirm_resets_failed_and_retries(monkeypatch, real_script): + monkeypatch.setattr(time, "sleep", lambda _s: None) + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) + pid = _last_pending_id() + monkeypatch.setattr( + submit, + "run_spark_submit", + lambda _cmd: (_ for _ in ()).throw(SparkSubmitError("boom")), + ) + with pytest.raises(SparkSubmitError): + submit.confirm_submit_job(pending_id=pid) + p = submit.pending_store.get(pid) + assert p.status == "FAILED" + assert "boom" in (p.error or "") + + monkeypatch.setattr(submit, "run_spark_submit", lambda _cmd: _spark_proc()) + result = submit.confirm_submit_job(pending_id=pid) + assert result.application_id == "application_17400000001" + p = submit.pending_store.get(pid) + assert p.status == "SUBMITTED" + assert p.error is None + + +def test_confirm_updates_pending_before_job_store(monkeypatch, real_script): + submit.prepare_submit_job( + connection="prod", + script_path=str(real_script), + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + app_name="test-app", + ) + pid = _last_pending_id() + monkeypatch.setattr(submit, "run_spark_submit", lambda _cmd: _spark_proc()) + + def _exploding_put(job): + raise RuntimeError("job_store unavailable") + monkeypatch.setattr(submit.job_store, "put", _exploding_put) + with pytest.raises(RuntimeError, match="job_store unavailable"): + submit.confirm_submit_job(pending_id=pid) + p = submit.pending_store.get(pid) + assert p.status == "SUBMITTED" + assert p.application_id == "application_17400000001" + assert p.tracking_url is not None + + # --- list_pending_jobs --- def test_list_pending_jobs_empty():