diff --git a/Makefile b/Makefile
index 977a094ad..fe6d63085 100644
--- a/Makefile
+++ b/Makefile
@@ -42,7 +42,7 @@ REGISTRY_NAME ?= $(shell echo $(REGISTRY) | sed "s/.azurecr.io//g")
 IMAGE_TAG = $(REGISTRY)/$(IMAGENAME):$(IMAGE_VERSION)
 IMAGE_TAG_LATEST = $(REGISTRY)/$(IMAGENAME):latest
 
-E2E_HELM_OPTIONS ?= --set image.nfs.repository=$(REGISTRY)/$(IMAGENAME) --set image.nfs.tag=$(IMAGE_VERSION) --set image.nfs.pullPolicy=Always --set feature.enableInlineVolume=true --set externalSnapshotter.enabled=true
+E2E_HELM_OPTIONS ?= --set image.nfs.repository=$(REGISTRY)/$(IMAGENAME) --set image.nfs.tag=$(IMAGE_VERSION) --set image.nfs.pullPolicy=Always --set feature.enableInlineVolume=true --set externalSnapshotter.enabled=true --set controller.runOnControlPlane=true
 E2E_HELM_OPTIONS += ${EXTRA_HELM_OPTIONS}
 
 # Output type of docker buildx build
diff --git a/charts/latest/csi-driver-nfs-v0.0.0.tgz b/charts/latest/csi-driver-nfs-v0.0.0.tgz
index 9bc8a1554..7473be9cb 100644
Binary files a/charts/latest/csi-driver-nfs-v0.0.0.tgz and b/charts/latest/csi-driver-nfs-v0.0.0.tgz differ
diff --git a/charts/latest/csi-driver-nfs/templates/csi-nfs-controller.yaml b/charts/latest/csi-driver-nfs/templates/csi-nfs-controller.yaml
index afea01ea7..c70d32f07 100644
--- a/charts/latest/csi-driver-nfs/templates/csi-nfs-controller.yaml
+++ b/charts/latest/csi-driver-nfs/templates/csi-nfs-controller.yaml
@@ -24,18 +24,29 @@ spec:
       hostNetwork: true  # controller also needs to mount nfs to create dir
       dnsPolicy: {{ .Values.controller.dnsPolicy }}
       serviceAccountName: {{ .Values.serviceAccount.controller }}
-{{- with .Values.controller.affinity }}
+    # runOnControlPlane=true or runOnMaster=true only takes effect if affinity is not set
+    {{- if contains (tpl "{{ .Values.controller.affinity }}" .) "nodeSelectorTerms" }}
+      {{- with .Values.controller.affinity }}
       affinity:
-{{ toYaml . | indent 8 }}
-{{- end }}
+      {{ toYaml . | indent 8 }}
+      {{- end }}
+    {{- else if or .Values.controller.runOnControlPlane .Values.controller.runOnMaster}}
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+              - matchExpressions:
+                {{- if .Values.controller.runOnControlPlane}}
+                - key: node-role.kubernetes.io/control-plane
+                  operator: Exists
+                {{- end}}
+                {{- if .Values.controller.runOnMaster}}
+                - key: node-role.kubernetes.io/master
+                  operator: Exists
+                {{- end}}
+    {{- end }}
       nodeSelector:
         kubernetes.io/os: linux
-        {{- if .Values.controller.runOnMaster}}
-        node-role.kubernetes.io/master: ""
-        {{- end}}
-        {{- if .Values.controller.runOnControlPlane}}
-        node-role.kubernetes.io/control-plane: ""
-        {{- end}}
 {{- with .Values.controller.nodeSelector }}
 {{ toYaml . | indent 8 }}
 {{- end }}
diff --git a/charts/v4.9.0/csi-driver-nfs-v4.9.0.tgz b/charts/v4.9.0/csi-driver-nfs-v4.9.0.tgz
index 294bfdd08..0e8cd5d0b 100644
Binary files a/charts/v4.9.0/csi-driver-nfs-v4.9.0.tgz and b/charts/v4.9.0/csi-driver-nfs-v4.9.0.tgz differ
diff --git a/charts/v4.9.0/csi-driver-nfs/templates/csi-nfs-controller.yaml b/charts/v4.9.0/csi-driver-nfs/templates/csi-nfs-controller.yaml
index 82332a0a9..78c1aac7b 100644
--- a/charts/v4.9.0/csi-driver-nfs/templates/csi-nfs-controller.yaml
+++ b/charts/v4.9.0/csi-driver-nfs/templates/csi-nfs-controller.yaml
@@ -24,18 +24,29 @@ spec:
       hostNetwork: true  # controller also needs to mount nfs to create dir
       dnsPolicy: {{ .Values.controller.dnsPolicy }}
       serviceAccountName: {{ .Values.serviceAccount.controller }}
-{{- with .Values.controller.affinity }}
+    # runOnControlPlane=true or runOnMaster=true only takes effect if affinity is not set
+    {{- if contains (tpl "{{ .Values.controller.affinity }}" .) "nodeSelectorTerms" }}
+      {{- with .Values.controller.affinity }}
       affinity:
-{{ toYaml . | indent 8 }}
-{{- end }}
+      {{ toYaml . | indent 8 }}
+      {{- end }}
+    {{- else if or .Values.controller.runOnControlPlane .Values.controller.runOnMaster}}
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+              - matchExpressions:
+                {{- if .Values.controller.runOnControlPlane}}
+                - key: node-role.kubernetes.io/control-plane
+                  operator: Exists
+                {{- end}}
+                {{- if .Values.controller.runOnMaster}}
+                - key: node-role.kubernetes.io/master
+                  operator: Exists
+                {{- end}}
+    {{- end }}
       nodeSelector:
         kubernetes.io/os: linux
-        {{- if .Values.controller.runOnMaster}}
-        node-role.kubernetes.io/master: ""
-        {{- end}}
-        {{- if .Values.controller.runOnControlPlane}}
-        node-role.kubernetes.io/control-plane: ""
-        {{- end}}
 {{- with .Values.controller.nodeSelector }}
 {{ toYaml . | indent 8 }}
 {{- end }}
@@ -171,4 +182,4 @@ spec:
             path: {{ .Values.kubeletDir }}/pods
             type: Directory
         - name: socket-dir
-          emptyDir: {}
+          emptyDir: {}
\ No newline at end of file