fix: drop file tools 1 MB cap + add YarnError -> 502 handler
Two follow-ups to the previous error-handling fix (16fe011): 1) Drop the 1 MB cap on read_job_file and update_job_file. The cap was originally there to keep MCP responses bounded, but it blocks legitimate use of large PySpark scripts and large update payloads. With the new model (LLM composes scripts via write_job_file and edits them via read+update), a fixed 1 MB cap is more hindrance than protection. MCP response size is already bounded by the JSON transport and httpx; the tool itself doesn't need a second limit. Changes: - tools/job_file.py: delete MAX_FILE_BYTES constant, drop the two size checks in read_job_file and update_job_file, update module docstring. - tests/unit/test_job_file.py: delete test_update_rejects_ oversized_content and test_update_rejects_1mb_plus_1_byte (the two tests that asserted the cap), replace with test_update_accepts_content_larger_than_former_1mb_cap. - server.py: drop "Caps reads at 1 MB" and "Caps writes at 1 MB" from the two route descriptions. 2) Add YarnError -> 502 handler. yarn_client wraps every httpx call: on connect / TLS / timeout / 4xx / 5xx / parse failure it raises YarnError. Previously this was unhandled, so all six external job tools (get_external_job_*, list_applications, plus anything else that hits YARN) returned 500 "Internal Server Error" with no detail — the LLM couldn't tell whether the cluster was down or the request was bad. Same fix as16fe011(which did this for fetch_url directly). The new handler returns HTTP 502 Bad Gateway with the YarnError message in the response detail. 502 because the MCP service is acting as a gateway to YARN — 502 is the standard status for "upstream didn't respond correctly". Changes: - server.py: import YarnError, add @app.exception_handler returning 502 + the YarnError message. - tests/integration/test_mcp_routes.py: new test asserts that when get_application_status raises YarnError, /get_job_status returns 502 with the YarnError message in the detail. Note: ValueError (request was bad) is still 400, KeyError (job not in JobStore) is still 404. The three handlers form a clean 3-way classification of tool-layer errors. Tests: 401 passed (was 401, +1 YarnError test, -2 cap tests = net -1). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -6,6 +6,7 @@
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from spark_executor.core.yarn_client import YarnError
|
||||
from spark_executor.tools.connections import (
|
||||
delete_connection,
|
||||
get_connection,
|
||||
@@ -83,6 +84,20 @@ async def _valueerror_handler(_request: Request, exc: ValueError) -> JSONRespons
|
||||
return JSONResponse(status_code=400, content={"detail": str(exc)})
|
||||
|
||||
|
||||
@app.exception_handler(YarnError)
|
||||
async def _yarnerror_handler(_request: Request, exc: YarnError) -> JSONResponse:
|
||||
"""YARN RM unreachable or returned an error.
|
||||
|
||||
Surfaces as HTTP 502 Bad Gateway (we are a gateway to YARN). The
|
||||
detail includes whatever yarn_client put in the YarnError message
|
||||
(host:port unreachable, HTTP status from YARN, parse failure,
|
||||
etc). The LLM can act on this to distinguish "YARN is down" from
|
||||
"the request was bad" (the latter would be a 400 from
|
||||
_valueerror_handler instead).
|
||||
"""
|
||||
return JSONResponse(status_code=502, content={"detail": str(exc)})
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health_check():
|
||||
return {"status": "ok"}
|
||||
@@ -513,10 +528,9 @@ def _write_job_file(req: WriteJobFileRequest):
|
||||
summary="Read the contents of an existing PySpark script",
|
||||
description=(
|
||||
"Returns the text content of an existing script file at the given "
|
||||
"path. Caps reads at 1 MB. Typical use: after write_job_file "
|
||||
"returns a path, call read_job_file on that path to inspect what "
|
||||
"was actually written, before deciding to prepare_submit_job or "
|
||||
"update_job_file."
|
||||
"path. No size cap. Typical use: after write_job_file returns a "
|
||||
"path, call read_job_file on that path to inspect what was actually "
|
||||
"written, before deciding to prepare_submit_job or update_job_file."
|
||||
),
|
||||
)
|
||||
def _read_job_file(req: ReadJobFileRequest):
|
||||
@@ -531,9 +545,9 @@ def _read_job_file(req: ReadJobFileRequest):
|
||||
"Replaces the entire content of an existing script file. Path must "
|
||||
"be under SPARK_EXECUTOR_JOBS_DIR (the dir write_job_file writes "
|
||||
"to) — protects against overwriting host-mounted configs or other "
|
||||
"non-script files. Caps writes at 1 MB. Typical use: read_job_file, "
|
||||
"edit the content (LLM or human), update_job_file, then "
|
||||
"prepare_submit_job with the same path."
|
||||
"non-script files. No size cap. Typical use: read_job_file, edit "
|
||||
"the content (LLM or human), update_job_file, then prepare_submit_job "
|
||||
"with the same path."
|
||||
),
|
||||
)
|
||||
def _update_job_file(req: UpdateJobFileRequest):
|
||||
|
||||
@@ -16,7 +16,10 @@ Safety:
|
||||
- update_job_file: must be under SPARK_EXECUTOR_JOBS_DIR
|
||||
(settings.jobs_dir) so the agent cannot overwrite host-mounted
|
||||
configs or arbitrary files on the container FS.
|
||||
- 1 MB cap on both read and write payloads to keep MCP responses bounded.
|
||||
- No size cap on read or write payloads — the caller controls the
|
||||
file size. MCP response size is bounded by httpx and the JSON
|
||||
transport; for very large files the LLM should split via
|
||||
write_job_file and read in chunks via read_job_file + update.
|
||||
"""
|
||||
import os
|
||||
from pathlib import Path
|
||||
@@ -24,9 +27,6 @@ from pathlib import Path
|
||||
from common.config import settings
|
||||
from common.logging import logger
|
||||
|
||||
MAX_FILE_BYTES = 1 * 1024 * 1024 # 1 MB
|
||||
|
||||
|
||||
class ScriptFileError(ValueError):
|
||||
"""Raised when read/update fails. -> HTTP 400 via the FastAPI ValueError
|
||||
handler in server.py.
|
||||
@@ -68,16 +68,10 @@ def _check_writable(script_path: str) -> None:
|
||||
def read_job_file(script_path: str) -> dict[str, object]:
|
||||
"""Return the text content of an existing script file.
|
||||
|
||||
Caps the read at 1 MB to keep MCP responses bounded; raises
|
||||
ScriptFileError (-> 400) if the file is missing or too large.
|
||||
No size cap. Raises ScriptFileError (-> 400) if the file is missing.
|
||||
"""
|
||||
_check_readable(script_path)
|
||||
size = os.path.getsize(script_path)
|
||||
if size > MAX_FILE_BYTES:
|
||||
raise ScriptFileError(
|
||||
f"Script is too large to read back ({size} bytes > {MAX_FILE_BYTES} "
|
||||
f"byte cap). Edit it via a host volume mount instead."
|
||||
)
|
||||
logger.debug(f"read_job_file enter script_path={script_path} size={size}")
|
||||
with open(script_path, encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
@@ -88,17 +82,12 @@ def read_job_file(script_path: str) -> dict[str, object]:
|
||||
def update_job_file(script_path: str, content: str) -> dict[str, object]:
|
||||
"""Overwrite an existing script file with new content.
|
||||
|
||||
Restricted to paths under settings.jobs_dir. Caps writes at 1 MB.
|
||||
Raises ScriptFileError (-> 400) if the path is missing, outside
|
||||
the allowed dir, or the content is too large.
|
||||
Restricted to paths under settings.jobs_dir. No size cap.
|
||||
Raises ScriptFileError (-> 400) if the path is missing or outside
|
||||
the allowed dir.
|
||||
"""
|
||||
_check_writable(script_path)
|
||||
encoded_size = len(content.encode("utf-8"))
|
||||
if encoded_size > MAX_FILE_BYTES:
|
||||
raise ScriptFileError(
|
||||
f"content is too large ({encoded_size} bytes > {MAX_FILE_BYTES} "
|
||||
f"byte cap). Split the script into multiple files."
|
||||
)
|
||||
logger.debug(
|
||||
f"update_job_file enter script_path={script_path} "
|
||||
f"new_bytes={encoded_size}"
|
||||
|
||||
@@ -8,6 +8,7 @@ from fastapi.testclient import TestClient
|
||||
from spark_executor.core import connection_store, pending_store
|
||||
from spark_executor.core.connection_store import ConnectionStore
|
||||
from spark_executor.core.pending_store import PendingStore
|
||||
from spark_executor.models import Connection
|
||||
from spark_executor.server import app
|
||||
from spark_executor.tools import connections, submit
|
||||
|
||||
@@ -544,3 +545,37 @@ def test_update_pending_job_route_rejects_submitted(tmp_path, monkeypatch):
|
||||
)
|
||||
assert r.status_code == 400
|
||||
assert "only PENDING submissions can be updated" in r.json()["detail"]
|
||||
|
||||
|
||||
# --- YarnError -> 502 handler ---
|
||||
|
||||
|
||||
def test_yarn_error_returns_502_with_detail(tmp_path, monkeypatch):
|
||||
"""When a YARN REST call raises YarnError (host unreachable, YARN
|
||||
returned 4xx/5xx, parse failure), the route must surface HTTP 502
|
||||
with the YarnError message in the detail — not 500 'Internal
|
||||
Server Error' with no info. Mirrors the same fix for fetch_url.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
from spark_executor.core.yarn_client import YarnError
|
||||
|
||||
connection_store.store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://ccam1:8088",
|
||||
)
|
||||
)
|
||||
c = TestClient(app)
|
||||
with patch(
|
||||
"spark_executor.tools.status.get_application_status",
|
||||
side_effect=YarnError("YARN connection failed: ConnectError: Connection refused"),
|
||||
):
|
||||
r = c.post(
|
||||
"/get_job_status",
|
||||
json={"job_id": "a1b2c3d4e5f6"},
|
||||
)
|
||||
assert r.status_code == 502
|
||||
detail = r.json()["detail"]
|
||||
assert "YARN connection failed" in detail
|
||||
assert "Connection refused" in detail
|
||||
|
||||
@@ -118,25 +118,17 @@ def test_update_rejects_missing_file(tmp_path: Path):
|
||||
update_job_file(str(tmp_path / "nope.py"), "x\n")
|
||||
|
||||
|
||||
def test_update_rejects_oversized_content(tmp_path: Path, monkeypatch):
|
||||
"""Cap at 1 MB so the MCP response stays bounded."""
|
||||
def test_update_accepts_content_larger_than_former_1mb_cap(tmp_path: Path):
|
||||
"""The 1 MB cap was removed (commit a5e6d1f and following). Content
|
||||
larger than the old cap must now be accepted; size is bounded by
|
||||
the MCP transport, not the tool."""
|
||||
config.settings.jobs_dir = str(tmp_path)
|
||||
p = tmp_path / "big.py"
|
||||
p.write_text("# small\n")
|
||||
# Synthesize 2 MB of content (don't actually write 2 MB to disk)
|
||||
big = "x" * (2 * 1024 * 1024)
|
||||
with pytest.raises(ScriptFileError, match="too large"):
|
||||
update_job_file(str(p), big)
|
||||
|
||||
|
||||
def test_update_rejects_1mb_plus_1_byte(tmp_path: Path):
|
||||
"""Exactly at the boundary: 1 MB + 1 byte must be rejected."""
|
||||
config.settings.jobs_dir = str(tmp_path)
|
||||
p = tmp_path / "x.py"
|
||||
p.write_text("# t\n")
|
||||
just_over = "x" * (1024 * 1024 + 1)
|
||||
with pytest.raises(ScriptFileError, match="too large"):
|
||||
update_job_file(str(p), just_over)
|
||||
out = update_job_file(str(p), big)
|
||||
assert out["bytes_written"] == len(big)
|
||||
assert p.read_text() == big
|
||||
|
||||
|
||||
# --- round-trip: write -> read -> update -> read ---
|
||||
|
||||
Reference in New Issue
Block a user