mcp-server/spark_executor/core/yarn_client.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen

YARN ResourceManager REST API client. Replaces the previous `yarn` CLI shell-out
so the runtime image does not need a Hadoop client installation — the
`httpx` library already in pyproject.toml is enough.

Endpoints used (YARN 2.6+):
  GET  /ws/v1/cluster/apps/{appid}             -> app status + state
  GET  /ws/v1/cluster/apps/{appid}/aggregated-logs  -> aggregated container logs
  PUT  /ws/v1/cluster/apps/{appid}/state       -> kill an app (body: {"state":"KILLED"})

The ResourceManager URL is passed in per call (snapshotted on the Job at
confirm_submit_job time) and falls back to the YARN_RESOURCE_MANAGER_URL env
var if unset. This matches the pattern the original Connection.yarn_rm_url
field was designed for, but no longer requires the `yarn` CLI to interpret it.
"""
import json

import httpx

from common.config import settings
from common.logging import logger


class YarnError(Exception):
    """Raised when a YARN REST API call fails."""


class YarnConfigError(YarnError):
    """Raised when the YARN ResourceManager URL is missing or malformed."""


def _base_url(yarn_rm_url: str | None) -> str:
    """Resolve and validate the RM URL. Raises YarnConfigError if unusable."""
    url = yarn_rm_url or settings.yarn_resource_manager_url
    if not url:
        raise YarnConfigError(
            "YARN ResourceManager URL is not configured. "
            "Set Connection.yarn_rm_url when saving the connection, "
            "or set the YARN_RESOURCE_MANAGER_URL environment variable."
        )
    base = url.rstrip("/")
    if not base.startswith(("http://", "https://")):
        raise YarnConfigError(
            f"YARN ResourceManager URL must start with http:// or https://: {url!r}"
        )
    return base


def _request(method: str, url: str, *, json_body: dict | None = None,
             timeout: float = 30.0) -> httpx.Response:
    logger.debug(f"YARN {method} {url}" + (f" body={json_body}" if json_body else ""))
    try:
        resp = httpx.request(method, url, json=json_body, timeout=timeout)
    except httpx.HTTPError as exc:
        logger.error(f"YARN {method} {url} failed: {exc}")
        raise YarnError(f"YARN connection failed: {exc}") from exc
    logger.debug(
        f"YARN {method} {url} -> {resp.status_code} "
        f"({len(resp.content)} bytes)"
    )
    return resp


def get_application_status(application_id: str, yarn_rm_url: str | None) -> tuple[str, str]:
    """Return (state, raw_json_text) for an application, or raise YarnError."""
    url = f"{_base_url(yarn_rm_url)}/ws/v1/cluster/apps/{application_id}"
    resp = _request("GET", url)
    if resp.status_code == 404:
        raise YarnError(f"YARN application {application_id!r} not found")
    if resp.status_code >= 400:
        logger.error(f"YARN GET {url} -> {resp.status_code}: {resp.text[:500]}")
        raise YarnError(f"YARN GET returned HTTP {resp.status_code}")
    data = resp.json()
    app = data.get("app", {})
    state = app.get("state")
    if not state:
        raise YarnError(f"Could not parse YARN state from response: {data!r}")
    logger.info(f"YARN status {application_id} -> {state}")
    return state, json.dumps(data, indent=2)


def _logs_unavailable_error(application_id: str) -> YarnError:
    """Consistent error when logs cannot be retrieved from either path."""
    return YarnError(
        f"YARN aggregated logs not available for {application_id!r}. "
        f"The application may not be in FINISHED state, or "
        f"yarn.log-aggregation-enable is false on the cluster."
    )


def _fetch_logs_via_nodemanager(application_id: str, yarn_rm_url: str | None) -> str:
    """
    Hadoop 2.x fallback: walk app attempts -> containers -> NodeManager
    container logs and concatenate the results.
    """
    base = _base_url(yarn_rm_url)
    app_url = f"{base}/ws/v1/cluster/apps/{application_id}"

    # Latest (or all) app attempts.
    attempts_url = f"{app_url}/appattempts"
    resp = _request("GET", attempts_url, timeout=30.0)
    if resp.status_code >= 400:
        raise _logs_unavailable_error(application_id)
    attempts = resp.json().get("appAttempts", {}).get("appAttempt", [])
    if not attempts:
        raise _logs_unavailable_error(application_id)

    # Collect containers across attempts.
    containers: list[dict] = []
    for attempt in attempts:
        attempt_id = attempt.get("id")
        if not attempt_id:
            continue
        containers_url = f"{app_url}/appattempts/{attempt_id}/containers"
        resp = _request("GET", containers_url, timeout=30.0)
        if resp.status_code >= 400:
            raise _logs_unavailable_error(application_id)
        containers.extend(resp.json().get("containers", {}).get("container", []))

    if not containers:
        raise _logs_unavailable_error(application_id)

    parts = []
    for container in containers:
        container_id = container.get("id")
        user = container.get("user")
        node_http_address = container.get("nodeHttpAddress")
        if not all((container_id, user, node_http_address)):
            continue
        nm_url = (
            f"http://{node_http_address}/node/containerlogs/{container_id}/{user}/"
        )
        resp = _request("GET", nm_url, timeout=60.0)
        if resp.status_code >= 400:
            raise _logs_unavailable_error(application_id)
        parts.append(f"=== container: {container_id} ===\n{resp.text}")

    if not parts:
        raise _logs_unavailable_error(application_id)
    return "\n\n".join(parts)


def get_application_logs(application_id: str, yarn_rm_url: str | None) -> str:
    """Return aggregated container logs for an application as text."""
    url = f"{_base_url(yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/aggregated-logs"
    resp = _request("GET", url, timeout=60.0)
    if resp.status_code in (404, 501):
        return _fetch_logs_via_nodemanager(application_id, yarn_rm_url)
    if resp.status_code >= 400:
        logger.error(f"YARN GET {url} -> {resp.status_code}: {resp.text[:500]}")
        raise YarnError(f"YARN GET logs returned HTTP {resp.status_code}")
    logger.info(f"YARN logs {application_id} -> {len(resp.text)} chars")
    return resp.text


def kill_application(application_id: str, yarn_rm_url: str | None) -> None:
    """PUT state=KILLED to /ws/v1/cluster/apps/{appid}/state."""
    url = f"{_base_url(yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/state"
    resp = _request("PUT", url, json_body={"state": "KILLED"})
    if resp.status_code >= 400:
        logger.error(f"YARN PUT {url} -> {resp.status_code}: {resp.text[:500]}")
        raise YarnError(f"YARN kill returned HTTP {resp.status_code}: {resp.text}")
    logger.info(f"YARN kill {application_id} -> ok")