diff --git a/main.py b/main.py index c55f764..323f4e6 100644 --- a/main.py +++ b/main.py @@ -1,29 +1,49 @@ # coding=utf-8 """ @Time :2026/6/24 -@Author :tao.chen +@Author :tao.chen """ from contextlib import asynccontextmanager + from fastapi import FastAPI -from spark_executor import app as spark_executor_app + from common.factory import init_mcp_server from common.logging import logger +from spark_executor import app as spark_executor_app + @asynccontextmanager async def lifespan(app: "FastAPI"): logger.info(f"Startup {app.title}") spark_executor_mcp = init_mcp_server(spark_executor_app) spark_executor_mcp.mount_http(app, mount_path="/spark-executor-mcp") - logger.info(f"include {spark_executor_app.title}, version {spark_executor_app.version}") + logger.info( + f"MCP server mounted at /spark-executor-mcp " + f"(source: {spark_executor_app.title} v{spark_executor_app.version})" + ) + + # Enumerate every MCP tool + its description so operators can see at a + # glance what the server exposes (useful for debugging / change review). + tools = spark_executor_mcp.tools + logger.info(f"MCP tools registered ({len(tools)}):") + for i, tool in enumerate(tools, 1): + desc = (tool.description or "").strip() + # fastapi-mcp appends an auto-generated "### Responses: ..." block to + # every tool description; drop it so the startup log stays scannable. + if "### Responses" in desc: + desc = desc.split("### Responses", 1)[0].rstrip() + # Collapse internal newlines / extra spaces into one line. + desc = " ".join(desc.split()) + logger.info(f" {i:>2}. {tool.name} - {desc}") + yield logger.info(f"Shutdown {app.title}") -app = FastAPI(title="Main App", lifespan=lifespan) +app = FastAPI(title="Main App", lifespan=lifespan) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000) - diff --git a/spark_executor/server.py b/spark_executor/server.py index 0b0e2cb..7734204 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -64,66 +64,140 @@ def health_check(): # --- Pending submission flow (two-step submit) --- -@app.post("/prepare_submit_job") +@app.post( + "/prepare_submit_job", + summary="Prepare a Spark job submission (no spark-submit yet)", + description=( + "Snapshot the named Connection's master / deploy_mode / spark_conf / " + "yarn_rm_url into a PendingSubmission record and persist it. " + "Does NOT invoke spark-submit. Returns pending_id for use with " + "confirm_submit_job (the user-second-confirmation step)." + ), +) def _prepare_submit_job(req: PrepareSubmitJobRequest): return prepare_submit_job(**req.model_dump()) -@app.post("/confirm_submit_job") +@app.post( + "/confirm_submit_job", + summary="Confirm and submit a previously-prepared job", + description=( + "Actually invoke spark-submit for the PendingSubmission identified " + "by pending_id. Requires status=PENDING. On success, transitions the " + "pending entry to SUBMITTED and creates a Job record. On failure, " + "marks the entry FAILED and re-raises." + ), +) def _confirm_submit_job(req: PendingIdRequest): return confirm_submit_job(pending_id=req.pending_id) -@app.post("/list_pending_jobs") +@app.post( + "/list_pending_jobs", + summary="List all pending submissions", + description="Return every PendingSubmission in any status (PENDING, SUBMITTED, CANCELLED, FAILED).", +) def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()): return list_pending_jobs() -@app.post("/get_pending_job") +@app.post( + "/get_pending_job", + summary="Get a single pending submission", + description="Return the PendingSubmission identified by pending_id, including its current status and outcome fields.", +) def _get_pending_job(req: PendingIdRequest): return get_pending_job(req.pending_id) -@app.post("/cancel_pending_job") +@app.post( + "/cancel_pending_job", + summary="Cancel a pending submission", + description=( + "Flip a PENDING (or already-CANCELLED) PendingSubmission to CANCELLED. " + "Refuses to cancel entries that are SUBMITTED or FAILED — those are " + "terminal and must be killed via kill_job instead." + ), +) def _cancel_pending_job(req: PendingIdRequest): return cancel_pending_job(req.pending_id) # --- Spark job tools --- -@app.post("/get_job_status") +@app.post( + "/get_job_status", + summary="Query YARN for a job's current status", + description=( + "Return the YARN application state (RUNNING / SUCCEEDED / FAILED / " + "KILLED / ACCEPTED / NEW / NEW_SAVING / SUBMITTED / etc.) plus the " + "raw YARN REST response body." + ), +) def _get_job_status(req: JobIdRequest): return get_job_status(req.job_id) -@app.post("/get_job_logs") +@app.post( + "/get_job_logs", + summary="Fetch aggregated container logs for a job", + description=( + "Pull aggregated logs from the YARN ResourceManager. Returns the last " + "tail_chars characters (default 5000). Requires yarn.log-aggregation-enable " + "to be true on the target cluster." + ), +) def _get_job_logs(req: GetJobLogsRequest): return get_job_logs(req.job_id, tail_chars=req.tail_chars) -@app.post("/kill_job") +@app.post( + "/kill_job", + summary="Kill a running job", + description="PUT state=KILLED to YARN REST API for the job's application_id.", +) def _kill_job(req: JobIdRequest): return kill_job(req.job_id) # --- Connection management tools --- -@app.post("/save_connection") +@app.post( + "/save_connection", + summary="Save or update a named Spark connection", + description=( + "Upsert a Connection record (master URL, deploy mode, optional YARN RM URL, " + "spark_conf K/V) keyed by name. Used by prepare_submit_job via the " + "connection parameter." + ), +) def _save_connection(req: SaveConnectionRequest): # exclude_none so we don't overwrite the function's default with explicit None return save_connection(**req.model_dump(exclude_none=True)) -@app.post("/list_connections") +@app.post( + "/list_connections", + summary="List all saved Spark connections", + description="Return every Connection in the registry (model_dump form).", +) def _list_connections(_req: EmptyRequest = EmptyRequest()): return list_connections() -@app.post("/get_connection") +@app.post( + "/get_connection", + summary="Get a single connection by name", + description="Return the Connection record, or 404 if not found.", +) def _get_connection(req: ConnectionNameRequest): return get_connection(req.name) -@app.post("/delete_connection") +@app.post( + "/delete_connection", + summary="Delete a saved connection", + description="Remove a Connection by name. 404 if not found.", +) def _delete_connection(req: ConnectionNameRequest): return delete_connection(req.name)