Add labels and annotations support for train client

Signed-off-by: Abhijeet Dhumal <[email protected]>

Author: abhijeet-dhumal

kubeflow/trainer/api/trainer_client.py

@@ -230,6 +230,10 @@ def train(
         runtime: Optional[types.Runtime] = None,
         initializer: Optional[types.Initializer] = None,
         trainer: Optional[Union[types.CustomTrainer, types.BuiltinTrainer]] = None,
+        labels: Optional[Dict[str, str]] = None,
+        annotations: Optional[Dict[str, str]] = None,
+        job_labels: Optional[Dict[str, str]] = None,
+        job_annotations: Optional[Dict[str, str]] = None,
     ) -> str:
         """
         Create the TrainJob. You can configure these types of training task:
@@ -246,6 +250,15 @@ def train(
                 Configuration for the dataset and model initializers.
             trainer:
                 Configuration for Custom Training Task or Config-driven Task with Builtin Trainer.
+            labels: Optional dictionary of labels to apply to the TrainJob metadata (.metadata.labels).
+                Used for TrainJob resource organization and filtering.
+            annotations: Optional dictionary of annotations to apply to the TrainJob metadata (.metadata.annotations).
+                Useful for storing additional metadata about the training job resource.
+            job_labels: Optional dictionary of labels to apply to the JobSet and Jobs (.spec.labels).
+                These labels are propagated to the derivative JobSet and Jobs. Use this for Kueue 
+                integration (e.g., {"kueue.x-k8s.io/queue-name": "ml-queue"}).
+            job_annotations: Optional dictionary of annotations to apply to the JobSet and Jobs (.spec.annotations).
+                These annotations are propagated to the derivative JobSet and Jobs.
 
         Returns:
             str: The unique name of the TrainJob that has been generated.
@@ -297,10 +310,14 @@ def train(
             apiVersion=constants.API_VERSION,
             kind=constants.TRAINJOB_KIND,
             metadata=models.IoK8sApimachineryPkgApisMetaV1ObjectMeta(
-                name=train_job_name
+                name=train_job_name,
+                labels=labels,
+                annotations=annotations
             ),
             spec=models.TrainerV1alpha1TrainJobSpec(
                 runtimeRef=models.TrainerV1alpha1RuntimeRef(name=runtime.name),
+                labels=job_labels,
+                annotations=job_annotations,
                 trainer=(
                     trainer_crd
                     if trainer_crd != models.TrainerV1alpha1Trainer()

kubeflow/trainer/api/trainer_client_test.py

@@ -253,16 +253,26 @@ def get_train_job(
     runtime_name: str,
     train_job_name: str = BASIC_TRAIN_JOB_NAME,
     train_job_trainer: Optional[models.TrainerV1alpha1Trainer] = None,
+    labels: Optional[Dict[str, str]] = None,
+    annotations: Optional[Dict[str, str]] = None,
+    job_labels: Optional[Dict[str, str]] = None,
+    job_annotations: Optional[Dict[str, str]] = None,
 ) -> models.TrainerV1alpha1TrainJob:
     """
     Create a mock TrainJob object with optional trainer configurations.
     """
     train_job = models.TrainerV1alpha1TrainJob(
         apiVersion=constants.API_VERSION,
         kind=constants.TRAINJOB_KIND,
-        metadata=models.IoK8sApimachineryPkgApisMetaV1ObjectMeta(name=train_job_name),
+        metadata=models.IoK8sApimachineryPkgApisMetaV1ObjectMeta(
+            name=train_job_name,
+            labels=labels,
+            annotations=annotations
+        ),
         spec=models.TrainerV1alpha1TrainJobSpec(
             runtimeRef=models.TrainerV1alpha1RuntimeRef(name=runtime_name),
+            labels=job_labels,
+            annotations=job_annotations,
             trainer=train_job_trainer,
         ),
     )
@@ -793,6 +803,86 @@ def test_get_runtime_packages(trainer_client, test_case):
             },
             expected_error=ValueError,
         ),
+        TestCase(
+            name="valid flow with labels and annotations",
+            expected_status=SUCCESS,
+            config={
+                "labels": {"kueue.x-k8s.io/queue-name": "ml-queue", "team": "ml-engineering"},
+                "annotations": {"experiment.id": "exp-001", "description": "Test training job"},
+            },
+            expected_output=get_train_job(
+                runtime_name=TORCH_RUNTIME,
+                train_job_name=BASIC_TRAIN_JOB_NAME,
+                labels={"kueue.x-k8s.io/queue-name": "ml-queue", "team": "ml-engineering"},
+                annotations={"experiment.id": "exp-001", "description": "Test training job"},
+            ),
+        ),
+        TestCase(
+            name="valid flow with only labels",
+            expected_status=SUCCESS,
+            config={
+                "labels": {"priority": "high"},
+            },
+            expected_output=get_train_job(
+                runtime_name=TORCH_RUNTIME,
+                train_job_name=BASIC_TRAIN_JOB_NAME,
+                labels={"priority": "high"},
+            ),
+        ),
+        TestCase(
+            name="valid flow with only annotations",
+            expected_status=SUCCESS,
+            config={
+                "annotations": {"created-by": "training-pipeline"},
+            },
+            expected_output=get_train_job(
+                runtime_name=TORCH_RUNTIME,
+                train_job_name=BASIC_TRAIN_JOB_NAME,
+                annotations={"created-by": "training-pipeline"},
+            ),
+        ),
+        TestCase(
+            name="valid flow with job_labels for Kueue",
+            expected_status=SUCCESS,
+            config={
+                "job_labels": {"kueue.x-k8s.io/queue-name": "ml-queue"},
+            },
+            expected_output=get_train_job(
+                runtime_name=TORCH_RUNTIME,
+                train_job_name=BASIC_TRAIN_JOB_NAME,
+                job_labels={"kueue.x-k8s.io/queue-name": "ml-queue"},
+            ),
+        ),
+        TestCase(
+            name="valid flow with job_annotations",
+            expected_status=SUCCESS,
+            config={
+                "job_annotations": {"experiment.id": "exp-001"},
+            },
+            expected_output=get_train_job(
+                runtime_name=TORCH_RUNTIME,
+                train_job_name=BASIC_TRAIN_JOB_NAME,
+                job_annotations={"experiment.id": "exp-001"},
+            ),
+        ),
+        TestCase(
+            name="valid flow with both resource and JobSet labels/annotations",
+            expected_status=SUCCESS,
+            config={
+                "labels": {"team": "ml-platform"},
+                "annotations": {"created-by": "sdk"},
+                "job_labels": {"kueue.x-k8s.io/queue-name": "gpu-queue"},
+                "job_annotations": {"experiment.id": "exp-001"},
+            },
+            expected_output=get_train_job(
+                runtime_name=TORCH_RUNTIME,
+                train_job_name=BASIC_TRAIN_JOB_NAME,
+                labels={"team": "ml-platform"},
+                annotations={"created-by": "sdk"},
+                job_labels={"kueue.x-k8s.io/queue-name": "gpu-queue"},
+                job_annotations={"experiment.id": "exp-001"},
+            ),
+        ),
     ],
 )
 def test_train(trainer_client, test_case):
@@ -805,7 +895,12 @@ def test_train(trainer_client, test_case):
         )
 
         train_job_name = trainer_client.train(
-            runtime=runtime, trainer=test_case.config.get("trainer", None)
+            runtime=runtime,
+            trainer=test_case.config.get("trainer", None),
+            labels=test_case.config.get("labels", None),
+            annotations=test_case.config.get("annotations", None),
+            job_labels=test_case.config.get("job_labels", None),
+            job_annotations=test_case.config.get("job_annotations", None)
         )
 
         assert test_case.expected_status == SUCCESS