orchestrator-kubernetes: make availability-zone node label configurable

misc/helm-charts/operator/README.md

@@ -107,6 +107,7 @@ The following table lists the configurable parameters of the Materialize operato
 
 | Parameter | Description | Default |
 |-----------|-------------|---------|
+| `availabilityZoneLabel` | The Kubernetes node label key used for availability zone affinity rules. When unset, defaults to topology.k8s.aws/zone-id for AWS cloud providers and topology.kubernetes.io/zone otherwise. | ``nil`` |
 | `balancerd.affinity` | Affinity to use for balancerd pods spawned by the operator | ``{}`` |
 | `balancerd.defaultResources.limits` | Default resource limits for balancerd's CPU and memory if not set in the Materialize CR | ``{"memory":"256Mi"}`` |
 | `balancerd.defaultResources.requests` | Default resources requested for balancerd's CPU and memory if not set in the Materialize CR | ``{"cpu":"500m","memory":"256Mi"}`` |

misc/helm-charts/operator/templates/deployment.yaml

@@ -237,6 +237,11 @@ spec:
         {{- if .Values.schedulerName }}
         - "--scheduler-name={{ .Values.schedulerName }}"
         {{- end }}
+        {{- if .Values.availabilityZoneLabel }}
+        - "--orchestrator-kubernetes-az-label={{ .Values.availabilityZoneLabel }}"
+        {{- else if eq .Values.operator.cloudProvider.type "aws" }}
+        - "--orchestrator-kubernetes-az-label=topology.k8s.aws/zone-id"
+        {{- end }}
 
         {{- if .Values.operator.additionalMaterializeCRDColumns }}
         - >

misc/helm-charts/operator/values.yaml

@@ -331,6 +331,11 @@ rbac:
 # -- Optionally use a non-default kubernetes scheduler.
 schedulerName:
 
+# -- The Kubernetes node label key used for availability zone affinity rules.
+# When unset, defaults to topology.k8s.aws/zone-id for AWS cloud providers
+# and topology.kubernetes.io/zone otherwise.
+availabilityZoneLabel:
+
 # Service account settings
 serviceAccount:
   # -- Whether to create a new service account for the operator

src/environmentd/src/environmentd/main.rs

@@ -249,6 +249,17 @@ pub struct Args {
         env = "ORCHESTRATOR_KUBERNETES_ENABLE_PROMETHEUS_SCRAPE_ANNOTATIONS"
     )]
     orchestrator_kubernetes_enable_prometheus_scrape_annotations: bool,
+    /// The Kubernetes node label key used to identify availability zones for
+    /// pod scheduling affinity rules.
+    ///
+    /// Defaults to `topology.kubernetes.io/zone`. Use
+    /// `topology.k8s.aws/zone-id` for AWS zone IDs, or any custom label.
+    #[clap(
+        long,
+        env = "ORCHESTRATOR_KUBERNETES_AZ_LABEL",
+        default_value = "topology.kubernetes.io/zone"
+    )]
+    orchestrator_kubernetes_az_label: String,
     #[clap(long, env = "ORCHESTRATOR_PROCESS_WRAPPER")]
     orchestrator_process_wrapper: Option<String>,
     /// Where the process orchestrator should store secrets.
@@ -834,6 +845,7 @@ fn run(mut args: Args) -> Result<(), anyhow::Error> {
                             .orchestrator_kubernetes_disable_pod_metrics_collection,
                         enable_prometheus_scrape_annotations: args
                             .orchestrator_kubernetes_enable_prometheus_scrape_annotations,
+                        availability_zone_label: args.orchestrator_kubernetes_az_label.clone(),
                     }))
                     .context("creating kubernetes orchestrator")?,
             );

src/orchestrator-kubernetes/src/lib.rs

@@ -109,6 +109,12 @@ pub struct KubernetesOrchestratorConfig {
     pub collect_pod_metrics: bool,
     /// Whether to annotate pods for prometheus service discovery.
     pub enable_prometheus_scrape_annotations: bool,
+    /// The Kubernetes node label key used to identify availability zones.
+    ///
+    /// Defaults to the standard `topology.kubernetes.io/zone`. Can be set to
+    /// `topology.k8s.aws/zone-id` for AWS-specific zone IDs, or any other
+    /// topology label.
+    pub availability_zone_label: String,
 }
 
 impl KubernetesOrchestratorConfig {
@@ -879,7 +885,7 @@ impl NamespacedOrchestrator for NamespacedKubernetesOrchestrator {
         let node_affinity = if let Some(availability_zones) = availability_zones {
             let selector = NodeSelectorTerm {
                 match_expressions: Some(vec![NodeSelectorRequirement {
-                    key: "materialize.cloud/availability-zone".to_string(),
+                    key: self.config.availability_zone_label.clone(),
                     operator: "In".to_string(),
                     values: Some(availability_zones),
                 }]),

src/orchestratord/src/bin/orchestratord.rs

@@ -106,6 +106,8 @@ pub struct Args {
     #[clap(long)]
     scheduler_name: Option<String>,
     #[clap(long)]
+    orchestrator_kubernetes_az_label: Option<String>,
+    #[clap(long)]
     enable_security_context: bool,
     #[clap(long)]
     enable_internal_statement_logging: bool,
@@ -336,6 +338,7 @@ async fn run(args: Args) -> Result<(), anyhow::Error> {
                     environmentd_availability_zones: args.environmentd_availability_zones,
                     ephemeral_volume_class: args.ephemeral_volume_class,
                     scheduler_name: args.scheduler_name.clone(),
+                    orchestrator_kubernetes_az_label: args.orchestrator_kubernetes_az_label.clone(),
                     enable_security_context: args.enable_security_context,
                     enable_internal_statement_logging: args.enable_internal_statement_logging,
                     disable_statement_logging: args.disable_statement_logging,

src/orchestratord/src/controller/materialize.rs

@@ -76,6 +76,7 @@ pub struct Config {
 
     pub ephemeral_volume_class: Option<String>,
     pub scheduler_name: Option<String>,
+    pub orchestrator_kubernetes_az_label: Option<String>,
     pub enable_security_context: bool,
     pub enable_internal_statement_logging: bool,
     pub disable_statement_logging: bool,

src/orchestratord/src/controller/materialize/generation.rs

@@ -785,6 +785,9 @@ fn create_environmentd_statefulset_object(
             scheduler_name
         ));
     }
+    if let Some(az_label) = &config.orchestrator_kubernetes_az_label {
+        args.push(format!("--orchestrator-kubernetes-az-label={}", az_label));
+    }
     if mz.meets_minimum_version(&V154_DEV0) {
         args.extend(
             mz.spec