feat: add confirm_submit_job (two-step submit flow)
This commit is contained in:
@@ -8,9 +8,17 @@ from datetime import datetime
|
||||
|
||||
from common.logging import logger
|
||||
from spark_executor.core.connection_store import store as conn_store
|
||||
from spark_executor.core.job_store import JobStore
|
||||
from spark_executor.core.log_parser import parse_spark_submit_output
|
||||
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
|
||||
from spark_executor.core.spark_submit import (
|
||||
SparkSubmitError,
|
||||
build_spark_submit_command,
|
||||
run_spark_submit,
|
||||
)
|
||||
from spark_executor.models import Job, PendingSubmission, SubmitResult
|
||||
import uuid
|
||||
from datetime import datetime as _datetime
|
||||
|
||||
|
||||
def _new_pending_id() -> str:
|
||||
@@ -55,3 +63,64 @@ def prepare_submit_job(
|
||||
"status": "PENDING",
|
||||
"parameters": pending.model_dump(),
|
||||
}
|
||||
|
||||
|
||||
# Module-level job store singleton; replaced in tests.
|
||||
job_store: JobStore = JobStore()
|
||||
|
||||
|
||||
def confirm_submit_job(*, pending_id: str) -> SubmitResult:
|
||||
"""Actually invoke spark-submit for a previously-prepared PendingSubmission."""
|
||||
pending = pending_store.get(pending_id)
|
||||
if pending is None:
|
||||
raise KeyError(f"Unknown pending_id: {pending_id}")
|
||||
if pending.status != "PENDING":
|
||||
raise ValueError(
|
||||
f"pending_id {pending_id} is in status {pending.status!r}, not PENDING"
|
||||
)
|
||||
|
||||
cmd = build_spark_submit_command(
|
||||
master=pending.master,
|
||||
deploy_mode=pending.deploy_mode,
|
||||
script_path=pending.script_path,
|
||||
queue=pending.queue,
|
||||
executor_memory=pending.executor_memory,
|
||||
executor_cores=pending.executor_cores,
|
||||
num_executors=pending.num_executors,
|
||||
spark_conf=pending.spark_conf,
|
||||
)
|
||||
logger.info(f"confirm_submit_job pending_id={pending_id} cmd={cmd}")
|
||||
try:
|
||||
result = run_spark_submit(cmd)
|
||||
except SparkSubmitError as exc:
|
||||
pending.status = "FAILED"
|
||||
pending.error = str(exc)
|
||||
pending_store.save(pending)
|
||||
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,
|
||||
)
|
||||
)
|
||||
|
||||
pending.status = "SUBMITTED"
|
||||
pending.job_id = job_id
|
||||
pending.application_id = application_id
|
||||
pending_store.save(pending)
|
||||
logger.info(
|
||||
f"confirm_submit_job pending_id={pending_id} job_id={job_id} application_id={application_id}"
|
||||
)
|
||||
return SubmitResult(
|
||||
job_id=job_id,
|
||||
application_id=application_id,
|
||||
tracking_url=tracking_url,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user