[Runtimes] Fix failing build as a result of wrong affinity type

mlrun/k8s_utils.py

@@ -626,7 +626,7 @@ def generate_preemptible_node_selector_requirements(
             kubernetes.client.V1NodeSelectorRequirement(
                 key=node_selector_key,
                 operator=node_selector_operator,
-                values=node_selector_value,
+                values=[node_selector_value],
             )
         )
     return match_expressions
@@ -648,7 +648,7 @@ def generate_preemptible_nodes_anti_affinity_terms() -> typing.List[
 
     # compile affinities with operator NotIn to make sure pods are not running on preemptible nodes.
     node_selector_requirements = generate_preemptible_node_selector_requirements(
-        NodeSelectorOperator.node_selector_op_not_in
+        NodeSelectorOperator.node_selector_op_not_in.value
     )
     return [
         kubernetes.client.V1NodeSelectorTerm(
@@ -673,7 +673,7 @@ def generate_preemptible_nodes_affinity_terms() -> typing.List[
 
     # compile affinities with operator In so pods could schedule on at least one of the preemptible nodes.
     node_selector_requirements = generate_preemptible_node_selector_requirements(
-        NodeSelectorOperator.node_selector_op_in
+        NodeSelectorOperator.node_selector_op_in.value
     )
     for expression in node_selector_requirements:
         node_selector_terms.append(

mlrun/runtimes/pod.py

@@ -523,15 +523,15 @@ def enrich_function_preemption_spec(self):
 
         # remove preemptible tolerations and remove preemption related configuration
         # and enrich with anti-affinity if preemptible tolerations configuration haven't been provided
-        if self.preemption_mode == PreemptionModes.prevent:
+        if self.preemption_mode == PreemptionModes.prevent.value:
             # ensure no preemptible node tolerations
             self._prune_tolerations(generate_preemptible_tolerations())
 
             # purge affinity preemption related configuration
             self._prune_affinity_node_selector_requirement(
                 (
                     generate_preemptible_node_selector_requirements(
-                        NodeSelectorOperator.node_selector_op_in
+                        NodeSelectorOperator.node_selector_op_in.value
                     )
                 )
             )
@@ -546,7 +546,7 @@ def enrich_function_preemption_spec(self):
                 )
 
         # enrich tolerations and override all node selector terms with preemptible node selector terms
-        elif self.preemption_mode == PreemptionModes.constrain:
+        elif self.preemption_mode == PreemptionModes.constrain.value:
             # enrich with tolerations
             self._merge_tolerations(generate_preemptible_tolerations())
 
@@ -560,18 +560,18 @@ def enrich_function_preemption_spec(self):
             )
 
         # purge any affinity / anti-affinity preemption related configuration and enrich with preemptible tolerations
-        elif self.preemption_mode == PreemptionModes.allow:
+        elif self.preemption_mode == PreemptionModes.allow.value:
 
             # remove anti-affinity
             self._prune_affinity_node_selector_requirement(
                 generate_preemptible_node_selector_requirements(
-                    NodeSelectorOperator.node_selector_op_not_in
+                    NodeSelectorOperator.node_selector_op_not_in.value
                 ),
             )
             # remove affinity
             self._prune_affinity_node_selector_requirement(
                 generate_preemptible_node_selector_requirements(
-                    NodeSelectorOperator.node_selector_op_in
+                    NodeSelectorOperator.node_selector_op_in.value
                 ),
             )
 

tests/api/runtimes/test_kubejob.py

@@ -151,25 +151,25 @@ def test_preemptible_modes_transitions(self, db: Session, client: TestClient):
             json.dumps(node_selector).encode("utf-8")
         )
         mlrun.mlconf.function_defaults.preemption_mode = (
-            mlrun.api.schemas.PreemptionModes.prevent
+            mlrun.api.schemas.PreemptionModes.prevent.value
         )
         runtime = self._generate_runtime()
         self._execute_run(runtime)
         self.assert_node_selection(affinity=self._generate_preemptible_anti_affinity())
 
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain.value)
         self._execute_run(runtime)
         self.assert_node_selection(affinity=self._generate_preemptible_affinity())
 
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.allow)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.allow.value)
         self._execute_run(runtime)
         self.assert_node_selection()
 
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain.value)
         self._execute_run(runtime)
         self.assert_node_selection(affinity=self._generate_preemptible_affinity())
 
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.prevent)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.prevent.value)
         self._execute_run(runtime)
         self.assert_node_selection(affinity=self._generate_preemptible_anti_affinity())
 
