Support setting "defined tags" in data science jobs. (#167)

Author: qiuosier

ads/jobs/ads_job.py

@@ -74,6 +74,9 @@ class Job(Builder):
                 .with_python_path("other_packages")
                 # Copy files in "code_dir/output" to object storage after job finishes.
                 .with_output("output", "oci://bucket_name@namespace/path/to/dir")
+                # Tags
+                .with_freeform_tag(my_tag="my_value")
+                .with_defined_tag(**{"Operations": {"CostCenter": "42"}})
             )
         )
         # Create and Run the job

ads/jobs/builders/infrastructure/base.py

@@ -61,6 +61,7 @@ def run(
         args: str = None,
         env_var: dict = None,
         freeform_tags: dict = None,
+        defined_tags: dict = None,
         wait: bool = False,
     ):
         """Runs a job on the infrastructure.
@@ -72,9 +73,11 @@ def run(
         args : str, optional
             Command line arguments for the job run, by default None.
         env_var : dict, optional
-            Environment variable for the job run, by default None
+            Environment variable for the job run, by default None.
         freeform_tags : dict, optional
-            Freeform tags for the job run, by default None
+            Freeform tags for the job run, by default None.
+        defined_tags : dict, optional
+            Defined tags for the job run, by default None.
         wait : bool, optional
             Indicate if this method should wait for the run to finish before it returns, by default False.
         """

ads/jobs/builders/infrastructure/dsc_job.py

@@ -35,10 +35,7 @@
 from ads.jobs.builders.runtimes.container_runtime import ContainerRuntime
 from ads.jobs.builders.runtimes.python_runtime import GitPythonRuntime
 
-from ads.common.dsc_file_system import (
-    OCIFileStorage,
-    DSCFileSystemManager
-)
+from ads.common.dsc_file_system import OCIFileStorage, DSCFileSystemManager
 
 logger = logging.getLogger(__name__)
 
@@ -445,6 +442,8 @@ def run(self, **kwargs) -> DataScienceJobRun:
             * command_line_arguments: str
             * maximum_runtime_in_minutes: int
             * display_name: str
+            * freeform_tags: dict(str, str)
+            * defined_tags: dict(str, dict(str, object))
 
         If display_name is not specified, it will be generated as "<JOB_NAME>-run-<TIMESTAMP>".
 
@@ -845,9 +844,12 @@ class DataScienceJob(Infrastructure):
             .with_storage_mount(
                 {
                     "src" : "<mount_target_ip_address>:<export_path>",
-	            "dest" : "<destination_directory_name>"
+                    "dest" : "<destination_directory_name>"
                 }
             )
+            # Tags
+            .with_freeform_tag(my_tag="my_value")
+            .with_defined_tag(**{"Operations": {"CostCenter": "42"}})
         )
 
     """
@@ -866,6 +868,8 @@ class DataScienceJob(Infrastructure):
     CONST_LOG_ID = "logId"
     CONST_LOG_GROUP_ID = "logGroupId"
     CONST_STORAGE_MOUNT = "storageMount"
+    CONST_FREEFORM_TAGS = "freeformTags"
+    CONST_DEFINED_TAGS = "definedTags"
 
     attribute_map = {
         CONST_PROJECT_ID: "project_id",
@@ -880,6 +884,8 @@ class DataScienceJob(Infrastructure):
         CONST_LOG_ID: "log_id",
         CONST_LOG_GROUP_ID: "log_group_id",
         CONST_STORAGE_MOUNT: "storage_mount",
+        CONST_FREEFORM_TAGS: "freeform_tags",
+        CONST_DEFINED_TAGS: "defined_tags",
     }
 
     shape_config_details_attribute_map = {
@@ -1231,9 +1237,7 @@ def log_group_id(self) -> str:
         """
         return self.get_spec(self.CONST_LOG_GROUP_ID)
 
-    def with_storage_mount(
-        self, *storage_mount: List[dict]
-    ) -> DataScienceJob:
+    def with_storage_mount(self, *storage_mount: List[dict]) -> DataScienceJob:
         """Sets the file systems to be mounted for the data science job.
         A maximum number of 5 file systems are allowed to be mounted for a single data science job.
 
@@ -1271,6 +1275,36 @@ def storage_mount(self) -> List[dict]:
         """
         return self.get_spec(self.CONST_STORAGE_MOUNT, [])
 
+    def with_freeform_tag(self, **kwargs) -> DataScienceJob:
+        """Sets freeform tags
+
+        Returns
+        -------
+        DataScienceJob
+            The DataScienceJob instance (self)
+        """
+        return self.set_spec(self.CONST_FREEFORM_TAGS, kwargs)
+
+    def with_defined_tag(self, **kwargs) -> DataScienceJob:
+        """Sets defined tags
+
+        Returns
+        -------
+        DataScienceJob
+            The DataScienceJob instance (self)
+        """
+        return self.set_spec(self.CONST_DEFINED_TAGS, kwargs)
+
+    @property
+    def freeform_tags(self) -> dict:
+        """Freeform tags"""
+        return self.get_spec(self.CONST_FREEFORM_TAGS, {})
+
+    @property
+    def defined_tags(self) -> dict:
+        """Defined tags"""
+        return self.get_spec(self.CONST_DEFINED_TAGS, {})
+
     def _prepare_log_config(self) -> dict:
         if not self.log_group_id and not self.log_id:
             return None
