# 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 == 404: # No apps match (or RM doesn't support the endpoint) return [] 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 []