diff --git a/spark_executor/server.py b/spark_executor/server.py index d9cacc5..6980877 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -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): diff --git a/spark_executor/tools/job_file.py b/spark_executor/tools/job_file.py index 860040d..b93a156 100644 --- a/spark_executor/tools/job_file.py +++ b/spark_executor/tools/job_file.py @@ -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}" diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index df69baf..87d9f2d 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -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 diff --git a/tests/unit/test_job_file.py b/tests/unit/test_job_file.py index 55819f3..4a0bce1 100644 --- a/tests/unit/test_job_file.py +++ b/tests/unit/test_job_file.py @@ -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 ---