Pass security context for Intel GPU based training jobs (#157)

Author: AlbertvanHouten

interactive_ai/workflows/geti_domain/common/jobs_common/k8s_helpers/trainer_image_info.py

@@ -48,6 +48,7 @@ class TrainerImageInfo:
 
     train_image_name: str
     sidecar_image_name: str
+    render_gid: int = 0  # Should be non-zero value when training with Intel GPUs
 
     @classmethod
     def create(cls, training_framework: TrainingFramework) -> "TrainerImageInfo":
@@ -64,6 +65,8 @@ def create(cls, training_framework: TrainingFramework) -> "TrainerImageInfo":
 
         configmap = asyncio.run(get_config_map(namespace=namespace, name=name))
 
+        render_gid = 0
+
         msg = "Cannot get `{0}` field from config map `{1}/{2}`"
 
         # This information is from `impt-configuration` config map in the namespace `impt`
@@ -73,7 +76,8 @@ def create(cls, training_framework: TrainingFramework) -> "TrainerImageInfo":
             raise ValueError(msg.format("ote_image", namespace, name))
         if (otx2_image := configmap.data.get("otx2_image")) is None:
             raise ValueError(msg.format("otx2_image", namespace, name))
-
+        if render_gid_value := configmap.data.get("render_gid"):
+            render_gid = int(render_gid_value)
         if FeatureFlagProvider.is_enabled(FeatureFlag.FEATURE_FLAG_OTX_VERSION_SELECTION):
             if parse_version(training_framework.version) < parse_version("2.0.0"):
                 image_name = ote_image
@@ -86,7 +90,7 @@ def create(cls, training_framework: TrainingFramework) -> "TrainerImageInfo":
             f"Trainer image has been selected {image_name}, where a model has trainer "
             f"identification for {training_framework.version}."
         )
-        return cls(train_image_name=image_name, sidecar_image_name=mlflow_sidecar_image)
+        return cls(train_image_name=image_name, sidecar_image_name=mlflow_sidecar_image, render_gid=render_gid)
 
     def to_primary_image_full_name(self) -> str:
         """Get primary image full name.

interactive_ai/workflows/geti_domain/common/jobs_common/k8s_helpers/trainer_pod_definition.py

@@ -7,7 +7,7 @@
 import os
 
 from flytekit import ContainerTask, PodTemplate, current_context
-from kubernetes.client import V1PodSpec
+from kubernetes.client import V1Capabilities, V1PodSpec, V1SecurityContext
 from kubernetes.client.models import (
     V1ConfigMapEnvSource,
     V1ConfigMapKeySelector,
@@ -200,6 +200,16 @@ def create_flyte_container_task(  # noqa: PLR0913
     runtime_class_name = "nvidia" if accelerator_name == "nvidia.com/gpu" else None
     logger.info(f"Create runtime_class_name={runtime_class_name}")
 
+    security_context = None
+    if trainer_image_info.render_gid != 0:
+        security_context = V1SecurityContext(
+            run_as_group=trainer_image_info.render_gid,
+            allow_privilege_escalation=False,
+            read_only_root_filesystem=False,
+            run_as_non_root=True,
+            run_as_user=10001,
+            capabilities=V1Capabilities(drop=["ALL"]),
+        )
     role = "flyte_workflows"
 
     pod_spec = V1PodSpec(
@@ -296,6 +306,7 @@ def create_flyte_container_task(  # noqa: PLR0913
                     V1VolumeMount(mount_path="/dev/shm", name="shared-memory"),  # noqa : S108 # nosec: B108
                     V1VolumeMount(mount_path="/shard_files", name="shard-files-dir"),
                 ],
+                security_context=security_context,
             )
         ],
         init_containers=[

interactive_ai/workflows/geti_domain/common/tests/unit/k8s_helpers/test_trainer_image_info.py

@@ -13,31 +13,36 @@
 @pytest.mark.JobsComponent
 class TestTrainerImageInfo:
     @pytest.mark.parametrize(
-        "training_framework, feature_flag_otx_version_selection, primary_image_full_name, sidecar_image_full_name",
+        "training_framework, feature_flag_otx_version_selection, "
+        "primary_image_full_name, sidecar_image_full_name, render_gid",
         [
             (
                 TrainingFramework(type=TrainingFrameworkType.OTX, version="2.1.0"),
                 "false",
                 "otx2_image",
                 "mlflow_sidecar_image",
+                0,
             ),
             (
                 TrainingFramework(type=TrainingFrameworkType.OTX, version="1.6.0"),
                 "false",
                 "otx2_image",
                 "mlflow_sidecar_image",
+                0,
             ),
             (
                 TrainingFramework(type=TrainingFrameworkType.OTX, version="2.1.0"),
                 "true",
                 "otx2_image",
                 "mlflow_sidecar_image",
+                992,
             ),
             (
                 TrainingFramework(type=TrainingFrameworkType.OTX, version="1.6.0"),
                 "true",
                 "ote_image",
                 "mlflow_sidecar_image",
+                992,
             ),
         ],
     )
@@ -49,6 +54,7 @@ def test_create(
         feature_flag_otx_version_selection,
         primary_image_full_name,
         sidecar_image_full_name,
+        render_gid,
     ):
         os.environ.update(
             {
@@ -62,6 +68,7 @@ def test_create(
             "mlflow_sidecar_image": "mlflow_sidecar_image",
             "ote_image": "ote_image",
             "otx2_image": "otx2_image",
+            "render_gid": str(render_gid),
         }
         mock_get_config_map.return_value = MagicMock()
         mock_get_config_map.return_value.data.get.side_effect = lambda key: configmap_data.get(key)
@@ -70,3 +77,4 @@ def test_create(
 
         assert trainer_image_info.to_primary_image_full_name() == primary_image_full_name
         assert trainer_image_info.to_sidecar_image_full_name() == sidecar_image_full_name
+        assert trainer_image_info.render_gid == render_gid