Ray TPU Webhook Reliability Improvements (#723)

* Generate hostnames at Pod creation Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Update tests and add error checking Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Make webhook stateless in between mutate calls Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Formatting changes Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Fix bug causing incorrect IDs * Add cluster role and log formatting changes Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Filter pods by Ray worker group label Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Vulnerability fixes Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Better names and add ServiceAccount Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Change version back to v1.1 Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Change implementation to use PodInformer Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Use PodLister Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * updateSliceToWorkerIDs returns error Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Use mutex lock in updateSliceToWorkerIDs Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Update unit tests and fix comments Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Remove global client var Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Just return err instead of logging Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * TODO comment Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Lock when reading from shared sliceToWorkerIDs mapping Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Switch to using RWMutex Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Ray TPU Webhook Autoscaling Support (#740) * Generate hostnames at Pod creation Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Should not fatal log in deletePod Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * deletePod admission always succeeds Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Remove unused tests make command Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Update tests and add error checking Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Just return an error instead of logging Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> --------- Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Generate hostnames at Pod creation Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Update tests and add error checking Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Close stop channel on webhook termination Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Refactor webhook to avoid using global vars Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Fix comments Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Change service account name Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Return BadRequest if invalid kind Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Fix comments Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Change error messages Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Fatal log in main Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Update function comments Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Refactor to minimize indentations Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Change sliceToWorkerIDs nil check to use len Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Write http.Error to header Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Don't fatal log in validateRayCluster Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Check for nil admission request Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Add doc comment Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Update expected errors Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Better getNextWorkerID logic Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Update replicaIndex and nextWorkerID tests Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Refactor webhook unit tests Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Create numOfHosts pods for Pod List Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Log admission request object name Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Fix nits and go vet output Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Initial cloudbuil commit Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Fix vet command Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Update cloudbuild Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Fix cloudbuild errors Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Add dir Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Remove arg Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Change to bash command Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * increase timeout time Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Fix validateRayCluster test Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Fix nits for cloudbuild Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Break early in validateRayCluster Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Remove unnecessary args from validateRayCluster test Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Change break to continue Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Remove unused vars from webhook tests and add edge cases Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Update helm chart Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> --------- Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
GoogleCloudPlatform · Jul 25, 2024 · 0ae82b1 · 0ae82b1
1 parent 465b226
commit 0ae82b1
Show file tree

Hide file tree

Showing 9 changed files with 1,027 additions and 1,112 deletions.
diff --git a/ray-on-gke/tpu/kuberay-tpu-webhook/Makefile b/ray-on-gke/tpu/kuberay-tpu-webhook/Makefile
@@ -16,7 +16,19 @@ webhook:
 
 # Run against the configured Kubernetes cluster in ~/.kube/config  
 run: webhook  
-	go run ./main.go  
+	go run ./main.go
+
+# Run go fmt against code.
+fmt:
+	go fmt ./...
+
+# Run go vet against code.
+vet:
+	go vet ./...
+
+# Run go test against code.
+test:
+	go test ./...
 
 uninstall:  
 	kubectl delete -f deployments/

diff --git a/ray-on-gke/tpu/kuberay-tpu-webhook/cloudbuild.yaml b/ray-on-gke/tpu/kuberay-tpu-webhook/cloudbuild.yaml
@@ -0,0 +1,45 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+steps:
+  - id: 'kuberay tpu webhook fmt'
+    name: 'golang:1.21.11'
+    dir: /workspace/ray-on-gke/tpu/kuberay-tpu-webhook
+    entrypoint: 'bash'
+    args:
+      - '-c'
+      - |
+        set -e
+        make fmt
+  - id: 'kuberay tpu webhook vet'
+    name: 'golang:1.21.11'
+    dir: /workspace/ray-on-gke/tpu/kuberay-tpu-webhook
+    entrypoint: 'bash'
+    args:
+      - '-c'
+      - |
+        set -e
+        make vet
+  - id: 'kuberay tpu webhook test'
+    name: 'golang:1.21.11'
+    dir: /workspace/ray-on-gke/tpu/kuberay-tpu-webhook
+    entrypoint: 'bash'
+    args:
+      - '-c'
+      - |
+        set -e
+        make test
+options:
+  machineType: 'E2_HIGHCPU_8'
+timeout: 600s
diff --git a/ray-on-gke/tpu/kuberay-tpu-webhook/deployments/deployment.yaml b/ray-on-gke/tpu/kuberay-tpu-webhook/deployments/deployment.yaml
@@ -17,6 +17,34 @@ kind: Namespace
 metadata:
   name: ray-system
 ---
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: kuberay-tpu-webhook
+  namespace: ray-system
+---
+kind: ClusterRole
+apiVersion: rbac.authorization.k8s.io/v1
+metadata:
+  name: kuberay-tpu-webhook-pod-reader
+rules:
+  - apiGroups: [""]
+    resources: ["pods"]
+    verbs: ["get", "list", "watch"]
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRoleBinding
+metadata:
+  name: kuberay-tpu-webhook-pod-reader
+subjects:
+  - kind: ServiceAccount
+    name: kuberay-tpu-webhook
+    namespace: ray-system
+roleRef:
+  kind: ClusterRole
+  name: kuberay-tpu-webhook-pod-reader
+  apiGroup: rbac.authorization.k8s.io
+---
 apiVersion: apps/v1
 kind: Deployment
 metadata:
@@ -34,12 +62,13 @@ spec:
       labels:
         app: kuberay-tpu-webhook
     spec:
+      serviceAccountName: kuberay-tpu-webhook
       containers:
         - image: us-docker.pkg.dev/ai-on-gke/kuberay-tpu-webhook/kuberay-tpu-webhook:v1.1
           imagePullPolicy: Always
           name: kuberay-tpu-webhook
           args:
-          - --v=0 # change this value to 1 for verbose logging 
+          - --v=0 # change this value to 1 for verbose logging
           ports:
           - name: https
             containerPort: 443

diff --git a/ray-on-gke/tpu/kuberay-tpu-webhook/deployments/validating-webhook-cfg.yaml b/ray-on-gke/tpu/kuberay-tpu-webhook/deployments/validating-webhook-cfg.yaml
@@ -19,30 +19,6 @@ metadata:
   annotations:
     cert-manager.io/inject-ca-from: ray-system/kuberay-tpu-webhook-certs
 webhooks:
-  - name: pods-kuberay-tpu-webhook.ray-system.svc
-    admissionReviewVersions: [v1]
-    sideEffects: NoneOnDryRun
-    namespaceSelector:
-      matchExpressions:
-      - key: kubernetes.io/metadata.name
-        operator: NotIn
-        values:
-        - kube-system
-        - kube-node-lease
-    clientConfig:
-      service:
-        name: kuberay-tpu-webhook
-        namespace: ray-system
-        path: /validate
-    rules:
-      - operations: ["DELETE"]
-        apiGroups: [""]
-        apiVersions: ["v1"]
-        resources: ["pods"]
-        scope: Namespaced
-    objectSelector:
-      matchLabels:
-        app.kubernetes.io/name: kuberay
   - name: raycluster-kuberay-tpu-webhook.ray-system.svc
     admissionReviewVersions: [v1]
     sideEffects: NoneOnDryRun

diff --git a/ray-on-gke/tpu/kuberay-tpu-webhook/go.mod b/ray-on-gke/tpu/kuberay-tpu-webhook/go.mod
@@ -1,12 +1,13 @@
 module github.com/GoogleCloudPlatform/kuberay-tpu-webhook
 
-go 1.21
+go 1.21.11
 
 require (
-	github.com/ray-project/kuberay/ray-operator v1.1.0-rc.0
+	github.com/ray-project/kuberay/ray-operator v1.1.1
 	github.com/stretchr/testify v1.8.4
 	k8s.io/api v0.29.1
 	k8s.io/apimachinery v0.29.1
+	k8s.io/client-go v0.29.0
 	k8s.io/klog/v2 v2.120.1
 	k8s.io/utils v0.0.0-20240102154912-e7106e64919e
 )
@@ -16,6 +17,7 @@ require (
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/emicklei/go-restful/v3 v3.11.0 // indirect
+	github.com/evanphx/json-patch v5.6.0+incompatible // indirect
 	github.com/evanphx/json-patch/v5 v5.8.0 // indirect
 	github.com/fsnotify/fsnotify v1.7.0 // indirect
 	github.com/go-logr/logr v1.4.1 // indirect
@@ -24,7 +26,7 @@ require (
 	github.com/go-openapi/swag v0.22.3 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
-	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/gnostic-models v0.6.8 // indirect
 	github.com/google/go-cmp v0.6.0 // indirect
 	github.com/google/gofuzz v1.2.0 // indirect
@@ -53,17 +55,15 @@ require (
 	golang.org/x/term v0.18.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
 	golang.org/x/time v0.3.0 // indirect
-	golang.org/x/tools v0.17.0 // indirect
 	gomodules.xyz/jsonpatch/v2 v2.4.0 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/protobuf v1.33.0 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	k8s.io/apiextensions-apiserver v0.29.0 // indirect
-	k8s.io/client-go v0.29.0 // indirect
 	k8s.io/component-base v0.29.0 // indirect
-	k8s.io/kube-openapi v0.0.0-20231010175941-2dd684a91f00 // indirect
+	k8s.io/kube-openapi v0.0.0-20240228011516-70dd3763d340 // indirect
 	sigs.k8s.io/controller-runtime v0.17.0 // indirect
 	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect

diff --git a/ray-on-gke/tpu/kuberay-tpu-webhook/go.sum b/ray-on-gke/tpu/kuberay-tpu-webhook/go.sum
@@ -34,12 +34,10 @@ github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69
 github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE=
 github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
 github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
-github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
-github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg=
-github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
+github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
+github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
 github.com/google/gnostic-models v0.6.8 h1:yo/ABAfM5IMRsS1VnXjTBvUb61tFIHozhlYvRgGre9I=
 github.com/google/gnostic-models v0.6.8/go.mod h1:5n7qKqH0f5wFt+aWF8CW6pZLLNOfYuF5OpfBSENuI8U=
-github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
@@ -95,8 +93,8 @@ github.com/prometheus/common v0.45.0 h1:2BGz0eBc2hdMDLnO/8n0jeB3oPrt2D08CekT0lne
 github.com/prometheus/common v0.45.0/go.mod h1:YJmSTw9BoKxJplESWWxlbyttQR4uaEcGyv9MZjVOJsY=
 github.com/prometheus/procfs v0.12.0 h1:jluTpSng7V9hY0O2R9DzzJHYb2xULk9VTR1V1R/k6Bo=
 github.com/prometheus/procfs v0.12.0/go.mod h1:pcuDEFsWDnvcgNzo4EEweacyhjeA9Zk3cnaOZAZEfOo=
-github.com/ray-project/kuberay/ray-operator v1.1.0-rc.0 h1:QZIFpSxxoTFyC64Z6NK+TUCbQqs6PKLj2lCKHsZpLOc=
-github.com/ray-project/kuberay/ray-operator v1.1.0-rc.0/go.mod h1:ZqyKKvMP5nKDldQoKmur+Wcx7wVlV9Q98phFqHzr+KY=
+github.com/ray-project/kuberay/ray-operator v1.1.1 h1:mVOA1ddS9aAsPvhhHrpf0ZXgTzccIAyTbeYeDqtcfAk=
+github.com/ray-project/kuberay/ray-operator v1.1.1/go.mod h1:ZqyKKvMP5nKDldQoKmur+Wcx7wVlV9Q98phFqHzr+KY=
 github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
 github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
@@ -156,8 +154,8 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
-golang.org/x/tools v0.17.0 h1:FvmRgNOcs3kOa+T20R1uhfP9F6HgG2mfxDv1vrx1Htc=
-golang.org/x/tools v0.17.0/go.mod h1:xsh6VxdV005rRVaS6SSAf9oiAqljS7UZUacMZ8Bnsps=
+golang.org/x/tools v0.18.0 h1:k8NLag8AGHnn+PHbl7g43CtqZAwG60vZkLqgyZgIHgQ=
+golang.org/x/tools v0.18.0/go.mod h1:GL7B4CwcLLeo59yx/9UWWuNOW1n3VZ4f5axWfML7Lcg=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
@@ -166,8 +164,6 @@ gomodules.xyz/jsonpatch/v2 v2.4.0 h1:Ci3iUJyx9UeRx7CeFN8ARgGbkESwJK+KB9lLcWxY/Zw
 gomodules.xyz/jsonpatch/v2 v2.4.0/go.mod h1:AH3dM2RI6uoBZxn3LVrfvJ3E0/9dG4cSrbuBJT4moAY=
 google.golang.org/appengine v1.6.7 h1:FZR1q0exgwxzPzp/aF+VccGrSfxfPpkBqjIIEq3ru6c=
 google.golang.org/appengine v1.6.7/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=
-google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
-google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
 google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
@@ -194,8 +190,8 @@ k8s.io/component-base v0.29.0 h1:T7rjd5wvLnPBV1vC4zWd/iWRbV8Mdxs+nGaoaFzGw3s=
 k8s.io/component-base v0.29.0/go.mod h1:sADonFTQ9Zc9yFLghpDpmNXEdHyQmFIGbiuZbqAXQ1M=
 k8s.io/klog/v2 v2.120.1 h1:QXU6cPEOIslTGvZaXvFWiP9VKyeet3sawzTOvdXb4Vw=
 k8s.io/klog/v2 v2.120.1/go.mod h1:3Jpz1GvMt720eyJH1ckRHK1EDfpxISzJ7I9OYgaDtPE=
-k8s.io/kube-openapi v0.0.0-20231010175941-2dd684a91f00 h1:aVUu9fTY98ivBPKR9Y5w/AuzbMm96cd3YHRTU83I780=
-k8s.io/kube-openapi v0.0.0-20231010175941-2dd684a91f00/go.mod h1:AsvuZPBlUDVuCdzJ87iajxtXuR9oktsTctW/R9wwouA=
+k8s.io/kube-openapi v0.0.0-20240228011516-70dd3763d340 h1:BZqlfIlq5YbRMFko6/PM7FjZpUb45WallggurYhKGag=
+k8s.io/kube-openapi v0.0.0-20240228011516-70dd3763d340/go.mod h1:yD4MZYeKMBwQKVht279WycxKyM84kkAx2DPrTXaeb98=
 k8s.io/utils v0.0.0-20240102154912-e7106e64919e h1:eQ/4ljkx21sObifjzXwlPKpdGLrCfRziVtos3ofG/sQ=
 k8s.io/utils v0.0.0-20240102154912-e7106e64919e/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0=
 sigs.k8s.io/controller-runtime v0.17.0 h1:fjJQf8Ukya+VjogLO6/bNX9HE6Y2xpsO5+fyS26ur/s=

diff --git a/ray-on-gke/tpu/kuberay-tpu-webhook/helm-chart/templates/deployment.yaml b/ray-on-gke/tpu/kuberay-tpu-webhook/helm-chart/templates/deployment.yaml
@@ -20,7 +20,7 @@ metadata:
 apiVersion: v1
 kind: ServiceAccount
 metadata:
-  name: kuberay-tpu-webhook-sa
+  name: kuberay-tpu-webhook
   namespace: {{ .Values.tpuWebhook.namespace.name }}
 ---
 kind: ClusterRole
@@ -38,7 +38,7 @@ metadata:
   name: kuberay-tpu-webhook-pod-reader
 subjects:
   - kind: ServiceAccount
-    name: kuberay-tpu-webhook-sa
+    name: kuberay-tpu-webhook
     namespace: {{ .Values.tpuWebhook.namespace.name }}
 roleRef:
   kind: ClusterRole
@@ -62,6 +62,7 @@ spec:
       labels:
         app:  kuberay-tpu-webhook
     spec:
+      serviceAccountName: kuberay-tpu-webhook
       containers:
         - image: "{{ .Values.tpuWebhook.image.repository }}:{{ .Values.tpuWebhook.image.tag }}"
           imagePullPolicy: {{ .Values.tpuWebhook.image.pullPolicy }}