[core][kuberay] Support configurable Kubernetes API authentication in the autoscaler

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,18 @@ 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
+            return self._headers, self._verify, self._cert
         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 +347,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 +384,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()

python/ray/tests/kuberay/test_kuberay_node_provider.py

@@ -1,3 +1,4 @@
+import contextlib
 import copy
 import sys
 from collections import defaultdict
@@ -9,12 +10,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 +374,177 @@ 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,
+    }
+
+
+@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__]))