From 38ac3c33d688666522a34f494fda29f60d666c37 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 30 Jun 2026 13:39:30 +0800 Subject: [PATCH] fix(yarn_client): manually follow 3xx Location redirects on log fetch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit httpx follows redirects by default, but it does NOT forward the Authorization header across host boundaries. In YARN deployments where the ResourceManager 307-redirects the NodeManager log fetch to a different host (load balancer, Knox, NM selection), the follow-up request lands unauthenticated and returns 401/403. Replace the _request call in _fetch_logs_via_am_container with _request_following_redirects, which: - Walks the 3xx Location chain (up to 3 hops). - Re-applies auth + verify on every hop. - Resolves relative Location URLs against the current request URL. - Raises YarnError on 3xx-without-Location (misconfigured server) and on hop count overflow (redirect loop protection). 301/302/303/307/308 are all treated as 'follow the Location' per RFC 7231 — the method/body handling for the rest is up to httpx when we re-issue the request. Tests: 4 new (cross-host 307, no-Location 307, two-hop chain, relative Location). Full suite 247 passed. Co-Authored-By: Claude Fable 5 --- spark_executor/core/yarn_client.py | 54 +++++++++++++- tests/unit/test_yarn_client.py | 113 ++++++++++++++++++++++++++++- 2 files changed, 160 insertions(+), 7 deletions(-) diff --git a/spark_executor/core/yarn_client.py b/spark_executor/core/yarn_client.py index 312f061..052626d 100644 --- a/spark_executor/core/yarn_client.py +++ b/spark_executor/core/yarn_client.py @@ -15,9 +15,11 @@ Endpoints used (YARN 2.6+): 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 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 @@ -26,6 +28,7 @@ field was designed for, but no longer requires the `yarn` CLI to interpret it. """ import json from dataclasses import dataclass +from urllib.parse import urljoin import httpx import httpx_kerberos @@ -137,6 +140,46 @@ def _request(method: str, url: str, *, json_body: dict | None = None, 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. + + Relative `Location` values are resolved against the current request + URL. A 3xx without a `Location` header is returned as-is (the caller + will treat it as a non-2xx response and surface a clean error). + """ + 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}" + ) + next_url = urljoin(url, location) + logger.debug(f"YARN {method} {url} -> {resp.status_code}, following to {next_url}") + url = next_url + 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}" @@ -185,7 +228,10 @@ def _fetch_logs_via_am_container(application_id: str, config: YarnClientConfig) 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()) + 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) diff --git a/tests/unit/test_yarn_client.py b/tests/unit/test_yarn_client.py index 6379e42..f767476 100644 --- a/tests/unit/test_yarn_client.py +++ b/tests/unit/test_yarn_client.py @@ -43,10 +43,13 @@ def _restore_settings(): config.settings.ssl_ca_bundle_default = snapshot.ssl_ca_bundle_default -def _resp(status: int, *, json_data=None, text: str | None = None) -> httpx.Response: +def _resp(status: int, *, json_data=None, text: str | None = None, headers: dict | None = None) -> httpx.Response: + kwargs: dict = {} + if headers: + kwargs["headers"] = headers if json_data is not None: - return httpx.Response(status, json=json_data) - return httpx.Response(status, text=text or "") + return httpx.Response(status, json=json_data, **kwargs) + return httpx.Response(status, text=text or "", **kwargs) # --- get_application_status --- @@ -191,6 +194,110 @@ def test_logs_returns_am_container_stdout(): ) +def test_logs_follows_307_redirect_to_other_host(): + """RM 307-redirects the log fetch to the actual NodeManager host. The + helper must follow Location and re-apply auth/verify on the new host.""" + responses = [ + _resp(404), # aggregated-logs missing + _resp( + 200, + json_data={ + "app": { + "amContainerLogs": "http://rm:8088/node/containerlogs/container_1/hdfs", + } + }, + ), + _resp(307, headers={"Location": "http://nm2:8042/node/containerlogs/container_1/hdfs/stdout"}), + _resp(200, text="driver stdout via redirect"), + ] + with patch( + "spark_executor.core.yarn_client.httpx.request", side_effect=responses + ) as m: + out = get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM)) + assert out == "driver stdout via redirect" + calls = m.call_args_list + assert calls[2].args == ("GET", "http://rm:8088/node/containerlogs/container_1/hdfs/stdout") + assert calls[3].args == ("GET", "http://nm2:8042/node/containerlogs/container_1/hdfs/stdout") + # Auth must be re-applied on the redirected request. + assert calls[3].kwargs["auth"] == calls[2].kwargs["auth"] + + +def test_logs_307_without_location_raises_unavailable(): + """3xx without a Location header is a misconfigured server — surface as + YarnError rather than silently returning empty text.""" + responses = [ + _resp(404), # aggregated-logs missing + _resp( + 200, + json_data={ + "app": { + "amContainerLogs": "http://rm:8088/node/containerlogs/container_1/hdfs", + } + }, + ), + _resp(307), # no Location header + ] + with patch( + "spark_executor.core.yarn_client.httpx.request", side_effect=responses + ): + with pytest.raises(YarnError, match="no Location header"): + get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM)) + + +def test_logs_follows_redirect_chain_then_200(): + """Two consecutive 307s followed by a 200 — the helper should walk the + chain and return the final body.""" + responses = [ + _resp(404), # aggregated-logs missing + _resp( + 200, + json_data={ + "app": { + "amContainerLogs": "http://lb:8088/node/containerlogs/container_1/hdfs", + } + }, + ), + _resp(307, headers={"Location": "http://rm:8088/node/containerlogs/container_1/hdfs/stdout"}), + _resp(307, headers={"Location": "http://nm2:8042/node/containerlogs/container_1/hdfs/stdout"}), + _resp(200, text="chained redirect body"), + ] + with patch( + "spark_executor.core.yarn_client.httpx.request", side_effect=responses + ) as m: + out = get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM)) + assert out == "chained redirect body" + urls = [c.args[1] for c in m.call_args_list[2:]] + assert urls == [ + "http://lb:8088/node/containerlogs/container_1/hdfs/stdout", + "http://rm:8088/node/containerlogs/container_1/hdfs/stdout", + "http://nm2:8042/node/containerlogs/container_1/hdfs/stdout", + ] + + +def test_logs_relative_location_resolved_against_request_url(): + """A relative Location header (e.g. '/redirected/stdout') must be + resolved against the request URL, not treated as an absolute path.""" + responses = [ + _resp(404), # aggregated-logs missing + _resp( + 200, + json_data={ + "app": { + "amContainerLogs": "http://nm1:8042/node/containerlogs/container_1/hdfs", + } + }, + ), + _resp(307, headers={"Location": "/redirected/stdout"}), + _resp(200, text="resolved body"), + ] + with patch( + "spark_executor.core.yarn_client.httpx.request", side_effect=responses + ) as m: + out = get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM)) + assert out == "resolved body" + assert m.call_args_list[3].args == ("GET", "http://nm1:8042/redirected/stdout") + + def test_logs_5xx_on_aggregated_endpoint_raises_immediately(): with patch( "spark_executor.core.yarn_client.httpx.request",