orchestratord: allow setting a PriorityClass on environmentd and clusterd pods

misc/helm-charts/operator/README.md

@@ -115,6 +115,7 @@ The following table lists the configurable parameters of the Materialize operato
 | `balancerd.tolerations` | Tolerations to use for balancerd pods spawned by the operator | ``{}`` |
 | `clusterd.affinity` | Affinity to use for clusterd pods spawned by the operator | ``{}`` |
 | `clusterd.nodeSelector` | Node selector to use for all clusterd pods spawned by the operator | ``{}`` |
+| `clusterd.priorityClassName` | PriorityClass to use for clusterd pods spawned by the operator. The PriorityClass must already exist. Kubernetes rejects a pod that names one it cannot resolve, so a typo here stops these pods being created at all. | ``nil`` |
 | `clusterd.scratchfsNodeSelector` | Additional node selector to use for clusterd pods when using an LVM scratch disk. This will be merged with the values in `nodeSelector`. | ``{"materialize.cloud/scratch-fs": "true"}`` |
 | `clusterd.swapNodeSelector` | Additional node selector to use for clusterd pods when using swap. This will be merged with the values in `nodeSelector`. | ``{"materialize.cloud/swap": "true"}`` |
 | `clusterd.tolerations` | Tolerations to use for clusterd pods spawned by the operator | ``{}`` |
@@ -129,6 +130,7 @@ The following table lists the configurable parameters of the Materialize operato
 | `environmentd.defaultResources.limits` | Default resource limits for environmentd's CPU and memory if not set in the Materialize CR | ``{"memory":"4Gi"}`` |
 | `environmentd.defaultResources.requests` | Default resources requested for environmentd's CPU and memory if not set in the Materialize CR | ``{"cpu":"1","memory":"4095Mi"}`` |
 | `environmentd.nodeSelector` | Node selector to use for environmentd pods spawned by the operator | ``{}`` |
+| `environmentd.priorityClassName` | PriorityClass to use for environmentd pods spawned by the operator. The PriorityClass must already exist. Kubernetes rejects a pod that names one it cannot resolve, so a typo here stops these pods being created at all. | ``nil`` |
 | `environmentd.tolerations` | Tolerations to use for environmentd pods spawned by the operator | ``{}`` |
 | `networkPolicies.egress.cidrs` | CIDR blocks to allow egress to | ``["0.0.0.0/0"]`` |
 | `networkPolicies.egress.enabled` | Whether to enable egress network policies to sources and sinks | ``false`` |

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

@@ -192,6 +192,9 @@ spec:
         - '--environmentd-toleration={{ toJson $toleration }}'
         {{- end }}
         {{- end }}
+        {{- if .Values.environmentd.priorityClassName }}
+        - "--environmentd-priority-class-name={{ .Values.environmentd.priorityClassName }}"
+        {{- end }}
         {{- if .Values.environmentd.defaultResources }}
         - '--environmentd-default-resources={{ toJson .Values.environmentd.defaultResources }}'
         {{- end }}
@@ -208,6 +211,9 @@ spec:
         - '--clusterd-toleration={{ toJson $toleration }}'
         {{- end }}
         {{- end }}
+        {{- if .Values.clusterd.priorityClassName }}
+        - "--clusterd-priority-class-name={{ .Values.clusterd.priorityClassName }}"
+        {{- end }}
         {{- if .Values.balancerd.nodeSelector }}
         {{- range $key, $value := .Values.balancerd.nodeSelector }}
         - "--balancerd-node-selector={{ $key }}={{ $value }}"

misc/helm-charts/operator/tests/deployment_test.yaml

@@ -274,9 +274,9 @@ tests:
 
 - it: should not pass the scheduler when not configured
   asserts:
-  - notContains:
-      path: spec.template.spec.containers[0].args
-      content: "--scheduler-name"
+  - notMatchRegex:
+      path: spec.template.spec.containers[0].args[*]
+      pattern: "^--scheduler-name"
 
 - it: should pass the scheduler when configured
   set:
@@ -286,6 +286,34 @@ tests:
       path: spec.template.spec.containers[0].args
       content: "--scheduler-name=my-scheduler"
 
