Files
mcp-server/spark_executor/tools/submit.py
T

58 lines
1.6 KiB
Python

# 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(),
}