feat: Add --headers/--verify to Ray State CLI and document State REST API#61958
feat: Add --headers/--verify to Ray State CLI and document State REST API#61958Anarion-zuo wants to merge 4 commits intoray-project:masterfrom
Conversation
… API Signed-off-by: aaronzuo <[email protected]>
Anarion-zuo
force-pushed
the
aaronzuo/state-cli-auth-headers-verify
branch
from
7243d9d to
0af48e9
Compare
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds authentication support to the Ray State CLI via --headers and --verify flags, and also documents the State REST API. The changes are well-implemented, with corresponding updates to documentation and new unit tests. I've identified a couple of minor opportunities for refactoring to reduce code duplication, which would improve maintainability. Overall, this is a solid contribution that enhances the usability of Ray in secure environments.
python/ray/util/state/api.py
Outdated
| request_headers: Dict[str, Any] = {} | ||
| if headers: | ||
| request_headers.update(headers) | ||
| request_headers.update(get_auth_headers_if_auth_enabled(request_headers)) |
There was a problem hiding this comment.
This logic for preparing request headers is duplicated in list_logs (lines 1369-1372). To improve maintainability and avoid code duplication, consider extracting this logic into a private helper function.
For example, you could add:
def _prepare_request_headers(headers: Optional[Dict[str, Any]]) -> Dict[str, Any]:
"""Prepares request headers by adding auth headers if not present."""
request_headers = headers.copy() if headers else {}
auth_headers = get_auth_headers_if_auth_enabled(request_headers)
request_headers.update(auth_headers)
return request_headersAnd then call request_headers = _prepare_request_headers(headers) in both get_log and list_logs.
python/ray/util/state/state_cli.py
Outdated
| headers=_handle_headers(headers), | ||
| verify=verify, |
There was a problem hiding this comment.
The _handle_headers function is called here and again on line 1133 when calling _print_log. This is slightly inefficient as it involves parsing the JSON string twice.
Consider calling _handle_headers once at the beginning of the log_cluster function and reusing the result:
def log_cluster(
ctx,
glob_filter: str,
address: Optional[str],
headers: Optional[str],
...
):
parsed_headers = _handle_headers(headers)
if node_id is None and node_ip is None:
node_ip = _get_head_node_ip(address)
logs = list_logs(
...,
headers=parsed_headers,
...
)
...
_print_log(
...,
headers=parsed_headers,
...
)![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Signed-off-by: aaronzuo <[email protected]>
Signed-off-by: aaronzuo <[email protected]>
Signed-off-by: aaronzuo <[email protected]>
ray-gardener
bot
added
core
1 participant
Summary
This PR makes Ray’s State CLI usable on authenticated / TLS-secured clusters by adding
--headersand--verifyflags (similar toray job), and documents the State REST API endpoints and query parameters.Resolved issue #50256.
Changes
--headersand--verifyto:ray list,ray get,ray summary (tasks|actors|objects), andray logs (...)RAY_STATE_HEADERSenv var as a default for--headersStateApiClient(..., headers=..., verify=...)get_log()/list_logs()acceptheadersandverifyTests
--headersparsing and propagationRAY_STATE_HEADERSenv var behaviorRun:
python -m pytest -q python/ray/tests/test_state_api.py -k "test_state_cli_"Usage example