Merge pull request #33670 from alex-hunt-materialize/configurable_balancerd_pods

support configuring replicas for balancerd and console in MZ CR

Author: alex-hunt-materialize

src/cloud-resources/src/crd/materialize.rs

@@ -99,6 +99,10 @@ pub mod v1alpha1 {
         pub balancerd_resource_requirements: Option<ResourceRequirements>,
         // Resource requirements for the console pod
         pub console_resource_requirements: Option<ResourceRequirements>,
+        // Number of balancerd pods to create
+        pub balancerd_replicas: Option<i32>,
+        // Number of console pods to create
+        pub console_replicas: Option<i32>,
 
         // Name of the kubernetes service account to use.
         // If not set, we will create one with the same name as this Materialize object.
@@ -268,6 +272,14 @@ pub mod v1alpha1 {
             self.name_prefixed("balancerd-external-tls")
         }
 
+        pub fn balancerd_replicas(&self) -> i32 {
+            self.spec.balancerd_replicas.unwrap_or(2)
+        }
+
+        pub fn console_replicas(&self) -> i32 {
+            self.spec.console_replicas.unwrap_or(2)
+        }
+
         pub fn console_configmap_name(&self) -> String {
             self.name_prefixed("console")
         }

src/orchestratord/src/controller/materialize/balancer.rs

@@ -324,7 +324,7 @@ fn create_balancerd_deployment_object(
     };
 
     let deployment_spec = DeploymentSpec {
-        replicas: Some(1),
+        replicas: Some(mz.balancerd_replicas()),
         selector: LabelSelector {
             match_labels: Some(pod_template_labels.clone()),
             ..Default::default()

src/orchestratord/src/controller/materialize/console.rs

@@ -399,7 +399,7 @@ ssl_certificate_key /nginx/tls/tls.key;",
     };
 
     let deployment_spec = DeploymentSpec {
-        replicas: Some(2),
+        replicas: Some(mz.console_replicas()),
         selector: LabelSelector {
             match_labels: Some(pod_template_labels.clone()),
             ..Default::default()

test/orchestratord/mzcompose.py

@@ -56,57 +56,48 @@ def get_image(image: str, tag: str | None) -> str:
     return f'{image.rsplit(":", 1)[0]}:{tag}'
 
 
-def get_orchestratord_data() -> dict[str, Any]:
+def get_pod_data(
+    labels: dict[str, str], namespace="materialize-environment"
+) -> dict[str, Any]:
     return json.loads(
         spawn.capture(
             [
                 "kubectl",
                 "get",
                 "pod",
                 "-l",
-                "app.kubernetes.io/instance=operator",
+                ",".join(f"{key}={value}" for key, value in labels.items()),
                 "-n",
-                "materialize",
+                namespace,
                 "-o",
                 "json",
             ]
         )
     )
 
 
+def get_orchestratord_data() -> dict[str, Any]:
+    return get_pod_data(
+        labels={"app.kubernetes.io/instance": "operator"},
+        namespace="materialize",
+    )
+
+
 def get_balancerd_data() -> dict[str, Any]:
-    return json.loads(
-        spawn.capture(
-            [
-                "kubectl",
-                "get",
-                "pod",
-                "-l",
-                "app=balancerd",
-                "-n",
-                "materialize-environment",
-                "-o",
-                "json",
-            ]
-        )
+    return get_pod_data(
+        labels={"materialize.cloud/app": "balancerd"},
+    )
+
+
+def get_console_data() -> dict[str, Any]:
+    return get_pod_data(
+        labels={"materialize.cloud/app": "console"},
     )
 
 
 def get_environmentd_data() -> dict[str, Any]:
-    return json.loads(
-        spawn.capture(
-            [
-                "kubectl",
-                "get",
-                "pod",
-                "-l",
-                "app=environmentd",
-                "-n",
-                "materialize-environment",
-                "-o",
-                "json",
-            ]
-        )
+    return get_pod_data(
+        labels={"materialize.cloud/app": "environmentd"},
     )
 
 
@@ -528,6 +519,63 @@ def validate(self, mods: dict[type[Modification], Any]) -> None:
             ), f"Expected no {expected} in environmentd args, but found it: {args}"
 
 
+class BalancerdReplicas(Modification):
+    @classmethod
+    def values(cls) -> list[Any]:
+        return [None, 1, 2]
+
+    @classmethod
+    def default(cls) -> Any:
+        return None
+
+    def modify(self, definition: dict[str, Any]) -> None:
+        if self.value is not None:
+            definition["materialize"]["spec"]["balancerdReplicas"] = self.value
+
+    def validate(self, mods: dict[type[Modification], Any]) -> None:
+        if not mods[BalancerdEnabled]:
+            return
+
+        def check_replicas():
+            balancerd = get_balancerd_data()
+            num_pods = len(balancerd["items"])
+            expected = self.value if self.value is not None else 2
+            assert (
+                num_pods == expected
+            ), f"Expected {expected} balancerd pods, but found {num_pods}"
+
+        retry(check_replicas, 120)
+
+
+class ConsoleReplicas(Modification):
+    @classmethod
+    def values(cls) -> list[Any]:
+        return [None, 1, 2]
+
+    @classmethod
+    def default(cls) -> Any:
+        return None
+
+    def modify(self, definition: dict[str, Any]) -> None:
+        if self.value is not None:
+            definition["materialize"]["spec"]["consoleReplicas"] = self.value
+
+    def validate(self, mods: dict[type[Modification], Any]) -> None:
+        if not mods[ConsoleEnabled]:
+            return
+
+        def check_replicas():
+            console = get_console_data()
+            num_pods = len(console["items"])
+            expected = self.value if self.value is not None else 2
+            assert (
+                num_pods == expected
+            ), f"Expected {expected} console pods, but found {num_pods}"
+
+        # console doesn't get launched until last
+        retry(check_replicas, 120)
+
+
 def validate_cluster_replica_size(
     size: dict[str, Any], swap_enabled: bool, storage_class_name_set: bool
 ):
@@ -571,24 +619,6 @@ def validate_container_resources(
         assert resources["requests"]["memory"] == resources["limits"]["memory"]
 
 
-def get_pod_data(labels: dict[str, str]) -> dict[str, Any]:
-    return json.loads(
-        spawn.capture(
-            [
-                "kubectl",
-                "get",
-                "pod",
-                "-l",
-                ",".join(f"{key}={value}" for key, value in labels.items()),
-                "-n",
-                "materialize-environment",
-                "-o",
-                "json",
-            ]
-        )
-    )
-
-
 class SwapEnabledGlobal(Modification):
     @classmethod
     def values(cls) -> list[Any]: