[Feature] Mechanism for injecting node affinities and tolerations to every Flyte task pod

[jeevb]

Motivation: Why do you think this is important?
Please consider adding more control for specifying the pod spec of the Flyte task pods - namely node affinities and tolerations. This is critical on shared K8S clusters to target task pods to dedicated nodes that are better suited to run these tasks. These nodes are also typically tainted to block other workloads from scheduling on them.

Currently, via Flytekit, users only have the ability to specify the resource requirements for a given python or dynamic task. sidecar tasks provide more flexibility in terms of defining the spec for the pod that will run that task - allowing for the specification of both node affinities and tolerations. While most python tasks can be adapted to sidecar tasks to afford this flexibility, there is no way to do so for dynamic tasks.

Goal: What should the final outcome look like, ideally?
As a workflow/task author, I would like to be able to specify a set of node affinities and tolerations to target Flyte task pods to dedicated/tainted pipeline node pools on shared K8S clusters.

Describe alternatives you've considered
We are currently considering using mutating admission controllers to patch Flyte task pods accordingly.

Flyte component

• Overall
• Flyte Setup and Installation scripts
• Flyte Documentation
• Flyte communication (slack/email etc)
• FlytePropeller
• FlyteIDL (Flyte specification language)
• Flytekit (Python SDK)
• FlyteAdmin (Control Plane service)
• FlytePlugins
• DataCatalog
• FlyteStdlib (common libraries)
• FlyteConsole (UI)
• Other

[Optional] Propose: Link/Inline
I see 3 levels of implementations that can be relevant for this problem:

• Specifying node affinities and tolerations within the plugin configuration. All task pods processed by this propeller plugin will have these applied.
• Specifying in the configuration used to register workflows/tasks via pyflyte register workflows -c. All tasks registered this way will have these applied.
• Specifying in a task-specific way via Flytekit, similar to what is possible with sidecar tasks. This will afford the most flexibility.

Additional context
Add any other context or screenshots about the feature request here.

Is this a blocker for you to adopt Flyte
This is NOT a blocker.

[kumare3]

@jeevb thank you for the excellent write up and explaining the problem. I completely understand in a shared K8s cluster it is desirable to run batch style jobs on certain nodes rather than others. I think we should support it (caveat). First let me dive into some concepts,
"python_task" or just a pure simple container task in Flyte is meant to be a no frills container execution. It can almost run in any container orchestration as it does not allow volume mounts/node affinity/tolerations etc - in short no K8s pod specific special treatement. This is because we can easily port a container to run in some other environment, which at times can be way more scalable - for example AWS ECS. ECS is very simple, just runs a container on some cluster of machines, but it does not have any additional capabilities. this also makes it way more scalable.
That being said, today we run most python tasks directly on K8s (the dynamic array tasks if you so configure can be run on AWS Batch). We have seen at our scale the K8s based simple implementation of "dynamic array tasks" does not scale well. We are going to release a new scheduler tweak later in the year that will allow us to potentially run arrays of size 5k+ (this is the scale at Lyft) natively on K8s - so stay tuned. the simplicity of containers, also in the future allows Ok As a solution,

1. I think as part 1, we will implement a simple backend configuration (already used for gpu based tolerations etc) that will be applied to every K8s pod launch. this will allow all k8s pods to have some node affinity / toleration.
   This is pretty trivial to implement and should help you

2. For the user SDK part to support taints and tolerations + affinity I want to still recommend using Sidecar tasks. Sidecar tasks can be used in dynamic tasks as well.
   One Action item @katrogan can you please help add a flytesnacks example of dynamic tasks with sidecar. As for running a very large number of array in sidecar, we would love to do that with the k8s native array implementation. Let me know if you guys are open to help us design that better. @migueltol22 is leading this effort.

Wdyt?

[jeevb]

I understand the design considerations around python_task and that makes a lot of sense. (1) sounds like a solid plan, especially if (2) is already possible. We're happy to continue to use sidecar_task if more pod spec customization is required, but I believe (1) would be sufficient for most of our use cases.

[kumare3]

@jeevb awesome and thank you for the quick response, I will create a ticket for 1, and lets keep this open, if we find a problem you can re-visit. We will use this ticket to write an example that shows how to use "sidecar" in dynamic. I forgot to mention we have users at Lyft doing that.

[kumare3]

Part 1: #439

[kumare3]

This can now be resolved. the way to add affinities, tolerations etc is to use flytekitplugins-pod. This is a first party and highly usable plugin and will be maintained. The only reason to keep it as a separate lib, is to keep the k8s python dependencies in the core to a minimum