Merge branch 'main' into feature/aquav1.0.3

VipulMascarenhas · VipulMascarenhas · commit 4b2b653c951e · 2024-07-22T13:05:02.000-07:00
diff --git a/ads/jobs/builders/infrastructure/dsc_job.py b/ads/jobs/builders/infrastructure/dsc_job.py
@@ -30,6 +30,7 @@
 from ads.common.oci_resource import ResourceNotFoundError
 from ads.jobs.builders.infrastructure.base import Infrastructure, RunInstance
 from ads.jobs.builders.infrastructure.dsc_job_runtime import (
+    ContainerRuntimeHandler,
     DataScienceJobRuntimeManager,
 )
 from ads.jobs.builders.infrastructure.utils import get_value
@@ -458,14 +459,19 @@ def run(self, **kwargs) -> DataScienceJobRun:
         ----------
         **kwargs :
             Keyword arguments for initializing a Data Science Job Run.
-            The keys can be any keys in supported by OCI JobConfigurationDetails and JobRun, including:
+            The keys can be any keys in supported by OCI JobConfigurationDetails, OcirContainerJobEnvironmentConfigurationDetails and JobRun, including:
             * hyperparameter_values: dict(str, str)
             * environment_variables: dict(str, str)
             * command_line_arguments: str
             * maximum_runtime_in_minutes: int
             * display_name: str
             * freeform_tags: dict(str, str)
             * defined_tags: dict(str, dict(str, object))
+            * image: str
+            * cmd: list[str]
+            * entrypoint: list[str]
+            * image_digest: str
+            * image_signature_id: str
 
         If display_name is not specified, it will be generated as "<JOB_NAME>-run-<TIMESTAMP>".
 
@@ -478,14 +484,28 @@ def run(self, **kwargs) -> DataScienceJobRun:
         if not self.id:
             self.create()
 
