[Core feature] Support overriding task `pod_template` via `with_overrides`

[swarup-stripe]

Motivation: Why do you think this is important?

We use pod_template to specify things like node selectors and tolerations. However, this is specific to certain workflows that share the core training tasks with other workflows. This means we have to duplicate the task function for each type of override we require. If we can support pod_template in with_overrides, this would reduce a lot of toil since we can supply pod templates in a central location and override downstream tasks, similar to how we can do so for resources.

Goal: What should the final outcome look like, ideally?

The with_overrides function here should be able to take in a pod_template argument, similar to how it can take in requests and limits.

Describe alternatives you've considered

We've duplicated each training function for every variation of pod template we have. This isn't very scalable though, so it restricts users from using many different specifications for different workflows.

Propose: Link/Inline OR Additional context

No response

Are you sure this issue hasn't been raised already?

• [X] Yes

Have you read the Code of Conduct?

• [X] Yes

[runllm[bot]]

Currently, with_overrides does not support overriding the pod_template parameter. It is designed to override resources like requests and limits (source). To achieve your goal, you would need to propose an enhancement to the Flyte project to include pod_template in the with_overrides function.

Would you like more details on how to propose this enhancement or any other specific part of this configuration?

Sources:

• Flyte Discussion on with_overrides

---

This is an AI-generated response and your feedback is appreciated! Please leave a 👍 if this is helpful and 👎 if it is not.

Mention @runllm to ask follow-up questions. See runllm.com for more!

[arbaobao]

#take

[swarup-stripe]

@arbaobao Thank you for picking this up!! Curious if there are any updates? 👀

[arbaobao]

@swarup-stripe I'm working on it, and it will be done in a few days.

[geekysuavo]

@arbaobao any updates?

[swarup-stripe]

Hey @arbaobao, wondering if you have any updates on this?

[arbaobao]

Hi @swarup-stripe @geekysuavo,

I think overriding task pod_template via with_overrides has been completed and merged.

Thanks for your patience.