diff --git a/python/ray/autoscaler/_private/kuberay/node_provider.py b/python/ray/autoscaler/_private/kuberay/node_provider.py index ab8ff9a50c3f..8f41b1f02357 100644 --- a/python/ray/autoscaler/_private/kuberay/node_provider.py +++ b/python/ray/autoscaler/_private/kuberay/node_provider.py @@ -53,6 +53,21 @@ ) KUBERNETES_SERVICE_PORT = os.getenv("KUBERNETES_SERVICE_PORT_HTTPS", "443") KUBERNETES_HOST = build_address(KUBERNETES_SERVICE_HOST, KUBERNETES_SERVICE_PORT) + +# Paths of the credentials projected into a Pod by Kubernetes. +IN_CLUSTER_TOKEN_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/token" +IN_CLUSTER_CA_CERT_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt" + +# Optional client TLS credentials, for Kubernetes distributions where API +# authentication is not the in-cluster ServiceAccount mechanism. For example, the +# API server may sit behind an authenticating proxy that identifies callers by a +# client certificate instead of a bearer token. Both must be set to take effect. +# When unset, the in-cluster ServiceAccount token is used, as before. +KUBERAY_CLIENT_CERT_PATH = os.getenv("RAY_KUBERAY_CLIENT_CERT_PATH") +KUBERAY_CLIENT_KEY_PATH = os.getenv("RAY_KUBERAY_CLIENT_KEY_PATH") +# Optional CA certificate used to verify the Kubernetes API server. When unset, +# the in-cluster ServiceAccount CA certificate is used, as before. +KUBERAY_CA_CERT_PATH = os.getenv("RAY_KUBERAY_CA_CERT_PATH") # Key for GKE label that identifies which multi-host replica a pod belongs to REPLICA_INDEX_KEY = "replicaIndex" @@ -165,23 +180,42 @@ def replace_patch(path: str, value: Any) -> Dict[str, Any]: return {"op": "replace", "path": path, "value": value} -def load_k8s_secrets() -> Tuple[Dict[str, str], str]: +def load_k8s_secrets() -> Tuple[Dict[str, str], str, Optional[Tuple[str, str]]]: """ Loads secrets needed to access K8s resources. + By default the in-cluster ServiceAccount credentials are used. Setting + $RAY_KUBERAY_CLIENT_CERT_PATH and $RAY_KUBERAY_CLIENT_KEY_PATH instead + authenticates with a client TLS certificate, and $RAY_KUBERAY_CA_CERT_PATH + overrides the CA certificate used to verify the API server. + Returns: - headers: Headers with K8s access token - verify: Path to certificate + headers: Headers with K8s access token, empty if no token is available + verify: Path to the CA certificate used to verify the API server + cert: Paths to the client certificate and key, None if not configured """ - with open("/var/run/secrets/kubernetes.io/serviceaccount/token") as secret: - token = secret.read() + if bool(KUBERAY_CLIENT_CERT_PATH) != bool(KUBERAY_CLIENT_KEY_PATH): + raise ValueError( + "$RAY_KUBERAY_CLIENT_CERT_PATH and $RAY_KUBERAY_CLIENT_KEY_PATH must be " + "set together to authenticate with a client TLS certificate." + ) + + cert = None + if KUBERAY_CLIENT_CERT_PATH and KUBERAY_CLIENT_KEY_PATH: + cert = (KUBERAY_CLIENT_CERT_PATH, KUBERAY_CLIENT_KEY_PATH) + + headers = {} + # When authenticating with a client certificate, a ServiceAccount token is not + # necessarily projected into the Pod, so treat it as optional. Without a client + # certificate the token is required, as before. + if cert is None or os.path.exists(IN_CLUSTER_TOKEN_PATH): + with open(IN_CLUSTER_TOKEN_PATH) as secret: + token = secret.read() + headers["Authorization"] = "Bearer " + token - headers = { - "Authorization": "Bearer " + token, - } - verify = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt" + verify = KUBERAY_CA_CERT_PATH or IN_CLUSTER_CA_CERT_PATH - return headers, verify + return headers, verify, cert def url_from_resource( @@ -282,18 +316,16 @@ def __init__(self, namespace: str, kuberay_crd_version: str = KUBERAY_CRD_VER): self._kuberay_crd_version = kuberay_crd_version self._namespace = namespace self._token_expires_at = datetime.datetime.now() + TOKEN_REFRESH_PERIOD - self._headers, self._verify = None, None + self._headers, self._verify, self._cert = None, None, None - def _get_refreshed_headers_and_verify(self): + def _get_refreshed_credentials(self): if (datetime.datetime.now() >= self._token_expires_at) or ( self._headers is None or self._verify is None ): logger.info("Refreshing K8s API client token and certs.") - self._headers, self._verify = load_k8s_secrets() + self._headers, self._verify, self._cert = load_k8s_secrets() self._token_expires_at = datetime.datetime.now() + TOKEN_REFRESH_PERIOD - return self._headers, self._verify - else: - return self._headers, self._verify + return self._headers, self._verify, self._cert def get(self, path: str) -> Dict[str, Any]: """Wrapper for REST GET of resource with proper headers. @@ -313,12 +345,13 @@ def get(self, path: str) -> Dict[str, Any]: kuberay_crd_version=self._kuberay_crd_version, ) - headers, verify = self._get_refreshed_headers_and_verify() + headers, verify, cert = self._get_refreshed_credentials() result = requests.get( url, headers=headers, timeout=KUBERAY_REQUEST_TIMEOUT_S, verify=verify, + cert=cert, ) if not result.status_code == 200: result.raise_for_status() @@ -349,13 +382,14 @@ def patch( path=path, kuberay_crd_version=self._kuberay_crd_version, ) - headers, verify = self._get_refreshed_headers_and_verify() + headers, verify, cert = self._get_refreshed_credentials() result = requests.patch( url, json.dumps(payload), headers={**headers, "Content-type": content_type}, timeout=KUBERAY_REQUEST_TIMEOUT_S, verify=verify, + cert=cert, ) if not result.status_code == 200: result.raise_for_status() diff --git a/python/ray/tests/kuberay/test_kuberay_node_provider.py b/python/ray/tests/kuberay/test_kuberay_node_provider.py index 6abf8c2372fb..29e05dcd093f 100644 --- a/python/ray/tests/kuberay/test_kuberay_node_provider.py +++ b/python/ray/tests/kuberay/test_kuberay_node_provider.py @@ -1,4 +1,6 @@ +import contextlib import copy +import datetime import sys from collections import defaultdict from pathlib import Path @@ -9,12 +11,18 @@ import pytest import yaml +from ray.autoscaler._private.kuberay import node_provider as node_provider_module from ray.autoscaler._private.kuberay.node_provider import ( + IN_CLUSTER_CA_CERT_PATH, + IN_CLUSTER_TOKEN_PATH, + KUBERAY_REQUEST_TIMEOUT_S, KubeRayNodeProvider, + KubernetesHttpApiClient, ScaleRequest, _worker_group_index, _worker_group_max_replicas, _worker_group_replicas, + load_k8s_secrets, ) from ray.autoscaler._private.util import NodeID from ray.autoscaler.batching_node_provider import NodeData @@ -367,5 +375,197 @@ def mock_patch(kuberay_provider, path, patch_payload): assert patched_tpu_workers_to_delete == tpu_workers_to_delete +CLIENT_CERT_PATH = "/etc/ray-k8s-auth/tls.crt" +CLIENT_KEY_PATH = "/etc/ray-k8s-auth/tls.key" +CA_CERT_PATH = "/etc/ray-k8s-auth/ca.crt" + + +def _mock_auth_config( + client_cert_path=None, client_key_path=None, ca_cert_path=None +) -> mock._patch: + """Patch the auth configuration read from the environment at import time.""" + return mock.patch.multiple( + node_provider_module, + KUBERAY_CLIENT_CERT_PATH=client_cert_path, + KUBERAY_CLIENT_KEY_PATH=client_key_path, + KUBERAY_CA_CERT_PATH=ca_cert_path, + ) + + +@contextlib.contextmanager +def _mock_token_file(token: str = "fake-token", exists: bool = True): + """Fake the ServiceAccount token file projected into a Pod by Kubernetes.""" + with mock.patch( + "builtins.open", mock.mock_open(read_data=token) + ) as mock_open, mock.patch.object( + node_provider_module.os.path, "exists", return_value=exists + ): + yield mock_open + + +def _mock_ok_response(payload): + response = mock.MagicMock() + response.status_code = 200 + response.json.return_value = payload + return response + + +def test_load_k8s_secrets_defaults_to_in_cluster_auth(): + """With no auth configured, the in-cluster ServiceAccount is used, as before.""" + with _mock_auth_config(), _mock_token_file() as mock_open: + headers, verify, cert = load_k8s_secrets() + + mock_open.assert_called_once_with(IN_CLUSTER_TOKEN_PATH) + assert headers == {"Authorization": "Bearer fake-token"} + assert verify == IN_CLUSTER_CA_CERT_PATH + assert cert is None + + +def test_load_k8s_secrets_with_client_cert(): + """A configured client certificate is returned for use as `requests`' `cert`.""" + with _mock_auth_config( + client_cert_path=CLIENT_CERT_PATH, + client_key_path=CLIENT_KEY_PATH, + ca_cert_path=CA_CERT_PATH, + ), _mock_token_file(): + headers, verify, cert = load_k8s_secrets() + + assert headers == {"Authorization": "Bearer fake-token"} + assert verify == CA_CERT_PATH + assert cert == (CLIENT_CERT_PATH, CLIENT_KEY_PATH) + + +def test_load_k8s_secrets_ca_cert_override_only(): + """The CA certificate can be overridden without configuring a client cert.""" + with _mock_auth_config(ca_cert_path=CA_CERT_PATH), _mock_token_file(): + headers, verify, cert = load_k8s_secrets() + + assert headers == {"Authorization": "Bearer fake-token"} + assert verify == CA_CERT_PATH + assert cert is None + + +def test_load_k8s_secrets_token_optional_with_client_cert(): + """A missing ServiceAccount token is tolerated when using a client cert.""" + with _mock_auth_config( + client_cert_path=CLIENT_CERT_PATH, client_key_path=CLIENT_KEY_PATH + ), _mock_token_file(exists=False): + headers, verify, cert = load_k8s_secrets() + + assert headers == {} + assert verify == IN_CLUSTER_CA_CERT_PATH + assert cert == (CLIENT_CERT_PATH, CLIENT_KEY_PATH) + + +def test_load_k8s_secrets_token_required_without_client_cert(): + """A missing ServiceAccount token still raises in the default configuration.""" + with _mock_auth_config(), mock.patch( + "builtins.open", side_effect=FileNotFoundError + ): + with pytest.raises(FileNotFoundError): + load_k8s_secrets() + + +@pytest.mark.parametrize( + "client_cert_path,client_key_path", + [(CLIENT_CERT_PATH, None), (None, CLIENT_KEY_PATH)], +) +def test_load_k8s_secrets_requires_cert_and_key_together( + client_cert_path, client_key_path +): + with _mock_auth_config( + client_cert_path=client_cert_path, client_key_path=client_key_path + ), _mock_token_file(): + with pytest.raises(ValueError, match="must be set together"): + load_k8s_secrets() + + +def test_client_get_defaults_unchanged(): + """`get` sends the in-cluster bearer token and no client cert, as before.""" + client = KubernetesHttpApiClient(namespace="default") + with _mock_auth_config(), _mock_token_file(), mock.patch.object( + node_provider_module.requests, "get" + ) as mock_get: + mock_get.return_value = _mock_ok_response({"kind": "PodList"}) + assert client.get("pods") == {"kind": "PodList"} + + args, kwargs = mock_get.call_args + assert args == ( + node_provider_module.url_from_resource(namespace="default", path="pods"), + ) + assert kwargs == { + "headers": {"Authorization": "Bearer fake-token"}, + "timeout": KUBERAY_REQUEST_TIMEOUT_S, + "verify": IN_CLUSTER_CA_CERT_PATH, + "cert": None, + } + + +def test_client_patch_defaults_unchanged(): + """`patch` sends the in-cluster bearer token and no client cert, as before.""" + client = KubernetesHttpApiClient(namespace="default") + with _mock_auth_config(), _mock_token_file(), mock.patch.object( + node_provider_module.requests, "patch" + ) as mock_patch: + mock_patch.return_value = _mock_ok_response({"kind": "RayCluster"}) + assert client.patch("rayclusters/fake", [{"op": "remove", "path": "/x"}]) == { + "kind": "RayCluster" + } + + _, kwargs = mock_patch.call_args + assert kwargs == { + "headers": { + "Authorization": "Bearer fake-token", + "Content-type": "application/json-patch+json", + }, + "timeout": KUBERAY_REQUEST_TIMEOUT_S, + "verify": IN_CLUSTER_CA_CERT_PATH, + "cert": None, + } + + +def test_client_caches_credentials_until_they_expire(): + """Credentials are loaded once and reused until the refresh period elapses.""" + client = KubernetesHttpApiClient(namespace="default") + with _mock_auth_config(), _mock_token_file(), mock.patch.object( + node_provider_module, "load_k8s_secrets", return_value=({}, "ca", None) + ) as mock_load, mock.patch.object(node_provider_module.requests, "get") as mock_get: + mock_get.return_value = _mock_ok_response({}) + + client.get("pods") + client.get("pods") + assert mock_load.call_count == 1 + + # Force the cached credentials to look expired. + client._token_expires_at = datetime.datetime.now() - datetime.timedelta( + seconds=1 + ) + client.get("pods") + assert mock_load.call_count == 2 + + +@pytest.mark.parametrize("method", ["get", "patch"]) +def test_client_uses_configured_client_cert(method: str): + """A configured client cert and CA are passed through to `requests`.""" + client = KubernetesHttpApiClient(namespace="default") + with _mock_auth_config( + client_cert_path=CLIENT_CERT_PATH, + client_key_path=CLIENT_KEY_PATH, + ca_cert_path=CA_CERT_PATH, + ), _mock_token_file(), mock.patch.object( + node_provider_module.requests, method + ) as mock_request: + mock_request.return_value = _mock_ok_response({}) + if method == "get": + client.get("pods") + else: + client.patch("rayclusters/fake", []) + + _, kwargs = mock_request.call_args + assert kwargs["cert"] == (CLIENT_CERT_PATH, CLIENT_KEY_PATH) + assert kwargs["verify"] == CA_CERT_PATH + assert kwargs["headers"]["Authorization"] == "Bearer fake-token" + + if __name__ == "__main__": sys.exit(pytest.main(["-v", __file__]))