[Core feature] Allow setting metadata (labels&annotations) individually on all K8s task types

[fg91]

Motivation: Why do you think this is important?

Today, as a Flyte user, I have the following options to set labels/annotations on the pods/CRD objects of K8s tasks in a flyte workflow execution:

1. Set via pod template:

   def task(
       pod_template=PodTemplate(annotations=..., labels=...)
   )

   This sets labels/annotations on the pods of individual tasks.
   For distributed tasks (like pytorch, ray, ...) this sets the metadata not on the CRD object but its pod template spec.

2. Set via pyflyte run --labels ... --annotations ...

   This applies the metadata to all K8s objects in a flyte workflow execution, including task pods and task CRD objects. However, this mechanism doesn't work on an individual task level.

As a Flyte user, I would like to be able to specify specific labels/annotations for individual k8s task CRD objects like pytorch jobs, ray job, ... (the same way I already can today for pods via the pod template):

@task(
    task_config=PyTorch(
        num_workers=...,
        ...
        # Proposed addition:
        metadata=ObjectMeta(
            annotations={"kueue.x-k8s.io/queue-name": "queue-name"},
            labels={...}
        )
    )
)

I propose to use the same syntax/flyteidl type for all K8s (non-pod) plugins like Elastic, TfJob, MpiJob, RayJobConfig, ...

---

In my concrete case, I would like to have this feature in order to leverage Kueue to gang schedule worker pods for distributed pytorch training tasks (e.g. as documented here).
This requires setting a queue name annotation on the underlying PytorchJob CRD object.

There have been previous asks from the community to enable such a feature/integration:

• Attempts to integrate Yunikorn and Kueue more deeply into flytepropeller which weren't accepted though.

  In contrast, the feature I propose allows users to choose to use Kueue while it isn't a Flyte-Kueue integration. Instead it is a very general feature that could be used for any other integration as well that makes use of annotations/labels to select workloads.

• There have been discussions in Slack about using Kueue, suggesting to use e.g. pyflyte run --labels/--annotations to set the required metadata. However, this is not good enough because this applies the metadata to all nodes in the graph while you might want to do queueing/gang scheduling only for a subset.

Describe alternatives you've considered

Add task kwargs for labels and annotations:

@task(
    # If we added these args ...
    labels={...},
    annotations={...),
    # ... for simple python function tasks this would conflict with this existing arg:
    pod_template=PodTemplate(annotations=...)
)

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

• Yes

Have you read the Code of Conduct?

• Yes

[fg91]

Notes contributors' sync:
Agree on precedence when metadata set in task_config conflicts with metadata set e.g. via pyflyte run.

[fg91]

@eapolinario there cannot be a conflict with metadata set via @task(pod_template=...) because that is applied to the CRD object's pod template spec and not the object itself.
But there can be a conflict with pyflyte run --annotations ... --labels ....

I would say when a user specifies metadata via pyflyte run, this should have precedence. But I'm open to implementing it either way. Do you have a strong opinion?

[eapolinario]

I would say when a user specifies metadata via pyflyte run, this should have precedence.

I agree.

[pingsutw]

I think if we could add the label and annotation to @task

• Add label and annotation to the individual task pod and task CRD object
• @task(label=) should take precedence if the user also specifies the label in the podTemplate.

@task(label=..., annotation=...)
def t1()

I would say when a user specifies metadata via pyflyte run, this should have precedence.

But currently, when you specify the --image flag in pyflyte run, it doesn't override the image for tasks that already have an image specified. It only overrides tasks that don't have an image set.

[fg91]

@pingsutw, implementing @task(labels=..., annotations=...) as opposed to @task(task_config=SomePlugin(metadata=...)) sounds good to me, I can do that 👍

I have some questions about how this info needs to be routed to the backend in a clean way:

Eventually in flytekit during the registration, the labels/annotations need to end up in flyteidl's TaskTemplate message, correct?

If I'm not mistaken TaskTemplate doesn't yet have a field we can reuse. It has

    // Extra metadata about the task.
    TaskMetadata metadata = 3;

but this is not metadata in the k8s sense.

And is has

    oneof target {
        Container container = 6;
        K8sPod k8s_pod = 17;
        Sql sql = 18;
    }

where K8sPod is

// Defines a pod spec and additional pod metadata that is created when a task is executed.
message K8sPod {
    // Contains additional metadata for building a kubernetes pod.
    K8sObjectMetadata metadata = 1;

But this metadata is applied to the pod spec of a distributed job CRD object, not the CRD object itself. So it feels wrong to put the annotations/labels I'd like to add to @task to K8sPod.

Do you have an opinion on where I should put them?

Update:

Maybe adding K8sObjectMetadata metadata to TaskMetadata is acceptable after all, it already has string pod_template_name = 12; which is also K8s specific.

[fg91]

Result of discussion with @pingsutw in contrib sync March 27th:

I'll add K8sObjectMetadata metadata to TaskMetadata in order to be able to expose @task(labels=..., annotations=...).

@eapolinario since we talked about this before (when the idea was still to add the metadata arg to the task configs), is there anything that speaks against this in your eyes? Thank you 🙏