Support adding artifacts for container jobs (#936)

Author: qiuosier

ads/jobs/builders/infrastructure/dsc_job.py

@@ -312,7 +312,7 @@ def _create_with_oci_api(self) -> None:
         logger.debug(oci_model)
         res = self.client.create_job(oci_model)
         self.update_from_oci_model(res.data)
-        if self.lifecycle_state == "ACTIVE":
+        if not self.artifact:
             return
         try:
             if issubclass(self.artifact.__class__, Artifact):
@@ -487,7 +487,9 @@ def run(self, **kwargs) -> DataScienceJobRun:
             oci.data_science.models.DefaultJobConfigurationDetails().swagger_types.keys()
         )
         env_config_swagger_types = {}
-        if hasattr(oci.data_science.models, "OcirContainerJobEnvironmentConfigurationDetails"):
+        if hasattr(
+            oci.data_science.models, "OcirContainerJobEnvironmentConfigurationDetails"
+        ):
             env_config_swagger_types = (
                 oci.data_science.models.OcirContainerJobEnvironmentConfigurationDetails().swagger_types.keys()
             )
@@ -501,7 +503,7 @@ def run(self, **kwargs) -> DataScienceJobRun:
                 value = kwargs.pop(key)
                 if key in [
                     ContainerRuntime.CONST_CMD,
-                    ContainerRuntime.CONST_ENTRYPOINT
+                    ContainerRuntime.CONST_ENTRYPOINT,
                 ] and isinstance(value, str):
                     value = ContainerRuntimeHandler.split_args(value)
                 env_config_kwargs[key] = value
@@ -535,9 +537,13 @@ def run(self, **kwargs) -> DataScienceJobRun:
 
         if env_config_kwargs:
             env_config_kwargs["jobEnvironmentType"] = "OCIR_CONTAINER"
-            env_config_override = kwargs.get("job_environment_configuration_override_details", {})
+            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
+            kwargs["job_environment_configuration_override_details"] = (
+                env_config_override
+            )
 
         wait = kwargs.pop("wait", False)
         run = DataScienceJobRun(**kwargs, **self.auth).create()

ads/jobs/builders/infrastructure/dsc_job_runtime.py

@@ -181,9 +181,9 @@ def _translate_config(self, runtime: Runtime) -> dict:
             "jobType": self.data_science_job.job_type,
         }
         if runtime.maximum_runtime_in_minutes:
-            job_configuration_details[
-                "maximum_runtime_in_minutes"
-            ] = runtime.maximum_runtime_in_minutes
+            job_configuration_details["maximum_runtime_in_minutes"] = (
+                runtime.maximum_runtime_in_minutes
+            )
         job_configuration_details["environment_variables"] = self._translate_env(
             runtime
         )
@@ -310,7 +310,7 @@ def extract(self, dsc_job):
         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.
 
@@ -968,23 +968,10 @@ def translate(self, runtime: Runtime) -> dict:
         payload["job_environment_configuration_details"] = job_env_config
         return payload
 
-    def _translate_artifact(self, runtime: Runtime):
-        """Specifies a dummy script as the job artifact.
-        runtime is not used in this method.
-
-        Parameters
-        ----------
-        runtime : Runtime
-            This is not used.
-
-        Returns
-        -------
-        str
-            Path to the dummy script.
-        """
-        return os.path.join(
-            os.path.dirname(__file__), "../../templates", "container.py"
-        )
+    def _translate_artifact(self, runtime: ContainerRuntime):
+        """Additional artifact for the container"""
+        if runtime.artifact_uri:
+            return ScriptArtifact(runtime.artifact_uri, runtime)
 
     def _translate_env_config(self, runtime: Runtime) -> dict:
         """Converts runtime properties to ``OcirContainerJobEnvironmentConfigurationDetails`` payload required by OCI Data Science job.
@@ -1007,7 +994,7 @@ def _translate_env_config(self, runtime: Runtime) -> dict:
             property = runtime.get_spec(key, None)
             if key in [
                 ContainerRuntime.CONST_CMD,
-                ContainerRuntime.CONST_ENTRYPOINT
+                ContainerRuntime.CONST_ENTRYPOINT,
             ] and isinstance(property, str):
                 property = self.split_args(property)
             if property is not None:
@@ -1063,7 +1050,7 @@ def _extract_envs(self, dsc_job):
             spec[ContainerRuntime.CONST_ENV_VAR] = envs
 
         return spec
-    
+
     def _extract_properties(self, dsc_job) -> dict:
         """Extract the runtime properties from data science job.
 
