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

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 _fetch_logs_via_nodemanager(application_id: str, config: YarnClientConfig) -> str:
    """
    Hadoop 2.x fallback: walk app attempts -> containers -> NodeManager
    container logs and concatenate the results.
    """
    base = _base_url(config.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, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
    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, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
        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
        scheme = "https" if config.yarn_rm_url and config.yarn_rm_url.startswith("https://") else "http"
        nm_url = f"{scheme}://{node_http_address}/node/containerlogs/{container_id}/{user}/"
        resp = _request("GET", nm_url, timeout=60.0, verify=config.verify_for_httpx(), auth=config.auth_for_httpx())
        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, 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_nodemanager(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")