feat: add container and pod security context support for kubernetes workloads

metaflow/metaflow_config.py

@@ -437,6 +437,10 @@
 KUBERNETES_DISK = from_conf("KUBERNETES_DISK", None)
 # Default kubernetes QoS class
 KUBERNETES_QOS = from_conf("KUBERNETES_QOS", "burstable")
+# Default container security context (JSON) for kubernetes pods
+KUBERNETES_SECURITY_CONTEXT = from_conf("KUBERNETES_SECURITY_CONTEXT", "")
+# Default pod security context (JSON) for kubernetes pods
+KUBERNETES_POD_SECURITY_CONTEXT = from_conf("KUBERNETES_POD_SECURITY_CONTEXT", "")
 
 # Architecture of kubernetes nodes - used for @conda/@pypi in metaflow-dev
 KUBERNETES_CONDA_ARCH = from_conf("KUBERNETES_CONDA_ARCH")

metaflow/plugins/argo/argo_workflows.py

@@ -2581,6 +2581,15 @@ def _container_templates(self):
                     )
                 }
 
+            pod_security_context = resources.get("pod_security_context", None)
+            _pod_security_context = {}
+            if pod_security_context is not None and len(pod_security_context) > 0:
+                _pod_security_context = {
+                    "security_context": kubernetes_sdk.V1PodSecurityContext(
+                        **pod_security_context
+                    )
+                }
+
             # Create a ContainerTemplate for this node. Ideally, we would have
             # liked to inline this ContainerTemplate and avoid scanning the workflow
             # twice, but due to issues with variable substitution, we will have to
@@ -2639,6 +2648,7 @@ def _container_templates(self):
                     port=port,
                     qos=resources["qos"],
                     security_context=security_context,
+                    pod_security_context=pod_security_context,
                 )
 
                 for k, v in env.items():
@@ -2804,6 +2814,16 @@ def _container_templates(self):
                         if resources["image_pull_secrets"]
                         else None
                     )
