Add provision to provide labels and annotations for the pytorchjob an… (kubeflow#2612)

Signed-off-by: Abhijeet Dhumal <abhijeetdhumal652@gmail.com>

Author: abhijeet-dhumal

sdk/python/kubeflow/training/api/training_client.py

@@ -96,6 +96,8 @@ def train(
         self,
         name: str,
         namespace: Optional[str] = None,
+        labels: Optional[Dict[str, str]] = None,
+        annotations: Optional[Dict[str, str]] = None,
         num_workers: int = 1,
         num_procs_per_worker: int = 1,
         resources_per_worker: Union[dict, client.V1ResourceRequirements, None] = None,
@@ -133,6 +135,10 @@ def train(
             name: Name of the PyTorchJob.
             namespace: Namespace for the PyTorchJob. By default namespace is taken from
                 `TrainingClient` object.
+            labels: Labels for the PyTorchJob. You can specify a dictionary as a mapping object
+                representing the labels.
+            annotations: Annotations for the PyTorchJob. You can specify a dictionary as a
+                mapping object representing the annotations.
             num_workers: Number of PyTorchJob workers.
             num_procs_per_worker: Number of processes per PyTorchJob worker for `torchrun` CLI. You
                 should use this parameter if you want to use more than 1 GPU per PyTorchJob worker.
@@ -327,6 +333,8 @@ def train(
         job = utils.get_pytorchjob_template(
             name=name,
             namespace=namespace,
+            labels=labels,
+            annotations=annotations,
             master_pod_template_spec=master_pod_template_spec,
             worker_pod_template_spec=worker_pod_template_spec,
             num_workers=num_workers,
@@ -340,6 +348,8 @@ def create_job(
         job: Optional[constants.JOB_MODELS_TYPE] = None,
         name: Optional[str] = None,
         namespace: Optional[str] = None,
+        labels: Optional[Dict[str, str]] = None,
+        annotations: Optional[Dict[str, str]] = None,
         job_kind: Optional[str] = None,
         base_image: Optional[str] = None,
         train_func: Optional[Callable] = None,
@@ -370,6 +380,10 @@ def create_job(
             name: Name for the Job. It must be set if `job` parameter is omitted.
             namespace: Namespace for the Job. By default namespace is taken from
                 `TrainingClient` object.
+            labels: Labels for the Job. You can specify a dictionary as a mapping
+                object representing the labels.
+            annotations: Annotations for the Job. You can specify a dictionary as
+                a mapping object representing the annotations.
             job_kind: Kind for the Job (e.g. `TFJob` or `PyTorchJob`). It must be set if
                 `job` parameter is omitted. By default Job kind is taken from
                 `TrainingClient` object.
@@ -429,16 +443,25 @@ def create_job(
             RuntimeError: Failed to create Job.
         """
 
-        # When Job is set, only namespace arg is allowed.
+        # When Job is set, only namespace, labels, and annotations args are allowed.
         if job is not None:
             for key, value in locals().items():
                 if (
                     key
-                    not in ["self", "job", "namespace", "pip_index_url", "num_workers"]
+                    not in [
+                        "self",
+                        "job",
+                        "namespace",
+                        "pip_index_url",
+                        "num_workers",
+                        "labels",
+                        "annotations",
+                    ]
                     and value is not None
                 ):
                     raise ValueError(
-                        "If `job` is set only `namespace` argument is allowed. "
+                        "If `job` is set only `namespace`, `labels`, and `annotations` "
+                        "arguments are allowed. "
                         f"Argument `{key}` must be None."
                     )
 
@@ -525,6 +548,8 @@ def create_job(
                 job = utils.get_tfjob_template(
                     name=name,
                     namespace=namespace,
+                    labels=labels,
+                    annotations=annotations,
                     pod_template_spec=pod_template_spec,
                     num_workers=num_workers,
                     num_chief_replicas=num_chief_replicas,
@@ -539,6 +564,8 @@ def create_job(
                 job = utils.get_pytorchjob_template(
                     name=name,
                     namespace=namespace,
+                    labels=labels,
+                    annotations=annotations,
                     worker_pod_template_spec=pod_template_spec,
                     num_workers=num_workers,
                     num_procs_per_worker=num_procs_per_worker,

sdk/python/kubeflow/training/api/training_client_test.py

@@ -143,6 +143,8 @@ def create_job(
     command=None,
     args=None,
     num_workers=2,
+    labels=None,
+    annotations=None,
     env_vars=None,
     volumes=None,
     volume_mounts=None,
@@ -195,10 +197,16 @@ def create_job(
             ),
         )
 
+    meta_kwargs = {"name": TEST_NAME, "namespace": TEST_NAME}
+    if labels is not None:
+        meta_kwargs["labels"] = labels
+    if annotations is not None:
+        meta_kwargs["annotations"] = annotations
+
     pytorchjob = KubeflowOrgV1PyTorchJob(
         api_version=constants.API_VERSION,
         kind=constants.PYTORCHJOB_KIND,
-        metadata=V1ObjectMeta(name=TEST_NAME, namespace=TEST_NAME),
+        metadata=V1ObjectMeta(**meta_kwargs),
         spec=KubeflowOrgV1PyTorchJobSpec(
             run_policy=KubeflowOrgV1RunPolicy(clean_pod_policy=None),
             pytorch_replica_specs=pytorch_replica_specs,
@@ -570,6 +578,38 @@ def __init__(self):
         ValueError,
         None,
     ),
+    (
+        "valid flow with labels",
+        {
+            "name": TEST_NAME,
+            "namespace": TEST_NAME,
+            "base_image": TEST_IMAGE,
+            "num_workers": 1,
+            "labels": {"upstream": "kubeflow", "component": "training"},
+        },
+        SUCCESS,
+        create_job(
+            num_workers=1,
+            labels={"upstream": "kubeflow", "component": "training"},
+            annotations=None,
+        ),
+    ),
+    (
+        "valid flow with annotations",
+        {
+            "name": TEST_NAME,
+            "namespace": TEST_NAME,
+            "base_image": TEST_IMAGE,
+            "num_workers": 1,
+            "annotations": {"purpose": "unit-test", "env": "test"},
+        },
+        SUCCESS,
+        create_job(
+            num_workers=1,
+            labels=None,
+            annotations={"purpose": "unit-test", "env": "test"},
+        ),
+    ),
 ]
 
 test_data_get_job_pods = [

sdk/python/kubeflow/training/utils/utils.py

@@ -273,6 +273,8 @@ def get_tfjob_template(
     namespace: str,
     pod_template_spec: models.V1PodTemplateSpec,
     num_workers: int,
+    labels: Optional[Dict[str, str]] = None,
+    annotations: Optional[Dict[str, str]] = None,
     num_chief_replicas: Optional[int] = None,
     num_ps_replicas: Optional[int] = None,
 ):
@@ -297,6 +299,12 @@ def get_tfjob_template(
             )
         )
 
+    if labels is not None:
+        tfjob.metadata.labels = labels
+
+    if annotations is not None:
+        tfjob.metadata.annotations = annotations
+
     if num_ps_replicas is not None:
         tfjob.spec.tf_replica_specs[constants.REPLICA_TYPE_PS] = (
             models.KubeflowOrgV1ReplicaSpec(
@@ -321,6 +329,8 @@ def get_pytorchjob_template(
     namespace: str,
     num_workers: int,
     worker_pod_template_spec: Optional[models.V1PodTemplateSpec],
+    labels: Optional[Dict[str, str]] = None,
+    annotations: Optional[Dict[str, str]] = None,
     master_pod_template_spec: Optional[models.V1PodTemplateSpec] = None,
     num_procs_per_worker: Optional[Union[int, str]] = None,
 ):
@@ -336,7 +346,13 @@ def get_pytorchjob_template(
         ),
     )
 
-    if num_procs_per_worker:
+    if labels is not None:
+        pytorchjob.metadata.labels = labels
+
+    if annotations is not None:
+        pytorchjob.metadata.annotations = annotations
+
+    if num_procs_per_worker is not None:
         pytorchjob.spec.nproc_per_node = str(num_procs_per_worker)
 
     # Create Master replica if that is set.