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; amContainerLogs
  GET  /ws/v1/cluster/apps/{appid}/aggregated-logs  -> aggregated container logs
  PUT  /ws/v1/cluster/apps/{appid}/state             -> kill an app (body: {"state":"KILLED"})

When aggregated logs are unavailable (HTTP 404/501, e.g. log aggregation
disabled), fall back to the amContainerLogs field reported by the app
endpoint and fetch the AM (driver) container's /stdout from the
NodeManager. Some clusters (Knox-proxied, wrapped YARN builds) wrap
/stdout in an HTML directory listing; in that case we fetch the bare
{amContainerLogs} URL, parse the listing, and pull each log file. This
returns driver logs only; executor container logs still require
yarn.log-aggregation-enable=true (the primary /aggregated-logs path).

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
from dataclasses import dataclass

import re

import httpx
import httpx_kerberos

from common.config import settings
from common.logging import logger
from spark_executor.models import Connection


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


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


@dataclass(frozen=True)
class YarnClientConfig:
    """Resolved per-connection + global-default configuration for YARN REST calls.

    Carries the ResourceManager URL and the SSL/TLS verification value that
    should be passed to httpx.request. Future authentication fields
    (basic/kerberos/SPNEGO) can be added here without renaming the class.
    """

    yarn_rm_url: str | None
    ssl_verify: bool | None = None
    ssl_ca_bundle: str | None = None
    auth_type: str = "none"
    auth_user: str | None = None
    auth_password: str | None = None
    auth_principal: str | None = None  # display/audit only for kerberos
    auth_keytab: str | None = None     # display/audit only for kerberos

    @classmethod
    def from_connection(cls, conn: Connection) -> "YarnClientConfig":
        return cls(
            yarn_rm_url=conn.yarn_rm_url,
            ssl_verify=conn.ssl_verify,
            ssl_ca_bundle=conn.ssl_ca_bundle,
            auth_type=conn.auth_type,
            auth_user=conn.auth_user,
            auth_password=conn.auth_password,
            auth_principal=conn.auth_principal,
            auth_keytab=conn.auth_keytab,
        )

    def verify_for_httpx(self) -> bool | str:
        """Return the value for httpx.request(verify=...).

        Per-connection CA bundle wins over the global default. If neither is
        set, fall back to the per-connection ssl_verify flag, then to the
        global ssl_verify_default.
        """
        ca = self.ssl_ca_bundle or settings.ssl_ca_bundle_default
        if ca:
            return ca
        if self.ssl_verify is not None:
            return self.ssl_verify
        return settings.ssl_verify_default

    def auth_for_httpx(self) -> httpx.Auth | None:
        """Return the httpx.Auth object implied by this config, if any."""
        if self.auth_type in ("none", "simple"):
            return None
        if self.auth_type == "basic":
            if not self.auth_user:
                raise YarnConfigError("auth_type='basic' requires auth_user")
            return httpx.BasicAuth(self.auth_user, self.auth_password or "")
        if self.auth_type == "kerberos":
            return httpx_kerberos.HTTPKerberosAuth()
        raise YarnConfigError(f"Unknown auth_type: {self.auth_type!r}")


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, verify: bool | str = True,
             auth: httpx.Auth | None = None) -> httpx.Response:
    headers = {"Accept": "application/json"}
    logger.debug(f"YARN {method} {url}" + (f" body={json_body}" if json_body else ""))
    try:
        resp = httpx.request(
            method, url, json=json_body, headers=headers, timeout=timeout, verify=verify, auth=auth
        )
    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, config: YarnClientConfig) -> tuple[str, str]:
    """Return (state, raw_json_text) for an application, or raise YarnError."""
    url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps/{application_id}"
    resp = _request("GET", url, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
    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 _looks_like_html(text: str) -> bool:
    """Cheap heuristic: True if the body starts with an HTML/XHTML prolog.

    Intentionally lightweight (no full parser) — used only to decide whether
    a NodeManager log response is raw text or an HTML directory listing.
    """
    head = text.lstrip()[:200].lower()
    return head.startswith(("<!doctype html", "<html", "<?xml"))


def _parse_log_listing_html(html: str) -> list[str]:
    """Extract plain log-file names from a NodeManager directory-listing page.

    Keeps simple filenames (stdout, stderr, syslog, prelaunch.err,
    launch_container.sh, ...); drops parent-dir links, absolute URLs, and
    sort-toggle query strings. Returns hrefs in document order, deduplicated.
    """
    hrefs = re.findall(r'href="([^"]*)"', html, re.IGNORECASE)
    seen: set[str] = set()
    out: list[str] = []
    for href in hrefs:
        if href in ("../", "/") or href.startswith("/"):
            continue
        if "://" in href or "?" in href:
            continue
        if href in seen:
            continue
        seen.add(href)
        out.append(href)
    return out


def _fetch_logs_via_am_container(application_id: str, config: YarnClientConfig) -> str:
    """
    Fetch ApplicationMaster (driver) container logs as a fallback when the
    aggregated-logs endpoint is unavailable.

    Fast path: GET {amContainerLogs}/stdout and return its text directly when
    the NodeManager serves raw log bytes (modern YARN).

    Fallback path: when /stdout returns an HTML directory listing (Knox /
    wrapped YARN), GET the bare {amContainerLogs} URL, parse the listing for
    log-file names, and fetch each one individually.

    Returns AM (driver) container logs only. Executor container logs require
    yarn.log-aggregation-enable=true, which is covered by the primary
    /aggregated-logs path.
    """
    base = _base_url(config.yarn_rm_url)
    app_url = f"{base}/ws/v1/cluster/apps/{application_id}"

    resp = _request("GET", app_url, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
    if resp.status_code >= 400:
        raise _logs_unavailable_error(application_id)

    am_container_logs = resp.json().get("app", {}).get("amContainerLogs")
    if not am_container_logs:
        raise _logs_unavailable_error(application_id)

    am_container_logs = am_container_logs.rstrip("/")
    verify = config.verify_for_httpx()
    auth = config.auth_for_httpx()

    # Fast path: modern YARN serves raw log bytes at /stdout.
    stdout_url = f"{am_container_logs}/stdout"
    resp = _request("GET", stdout_url, timeout=60.0, verify=verify, auth=auth)
    if resp.status_code < 400 and not _looks_like_html(resp.text):
        return resp.text

    # Fallback path: Knox / wrapped YARN serves an HTML directory listing for
    # every URL on the NodeManager log endpoint. Parse it and fetch each file.
    resp = _request("GET", am_container_logs, timeout=60.0, verify=verify, auth=auth)
    if resp.status_code < 400:
        parts: list[str] = []
        for filename in _parse_log_listing_html(resp.text):
            file_url = f"{am_container_logs}/{filename}"
            file_resp = _request("GET", file_url, timeout=60.0, verify=verify, auth=auth)
            if file_resp.status_code < 400 and not _looks_like_html(file_resp.text):
                parts.append(f"=== {filename} ===\n{file_resp.text}")
        if parts:
            return "\n\n".join(parts)

    raise _logs_unavailable_error(application_id)


def get_application_logs(application_id: str, config: YarnClientConfig) -> str:
    """Return aggregated container logs for an application as text."""
    url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/aggregated-logs"
    resp = _request("GET", url, timeout=60.0, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
    if resp.status_code in (404, 501):
        return _fetch_logs_via_am_container(application_id, config)
    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, config: YarnClientConfig) -> None:
    """PUT state=KILLED to /ws/v1/cluster/apps/{appid}/state."""
    url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps/{application_id}/state"
    resp = _request("PUT", url, json_body={"state": "KILLED"}, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
    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")