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.
2026-06-26 15:51:51 +08:00
parent 9e7b425f09
commit a613c7bdb8
5 changed files with 240 additions and 29 deletions
@@ -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


@@ -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
@@ -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]]:
@@ -1,4 +1,5 @@
 # coding=utf-8
+import json
 from datetime import datetime
 from pathlib import Path

@@ -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():