[API] Add preemption nodes info to frontend spec (#1880)

mlrun/api/api/endpoints/frontend_spec.py

@@ -65,6 +65,7 @@ def get_frontend_spec(
         auto_mount_params=config.get_storage_auto_mount_params(),
         default_artifact_path=config.artifact_path,
         default_function_pod_resources=mlrun.mlconf.default_function_pod_resources.to_dict(),
+        default_function_preemption_mode=mlrun.mlconf.function_defaults.preemption_mode,
     )
 
 
@@ -93,8 +94,14 @@ def _resolve_feature_flags() -> mlrun.api.schemas.FeatureFlags:
         mlrun.runtimes.utils.resolve_nuclio_version()
     ) >= semver.VersionInfo.parse("1.7.8"):
         nuclio_streams = mlrun.api.schemas.NuclioStreamsFeatureFlag.enabled
+
+    preemption_nodes = mlrun.api.schemas.PreemptionNodesFeatureFlag.disabled
+    if mlrun.mlconf.is_preemption_nodes_configured():
+        preemption_nodes = mlrun.api.schemas.PreemptionNodesFeatureFlag.enabled
+
     return mlrun.api.schemas.FeatureFlags(
         project_membership=project_membership,
         authentication=authentication,
         nuclio_streams=nuclio_streams,
+        preemption_nodes=preemption_nodes,
     )

mlrun/api/schemas/__init__.py

@@ -55,6 +55,7 @@
     FeatureFlags,
     FrontendSpec,
     NuclioStreamsFeatureFlag,
+    PreemptionNodesFeatureFlag,
     ProjectMembershipFeatureFlag,
 )
 from .function import FunctionState, PreemptionModes

mlrun/api/schemas/frontend_spec.py

@@ -11,6 +11,11 @@ class ProjectMembershipFeatureFlag(str, enum.Enum):
     disabled = "disabled"
 
 
+class PreemptionNodesFeatureFlag(str, enum.Enum):
+    enabled = "enabled"
+    disabled = "disabled"
+
+
 class AuthenticationFeatureFlag(str, enum.Enum):
     none = "none"
     basic = "basic"
@@ -27,6 +32,7 @@ class FeatureFlags(pydantic.BaseModel):
     project_membership: ProjectMembershipFeatureFlag
     authentication: AuthenticationFeatureFlag
     nuclio_streams: NuclioStreamsFeatureFlag
+    preemption_nodes: PreemptionNodesFeatureFlag
 
 
 class FrontendSpec(pydantic.BaseModel):
@@ -44,3 +50,4 @@ class FrontendSpec(pydantic.BaseModel):
     auto_mount_params: typing.Dict[str, str] = {}
     default_artifact_path: str
     default_function_pod_resources: Resources = Resources()
+    default_function_preemption_mode: str

mlrun/config.py

@@ -467,6 +467,14 @@ def get_preemptible_tolerations(self) -> list:
             "preemptible_nodes.tolerations", list
         )
 
+    def is_preemption_nodes_configured(self):
+        if (
+            not self.get_preemptible_tolerations()
+            and not self.get_preemptible_node_selector()
+        ):
+            return False
+        return True
+
     @staticmethod
     def get_valid_function_priority_class_names():
         valid_function_priority_class_names = []

mlrun/runtimes/pod.py

@@ -506,10 +506,7 @@ def enrich_function_preemption_spec(self):
                         > Adds anti-affinity IF no tolerations were configured
         """
         # nothing to do here, configuration is not populated
-        if (
-            not mlconf.get_preemptible_tolerations()
-            and not mlconf.get_preemptible_node_selector()
-        ):
+        if not mlconf.is_preemption_nodes_configured():
             return
 
         if not self.preemption_mode:

tests/api/api/test_frontend_spec.py

@@ -46,6 +46,10 @@ def test_get_frontend_spec(
         frontend_spec.feature_flags.nuclio_streams
         == mlrun.api.schemas.NuclioStreamsFeatureFlag.disabled
     )
+    assert (
+        frontend_spec.feature_flags.preemption_nodes
+        == mlrun.api.schemas.PreemptionNodesFeatureFlag.disabled
+    )
     assert frontend_spec.default_function_image_by_kind is not None
     assert frontend_spec.function_deployment_mlrun_command is not None
     assert frontend_spec.default_artifact_path is not None
@@ -70,6 +74,11 @@ def test_get_frontend_spec(
         == mlrun.runtimes.utils.resolve_function_target_image_registries_to_enforce_prefix()
     )
 
+    assert (
+        frontend_spec.default_function_preemption_mode
+        == mlrun.api.schemas.PreemptionModes.prevent.value
+    )
+
 
 def test_get_frontend_spec_jobs_dashboard_url_resolution(
     db: sqlalchemy.orm.Session, client: fastapi.testclient.TestClient