Feat: Inject user identity as pod label in K8s plugin

[fg91]

Why are the changes needed?

@ByronHsu recently implemented middleware to inject the user identity into the flyte workflow's ExecutionSpec: flyteorg/flyteadmin#549

Building on this, we would like to inject the user identity (not credentials, just the information who started the execution) into the task pod as an additional label in order to enable the following user stories:

• As a platform engineer, I would like to know which user started which Pod.
• As a platform engineer, I would like attribute compute cost to users (e.g. by activating GKE cost export into big query which associates compute cost with pod labels).
• As a platform user/ML engineer, I would like my task executions to automatically log training progress under my account in an experiment tracking server.

What changes were proposed in this pull request?

• I propose to add a new flag to the k8s plugin config, specifying whether the execution identity is injected as a label. The default is false to make this opt in.

• I propose to inject the additional execution identity label in newTaskExecutionMetadata where we also inject labels and annotations for secrets.

  Because of the following reasons I believe that this is a good place but I'm looking for feedback on this:

  • The TaskExecutionMetadata type returned by this function has the following doc string

    TaskExecutionMetadata provides a layer on top of the core TaskExecutionMetadata with customized annotations and labels for k8s plugins.

    I find this fitting the proposed use case.

  • Injecting the label here means that it is automatically injected not only for normal python tasks but also for plugins like kubeflow distributed training tasks.

How was this patch tested?

• Added unit tests.
• Ran propeller with the proposed changes in our staging cluster and ensured that the execution identity is correctly injected for python tasks as well as pytorch tasks.
  • I also tested the case where the client authenticates with a client credentials flow instead of pkce flow as in this case the execution identity is not set in the flyte workflow.

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (513c3e1) 58.10% compared to head (988a943) 58.08%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #4637      +/-   ##
==========================================
- Coverage   58.10%   58.08%   -0.03%     
==========================================
  Files         626      626              
  Lines       53834    53836       +2     
==========================================
- Hits        31280    31270      -10     
- Misses      20048    20058      +10     
- Partials     2506     2508       +2     

Flag	Coverage Δ
unittests	58.08% <100.00%> (-0.03%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[fg91]

For documentation purposes:

One can set the following e.g. in the default flyte pod template in order to use the execution identity as a pod template:

apiVersion: v1
kind: PodTemplate
template:
  spec:
    containers:
    - env:
      - name: EXECUTION_IDENTITY
        valueFrom:
          fieldRef:
            apiVersion: v1
            fieldPath: metadata.labels['execution_identity']

[hamersaw]

@fg91 this looks great. I agree, the metadata is the correct place to inject this. My first question here is whether this should be behind a configuration flag? Is there a scenario where we would not want the execution identity available? I'm assuming that if a user has Flyte configured to inject user IDs from middleware, it would be beneficial to include this information in k8s resources, thoughts? cc. @ByronHsu

[fg91]

My first question here is whether this should be behind a configuration flag? Is there a scenario where we would not want the execution identity available?

I can't come up with such a scenario tbh. I first proposed a configuration flag as the most conservative change because I wasn't sure how others feel about this. I'm happy to remove the flag though @hamersaw! a9652e8 115c23d

I'm assuming that if a user has Flyte configured to inject user IDs from middleware, it would be beneficial to include this information in k8s resources, thoughts? cc. @ByronHsu

Injecting user ids from middleware in flyteadmin is not behind a config flag (see here) so users don't have to opt in.

[hamersaw]

Very useful, can we clean up the unit tests?

[fg91]

[...] can we clean up the unit tests?

Done ✅