-        swagger_types = (
+        config_swagger_types = (
             oci.data_science.models.DefaultJobConfigurationDetails().swagger_types.keys()
         )
+        env_config_swagger_types = {}
+        if hasattr(oci.data_science.models, "OcirContainerJobEnvironmentConfigurationDetails"):
+            env_config_swagger_types = (
+                oci.data_science.models.OcirContainerJobEnvironmentConfigurationDetails().swagger_types.keys()
+            )
         config_kwargs = {}
+        env_config_kwargs = {}
         keys = list(kwargs.keys())
         for key in keys:
-            if key in swagger_types:
+            if key in config_swagger_types:
                 config_kwargs[key] = kwargs.pop(key)
+            elif key in env_config_swagger_types:
+                value = kwargs.pop(key)
+                if key in [
+                    ContainerRuntime.CONST_CMD,
+                    ContainerRuntime.CONST_ENTRYPOINT
+                ] and isinstance(value, str):
+                    value = ContainerRuntimeHandler.split_args(value)
+                env_config_kwargs[key] = value
 
         # remove timestamp from the job name (added in default names, when display_name not specified by user)
         if self.display_name:
@@ -514,6 +534,12 @@ def run(self, **kwargs) -> DataScienceJobRun:
             config_override.update(config_kwargs)
             kwargs["job_configuration_override_details"] = config_override
 
+        if env_config_kwargs:
+            env_config_kwargs["jobEnvironmentType"] = "OCIR_CONTAINER"
+            env_config_override = kwargs.get("job_environment_configuration_override_details", {})
+            env_config_override.update(env_config_kwargs)
+            kwargs["job_environment_configuration_override_details"] = env_config_override
+
         wait = kwargs.pop("wait", False)
         run = DataScienceJobRun(**kwargs, **self.auth).create()
         if wait:
diff --git a/ads/jobs/builders/infrastructure/dsc_job_runtime.py b/ads/jobs/builders/infrastructure/dsc_job_runtime.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2021, 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 """Contains classes for conversion between ADS runtime and OCI Data Science Job implementation.
 This module is for ADS developers only.
@@ -305,10 +305,29 @@ def extract(self, dsc_job):
             self._extract_envs,
             self._extract_artifact,
             self._extract_runtime_minutes,
+            self._extract_properties,
         ]
         for extraction in extractions:
             runtime_spec.update(extraction(dsc_job))
         return self.RUNTIME_CLASS(self._format_env_var(runtime_spec))
+    
+    def _extract_properties(self, dsc_job) -> dict:
+        """Extract the job runtime properties from data science job.
+
+        This is the base method which does not extract the job runtime properties.
+        Sub-class should implement the extraction if needed.
+
+        Parameters
+        ----------
+        dsc_job : DSCJob or oci.datascience.models.Job
+            The data science job containing runtime information.
+
+        Returns
+        -------
+        dict
+            A runtime specification dictionary for initializing a runtime.
+        """
+        return {}
 
     def _extract_args(self, dsc_job) -> dict:
         """Extracts the command line arguments from data science job.
@@ -942,9 +961,12 @@ def _extract_artifact(self, dsc_job):
 class ContainerRuntimeHandler(RuntimeHandler):
     RUNTIME_CLASS = ContainerRuntime
     CMD_DELIMITER = ","
-    CONST_CONTAINER_IMAGE = "CONTAINER_CUSTOM_IMAGE"
-    CONST_CONTAINER_ENTRYPOINT = "CONTAINER_ENTRYPOINT"
-    CONST_CONTAINER_CMD = "CONTAINER_CMD"
+
+    def translate(self, runtime: Runtime) -> dict:
+        payload = super().translate(runtime)
+        job_env_config = self._translate_env_config(runtime)
+        payload["job_environment_configuration_details"] = job_env_config
+        return payload
 
     def _translate_artifact(self, runtime: Runtime):
         """Specifies a dummy script as the job artifact.
@@ -964,29 +986,34 @@ def _translate_artifact(self, runtime: Runtime):
             os.path.dirname(__file__), "../../templates", "container.py"
         )
 
-    def _translate_env(self, runtime: ContainerRuntime) -> dict:
-        """Translate the environment variable.
+    def _translate_env_config(self, runtime: Runtime) -> dict:
+        """Converts runtime properties to ``OcirContainerJobEnvironmentConfigurationDetails`` payload required by OCI Data Science job.
 
         Parameters
         ----------
-        runtime : GitPythonRuntime
-            An instance of GitPythonRuntime
+        runtime : Runtime
+            The runtime containing the properties to be converted.
 
         Returns
         -------
         dict
-            A dictionary containing environment variables for OCI data science job.
+            A dictionary storing the ``OcirContainerJobEnvironmentConfigurationDetails`` payload for OCI data science job.
         """
-        if not runtime.image:
-            raise ValueError("Specify container image for ContainerRuntime.")
-        envs = super()._translate_env(runtime)
-        spec_mappings = {
-            ContainerRuntime.CONST_IMAGE: self.CONST_CONTAINER_IMAGE,
-            ContainerRuntime.CONST_ENTRYPOINT: self.CONST_CONTAINER_ENTRYPOINT,
-            ContainerRuntime.CONST_CMD: self.CONST_CONTAINER_CMD,
+        job_environment_configuration_details = {
+            "job_environment_type": runtime.job_env_type
         }
-        envs.update(self._translate_specs(runtime, spec_mappings, self.CMD_DELIMITER))
-        return envs
+
+        for key, value in ContainerRuntime.attribute_map.items():
+            property = runtime.get_spec(key, None)
+            if key in [
+                ContainerRuntime.CONST_CMD,
+                ContainerRuntime.CONST_ENTRYPOINT
+            ] and isinstance(property, str):
+                property = self.split_args(property)
+            if property is not None:
+                job_environment_configuration_details[value] = property
+
+        return job_environment_configuration_details
 
     @staticmethod
     def split_args(args: str) -> list:
@@ -1031,17 +1058,37 @@ def _extract_envs(self, dsc_job):
         """
         spec = super()._extract_envs(dsc_job)
         envs = spec.pop(ContainerRuntime.CONST_ENV_VAR, {})
-        if self.CONST_CONTAINER_IMAGE not in envs:
-            raise IncompatibleRuntime()
-        spec[ContainerRuntime.CONST_IMAGE] = envs.pop(self.CONST_CONTAINER_IMAGE)
-        cmd = self.split_args(envs.pop(self.CONST_CONTAINER_CMD, ""))
-        if cmd:
-            spec[ContainerRuntime.CONST_CMD] = cmd
-        entrypoint = self.split_args(envs.pop(self.CONST_CONTAINER_ENTRYPOINT, ""))
-        if entrypoint:
-            spec[ContainerRuntime.CONST_ENTRYPOINT] = entrypoint
+
         if envs:
             spec[ContainerRuntime.CONST_ENV_VAR] = envs
+
+        return spec
+    
+    def _extract_properties(self, dsc_job) -> dict:
+        """Extract the runtime properties from data science job.
+
+        Parameters
+        ----------
+        dsc_job : DSCJob or oci.datascience.models.Job
+            The data science job containing runtime information.
+
+        Returns
+        -------
+        dict
+            A runtime specification dictionary for initializing a runtime.
+        """
+        spec = super()._extract_envs(dsc_job)
+        
+        job_env_config = getattr(dsc_job, "job_environment_configuration_details", None)
+        job_env_type = getattr(job_env_config, "job_environment_type", None)
+        
+        if not (job_env_config and job_env_type == "OCIR_CONTAINER"):
+            raise IncompatibleRuntime()
+
+        for key, value in ContainerRuntime.attribute_map.items():
+            property = getattr(job_env_config, value, None)
+            if property is not None:
+                spec[key] = property
         return spec
 
 
diff --git a/ads/jobs/builders/runtimes/container_runtime.py b/ads/jobs/builders/runtimes/container_runtime.py
@@ -3,28 +3,36 @@
 
 # Copyright (c) 2021, 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+import logging
 from typing import Union
 from ads.jobs.builders.runtimes.base import MultiNodeRuntime
 
+logger = logging.getLogger(__name__)
+
 
 class ContainerRuntime(MultiNodeRuntime):
     """Represents a container job runtime
 
     To define container runtime:
 
     >>> ContainerRuntime()
-    >>> .with_image("iad.ocir.io/<your_tenancy>/<your_image>")
+    >>> .with_image("iad.ocir.io/<your_tenancy>/<your_image>:<tag>")
     >>> .with_cmd("sleep 5 && echo Hello World")
     >>> .with_entrypoint(["/bin/sh", "-c"])
+    >>> .with_image_digest("<image_digest>")
+    >>> .with_image_signature_id("<image_signature_id>")
     >>> .with_environment_variable(MY_ENV="MY_VALUE")
 
-    Alternatively, you can define the ``entrypoint`` and ``cmd`` along with the image.
+    Alternatively, you can define the ``entrypoint``, ``cmd``,
+    ``image_digest``and ``image_signature_id`` along with the image.
 
     >>> ContainerRuntime()
     >>> .with_image(
-    >>>     "iad.ocir.io/<your_tenancy>/<your_image>",
+    >>>     "iad.ocir.io/<your_tenancy>/<your_image>:<tag>",
     >>>     entrypoint=["/bin/sh", "-c"],
     >>>     cmd="sleep 5 && echo Hello World",
+    >>>     image_digest="<image_digest>",
+    >>>     image_signature_id="<image_signature_id>",
     >>> )
     >>> .with_environment_variable(MY_ENV="MY_VALUE")
 
@@ -46,20 +54,34 @@ class ContainerRuntime(MultiNodeRuntime):
     CONST_IMAGE = "image"
     CONST_ENTRYPOINT = "entrypoint"
     CONST_CMD = "cmd"
+    CONST_IMAGE_DIGEST = "imageDigest"
+    CONST_IMAGE_SIGNATURE_ID = "imageSignatureId"
     attribute_map = {
         CONST_IMAGE: CONST_IMAGE,
         CONST_ENTRYPOINT: CONST_ENTRYPOINT,
         CONST_CMD: CONST_CMD,
+        CONST_IMAGE_DIGEST: "image_digest",
+        CONST_IMAGE_SIGNATURE_ID: "image_signature_id",
     }
     attribute_map.update(MultiNodeRuntime.attribute_map)
 
+    @property
+    def job_env_type(self) -> str:
+        """The container type"""
+        return "OCIR_CONTAINER"
+
     @property
     def image(self) -> str:
         """The container image"""
         return self.get_spec(self.CONST_IMAGE)
 
     def with_image(
-        self, image: str, entrypoint: Union[str, list, None] = None, cmd: str = None
+        self,
+        image: str,
+        entrypoint: Union[str, list, None] = None, 
+        cmd: str = None,
+        image_digest: str = None,
+        image_signature_id: str = None,
     ) -> "ContainerRuntime":
         """Specify the image for the container job.
 
@@ -71,16 +93,73 @@ def with_image(
             Entrypoint for the job, by default None (the entrypoint defined in the image will be used).
         cmd : str, optional
             Command for the job, by default None.
+        image_digest: str, optional
+            The image digest, by default None.
+        image_signature_id: str, optional
+            The image signature id, by default None.
 
         Returns
         -------
         ContainerRuntime
             The runtime instance.
         """
+        if not isinstance(image, str):
+            raise ValueError(
+                "Custom image must be provided as a string."
+            )
+        if image.find(":") < 0:
+            logger.warning(
+                "Tag is required for custom image. Accepted format: iad.ocir.io/<tenancy>/<image>:<tag>."
+            )
         self.with_entrypoint(entrypoint)
         self.set_spec(self.CONST_CMD, cmd)
+        self.with_image_digest(image_digest)
+        self.with_image_signature_id(image_signature_id)
         return self.set_spec(self.CONST_IMAGE, image)
 
+    @property
+    def image_digest(self) -> str:
+        """The container image digest."""
+        return self.get_spec(self.CONST_IMAGE_DIGEST)
+    
+    def with_image_digest(self, image_digest: str) -> "ContainerRuntime":
+        """Sets the digest of custom image.
+
+        Parameters
+        ----------
+        image_digest: str
+            The image digest.
+
+        Returns
+        -------
+        ContainerRuntime
+            The runtime instance.
+        """
+        return self.set_spec(self.CONST_IMAGE_DIGEST, image_digest)
+    
+    @property
+    def image_signature_id(self) -> str:
+        """The container image signature id."""
+        return self.get_spec(self.CONST_IMAGE_SIGNATURE_ID)
+    
+    def with_image_signature_id(self, image_signature_id: str) -> "ContainerRuntime":
+        """Sets the signature id of custom image.
+
+        Parameters
+        ----------
+        image_signature_id: str
+            The image signature id.
+
+        Returns
+        -------
+        ContainerRuntime
+            The runtime instance.
+        """
+        return self.set_spec(
+            self.CONST_IMAGE_SIGNATURE_ID,
+            image_signature_id
+        )
+
     @property
     def entrypoint(self) -> str:
         """Entrypoint of the container job"""
diff --git a/docs/source/user_guide/jobs/run_container.rst b/docs/source/user_guide/jobs/run_container.rst
@@ -22,7 +22,7 @@ Here is an example to create and run a container job:
 
 To configure ``ContainerRuntime``, you must specify the container ``image``.
 Similar to other runtime, you can add environment variables.
-You can optionally specify the `entrypoint` and `cmd` for running the container.
+You can optionally specify the `entrypoint`, `cmd`, `image_digest` and `image_signature_id` for running the container.
 
 See also:
 
diff --git a/docs/source/user_guide/jobs/tabs/container_runtime.rst b/docs/source/user_guide/jobs/tabs/container_runtime.rst
@@ -27,7 +27,9 @@
         )
         .with_runtime(
             ContainerRuntime()
-            .with_image("<region>.ocir.io/<your_tenancy>/<your_image>")
+            .with_image("<region>.ocir.io/<tenancy>/<image>:<tag>")
+            .with_image_digest("<image_digest>")
+            .with_image_signature_id("<image_signature_id>")
             .with_environment_variable(GREETINGS="Welcome to OCI Data Science")
             .with_entrypoint(["/bin/sh", "-c"])
             .with_cmd("sleep 5 && echo $GREETINGS")
diff --git a/tests/unitary/default_setup/jobs/test_jobs_base.py b/tests/unitary/default_setup/jobs/test_jobs_base.py
diff --git a/tests/unitary/default_setup/jobs/test_jobs_container_runtime.py b/tests/unitary/default_setup/jobs/test_jobs_container_runtime.py
