feat: add prepare_submit_job (no spark-submit invocation)
This commit is contained in:
@@ -0,0 +1,57 @@
|
|||||||
|
# coding=utf-8
|
||||||
|
"""
|
||||||
|
@Time :2026/6/24
|
||||||
|
@Author :tao.chen
|
||||||
|
"""
|
||||||
|
import secrets
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
from common.logging import logger
|
||||||
|
from spark_executor.core.connection_store import store as conn_store
|
||||||
|
from spark_executor.core.pending_store import store as pending_store
|
||||||
|
from spark_executor.core.spark_submit import run_spark_submit # re-exported for monkeypatch in tests
|
||||||
|
from spark_executor.models import PendingSubmission
|
||||||
|
|
||||||
|
|
||||||
|
def _new_pending_id() -> str:
|
||||||
|
return "p_" + secrets.token_hex(6)
|
||||||
|
|
||||||
|
|
||||||
|
def prepare_submit_job(
|
||||||
|
*,
|
||||||
|
connection: str,
|
||||||
|
script_path: str,
|
||||||
|
queue: str = "default",
|
||||||
|
executor_memory: str = "4G",
|
||||||
|
executor_cores: int = 2,
|
||||||
|
num_executors: int = 2,
|
||||||
|
) -> dict[str, object]:
|
||||||
|
"""Snapshot connection params and persist a PendingSubmission. Does NOT submit."""
|
||||||
|
conn = conn_store.get(connection)
|
||||||
|
if conn is None:
|
||||||
|
raise KeyError(f"Unknown connection: {connection}")
|
||||||
|
|
||||||
|
pending_id = _new_pending_id()
|
||||||
|
pending = PendingSubmission(
|
||||||
|
pending_id=pending_id,
|
||||||
|
connection=connection,
|
||||||
|
master=conn.master,
|
||||||
|
deploy_mode=conn.deploy_mode,
|
||||||
|
script_path=script_path,
|
||||||
|
queue=queue,
|
||||||
|
executor_memory=executor_memory,
|
||||||
|
executor_cores=executor_cores,
|
||||||
|
num_executors=num_executors,
|
||||||
|
spark_conf=dict(conn.spark_conf),
|
||||||
|
created_at=datetime.utcnow(),
|
||||||
|
status="PENDING",
|
||||||
|
)
|
||||||
|
pending_store.save(pending)
|
||||||
|
logger.info(
|
||||||
|
f"prepare_submit_job pending_id={pending_id} connection={connection} master={conn.master}"
|
||||||
|
)
|
||||||
|
return {
|
||||||
|
"pending_id": pending_id,
|
||||||
|
"status": "PENDING",
|
||||||
|
"parameters": pending.model_dump(),
|
||||||
|
}
|
||||||
@@ -0,0 +1,79 @@
|
|||||||
|
# coding=utf-8
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(autouse=True)
|
||||||
|
def _fresh(tmp_path: Path, monkeypatch):
|
||||||
|
monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
|
||||||
|
monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
|
||||||
|
monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path))
|
||||||
|
monkeypatch.setattr(pending_store, "store", PendingStore())
|
||||||
|
submit.conn_store = connection_store.store
|
||||||
|
submit.pending_store = pending_store.store
|
||||||
|
connection_store.store.save(Connection(name="prod", master="yarn", deploy_mode="cluster"))
|
||||||
|
|
||||||
|
|
||||||
|
def _last_pending_id() -> str:
|
||||||
|
pending = submit.pending_store.list_all()
|
||||||
|
assert pending, "no pending submission was created"
|
||||||
|
return pending[-1].pending_id
|
||||||
|
|
||||||
|
|
||||||
|
# --- prepare_submit_job ---
|
||||||
|
|
||||||
|
def test_prepare_does_not_invoke_spark_submit(monkeypatch):
|
||||||
|
called = []
|
||||||
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd))
|
||||||
|
out = submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
||||||
|
assert called == [] # never called
|
||||||
|
assert out["status"] == "PENDING"
|
||||||
|
assert out["pending_id"].startswith("p_")
|
||||||
|
|
||||||
|
|
||||||
|
def test_prepare_persists_pending_with_snapshot(monkeypatch):
|
||||||
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
||||||
|
connection_store.store.save(
|
||||||
|
Connection(
|
||||||
|
name="prod",
|
||||||
|
master="yarn",
|
||||||
|
deploy_mode="cluster",
|
||||||
|
spark_conf={"spark.sql.shuffle.partitions": "200"},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py", queue="research")
|
||||||
|
p = submit.pending_store.get(_last_pending_id())
|
||||||
|
assert p is not None
|
||||||
|
assert p.connection == "prod"
|
||||||
|
assert p.master == "yarn"
|
||||||
|
assert p.deploy_mode == "cluster"
|
||||||
|
assert p.spark_conf == {"spark.sql.shuffle.partitions": "200"}
|
||||||
|
assert p.queue == "research"
|
||||||
|
assert p.script_path == "/tmp/j.py"
|
||||||
|
|
||||||
|
|
||||||
|
def test_prepare_raises_for_unknown_connection():
|
||||||
|
with pytest.raises(KeyError, match="missing"):
|
||||||
|
submit.prepare_submit_job(connection="missing", script_path="/tmp/j.py")
|
||||||
|
|
||||||
|
|
||||||
|
def test_prepare_snapshots_connection_at_prepare_time(monkeypatch):
|
||||||
|
"""Editing the connection between prepare and confirm must NOT silently retarget."""
|
||||||
|
monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
|
||||||
|
submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
|
||||||
|
p = submit.pending_store.get(_last_pending_id())
|
||||||
|
# Now mutate the connection
|
||||||
|
connection_store.store.save(
|
||||||
|
Connection(name="prod", master="spark://attacker:7077", deploy_mode="client")
|
||||||
|
)
|
||||||
|
# Snapshot is unchanged
|
||||||
|
p2 = submit.pending_store.get(p.pending_id)
|
||||||
|
assert p2.master == "yarn"
|
||||||
|
assert p2.deploy_mode == "cluster"
|
||||||
Reference in New Issue
Block a user