@@ -179,13 +179,13 @@ def test_preemptible_modes_transitions(self, db: Session, client: TestClient):
         mlrun.mlconf.preemptible_nodes.tolerations = base64.b64encode(
             json.dumps(serialized_tolerations).encode("utf-8")
         )
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain.value)
         self._execute_run(runtime)
         self.assert_node_selection(
             affinity=self._generate_preemptible_affinity(), tolerations=tolerations
         )
 
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.prevent)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.prevent.value)
         self._execute_run(runtime)
         self.assert_node_selection()
 
@@ -239,15 +239,15 @@ def test_run_with_prevent_preemptible_mode(self, db: Session, client: TestClient
             json.dumps(node_selector).encode("utf-8")
         )
         runtime = self._generate_runtime()
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.prevent)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.prevent.value)
         self._execute_run(runtime)
 
         self.assert_node_selection(affinity=self._generate_preemptible_anti_affinity())
 
         # tolerations are set, but expect to stay because those tolerations aren't configured as preemptible tolerations
         runtime = self._generate_runtime()
         runtime.with_node_selection(tolerations=self._generate_tolerations())
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.prevent)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.prevent.value)
         self._execute_run(runtime)
         self.assert_node_selection(
             affinity=self._generate_preemptible_anti_affinity(),
@@ -265,7 +265,7 @@ def test_run_with_prevent_preemptible_mode(self, db: Session, client: TestClient
         runtime.with_node_selection(
             tolerations=self._generate_preemptible_tolerations()
         )
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.prevent)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.prevent.value)
         self._execute_run(runtime)
         self.assert_node_selection()
 
@@ -275,7 +275,7 @@ def test_run_with_constrain_preemptible_mode(self, db: Session, client: TestClie
             json.dumps(node_selector).encode("utf-8")
         )
         runtime = self._generate_runtime()
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain.value)
         self._execute_run(runtime)
         self.assert_node_selection(affinity=self._generate_preemptible_affinity())
         # set default preemptible tolerations
@@ -285,7 +285,7 @@ def test_run_with_constrain_preemptible_mode(self, db: Session, client: TestClie
             json.dumps(serialized_tolerations).encode("utf-8")
         )
         runtime = self._generate_runtime()
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain.value)
         self._execute_run(runtime)
         self.assert_node_selection(
             affinity=self._generate_preemptible_affinity(),
@@ -294,7 +294,7 @@ def test_run_with_constrain_preemptible_mode(self, db: Session, client: TestClie
         # sets different affinity before, expects to override the required_during_scheduling_ignored_during_execution
         runtime = self._generate_runtime()
         runtime.with_node_selection(affinity=self._generate_affinity())
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.constrain.value)
         self._execute_run(runtime)
         expected_affinity = self._generate_affinity()
         expected_affinity.node_affinity.required_during_scheduling_ignored_during_execution = k8s_client.V1NodeSelector(
@@ -312,7 +312,7 @@ def test_run_with_allow_preemptible_mode(self, db: Session, client: TestClient):
         )
         # without default preemptible tolerations, expecting default to apply
         runtime = self._generate_runtime()
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.allow)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.allow.value)
         self._execute_run(runtime)
         self.assert_node_selection()
 
@@ -326,7 +326,7 @@ def test_run_with_allow_preemptible_mode(self, db: Session, client: TestClient):
         self.assert_node_selection()
 
         runtime = self._generate_runtime()
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.allow)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.allow.value)
         self._execute_run(runtime)
         self.assert_node_selection(tolerations=self._generate_preemptible_tolerations())
 
@@ -346,7 +346,7 @@ def test_run_with_allow_preemptible_mode(self, db: Session, client: TestClient):
         )
         runtime = self._generate_runtime()
         runtime.with_node_selection(affinity=affinity)
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.allow)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.allow.value)
         self._execute_run(runtime)
         self.assert_node_selection(
             tolerations=self._generate_preemptible_tolerations(),
@@ -366,7 +366,7 @@ def test_run_with_allow_preemptible_mode(self, db: Session, client: TestClient):
         expected_affinity = self._generate_affinity()
         runtime = self._generate_runtime()
         runtime.with_node_selection(affinity=affinity)
-        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.allow)
+        runtime.with_preemption_mode(mlrun.api.schemas.PreemptionModes.allow.value)
         self._execute_run(runtime)
         self.assert_node_selection(
             tolerations=self._generate_preemptible_tolerations(),