+- it: should not pass the environmentd priority class when not configured
+  asserts:
+  - notMatchRegex:
+      path: spec.template.spec.containers[0].args[*]
+      pattern: "^--environmentd-priority-class-name"
+
+- it: should pass the environmentd priority class when configured
+  set:
+    environmentd.priorityClassName: mz-environmentd-high
+  asserts:
+  - contains:
+      path: spec.template.spec.containers[0].args
+      content: "--environmentd-priority-class-name=mz-environmentd-high"
+
+- it: should not pass the clusterd priority class when not configured
+  asserts:
+  - notMatchRegex:
+      path: spec.template.spec.containers[0].args[*]
+      pattern: "^--clusterd-priority-class-name"
+
+- it: should pass the clusterd priority class when configured
+  set:
+    clusterd.priorityClassName: mz-clusterd-high
+  asserts:
+  - contains:
+      path: spec.template.spec.containers[0].args
+      content: "--clusterd-priority-class-name=mz-clusterd-high"
+
 - it: should have license key checks disabled by default
   asserts:
   - contains:

misc/helm-charts/operator/values.yaml

@@ -306,6 +306,10 @@ environmentd:
   # @notationType -- k8s/tolerations
   # @default -- ``{}``
   tolerations: {}
+  # -- PriorityClass to use for environmentd pods spawned by the operator. The
+  # PriorityClass must already exist. Kubernetes rejects a pod that names one
+  # it cannot resolve, so a typo here stops these pods being created at all.
+  priorityClassName:
   # @notationType -- k8s/resources
   defaultResources:
     # -- Default resource limits for environmentd's CPU and memory
@@ -345,6 +349,10 @@ clusterd:
   # @notationType -- k8s/tolerations
   # @default -- ``{}``
   tolerations: {}
+  # -- PriorityClass to use for clusterd pods spawned by the operator. The
+  # PriorityClass must already exist. Kubernetes rejects a pod that names one
+  # it cannot resolve, so a typo here stops these pods being created at all.
+  priorityClassName:
 
 balancerd:
   # -- Flag to indicate whether to create balancerd pods for the environments

src/environmentd/src/environmentd/main.rs