+                    # Set pod security context via pod_spec_patch
+                    .pod_spec_patch(
+                        {
+                            "securityContext": kubernetes_sdk.V1PodSecurityContext(
+                                **pod_security_context
+                            ).to_dict()
+                        }
+                        if pod_security_context
+                        else None
+                    )
                     # Set container
                     .container(
                         # TODO: Unify the logic with kubernetes.py
@@ -4503,7 +4523,12 @@ def pod_spec_patch(self, pod_spec_patch=None):
         if pod_spec_patch is None:
             return self
 
-        self.payload["podSpecPatch"] = json.dumps(pod_spec_patch)
+        if "podSpecPatch" in self.payload:
+            existing = json.loads(self.payload["podSpecPatch"])
+            existing.update(pod_spec_patch)
+            self.payload["podSpecPatch"] = json.dumps(existing)
+        else:
+            self.payload["podSpecPatch"] = json.dumps(pod_spec_patch)
 
         return self
 

metaflow/plugins/kubernetes/kubernetes.py

@@ -198,6 +198,7 @@ def create_jobset(
         num_parallel=None,
         qos=None,
         security_context=None,
+        pod_security_context=None,
     ):
         name = "js-%s" % str(uuid4())[:6]
         jobset = (
@@ -233,6 +234,7 @@ def create_jobset(
                 num_parallel=num_parallel,
                 qos=qos,
                 security_context=security_context,
+                pod_security_context=pod_security_context,
             )
             .environment_variable("METAFLOW_CODE_METADATA", code_package_metadata)
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
@@ -499,6 +501,7 @@ def create_job_object(
         qos=None,
         annotations=None,
         security_context=None,
+        pod_security_context=None,
     ):
         if env is None:
             env = {}
@@ -544,6 +547,7 @@ def create_job_object(
                 port=port,
                 qos=qos,
                 security_context=security_context,
+                pod_security_context=pod_security_context,
             )
             .environment_variable("METAFLOW_CODE_METADATA", code_package_metadata)
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -158,6 +158,12 @@ def kubernetes():
     type=JSONTypeClass(),
     multiple=False,
 )
+@click.option(
+    "--pod-security-context",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
 @click.pass_context
 def step(
     ctx,
@@ -192,6 +198,7 @@ def step(
     labels=None,
     annotations=None,
     security_context=None,
+    pod_security_context=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -338,6 +345,7 @@ def _sync_metadata():
                 labels=labels,
                 annotations=annotations,
                 security_context=security_context,
+                pod_security_context=pod_security_context,
             )
     except Exception:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -27,6 +27,8 @@
     KUBERNETES_NODE_SELECTOR,
     KUBERNETES_PERSISTENT_VOLUME_CLAIMS,
     KUBERNETES_PORT,
+    KUBERNETES_SECURITY_CONTEXT,
+    KUBERNETES_POD_SECURITY_CONTEXT,
     KUBERNETES_SERVICE_ACCOUNT,
     KUBERNETES_SHARED_MEMORY,
     KUBERNETES_TOLERATIONS,
@@ -136,6 +138,17 @@ class KubernetesDecorator(StepDecorator):
         - run_as_user: int, optional, default None
         - run_as_group: int, optional, default None
         - run_as_non_root: bool, optional, default None
+        - read_only_root_filesystem: bool, optional, default None
+        - capabilities: Dict[str, List[str]], optional, default None
+        Can also be set via METAFLOW_KUBERNETES_SECURITY_CONTEXT (JSON).
+    pod_security_context: Dict[str, Any], optional, default None
+        Pod-level security context. Applies to all containers in the pod. Allows the following keys:
+        - run_as_user: int, optional, default None
+        - run_as_group: int, optional, default None
+        - run_as_non_root: bool, optional, default None
+        - fs_group: int, optional, default None
+        - supplemental_groups: List[int], optional, default None
+        Can also be set via METAFLOW_KUBERNETES_POD_SECURITY_CONTEXT (JSON).
     """
 
     name = "kubernetes"
@@ -168,6 +181,7 @@ class KubernetesDecorator(StepDecorator):
         "hostname_resolution_timeout": 10 * 60,
         "qos": KUBERNETES_QOS,
         "security_context": None,
+        "pod_security_context": None,
     }
     package_metadata = None
     package_url = None
@@ -310,6 +324,19 @@ def init(self):
         if not self.attributes["port"]:
             self.attributes["port"] = KUBERNETES_PORT
 
+        # Security context: decorator takes precedence over env var
+        if not self.attributes["security_context"] and KUBERNETES_SECURITY_CONTEXT:
+            self.attributes["security_context"] = json.loads(
+                KUBERNETES_SECURITY_CONTEXT
+            )
+        if (
+            not self.attributes["pod_security_context"]
+            and KUBERNETES_POD_SECURITY_CONTEXT
+        ):
+            self.attributes["pod_security_context"] = json.loads(
+                KUBERNETES_POD_SECURITY_CONTEXT
+            )
+
     # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
     def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
         # Executing Kubernetes jobs requires a non-local datastore.
@@ -500,6 +527,7 @@ def runtime_step_cli(
                     "labels",
                     "annotations",
                     "security_context",
+                    "pod_security_context",
                 ]:
                     cli_args.command_options[k] = json.dumps(v)
                 else:

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -93,6 +93,13 @@ def create_job_spec(self):
                 "security_context": client.V1SecurityContext(**security_context)
             }
 
+        pod_security_context = self._kwargs.get("pod_security_context", {})
+        _pod_security_context = {}
+        if pod_security_context is not None and len(pod_security_context) > 0:
+            _pod_security_context = {
+                "security_context": client.V1PodSecurityContext(**pod_security_context)
+            }
+
         return client.V1JobSpec(
             # Retries are handled by Metaflow when it is responsible for
             # executing the flow. The responsibility is moved to Kubernetes
@@ -277,6 +284,7 @@ def create_job_spec(self):
                         if self._kwargs["persistent_volume_claims"] is not None
                         else []
                     ),
+                    **_pod_security_context,
                 ),
             ),
         )

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -569,6 +569,13 @@ def dump(self):
             _security_context = {
                 "security_context": client.V1SecurityContext(**security_context)
             }
+
+        pod_security_context = self._kwargs.get("pod_security_context", {})
+        _pod_security_context = {}
+        if pod_security_context is not None and len(pod_security_context) > 0:
+            _pod_security_context = {
+                "security_context": client.V1PodSecurityContext(**pod_security_context)
+            }
         return dict(
             name=self.name,
             template=client.api_client.ApiClient().sanitize_for_serialization(
@@ -784,6 +791,7 @@ def dump(self):
                                     is not None
                                     else []
                                 ),
+                                **_pod_security_context,
                             ),
                         ),
                     ),