tao.chen/mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(yarn_client): manually follow 3xx Location redirects on log fetch

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-06-30 13:39:30 +08:00
co-authored by Claude Fable 5
parent b10fae5fd1
commit 38ac3c33d6
2 changed files with 160 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spark_executor/core/yarn_client.py
+50 -4
View File
@@ -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)
tests/unit/test_yarn_client.py
+110 -3
View File
@@ -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",