[Argo-Workflows]: Copy Kubernetes Security Context into error and heartbeat templates

[kevslinger]

The functions _lifecycle_hook_from_deco, _error_msg_capture_hook_templates, and _heartbeat_daemon_template

metaflow/metaflow/plugins/argo/argo_workflows.py

Line 3041 in d2dd7f9

def _lifecycle_hook_from_deco(self, deco):

metaflow/metaflow/plugins/argo/argo_workflows.py

Line 3188 in d2dd7f9

def _error_msg_capture_hook_templates(self):

metaflow/metaflow/plugins/argo/argo_workflows.py

Line 3650 in d2dd7f9

def _heartbeat_daemon_template(self):

all include the same pattern of copying the start step's @kubernetes attributes

metaflow/metaflow/plugins/argo/argo_workflows.py

Lines 3192 to 3198 in d2dd7f9

	# We want to grab the base image used by the start step, as this is known to be pullable from within the cluster,
	# and it might contain the required libraries, allowing us to start up faster.
	resources = dict(
	[deco for deco in start_step.decorators if deco.name == "kubernetes"][
	0
	].attributes
	)

and then use this resources dict to fill in attributes such as the image and secrets of the container. However, it does not copy over the security context defined from the original @kubernetes decorator, which can result in, for instance, the resulting pod being labeled as insecure or even rejected.

To resolve this issue, I think it should be sufficient to specify the security_context in the generated container using the same security context from resources (resources["security_container"]).

Reproducing

I used the following flow and ran the comment python flow.py argo-workflows create --only-json (note: datastore required).

from metaflow import FlowSpec, kubernetes, step

_SECURITY_CONTEXT = {"allow_privilege_escalation": False}

class TestFlow(FlowSpec):
    """
    Minimal example to show error-msg-capture-hook does not inherit
    the same security context as all the steps
    """

    @kubernetes(security_context=_SECURITY_CONTEXT)
    @step
    def start(self):
        print("Start")
        self.next(self.end)
    
    @kubernetes(security_context=_SECURITY_CONTEXT)
    @step
    def end(self):
        print("End")

if __name__ == '__main__':
    TestFlow()

This will produce a start, an end, and an error-msg-capture-hook. The start and end steps have the following security context:

"securityContext": {
    "allowPrivilegeEscalation": false,
    "appArmorProfile": null,
    "capabilities": null,
    "privileged": null,
    "procMount": null,
    "readOnlyRootFilesystem": null,
    "runAsGroup": null,
    "runAsNonRoot": null,
    "runAsUser": null,
    "seLinuxOptions": null,
    "seccompProfile": null,
    "windowsOptions": null
},

meanwhile the error hook has "securityContext": null.

[goingforstudying-ctrl]

I'd like to pick this up.

The three template builders all rebuild a resources dict from the start step's @kubernetes decorator attributes and pass image, secrets, service account and so on into the generated V1Container, but they never read resources["security_context"]. Meanwhile the regular step containers build a V1SecurityContext from it in _container_templates (the security_context / _security_context block around line 2637). So the hook and heartbeat pods render with securityContext: null and can get rejected by a Pod Security Admission restricted namespace while every step pod passes.

Plan: replicate that same block in _lifecycle_hook_from_deco, _error_msg_capture_hook_templates, and _heartbeat_daemon_template. Read resources.get("security_context"), build a V1SecurityContext when it is a non-empty dict, and pass it into the V1Container kwargs so the generated template omits securityContext entirely when the step did not set one.

For tests I would render a flow like the one in your repro and assert the error-msg-capture-hook and heartbeat daemon templates carry the same securityContext as the step pods, plus a case without any security_context asserting the field stays absent from the rendered container.

One question before I start: these templates inherit everything from the start step's @kubernetes decorator (image, secrets, service account), and I would follow the same rule for security_context. Does that match your expectation, or should the error hook inherit from the failed step instead? The templates are generated up front from the start step, so start-step inheritance is the consistent choice, but I want to confirm the direction before opening a PR.

[kevslinger]

Hey @goingforstudying-ctrl thanks for your interest in this issue! I tackled it and opened PR #3354, is that also how you would solve it?

[goingforstudying-ctrl]

Yes, that's how I would have solved it. Start step inheritance is what these templates already do for image, secrets and service account, and your three hunks mirror the security context block from _container_templates. One nit: when security_context is unset you pass None explicitly while the step containers omit the kwarg entirely. Both render the same, so it's cosmetic at most.

Since you've got the PR, I'll leave this one to you. I'll pick up #3309 instead, same class of bug in the same templates, just nodeSelector and tolerations instead of the security context. Will comment over there.