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 directly from
the NodeManager. The log fetch follows 3xx Location redirects manually
(httpx's default redirect-following drops Authorization across host
boundaries, which breaks RM→NM 307 redirects in many deployments).
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 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,
             params: dict[str, str] | 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, params=params, 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


# Status codes that indicate "follow the Location header". 307/308 preserve
# the method+body; the rest collapse to GET for the redirect target per
# RFC 7231 — httpx will do the right thing for us on the second request.
_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308})
_MAX_LOG_REDIRECTS = 3


def _request_following_redirects(
    method: str, url: str, *, max_redirects: int = _MAX_LOG_REDIRECTS,
    **kwargs,
) -> httpx.Response:
    """Like `_request` but follows 3xx Location redirects up to `max_redirects`.

    httpx follows redirects by default, but it does NOT forward the
    `Authorization` header across host boundaries — so when the YARN RM
    307-redirects us to a different NodeManager host, the follow-up
    request lands unauthenticated and gets 401/403. Doing it ourselves
    keeps the auth + verify config on every hop.

    YARN always returns a fully-qualified absolute URL in `Location`, so
    we use it verbatim. A 3xx without a `Location` header is treated as
    a misconfigured server and raises YarnError.
    """
    for hop in range(max_redirects + 1):
        resp = _request(method, url, **kwargs)
        if resp.status_code not in _REDIRECT_STATUSES:
            return resp
        location = resp.headers.get("Location") or resp.headers.get("location")
        if not location:
            raise YarnError(
                f"YARN {method} returned {resp.status_code} with no Location header at {url}"
            )
        logger.debug(f"YARN {method} {url} -> {resp.status_code}, following to {location}")
        url = location
    raise YarnError(
        f"YARN {method} exceeded {max_redirects} redirects (last url: {url})"
    )


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_am_container(application_id: str, config: YarnClientConfig) -> str:
    """
    Fetch ApplicationMaster (driver) container stdout as a fallback when the
    aggregated-logs endpoint is unavailable.

    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)

    log_url = f"{am_container_logs.rstrip('/')}/stdout"
    resp = _request_following_redirects(
        "GET", log_url, timeout=60.0,
        verify=config.verify_for_httpx(), auth=config.auth_for_httpx(),
    )
    if resp.status_code >= 400:
        logger.error(f"YARN GET {log_url} -> {resp.status_code}: {resp.text[:500]}")
        raise _logs_unavailable_error(application_id)
    return resp.text


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")


def list_applications(
    config: YarnClientConfig,
    *,
    state: str | None = None,
    queue: str | None = None,
    limit: int | None = None,
) -> list[dict]:
    """List YARN applications, optionally filtered.

    YARN endpoint: GET /ws/v1/cluster/apps?state=...&queue=...&limit=...

    Filters:
      - state: YARN application state. Common values:
        "NEW", "NEW_SAVING", "SUBMITTED", "ACCEPTED", "RUNNING",
        "FINISHED", "FAILED", "KILLED".
        Note: "FINISHED" is the umbrella state covering SUCCEEDED/FAILED/KILLED.
      - queue: YARN queue name
      - limit: cap on number of returned apps (YARN has no pagination;
        callers that need a full enumeration should make multiple
        calls with state=... filters or accept the cap)

    Returns a list of YARN app dicts (each with id, name, user, queue,
    state, finalStatus, applicationType, startedTime, finishedTime,
    trackingUrl, progress, etc). Empty list if no apps match.

    Raises YarnError on transport / 4xx / 5xx.
    """
    params: dict[str, str] = {}
    if state is not None:
        params["state"] = state
    if queue is not None:
        params["queue"] = queue
    if limit is not None:
        params["limit"] = str(limit)

    url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps"
    resp = _request("GET", url, params=params,
                    verify=config.verify_for_httpx(),
                    auth=config.auth_for_httpx())
    if resp.status_code >= 400:
        raise YarnError(
            f"YARN list applications failed: {resp.status_code} {resp.text[:200]}"
        )
    data = resp.json()
    apps_container = data.get("apps") or {}
    return apps_container.get("app", []) or []