@@ -1425,7 +1459,8 @@ def _update_job_infra(self, dsc_job: DSCJob) -> DataScienceJob:
                     "Storage mount hasn't been supported in the current OCI SDK installed."
                 )
             dsc_job.job_storage_mount_configuration_details_list = [
-                DSCFileSystemManager.initialize(file_system) for file_system in self.storage_mount
+                DSCFileSystemManager.initialize(file_system)
+                for file_system in self.storage_mount
             ]
         return self
 
@@ -1467,6 +1502,10 @@ def create(self, runtime, **kwargs) -> DataScienceJob:
 
         payload["display_name"] = display_name
         payload["job_log_configuration_details"] = self._prepare_log_config()
+        if not payload.get("freeform_tags"):
+            payload["freeform_tags"] = self.freeform_tags
+        if not payload.get("defined_tags"):
+            payload["defined_tags"] = self.defined_tags
 
         self.dsc_job = DSCJob(**payload)
         # Set Job infra to user values after DSCJob initialized the defaults
@@ -1477,7 +1516,13 @@ def create(self, runtime, **kwargs) -> DataScienceJob:
         return self
 
     def run(
-        self, name=None, args=None, env_var=None, freeform_tags=None, wait=False
+        self,
+        name=None,
+        args=None,
+        env_var=None,
+        freeform_tags=None,
+        defined_tags=None,
+        wait=False,
     ) -> DataScienceJobRun:
         """Runs a job on OCI Data Science job
 
@@ -1491,6 +1536,8 @@ def run(
             Environment variable for the job run, by default None
         freeform_tags : dict, optional
             Freeform tags for the job run, by default None
+        defined_tags : dict, optional
+            Defined tags for the job run, by default None
         wait : bool, optional
             Indicate if this method should wait for the run to finish before it returns, by default False.
 
@@ -1505,11 +1552,18 @@ def run(
             raise RuntimeError(
                 "Job is not created. Call create() to create the job first."
             )
-        tags = self.runtime.freeform_tags
-        if not tags:
-            tags = {}
-        if freeform_tags:
-            tags.update(freeform_tags)
+
+        if not freeform_tags:
+            freeform_tags = {}
+        runtime_freeform_tags = self.runtime.freeform_tags
+        if runtime_freeform_tags:
+            freeform_tags.update(runtime_freeform_tags)
+
+        if not defined_tags:
+            defined_tags = {}
+        runtime_defined_tags = self.runtime.defined_tags
+        if runtime_defined_tags:
+            defined_tags.update(runtime_defined_tags)
 
         if name:
             envs = self.runtime.envs
@@ -1521,7 +1575,8 @@ def run(
             display_name=name,
             command_line_arguments=args,
             environment_variables=env_var,
-            freeform_tags=tags,
+            freeform_tags=freeform_tags,
+            defined_tags=defined_tags,
             wait=wait,
         )
 

ads/jobs/builders/infrastructure/dsc_job_runtime.py

@@ -38,7 +38,7 @@
 from ads.jobs.builders.infrastructure.utils import get_value
 
 
-class IncompatibleRuntime(Exception):   # pragma: no cover
+class IncompatibleRuntime(Exception):  # pragma: no cover
     """Represents an exception when runtime is not compatible with the OCI data science job configuration.
     This exception is designed to be raised during the extraction of a runtime from OCI data science job.
     The data science job does not explicitly contain information of the type of the ADS runtime.
@@ -104,6 +104,8 @@ def translate(self, runtime: Runtime) -> dict:
         payload["job_configuration_details"] = self._translate_config(runtime)
         if runtime.freeform_tags:
             payload["freeform_tags"] = runtime.freeform_tags
+        if runtime.defined_tags:
+            payload["defined_tags"] = runtime.defined_tags
         self.data_science_job.runtime = runtime
         return payload
 
@@ -353,10 +355,14 @@ def _extract_tags(self, dsc_job):
         dict
             A runtime specification dictionary for initializing a runtime.
         """
+        tags = {}
         value = get_value(dsc_job, "freeform_tags")
         if value:
-            return {Runtime.CONST_TAG: value}
-        return {}
+            tags[Runtime.CONST_FREEFORM_TAGS] = value
+        value = get_value(dsc_job, "defined_tags")
+        if value:
+            tags[Runtime.CONST_DEFINED_TAGS] = value
+        return tags
 
     def _extract_artifact(self, dsc_job):
         """Extract the job artifact from data science job.

ads/jobs/builders/runtimes/base.py

@@ -24,10 +24,12 @@ class Runtime(Builder):
     CONST_ENV_VAR = "env"
     CONST_ARGS = "args"
     CONST_MAXIMUM_RUNTIME_IN_MINUTES = "maximumRuntimeInMinutes"