@@ -1078,10 +1065,10 @@ def _extract_properties(self, dsc_job) -> 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()
 

ads/jobs/builders/runtimes/artifact.py

@@ -183,11 +183,6 @@ def build(self):
         if os.path.isdir(source):
             basename = os.path.basename(str(source).rstrip("/"))
             source = str(source).rstrip("/")
-            # Runtime must have entrypoint if the source is a directory
-            if self.runtime and not self.runtime.entrypoint:
-                raise ValueError(
-                    "Please specify entrypoint when script source is a directory."
-                )
             output = os.path.join(self.temp_dir.name, basename)
             shutil.make_archive(
                 output, "zip", os.path.dirname(source), base_dir=basename

ads/jobs/builders/runtimes/container_runtime.py

@@ -56,6 +56,7 @@ class ContainerRuntime(MultiNodeRuntime):
     CONST_CMD = "cmd"
     CONST_IMAGE_DIGEST = "imageDigest"
     CONST_IMAGE_SIGNATURE_ID = "imageSignatureId"
+    CONST_SCRIPT_PATH = "scriptPathURI"
     attribute_map = {
         CONST_IMAGE: CONST_IMAGE,
         CONST_ENTRYPOINT: CONST_ENTRYPOINT,
@@ -121,7 +122,7 @@ def with_image(
     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.
 
@@ -136,12 +137,12 @@ def with_image_digest(self, image_digest: str) -> "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.
 
@@ -217,3 +218,25 @@ def init(self, **kwargs) -> "ContainerRuntime":
             entrypoint=["bash", "--login", "-c"],
             cmd="{Container CMD. For MLflow and Operator will be auto generated}",
         )
+
+    @property
+    def artifact_uri(self) -> str:
+        """The URI of the source code"""
+        return self.get_spec(self.CONST_SCRIPT_PATH)
+
+    def with_artifact(self, uri: str):
+        """Specifies the artifact to be added to the container.
+
+        Parameters
+        ----------
+        uri : str
+            URI to the source code script, which can be any URI supported by fsspec,
+            including http://, https:// and OCI object storage.
+            For example: oci://your_bucket@your_namespace/path/to/script.py
+
+        Returns
+        -------
+        self
+            The runtime instance.
+        """
+        return self.set_spec(self.CONST_SCRIPT_PATH, uri)

ads/opctl/conda/cmds.py

@@ -213,7 +213,11 @@ def _create(
 
     if is_in_notebook_session() or NO_CONTAINER:
         command = f"conda env create --prefix {pack_folder_path} --file {os.path.abspath(os.path.expanduser(env_file))}"
-        run_command(command, shell=True)
+        proc = run_command(command, shell=True)
+        if proc.returncode != 0:
+            raise RuntimeError(
+                f"Failed to create conda environment. (exit code {proc.returncode})"
+            )
     else:
         _check_job_image_exists(gpu)
         docker_pack_folder_path = os.path.join(DEFAULT_IMAGE_HOME_DIR, slug)
@@ -660,7 +664,11 @@ def _publish(
             command = f"python {pack_script} --conda-path {pack_folder_path}"
             if publish_type:
                 command = f"CONDA_PUBLISH_TYPE={publish_type} {command}"
-            run_command(command, shell=True)
+            proc = run_command(command, shell=True)
+            if proc.returncode != 0:
+                raise RuntimeError(
+                    f"Failed to archive the conda environment. (exit code {proc.returncode})"
+                )
         else:
             volumes = {
                 pack_folder_path: {

tests/unitary/default_setup/jobs/test_jobs_dscjob.py

@@ -41,19 +41,19 @@ def setup_method(self):
 
     @property
     def sample_create_job_response(self):
-        self.payload[
-            "lifecycle_state"
-        ] = oci.data_science.models.Job.LIFECYCLE_STATE_ACTIVE
+        self.payload["lifecycle_state"] = (
+            oci.data_science.models.Job.LIFECYCLE_STATE_ACTIVE
+        )
         self.payload["id"] = "ocid1.datasciencejob.oc1.iad.<unique_ocid>"
         return Response(
             data=Job(**self.payload), status=None, headers=None, request=None
         )
 
     @property
     def sample_create_job_response_with_default_display_name(self):
-        self.payload[
-            "lifecycle_state"
-        ] = oci.data_science.models.Job.LIFECYCLE_STATE_ACTIVE
+        self.payload["lifecycle_state"] = (
+            oci.data_science.models.Job.LIFECYCLE_STATE_ACTIVE
+        )
         random.seed(random_seed)
         self.payload["display_name"] = utils.get_random_name_for_resource()
         return Response(
@@ -62,9 +62,9 @@ def sample_create_job_response_with_default_display_name(self):
 
     @property
     def sample_create_job_response_with_default_display_name_with_artifact(self):
-        self.payload[
-            "lifecycle_state"
-        ] = oci.data_science.models.Job.LIFECYCLE_STATE_ACTIVE
+        self.payload["lifecycle_state"] = (
+            oci.data_science.models.Job.LIFECYCLE_STATE_ACTIVE
+        )
         random.seed(random_seed)
         self.payload["display_name"] = "my_script_name"
         return Response(
@@ -155,14 +155,15 @@ def test_create_job_with_default_display_name_with_artifact(
         job = DSCJob(artifact=artifact, **self.payload)
         with patch.object(
             DSCJob, "client", mock_client_with_default_display_name_with_artifact
+        ), patch.object(
+            OCIModelMixin,
+            "to_oci_model",
+            mock_details_with_default_display_name_with_artifact,
+        ), patch.object(
+            DSCJob, "upload_artifact"
         ):
-            with patch.object(
-                OCIModelMixin,
-                "to_oci_model",
-                mock_details_with_default_display_name_with_artifact,
-            ):
-                job.create()
-                assert job.display_name == os.path.basename(str(artifact)).split(".")[0]
+            job.create()
+            assert job.display_name == os.path.basename(str(artifact)).split(".")[0]
 
     @pytest.mark.parametrize(
         "test_config_details, expected_result",

tests/unitary/with_extras/opctl/test_opctl_conda.py

@@ -4,7 +4,7 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import pytest
-from mock import patch, MagicMock
+from unittest.mock import patch, MagicMock, PropertyMock
 import tempfile
 import os
 from pathlib import Path
@@ -24,6 +24,8 @@ class TestOpctlConda:
     def test_conda_create(
         self, mock_run_container, mock_docker, mock_run_cmd, monkeypatch
     ):
+        type(mock_run_cmd.return_value).returncode = PropertyMock(return_value=0)
+        mock_run_cmd.returncode = 0
         with pytest.raises(FileNotFoundError):
             create(slug="test Abc", environment_file="environment.yaml")
         with tempfile.TemporaryDirectory() as td:
@@ -101,6 +103,7 @@ def test_conda_publish(
         mock_run_cmd,
         monkeypatch,
     ):
+        type(mock_run_cmd.return_value).returncode = PropertyMock(return_value=0)
         with tempfile.TemporaryDirectory() as td:
             with pytest.raises(FileNotFoundError):
                 publish(