116 lines
2.8 KiB
Python
116 lines
2.8 KiB
Python
# coding=utf-8
|
|
"""
|
|
@Time :2026/6/24
|
|
@Author :tao.chen
|
|
"""
|
|
from fastapi import FastAPI
|
|
|
|
from spark_executor.tools.connections import (
|
|
delete_connection,
|
|
get_connection,
|
|
list_connections,
|
|
save_connection,
|
|
)
|
|
from spark_executor.tools.kill import kill_job
|
|
from spark_executor.tools.logs import get_job_logs
|
|
from spark_executor.tools.status import get_job_status
|
|
from spark_executor.tools.submit import (
|
|
cancel_pending_job,
|
|
confirm_submit_job,
|
|
get_pending_job,
|
|
list_pending_jobs,
|
|
prepare_submit_job,
|
|
)
|
|
|
|
app = FastAPI(title="Spark Executor", version="0.0.1", description="Spark Executor MCP Server")
|
|
|
|
|
|
@app.get("/health")
|
|
def health_check():
|
|
return {"status": "ok"}
|
|
|
|
|
|
# MCP tool routes. fastapi-mcp discovers these and registers them as MCP tools.
|
|
|
|
# --- Pending submission flow (two-step submit) ---
|
|
|
|
@app.post("/prepare_submit_job")
|
|
def _prepare_submit_job(connection: str, script_path: str, queue: str = "default",
|
|
executor_memory: str = "4G", executor_cores: int = 2,
|
|
num_executors: int = 2):
|
|
return prepare_submit_job(
|
|
connection=connection,
|
|
script_path=script_path,
|
|
queue=queue,
|
|
executor_memory=executor_memory,
|
|
executor_cores=executor_cores,
|
|
num_executors=num_executors,
|
|
)
|
|
|
|
|
|
@app.post("/confirm_submit_job")
|
|
def _confirm_submit_job(pending_id: str):
|
|
return confirm_submit_job(pending_id=pending_id)
|
|
|
|
|
|
@app.post("/list_pending_jobs")
|
|
def _list_pending_jobs():
|
|
return list_pending_jobs()
|
|
|
|
|
|
@app.post("/get_pending_job")
|
|
def _get_pending_job(pending_id: str):
|
|
return get_pending_job(pending_id)
|
|
|
|
|
|
@app.post("/cancel_pending_job")
|
|
def _cancel_pending_job(pending_id: str):
|
|
return cancel_pending_job(pending_id)
|
|
|
|
|
|
# --- Spark job tools ---
|
|
|
|
@app.post("/get_job_status")
|
|
def _get_job_status(job_id: str):
|
|
return get_job_status(job_id)
|
|
|
|
|
|
@app.post("/get_job_logs")
|
|
def _get_job_logs(job_id: str, tail_chars: int = 5000):
|
|
return get_job_logs(job_id, tail_chars=tail_chars)
|
|
|
|
|
|
@app.post("/kill_job")
|
|
def _kill_job(job_id: str):
|
|
return kill_job(job_id)
|
|
|
|
|
|
# --- Connection management tools ---
|
|
|
|
@app.post("/save_connection")
|
|
def _save_connection(name: str, master: str, deploy_mode: str = "cluster",
|
|
yarn_rm_url: str | None = None,
|
|
spark_conf: dict[str, str] | None = None):
|
|
return save_connection(
|
|
name=name,
|
|
master=master,
|
|
deploy_mode=deploy_mode,
|
|
yarn_rm_url=yarn_rm_url,
|
|
spark_conf=spark_conf,
|
|
)
|
|
|
|
|
|
@app.post("/list_connections")
|
|
def _list_connections():
|
|
return list_connections()
|
|
|
|
|
|
@app.post("/get_connection")
|
|
def _get_connection(name: str):
|
|
return get_connection(name)
|
|
|
|
|
|
@app.post("/delete_connection")
|
|
def _delete_connection(name: str):
|
|
return delete_connection(name)
|