-    CONST_TAG = "freeformTags"
+    CONST_FREEFORM_TAGS = "freeformTags"
+    CONST_DEFINED_TAGS = "definedTags"
 
     attribute_map = {
-        CONST_TAG: "freeform_tags",
+        CONST_FREEFORM_TAGS: "freeform_tags",
+        CONST_DEFINED_TAGS: "defined_tags",
         CONST_ENV_VAR: CONST_ENV_VAR,
     }
 
@@ -164,14 +166,24 @@ def with_environment_variable(self: Self, **kwargs) -> Self:
         return self.set_spec(self.CONST_ENV_VAR, envs)
 
     def with_freeform_tag(self: Self, **kwargs) -> Self:
-        """Sets freeform tag
+        """Sets freeform tags
 
         Returns
         -------
         Self
             This method returns self to support chaining methods.
         """
-        return self.set_spec(self.CONST_TAG, kwargs)
+        return self.set_spec(self.CONST_FREEFORM_TAGS, kwargs)
+
+    def with_defined_tag(self: Self, **kwargs) -> Self:
+        """Sets defined tags
+
+        Returns
+        -------
+        Self
+            This method returns self to support chaining methods.
+        """
+        return self.set_spec(self.CONST_DEFINED_TAGS, kwargs)
 
     def with_maximum_runtime_in_minutes(
         self: Self, maximum_runtime_in_minutes: int
@@ -209,8 +221,13 @@ def envs(self) -> dict:
 
     @property
     def freeform_tags(self) -> dict:
-        """freeform_tags"""
-        return self.get_spec(self.CONST_TAG, {})
+        """Freeform tags"""
+        return self.get_spec(self.CONST_FREEFORM_TAGS, {})
+
+    @property
+    def defined_tags(self) -> dict:
+        """Defined tags"""
+        return self.get_spec(self.CONST_DEFINED_TAGS, {})
 
     @property
     def args(self) -> list:

tests/unitary/default_setup/jobs/test_jobs_base.py

@@ -10,7 +10,7 @@
 import unittest
 from unittest import mock
 from zipfile import ZipFile
-from unittest.mock import PropertyMock, patch
+from unittest.mock import patch
 import pytest
 
 

tests/unitary/default_setup/jobs/test_jobs_tags.py

@@ -0,0 +1,65 @@
+from ads.jobs import Job, DataScienceJob, ContainerRuntime
+from tests.unitary.default_setup.jobs.test_jobs_base import DataScienceJobPayloadTest
+
+
+class JobTagTestCase(DataScienceJobPayloadTest):
+    @staticmethod
+    def runtime() -> ContainerRuntime:
+        return ContainerRuntime().with_image(
+            "iad.ocir.io/my_namespace/my_ubuntu_image",
+            entrypoint="/bin/sh",
+            cmd="-c,echo Hello World",
+        )
+
+    @staticmethod
+    def infra() -> DataScienceJob:
+        return (
+            DataScienceJob()
+            .with_compartment_id("ocid1.compartment.oc1..<unique_ocid>")
+            .with_project_id("ocid1.datascienceproject.oc1.iad.<unique_ocid>")
+        )
+
+    def create_job(self, infra, runtime) -> dict:
+        job = (
+            Job(name=self.__class__.__name__)
+            .with_infrastructure(infra)
+            .with_runtime(runtime)
+        )
+        job = self.mock_create_job(job)
+        return job.infrastructure.dsc_job.to_dict()
+
+    def test_create_job_with_runtime_tags(self):
+        runtime = (
+            self.runtime()
+            .with_freeform_tag(freeform_tag="freeform_tag_val")
+            .with_defined_tag(Operations={"CostCenter": "42"})
+        )
+        payload = self.create_job(self.infra(), runtime)
+        self.assertEqual(payload["freeformTags"], dict(freeform_tag="freeform_tag_val"))
+        self.assertEqual(payload["definedTags"], {"Operations": {"CostCenter": "42"}})
+
+    def test_create_job_with_infra_tags(self):
+        infra = (
+            self.infra()
+            .with_freeform_tag(freeform_tag="freeform_tag_val")
+            .with_defined_tag(Operations={"CostCenter": "42"})
+        )
+        payload = self.create_job(infra, self.runtime())
+        self.assertEqual(payload["freeformTags"], dict(freeform_tag="freeform_tag_val"))
+        self.assertEqual(payload["definedTags"], {"Operations": {"CostCenter": "42"}})
+
+    def test_create_job_with_infra_and_runtime_tags(self):
+        # Tags defined in runtime will have higher priority
+        infra = (
+            self.infra()
+            .with_freeform_tag(freeform_tag="freeform_tag_val")
+            .with_defined_tag(Operations={"CostCenter": "41"})
+        )
+        runtime = (
+            self.runtime()
+            .with_freeform_tag(freeform_tag="freeform_tag_val")
+            .with_defined_tag(Operations={"CostCenter": "42"})
+        )
+        payload = self.create_job(infra, runtime)
+        self.assertEqual(payload["freeformTags"], dict(freeform_tag="freeform_tag_val"))
+        self.assertEqual(payload["definedTags"], {"Operations": {"CostCenter": "42"}})