Make it possible to specify topology_spread_constraints per service/i…

…nstance. PAASTA-18132

paasta_tools/cli/schemas/kubernetes_schema.json

@@ -866,6 +866,29 @@
                     "type": "string",
                     "pattern": "^[a-z0-9]([-a-z0-9]*[a-z0-9])?$",
                     "maxLength": 63
+                },
+                "topology_spread_constraints": {
+                    "type": "array",
+                    "items": {
+                        "type": "object",
+                        "properties": {
+                            "topology_key": {
+                                "type": "string"
+                            },
+                            "when_unsatisfiable": {
+                                "type": "string",
+                                "enum": [
+                                    "ScheduleAnyway",
+                                    "DoNotSchedule"
+                                ]
+                            },
+                            "max_skew": {
+                                "type": "integer"
+                            }
+                        },
+                        "required": []
+                    },
+                    "uniqueItems": true
                 }
             }
         }

paasta_tools/kubernetes_tools.py

@@ -161,6 +161,7 @@
 from paasta_tools.utils import SecretVolume
 from paasta_tools.utils import SystemPaastaConfig
 from paasta_tools.utils import time_cache
+from paasta_tools.utils import TopologySpreadConstraintDict
 from paasta_tools.utils import VolumeWithMode
 
 
@@ -380,6 +381,7 @@ class KubernetesDeploymentConfigDict(LongRunningServiceConfigDict, total=False):
     boto_keys: List[str]
     crypto_keys: CryptoKeyConfig
     datastore_credentials: DatastoreCredentialsConfig
+    topology_spread_constraints: List[TopologySpreadConstraintDict]
 
 
 def load_kubernetes_service_config_no_cache(
@@ -2167,7 +2169,9 @@ def get_pod_template_spec(
         pod_topology_spread_constraints = create_pod_topology_spread_constraints(
             service=self.get_service(),
             instance=self.get_instance(),
-            topology_spread_constraints=system_paasta_config.get_topology_spread_constraints(),
+            topology_spread_constraints=self.get_topology_spread_constraints(
+                system_paasta_config.get_topology_spread_constraints()
+            ),
         )
         if pod_topology_spread_constraints:
             constraints = pod_spec_kwargs.get("topology_spread_constraints", [])
@@ -2466,6 +2470,14 @@ def get_prometheus_path(self) -> Optional[str]:
     def get_prometheus_port(self) -> Optional[int]:
         return self.config_dict.get("prometheus_port")
 
+    def get_topology_spread_constraints(
+        self,
+        default_pod_topology_spread_constraints: List[TopologySpreadConstraintDict],
+    ) -> List[TopologySpreadConstraintDict]:
+        return self.config_dict.get(
+            "topology_spread_constraints", default_pod_topology_spread_constraints
+        )
+
 
 def get_kubernetes_secret_hashes(
     environment_variables: Mapping[str, str], service: str, namespace: str
@@ -3812,7 +3824,7 @@ def load_custom_resource_definitions(
 def create_pod_topology_spread_constraints(
     service: str,
     instance: str,
-    topology_spread_constraints: List[Dict[str, Any]],
+    topology_spread_constraints: List[TopologySpreadConstraintDict],
 ) -> List[V1TopologySpreadConstraint]:
     """
     Applies cluster-level topology spread constraints to every Pod template.

paasta_tools/utils.py

@@ -57,6 +57,7 @@
 from typing import Iterable
 from typing import Iterator
 from typing import List
+from typing import Literal
 from typing import Mapping
 from typing import NamedTuple
 from typing import Optional
@@ -1929,6 +1930,12 @@ class KubeStateMetricsCollectorConfigDict(TypedDict, total=False):
     label_renames: Dict[str, str]
 
 
+class TopologySpreadConstraintDict(TypedDict, total=False):
+    topology_key: str
+    when_unsatisfiable: Literal["ScheduleAnyway", "DoNotSchedule"]
+    max_skew: int
+
+
 class SystemPaastaConfigDict(TypedDict, total=False):
     allowed_pools: Dict[str, List[str]]
     api_endpoints: Dict[str, str]
@@ -2008,7 +2015,7 @@ class SystemPaastaConfigDict(TypedDict, total=False):
     pdb_max_unavailable: Union[str, int]
     pki_backend: str
     pod_defaults: Dict[str, Any]
-    topology_spread_constraints: List[Dict[str, Any]]
+    topology_spread_constraints: List[TopologySpreadConstraintDict]
     previous_marathon_servers: List[MarathonConfigDict]
     readiness_check_prefix_template: List[str]
     register_k8s_pods: bool
@@ -2610,7 +2617,7 @@ def get_taskproc(self) -> Dict:
     def get_disabled_watchers(self) -> List:
         return self.config_dict.get("disabled_watchers", [])
 
-    def get_topology_spread_constraints(self) -> List[Dict[str, Any]]:
+    def get_topology_spread_constraints(self) -> List[TopologySpreadConstraintDict]:
         """List of TopologySpreadConstraints that will be applied to all Pods in the cluster"""
         return self.config_dict.get("topology_spread_constraints", [])
 

tests/test_kubernetes_tools.py

@@ -2,6 +2,7 @@
 from base64 import b64encode
 from typing import Any
 from typing import Dict
+from typing import List
 from typing import Sequence
 
 import asynctest
@@ -150,6 +151,7 @@
 from paasta_tools.utils import SecretVolume
 from paasta_tools.utils import SecretVolumeItem
 from paasta_tools.utils import SystemPaastaConfig
+from paasta_tools.utils import TopologySpreadConstraintDict
 
 
 def test_force_delete_pods():
@@ -1991,7 +1993,7 @@ def test_routable_ip(
         assert ret == expected
 
     def test_create_pod_topology_spread_constraints(self):
-        configured_constraints = [
+        configured_constraints: List[TopologySpreadConstraintDict] = [
             {
                 "topology_key": "kubernetes.io/hostname",
                 "max_skew": 1,