# 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. 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, 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_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("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")