Expose TrainJob labels and annotations in the SDK

[kramaranya]

What you would like to be added?

We want the SDK to let users set labels and annotations when creating a TrainJob. Add two optional fields to the train method so users can set labels and annotations.

Why is this needed?

Right now, users can't set labels and annotations via the SDK and need YAML for that. It also keeps the SDK aligned with the controller API and KEP Trainer V2

Love this feature?

Give it a 👍 We prioritize the features with most 👍

[kramaranya]

/remove-lifecycle stale
/remove-label lifecycle/needs-triage
/good-first-issue

cc @kubeflow/kubeflow-sdk-team @abhijeet-dhumal

[google-oss-prow]

@kramaranya:
This request has been marked as suitable for new contributors.

Please ensure the request meets the requirements listed here.

If this request no longer meets these requirements, the label can be removed
by commenting with the /remove-good-first-issue command.

In response to this:

/remove-lifecycle stale
/remove-label lifecycle/needs-triage
/good-first-issue

cc @kubeflow/kubeflow-sdk-team @abhijeet-dhumal

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

[abhijeet-dhumal]

Thanks @kramaranya for raising issue on my behalf 🙌

Just to add in description :
For the uscases such as Kueue-integration support for Trainer V2 (in-progress) this fix will be essential:
kubernetes-sigs/kueue#6597

Can take reference from the kubeflow-training SDK based implementation done here :
kubeflow/trainer#2612
kubeflow/trainer#2636

Note: The labels and annotations are supported for TrainJob in trainer V2 operator side.

[andreyvelich]

Thanks for opening this @kramaranya!
In the TrainJob API we have many labels that users can set (.metadata.labels, .spec.trainer.labels, potentially podTemplateOverrides labels after this PR: kubeflow/trainer#2785).

What kind of labels and annotations we want to allow users to configure with train() API ?

[kramaranya]

In the TrainJob API we have many labels that users can set (.metadata.labels, .spec.trainer.labels, potentially podTemplateOverrides labels after this PR: kubeflow/trainer#2785).

What kind of labels and annotations we want to allow users to configure with train() API ?

I'm thinking about Kueue labels. If we put those into podSpecOverrides, the labels will go only on pod templates, right? But I thought Kueue operates at the JobSet level instead, which means we should add them to .spec.labels, so Kueue will see them at JobSet creation time. Please correct me if I'm wrong
cc @astefanutti

[abhijeet-dhumal]

@andreyvelich @astefanutti @szaher @kramaranya :

IIUC Labels like kueue.x-k8s.io/queue-name should go in .spec.labels so they're applied to the JobSet that Kueue manages, not just the TrainJob resource metadata, right?

For Kueue specifically:
If I see here

if trainJobQueueName := jobframework.QueueNameForObject(trainJob.Object()); trainJobQueueName != "" {
    if trainJob.Spec.Labels == nil {
        trainJob.Spec.Labels = make(map[string]string, 1)
    }
    trainJob.Spec.Labels[controllerconstants.QueueLabel] = string(trainJobQueueName)
}

according to this Kueue expects allkueue.x-k8s.io/queue-name label in .spec.labels!

The SDK needs to support both label locations, So we have 3 scenarios here :

1. .metadata.labels/.metadata.annotations - TrainJob resource level
2. .spec.labels/.spec.annotations - JobSet/Jobs level (Kueue integration)
3. .podSpecOverride.labels/.podSpecOverride.annotations - once we get this fix merged in Trainer V2 : feat(api): Add PodTemplateOverrides API into TrainJob trainer#2785

def train(
    self,
    runtime: Optional[types.Runtime] = None,
    initializer: Optional[types.Initializer] = None,
    trainer: Optional[Union[types.CustomTrainer, types.BuiltinTrainer]] = None,
    # Resource-level metadata (for TrainJob CRD organization)
    labels: Optional[Dict[str, str]] = None,
    annotations: Optional[Dict[str, str]] = None,
    # JobSet/Jobs-level metadata (for Kueue integration!)
    job_labels: Optional[Dict[str, str]] = None,
    job_annotations: Optional[Dict[str, str]] = None,
)
train_job = models.TrainerV1alpha1TrainJob(
    metadata=models.IoK8sApimachineryPkgApisMetaV1ObjectMeta(
        name=train_job_name,
        labels=labels,           # TrainJob resource level
        annotations=annotations  # TrainJob resource level
    ),
    spec=models.TrainerV1alpha1TrainJobSpec(
        labels=job_labels,       # JobSet level (Kueue expects this!)
        annotations=job_annotations,  # JobSet level
        runtimeRef=models.TrainerV1alpha1RuntimeRef(name=runtime.name),
        # ... rest of spec
    ),
)

job_name = client.train(
    runtime=runtime,
    trainer=trainer,
     labels={ "project": "training"},   # Goes to .metadata.labels
    job_labels={"kueue.x-k8s.io/queue-name": "ml-queue"},  # Goes to .spec.labels (Kueue reads this!)
   # Pod level (future)
    pod_overrides=[
        types.PodSpecOverride(
            target_jobs=["node"],
            labels={"gpu-type": "a100", "workload": "training"},
            annotations={"prometheus.io/scrape": "true"},
            node_selector={"accelerator": "nvidia-tesla-a100"}
        )
    ]
)

Is this approach seems correct?

[abhijeet-dhumal]

/assign