@@ -181,6 +181,9 @@ pub struct Args {
     /// Name of a non-default Kubernetes scheduler, if any.
     #[structopt(long, env = "ORCHESTRATOR_KUBERNETES_SCHEDULER_NAME")]
     orchestrator_kubernetes_scheduler_name: Option<String>,
+    /// Name of a `PriorityClass` to assign to services, if any.
+    #[structopt(long, env = "ORCHESTRATOR_KUBERNETES_PRIORITY_CLASS_NAME")]
+    orchestrator_kubernetes_priority_class_name: Option<String>,
     /// Annotations to apply to all services created by the Kubernetes orchestrator
     /// in the form `KEY=VALUE`.
     #[structopt(long, env = "ORCHESTRATOR_KUBERNETES_SERVICE_ANNOTATION")]
@@ -822,6 +825,7 @@ fn run(mut args: Args) -> Result<(), anyhow::Error> {
                     .block_on(KubernetesOrchestrator::new(KubernetesOrchestratorConfig {
                         context: args.orchestrator_kubernetes_context.clone(),
                         scheduler_name: args.orchestrator_kubernetes_scheduler_name,
+                        priority_class_name: args.orchestrator_kubernetes_priority_class_name,
                         service_annotations: args
                             .orchestrator_kubernetes_service_annotation
                             .into_iter()

src/orchestrator-kubernetes/src/lib.rs

@@ -76,6 +76,8 @@ pub struct KubernetesOrchestratorConfig {
     pub context: String,
     /// The name of a non-default Kubernetes scheduler to use, if any.
     pub scheduler_name: Option<String>,
+    /// The name of a `PriorityClass` to assign to services, if any.
+    pub priority_class_name: Option<String>,
     /// Annotations to install on every service created by the orchestrator.
     pub service_annotations: BTreeMap<String, String>,
     /// Labels to install on every service created by the orchestrator.
@@ -1200,6 +1202,7 @@ impl NamespacedOrchestrator for NamespacedKubernetesOrchestrator {
                 security_context,
                 node_selector: Some(node_selector),
                 scheduler_name: self.config.scheduler_name.clone(),
+                priority_class_name: self.config.priority_class_name.clone(),
                 service_account: self.config.service_account.clone(),
                 affinity: Some(affinity),
                 topology_spread_constraints: topology_spread,

src/orchestratord/src/bin/orchestratord.rs

@@ -236,12 +236,16 @@ pub struct Args {
     #[clap(long, value_parser = parse_resources)]
     environmentd_default_resources: Option<ResourceRequirements>,
     #[clap(long)]
+    environmentd_priority_class_name: Option<String>,
+    #[clap(long)]
     clusterd_node_selector: Vec<KeyValueArg<String, String>>,
     #[clap(long, value_parser = parse_affinity)]
     clusterd_affinity: Option<Affinity>,
     #[clap(long = "clusterd-toleration", value_parser = parse_tolerations)]
     clusterd_tolerations: Option<Vec<Toleration>>,
     #[clap(long)]
+    clusterd_priority_class_name: Option<String>,
+    #[clap(long)]
     balancerd_node_selector: Vec<KeyValueArg<String, String>>,
     #[clap(long, value_parser = parse_affinity)]
     balancerd_affinity: Option<Affinity>,
@@ -666,9 +670,11 @@ async fn run(args: Args) -> Result<(), anyhow::Error> {
             environmentd_affinity: args.environmentd_affinity,
             environmentd_tolerations: args.environmentd_tolerations,
             environmentd_default_resources: args.environmentd_default_resources,
+            environmentd_priority_class_name: args.environmentd_priority_class_name,
             clusterd_node_selector: args.clusterd_node_selector,
             clusterd_affinity: args.clusterd_affinity,
             clusterd_tolerations: args.clusterd_tolerations,
+            clusterd_priority_class_name: args.clusterd_priority_class_name,
             image_pull_policy: args.image_pull_policy,
             network_policies_internal_enabled: args.network_policies_internal_enabled,
             network_policies_ingress_enabled: args.network_policies_ingress_enabled,

src/orchestratord/src/controller/materialize.rs

@@ -89,9 +89,13 @@ pub struct Config {
     pub environmentd_affinity: Option<Affinity>,
     pub environmentd_tolerations: Option<Vec<Toleration>>,
     pub environmentd_default_resources: Option<ResourceRequirements>,
+    /// The name of a `PriorityClass` to assign to environmentd pods, if any.
+    pub environmentd_priority_class_name: Option<String>,
     pub clusterd_node_selector: Vec<KeyValueArg<String, String>>,
     pub clusterd_affinity: Option<Affinity>,
     pub clusterd_tolerations: Option<Vec<Toleration>>,
+    /// The name of a `PriorityClass` to assign to clusterd pods, if any.
+    pub clusterd_priority_class_name: Option<String>,
     pub image_pull_policy: KubernetesImagePullPolicy,
     pub network_policies_internal_enabled: bool,
     pub network_policies_ingress_enabled: bool,

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

@@ -73,6 +73,18 @@ static V154_DEV0: LazyLock<Version> = LazyLock::new(|| Version {
 });
 pub const V161: Version = Version::new(0, 161, 0);
 
+/// Version at which `environmentd` learned
+/// `--orchestrator-kubernetes-priority-class-name`. Older images reject the
+/// unknown argument and fail to start, so the flag is only forwarded at or
+/// above this version.
+static V26_40_0: LazyLock<Version> = LazyLock::new(|| Version {
+    major: 26,
+    minor: 40,
+    patch: 0,
+    pre: Prerelease::new("dev.0").expect("dev.0 is valid prerelease"),
+    build: BuildMetadata::new("").expect("empty string is valid buildmetadata"),
+});
+
 static V26_1_0: LazyLock<Version> = LazyLock::new(|| Version {
     major: 26,
     minor: 1,
@@ -804,6 +816,13 @@ fn create_environmentd_statefulset_object(
             scheduler_name
         ));
     }
+    if mz.meets_minimum_version(&V26_40_0) {
+        if let Some(priority_class_name) = &config.clusterd_priority_class_name {
+            args.push(format!(
+                "--orchestrator-kubernetes-priority-class-name={priority_class_name}"
+            ));
+        }
+    }
     if mz.meets_minimum_version(&V154_DEV0) {
         args.extend(
             mz.spec
@@ -1246,6 +1265,7 @@ fn create_environmentd_statefulset_object(
             ),
             affinity: config.environmentd_affinity.clone(),
             scheduler_name: config.scheduler_name.clone(),
+            priority_class_name: config.environmentd_priority_class_name.clone(),
             service_account_name: Some(mz.service_account_name()),
             volumes: Some(volumes),
             security_context: Some(PodSecurityContext {

test/orchestratord/mzcompose.py

@@ -768,6 +768,84 @@ def check() -> None:
         retry(check, 240)
 
 
+# Must match test/orchestratord/priorityclass.yaml.
+PRIORITY_CLASS_NAME = "mz-test-priority"
+PRIORITY_CLASS_VALUE = 1000000000
+# Release in which environmentd learned
+# --orchestrator-kubernetes-priority-class-name.
+PRIORITY_CLASS_VERSION = "v26.40.0"
+
+
+def assert_priority_class(pod: dict[str, Any], expected: str | None) -> None:
+    """Assert a pod's priority class by name and by resolved value.
+
+    `spec.priority` is filled in by the API server from the named class, so
+    checking it is what distinguishes a class Kubernetes actually resolved from
+    a string the operator copied into the spec.
+    """
+    spec = pod["spec"]
+    name = pod["metadata"]["name"]
+    actual = spec.get("priorityClassName")
+    priority = spec.get("priority")
+    if expected is None:
+        assert not actual, f"{name}: unexpected priorityClassName {actual}"
+        assert not priority, f"{name}: unexpected priority {priority}"
+    else:
+        assert actual == expected, f"{name}: expected {expected}, got {actual}"
+        assert (
+            priority == PRIORITY_CLASS_VALUE
+        ), f"{name}: expected priority {PRIORITY_CLASS_VALUE}, got {priority}"
+
+
+class PriorityClassName(Modification):
+    @classmethod
+    def values(cls, version: MzVersion) -> list[Any]:
+        return [None, PRIORITY_CLASS_NAME]
+
+    @classmethod
+    def default(cls) -> Any:
+        return None
+
+    def modify(self, definition: dict[str, Any]) -> None:
+        if not self.value:
+            return
+        definition["operator"]["environmentd"]["priorityClassName"] = self.value
+        definition["operator"]["clusterd"]["priorityClassName"] = self.value
+        # environmentd and clusterd only pick up generation-affecting changes on
+        # a requested rollout, so without this the pods keep the spec they were
+        # created with and the assertions below check the wrong generation.
+        request = str(uuid.uuid4())
+        if definition["materialize"]["apiVersion"] == "materialize.cloud/v1alpha1":
+            definition["materialize"]["spec"]["requestRollout"] = request
+        definition["materialize"]["spec"]["forceRollout"] = request
+
+    def validate(self, mods: dict[type[Modification], Any]) -> None:
+        version = MzVersion.parse_mz(mods[EnvironmentdImageRef])
+        # The operator sets environmentd's class directly, but forwards
+        # clusterd's to environmentd behind a version gate, so an older image
+        # gets one and not the other.
+        clusterd_expected = (
+            self.value
+            if version >= MzVersion.parse_mz(PRIORITY_CLASS_VERSION)
+            else None
+        )
+
+        def check() -> None:
+            environmentd = get_environmentd_data()["items"]
+            clusterd = get_clusterd_data()["items"]
+            # A class the API server cannot resolve yields no pod at all, which
+            # an assertion over an empty list would otherwise call a pass.
+            assert environmentd, "no environmentd pods"
+            assert clusterd, "no clusterd pods"
+            for pod in environmentd:
+                assert_priority_class(pod, self.value)
+            for pod in clusterd:
+                assert_priority_class(pod, clusterd_expected)
+
+        # Clusterd is recreated by the rollout and can take a while
+        retry(check, 240)
+
+
 class ConsoleEnabled(Modification):
     @classmethod
     def values(cls, version: MzVersion) -> list[Any]:
@@ -4571,6 +4649,14 @@ def setup(c: Composition, args) -> dict[str, Any]:
                 MZ_ROOT / "test" / "orchestratord" / "storageclass.yaml",
             ]
         )
+        spawn.runv(
+            [
+                "kubectl",
+                "apply",
+                "-f",
+                MZ_ROOT / "test" / "orchestratord" / "priorityclass.yaml",
+            ]
+        )
 
     if not args.tag:
         services = [

test/orchestratord/priorityclass.yaml

@@ -0,0 +1,12 @@
+# test/orchestratord/priorityclass.yaml, applied beside storageclass.yaml (:4571).
+# `value` is required. `globalDefault: false` keeps the unset baseline testable:
+# a global default would put a class on every pod, including the None case.
+# `preemptionPolicy: Never` keeps the test from evicting anything in the cluster.
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: mz-test-priority
+value: 1000000000
+globalDefault: false
+preemptionPolicy: Never
+description: "Test fixture for the priorityClassName chart values."