Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
307da276bf | ||
|
|
a460b750a8 | ||
|
|
6109c6d11a | ||
|
|
3b698e1bdc | ||
|
|
a5e6d1ffe1 | ||
|
|
16fe011fa0 | ||
|
|
59cab3346e | ||
|
|
8ededc50a9 | ||
|
|
f43e5b6403 | ||
|
|
585a1f6e43 |
@@ -20,3 +20,6 @@ data/
|
||||
.pytest_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
|
||||
# macOS
|
||||
.DS_Store
|
||||
|
||||
@@ -40,7 +40,7 @@
|
||||
| MCP 暴露 | fastapi-mcp |
|
||||
| 数据模型 | Pydantic v2 |
|
||||
| HTTP 客户端 | httpx、httpx-kerberos |
|
||||
| Spark 集成 | `spark-submit`、YARN REST API、`yarn logs` |
|
||||
| Spark 集成 | `spark-submit`、YARN REST API (httpx 直连) |
|
||||
| 日志 | loguru |
|
||||
| 进程管理 | uvicorn、gunicorn |
|
||||
| 包管理 | uv |
|
||||
@@ -221,6 +221,30 @@ get_job_logs
|
||||
kill_job
|
||||
```
|
||||
|
||||
## RM / SHS API 端点
|
||||
|
||||
本服务**不再 shell-out `yarn` CLI**,所有集群交互都通过 `httpx` 直接打 REST(由 `spark_executor/core/yarn_client.py` 统一封装)。鉴权走 `Connection.auth_type`(`none` / `simple` / `basic` / `kerberos` + SPNEGO),SSL 走 `Connection.ssl_verify` / `ssl_ca_bundle`。
|
||||
|
||||
### YARN ResourceManager
|
||||
|
||||
| 方法 | 端点 | 用途 |
|
||||
| --- | --- | --- |
|
||||
| `GET` | `/ws/v1/cluster/apps?state=&queue=&limit=` | 列出 YARN 应用(支持 `state` / `queue` / `limit` 过滤),由 `list_applications` 触发 |
|
||||
| `GET` | `/ws/v1/cluster/apps/{appid}` | 取单个 app 的 state + amContainerLogs;日志降级路径会复用 |
|
||||
| `GET` | `/ws/v1/cluster/apps/{appid}/aggregated-logs` | 拉聚合后的 container 日志,404/501 时自动降级到 amContainerLogs 路径 |
|
||||
| `PUT` | `/ws/v1/cluster/apps/{appid}/state`,body `{"state":"KILLED"}` | 杀应用,由 `kill_job` 触发 |
|
||||
|
||||
> **日志降级路径**会**手动**跟随 3xx Location 跳到 NodeManager(最多 3 次)抓 AM driver 的 `/stdout`。原因:httpx 默认跨主机重定向会丢掉 `Authorization`,RM→NM 的 307 在很多集群会因此 401/403。driver 日志可以这样取,**executor 日志仍然依赖 `yarn.log-aggregation-enable=true`**,这个不走降级。
|
||||
|
||||
### Spark History Server
|
||||
|
||||
**当前没有任何内置调用。** `Connection.history_server_url` 字段只是原样存储,留作未来 SHS 专用工具消费(`models.py` 的 docstring 明确写明)。
|
||||
|
||||
要走 SHS API(典型路径 `/api/v1/applications/{id}/jobs`、`/api/v1/applications/{id}/stages` 等),**唯一**方式是用 `fetch_url` 工具:
|
||||
|
||||
1. 把 SHS 主机 glob 加到对应 `Connection.url_allowlist`(单标签主机如 `ccam*`,或多级域名如 `*.hadoop.internal`)。
|
||||
2. 调用 `fetch_url` 时传入完整 SHS URL,鉴权和 SSL 复用同一条 Connection 的配置。
|
||||
|
||||
## 配置项
|
||||
|
||||
### 应用配置
|
||||
|
||||
@@ -96,6 +96,8 @@ class Connection(BaseModel):
|
||||
"List of fnmatch glob patterns for hosts the fetch_url tool may access. "
|
||||
"The list is mandatory-opt-in: an empty list (the default) denies all "
|
||||
"hosts, so you must populate it before fetch_url can access any URL. "
|
||||
"Set or change via save_connection (pass url_allowlist on create) or "
|
||||
"update_connection (PATCH the field on an existing connection). "
|
||||
"Useful for clusters whose hostnames do NOT share a common suffix — "
|
||||
"e.g. single-label hosts like 'ccam1'-'ccam99' (configure ['ccam*']) "
|
||||
"or HDFS namenode on a different subdomain ('*.hadoop.internal'). "
|
||||
@@ -104,6 +106,18 @@ class Connection(BaseModel):
|
||||
"'ccam50' but not 'ccam50.evil.com'."
|
||||
),
|
||||
)
|
||||
history_server_url: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional Spark History Server (SHS) base URL, e.g. "
|
||||
"'http://history.prod.internal:18080'. Stored as part of the "
|
||||
"connection for reference and to be consumed by future SHS-"
|
||||
"specific tools. Currently **not consumed by any tool** — to "
|
||||
"fetch SHS endpoints today, use fetch_url with the SHS host "
|
||||
"added to url_allowlist. The SHS host and the YARN RM host "
|
||||
"are usually different, so this is independent of yarn_rm_url."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("master")
|
||||
@classmethod
|
||||
|
||||
+41
-12
@@ -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"}
|
||||
@@ -426,10 +441,17 @@ def _list_connections(_req: EmptyRequest = EmptyRequest()):
|
||||
operation_id="get_connection",
|
||||
summary="Get a single connection by name",
|
||||
description=(
|
||||
"Return the Connection record, or 404 if not found. Useful to "
|
||||
"verify a connection was saved correctly, or to inspect the "
|
||||
"resolved yarn_rm_url / auth_type before submitting a job or "
|
||||
"calling one of the get_external_* tools."
|
||||
"Return the full Connection record, or 404 if not found. **Call this "
|
||||
"whenever you need any cluster-level config** — common lookups: "
|
||||
"yarn_rm_url (YARN RM endpoint), history_server_url (Spark History "
|
||||
"Server), url_allowlist (which hosts fetch_url may access), "
|
||||
"auth_type / auth_user / ssl_verify (for any YARN REST or HTTP call), "
|
||||
"master / deploy_mode / spark_conf (for prepare_submit_job).\n\n"
|
||||
"The response is a full Pydantic model dump — all fields including "
|
||||
"secrets (auth_password, auth_keytab). Treat it as sensitive. "
|
||||
"Useful to verify a connection was saved correctly, or to discover "
|
||||
"the right endpoint to call before invoking get_external_*, "
|
||||
"fetch_url, or list_applications."
|
||||
),
|
||||
)
|
||||
def _get_connection(req: ConnectionNameRequest):
|
||||
@@ -506,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):
|
||||
@@ -524,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):
|
||||
@@ -549,7 +570,15 @@ def _update_job_file(req: UpdateJobFileRequest):
|
||||
"guardrails — the allowlist is the only gate — so keep it tight. The "
|
||||
"Connection's saved auth is reused, so the agent does not need cluster "
|
||||
"credentials.\n\n"
|
||||
"30s timeout, redirects followed."
|
||||
"**Limits:** 30s timeout, redirects followed.\n\n"
|
||||
"**Errors:** when the host is unreachable (connect refused, DNS "
|
||||
"failure, TLS handshake error, timeout, etc.) this tool returns "
|
||||
"HTTP 400 with the exception class and message in the response "
|
||||
"detail — e.g. `fetch_url could not reach ...: ConnectError: "
|
||||
"Connection refused`. Upstream HTTP 4xx/5xx responses that DID "
|
||||
"come back are returned in the result body with their status code "
|
||||
"preserved (not translated to an error) so you can see what the "
|
||||
"server actually said."
|
||||
),
|
||||
)
|
||||
def _fetch_url(req: FetchUrlRequest):
|
||||
|
||||
@@ -26,6 +26,7 @@ def save_connection(
|
||||
auth_principal: str | None = _UNSET, # type: ignore[assignment]
|
||||
auth_keytab: str | None = _UNSET, # type: ignore[assignment]
|
||||
url_allowlist: list[str] | None = _UNSET, # type: ignore[assignment]
|
||||
history_server_url: str | None = _UNSET, # type: ignore[assignment]
|
||||
) -> dict[str, object]:
|
||||
logger.debug(
|
||||
f"save_connection enter name={name} master={master} "
|
||||
@@ -56,6 +57,8 @@ def save_connection(
|
||||
fields["auth_keytab"] = auth_keytab
|
||||
if url_allowlist is not _UNSET:
|
||||
fields["url_allowlist"] = url_allowlist or []
|
||||
if history_server_url is not _UNSET:
|
||||
fields["history_server_url"] = history_server_url
|
||||
return update_connection(name, **fields)
|
||||
new_fields: dict[str, object] = {
|
||||
"name": name,
|
||||
@@ -71,6 +74,7 @@ def save_connection(
|
||||
"auth_principal": auth_principal if auth_principal is not _UNSET else None,
|
||||
"auth_keytab": auth_keytab if auth_keytab is not _UNSET else None,
|
||||
"url_allowlist": (url_allowlist if url_allowlist is not _UNSET else None) or [],
|
||||
"history_server_url": history_server_url if history_server_url is not _UNSET else None,
|
||||
}
|
||||
conn = Connection(**new_fields)
|
||||
store.save(conn)
|
||||
@@ -86,7 +90,7 @@ def update_connection(name: str, **fields) -> dict[str, object]:
|
||||
|
||||
Mutable fields: master, deploy_mode, yarn_rm_url, spark_conf,
|
||||
ssl_verify, ssl_ca_bundle, auth_type, auth_user, auth_password,
|
||||
auth_principal, auth_keytab, url_allowlist.
|
||||
auth_principal, auth_keytab, url_allowlist, history_server_url.
|
||||
"""
|
||||
logger.debug(f"update_connection enter name={name} fields={sorted(fields.keys())}")
|
||||
return store.update(name, **fields).model_dump()
|
||||
|
||||
@@ -83,6 +83,7 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult:
|
||||
verify = config.verify_for_httpx()
|
||||
|
||||
for hop in range(_MAX_REDIRECTS + 1):
|
||||
try:
|
||||
resp = httpx.get(
|
||||
url,
|
||||
auth=auth,
|
||||
@@ -90,6 +91,19 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult:
|
||||
timeout=_REQUEST_TIMEOUT_SECONDS,
|
||||
follow_redirects=False,
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
# Network-level failure (no response received). Surface the
|
||||
# exception class + message so the LLM can act on it.
|
||||
# ValueError -> 400 via the existing handler in server.py.
|
||||
# The cause chain (`from exc`) preserves the original
|
||||
# exception for loguru.
|
||||
raise ValueError(
|
||||
f"fetch_url could not reach {url!r}: "
|
||||
f"{type(exc).__name__}: {exc}. "
|
||||
f"Check that the URL is reachable from the MCP service, "
|
||||
f"the host is in Connection.url_allowlist, and the "
|
||||
f"connection's auth/SSL settings are correct."
|
||||
) from exc
|
||||
if resp.status_code not in _REDIRECT_STATUSES:
|
||||
logger.info(
|
||||
f"fetch_url ok url={url} connection_name={connection_name} "
|
||||
|
||||
@@ -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}"
|
||||
|
||||
@@ -134,6 +134,14 @@ class SaveConnectionRequest(BaseModel):
|
||||
"populated via update_connection."
|
||||
),
|
||||
)
|
||||
history_server_url: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional Spark History Server (SHS) base URL, e.g. "
|
||||
"'http://history.prod.internal:18080'. Stored for reference. "
|
||||
"Currently not consumed by any tool — see Connection.history_server_url."
|
||||
),
|
||||
)
|
||||
|
||||
class PrepareSubmitJobRequest(BaseModel):
|
||||
connection: str = Field(
|
||||
@@ -408,8 +416,10 @@ class FetchUrlRequest(BaseModel):
|
||||
...,
|
||||
description=(
|
||||
"Name of a saved Connection (see list_connections). The "
|
||||
"Connection's yarn_rm_url defines the allowed host domain. "
|
||||
"The Connection's auth_type / auth_user / auth_password / "
|
||||
"Connection's url_allowlist is the host gate for this tool — "
|
||||
"**not yarn_rm_url** (yarn_rm_url is only used by the "
|
||||
"get_external_* tools to locate the YARN RM). The "
|
||||
"Connection's auth_type / auth_user / auth_password / "
|
||||
"ssl_verify / ssl_ca_bundle are reused for the request."
|
||||
),
|
||||
)
|
||||
@@ -475,6 +485,15 @@ class UpdateConnectionRequest(BaseModel):
|
||||
"Example: ['ccam*'] allows ccam1-ccam99."
|
||||
),
|
||||
)
|
||||
history_server_url: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"New Spark History Server base URL. Omit to keep current. "
|
||||
"Pass None to clear (use delete_connection + save_connection "
|
||||
"if you need explicit clear semantics; same caveat as other "
|
||||
"Optional fields in this tool)."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ListApplicationsRequest(BaseModel):
|
||||
|
||||
@@ -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,147 @@ 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
|
||||
|
||||
|
||||
def test_external_job_status_returns_502_on_yarn_error(tmp_path):
|
||||
"""get_external_job_status uses get_application_status under the
|
||||
hood — same handler, same 502, same detail."""
|
||||
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.external_jobs.get_application_status",
|
||||
side_effect=YarnError("YARN application 'application_xxx' not found"),
|
||||
):
|
||||
r = c.post(
|
||||
"/get_external_job_status",
|
||||
json={"application_id": "application_1740000000001_0001", "connection_name": "prod"},
|
||||
)
|
||||
assert r.status_code == 502
|
||||
assert "not found" in r.json()["detail"]
|
||||
|
||||
|
||||
def test_external_job_result_returns_502_on_yarn_error(tmp_path):
|
||||
"""get_external_job_result also uses get_application_status (same
|
||||
underlying YARN endpoint), and exercises the same YarnError path."""
|
||||
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.external_jobs.get_application_status",
|
||||
side_effect=YarnError("YARN connection failed: ConnectError: Connection refused"),
|
||||
):
|
||||
r = c.post(
|
||||
"/get_external_job_result",
|
||||
json={"application_id": "application_1740000000001_0001", "connection_name": "prod"},
|
||||
)
|
||||
assert r.status_code == 502
|
||||
assert "YARN connection failed" in r.json()["detail"]
|
||||
|
||||
|
||||
def test_external_job_logs_returns_502_on_yarn_error(tmp_path):
|
||||
"""get_external_job_logs uses get_application_logs — different
|
||||
function, but raises YarnError the same way. Handler must catch
|
||||
it the same way."""
|
||||
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.external_jobs.get_application_logs",
|
||||
side_effect=YarnError("YARN GET logs returned HTTP 500: cluster overloaded"),
|
||||
):
|
||||
r = c.post(
|
||||
"/get_external_job_logs",
|
||||
json={
|
||||
"application_id": "application_1740000000001_0001",
|
||||
"connection_name": "prod",
|
||||
"tail_chars": 5000,
|
||||
},
|
||||
)
|
||||
assert r.status_code == 502
|
||||
assert "YARN GET logs returned HTTP 500" in r.json()["detail"]
|
||||
assert "cluster overloaded" in r.json()["detail"]
|
||||
|
||||
|
||||
def test_list_applications_returns_502_on_yarn_error(tmp_path):
|
||||
"""list_applications uses list_applications_yarn — yet another
|
||||
YARN-touching tool. Same YarnError -> 502 contract."""
|
||||
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.external_jobs.list_applications_yarn",
|
||||
side_effect=YarnError("YARN list applications failed: 503 Service Unavailable"),
|
||||
):
|
||||
r = c.post(
|
||||
"/list_applications",
|
||||
json={"connection_name": "prod"},
|
||||
)
|
||||
assert r.status_code == 502
|
||||
assert "YARN list applications failed" in r.json()["detail"]
|
||||
|
||||
@@ -195,3 +195,39 @@ def test_update_connection_persists_to_disk(_fresh_store, tmp_path):
|
||||
connections.update_connection(name="prod", master="spark://new:7077")
|
||||
raw = json.loads((tmp_path / "connections.json").read_text())
|
||||
assert raw["prod"]["master"] == "spark://new:7077"
|
||||
|
||||
|
||||
# --- history_server_url ---
|
||||
|
||||
def test_save_connection_with_history_server_url(_fresh_store):
|
||||
connections.save_connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
history_server_url="http://history.prod.internal:18080",
|
||||
)
|
||||
assert _fresh_store.get("prod").history_server_url == "http://history.prod.internal:18080"
|
||||
|
||||
|
||||
def test_save_connection_history_server_url_defaults_to_none(_fresh_store):
|
||||
connections.save_connection(name="prod", master="yarn")
|
||||
assert _fresh_store.get("prod").history_server_url is None
|
||||
|
||||
|
||||
def test_update_connection_changes_history_server_url(_fresh_store):
|
||||
connections.save_connection(name="prod", master="yarn")
|
||||
out = connections.update_connection(
|
||||
name="prod",
|
||||
history_server_url="http://history2.prod.internal:18080",
|
||||
)
|
||||
assert out["history_server_url"] == "http://history2.prod.internal:18080"
|
||||
assert _fresh_store.get("prod").history_server_url == "http://history2.prod.internal:18080"
|
||||
|
||||
|
||||
def test_update_connection_keeps_history_server_url_when_omitted(_fresh_store):
|
||||
connections.save_connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
history_server_url="http://history.prod.internal:18080",
|
||||
)
|
||||
out = connections.update_connection(name="prod", master="spark://new:7077")
|
||||
assert out["history_server_url"] == "http://history.prod.internal:18080"
|
||||
|
||||
@@ -327,3 +327,76 @@ def test_list_applications_request_limit_bounds():
|
||||
ListApplicationsRequest(connection_name="prod", limit=0)
|
||||
with pytest.raises(ValidationError):
|
||||
ListApplicationsRequest(connection_name="prod", limit=10001)
|
||||
|
||||
|
||||
# --- Network error handling (httpx.HTTPError -> 400 with detail) ---
|
||||
|
||||
|
||||
def test_fetch_url_raises_400_with_detail_on_connect_error(fresh_stores):
|
||||
"""When httpx.get raises ConnectError (host unreachable / port closed),
|
||||
fetch_url must raise ValueError (-> 400) with the exception class
|
||||
and message in the detail. Previously this bubbled up as a bare
|
||||
500 Internal Server Error with no info."""
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
url_allowlist=["*.prod.internal"],
|
||||
)
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get",
|
||||
side_effect=httpx.ConnectError("Connection refused"),
|
||||
):
|
||||
with pytest.raises(ValueError) as ei:
|
||||
fetch_url.fetch_url("http://nm01.prod.internal:8042/", "prod")
|
||||
msg = str(ei.value)
|
||||
assert "ConnectError" in msg
|
||||
assert "Connection refused" in msg
|
||||
assert "nm01.prod.internal" in msg
|
||||
# Cause chain preserved for loguru
|
||||
assert isinstance(ei.value.__cause__, httpx.ConnectError)
|
||||
|
||||
|
||||
def test_fetch_url_raises_400_with_detail_on_timeout(fresh_stores):
|
||||
"""When httpx.get raises TimeoutException (request exceeded 30s),
|
||||
fetch_url must raise ValueError with the exception details surfaced."""
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
url_allowlist=["*.prod.internal"],
|
||||
)
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.fetch_url.httpx.get",
|
||||
side_effect=httpx.ReadTimeout("Timed out reading"),
|
||||
):
|
||||
with pytest.raises(ValueError) as ei:
|
||||
fetch_url.fetch_url("http://nm01.prod.internal:8042/", "prod")
|
||||
msg = str(ei.value)
|
||||
assert "ReadTimeout" in msg
|
||||
assert "Timed out reading" in msg
|
||||
assert isinstance(ei.value.__cause__, httpx.ReadTimeout)
|
||||
|
||||
|
||||
def test_fetch_url_returns_body_for_4xx_5xx_upstream(fresh_stores):
|
||||
"""Upstream HTTP errors (4xx/5xx responses that DID come back) are
|
||||
NOT translated to ValueError — the FetchUrlResult carries the status
|
||||
code and body so the LLM can see what the server actually said. This
|
||||
is the intentional contrast with the no-response-at-all case."""
|
||||
fetch_url.conn_store.save(
|
||||
Connection(
|
||||
name="prod",
|
||||
master="yarn",
|
||||
yarn_rm_url="http://rm.prod.internal:8088",
|
||||
url_allowlist=["*.prod.internal"],
|
||||
)
|
||||
)
|
||||
resp = httpx.Response(503, text="Service Unavailable - try again later")
|
||||
with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp):
|
||||
out = fetch_url.fetch_url("http://nm01.prod.internal:8042/", "prod")
|
||||
assert out.status_code == 503
|
||||
assert "Service Unavailable" in out.body
|
||||
|
||||
@@ -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