add test

Author: Ziqun Ye

ads/opctl/backend/local.py

@@ -642,9 +642,25 @@ def _log_orchestration_message(self, str: str) -> None:
 
 class LocalModelDeploymentBackend(LocalBackend):
     def __init__(self, config: Dict) -> None:
+        """
+        Initialize a LocalModelDeploymentBackend object with given config.
+
+        Parameters
+        ----------
+        config: dict
+            dictionary of configurations
+        """
         super().__init__(config)
 
     def predict(self) -> None:
+        """
+        Conducts local verify.
+
+        Returns
+        -------
+        None
+            Nothing.
+        """
         artifact_directory = self.config["execution"].get("artifact_directory")
         ocid = self.config["execution"].get("ocid")
         data = self.config["execution"].get("payload")
@@ -658,8 +674,8 @@ def predict(self) -> None:
             _download_model(ocid=ocid, artifact_directory=artifact_directory, region=region, bucket_uri=bucket_uri, timeout=timeout)
 
         if ocid:
-            conda_slug, conda_path = self._get_conda_info_from_catalog(ocid)
-        elif artifact_directory:
+            conda_slug, conda_path = self._get_conda_info_from_custom_metadata(ocid)
+        if artifact_directory or not conda_path:
             if not os.path.exists(artifact_directory) or len(os.listdir(artifact_directory)) == 0:
                 raise ValueError(f"`artifact_directory` {artifact_directory} does not exist or is empty.")
             conda_slug, conda_path = self._get_conda_info_from_runtime(artifact_dir=artifact_directory)
@@ -699,14 +715,33 @@ def predict(self) -> None:
                 f"Run with the --debug argument to view container logs."
             )
 
-    def _get_conda_info_from_catalog(self, ocid):
+    def _get_conda_info_from_custom_metadata(self, ocid):
+        """
+        Get conda env info from custom metadata from model catalog.
+
+        Returns
+        -------
+        (str, str)
+            conda slug and conda path.
+        """
         response = self.client.get_model(ocid)
         custom_metadata = ModelCustomMetadata._from_oci_metadata(response.data.custom_metadata_list)
-        conda_path = custom_metadata['CondaEnvironmentPath'].value
-        conda_slug = custom_metadata['SlugName'].value
+        conda_slug, conda_path = None, None
+        if "CondaEnvironmentPath" in custom_metadata:
+            conda_path = custom_metadata['CondaEnvironmentPath'].value
+        if "SlugName" in custom_metadata:
+            conda_slug = custom_metadata['SlugName'].value
         return conda_slug, conda_path
 
     def _get_conda_info_from_runtime(self, artifact_dir):
+        """
+        Get conda env info from runtime yaml file.
+
+        Returns
+        -------
+        (str, str)
+            conda slug and conda path.
+        """
         runtime_yaml_file = os.path.join(artifact_dir, "runtime.yaml")
         runtime_info = RuntimeInfo.from_yaml(uri=runtime_yaml_file)
         conda_slug = runtime_info.model_deployment.inference_conda_env.inference_env_slug

ads/opctl/cli.py

@@ -28,7 +28,7 @@
 from ads.opctl.utils import suppress_traceback
 from ads.opctl.config.merger import ConfigMerger
 from ads.opctl.constants import BACKEND_NAME
-from ads.opctl.backend.local import DEFAULT_MODEL_FOLDER
+from ads.opctl.constants import DEFAULT_MODEL_FOLDER
 
 import ads.opctl.conda.cli
 import ads.opctl.spark.cli

ads/opctl/model/cli.py

@@ -8,7 +8,7 @@
 from ads.common.auth import AuthType
 from ads.opctl.utils import suppress_traceback
 from ads.opctl.model.cmds import download_model as download_model_cmd
-from ads.opctl.backend.local import DEFAULT_MODEL_FOLDER
+from ads.opctl.constants import DEFAULT_MODEL_FOLDER
 
 
 @click.group("model")

ads/opctl/model/cmds.py

@@ -4,7 +4,7 @@
 from ads.common.auth import create_signer
 from ads.model.datascience_model import DataScienceModel
 from ads.opctl import logger
-from ads.opctl.backend.local import DEFAULT_MODEL_FOLDER
+from ads.opctl.constants import DEFAULT_MODEL_FOLDER
 from ads.opctl.config.base import ConfigProcessor
 from ads.opctl.config.merger import ConfigMerger
 
@@ -55,4 +55,5 @@ def _download_model(ocid, artifact_directory, oci_auth, region, bucket_uri, time
         )
     except Exception as e:
         print(type(e))
-        shutil.rmtree(artifact_directory, ignore_errors=True)
+        shutil.rmtree(artifact_directory, ignore_errors=True)
+        raise e

tests/unitary/with_extras/opctl/test_opctl_model.py

@@ -0,0 +1,32 @@
+from ads.opctl.model.cmds import _download_model, download_model
+import pytest 
+from unittest.mock import ANY, call, patch
+from ads.model.datascience_model import DataScienceModel
+from unittest.mock import MagicMock, Mock
+from ads.opctl.model.cmds import create_signer
+
+
+@patch.object(DataScienceModel, "from_id")
+def test_model__download_model(mock_from_id):
+    mock_datascience_model = MagicMock()
+    mock_from_id.return_value = mock_datascience_model
+    _download_model("fake_model_id", "fake_dir", "fake_auth", "region", "bucket_uri", 36, False)
+    mock_from_id.assert_called_with("fake_model_id")
+    mock_datascience_model.download_artifact.assert_called_with(target_dir='fake_dir', force_overwrite=False, overwrite_existing_artifact=True, remove_existing_artifact=True, auth='fake_auth', region='region', timeout=36, bucket_uri='bucket_uri')
+
+
+@patch.object(DataScienceModel, "from_id", side_effect=Exception("Fake error."))
+def test_model__download_model_error(mock_from_id):
+    with pytest.raises(Exception, match="Fake error."):
+        _download_model("fake_model_id", "fake_dir", "fake_auth", "region", "bucket_uri", 36, False)
+
+
+@patch("ads.opctl.model.cmds._download_model")
+@patch("ads.opctl.model.cmds.create_signer")
+def test_download_model(mock_create_signer, mock__download_model):
+    auth_mock = MagicMock()
+    mock_create_signer.return_value = auth_mock
+    download_model(ocid = "fake_model_id")
+    mock_create_signer.assert_called_once()
+    mock__download_model.assert_called_once_with(ocid='fake_model_id', artifact_directory='/Users/ziye/.ads_ops/models/fake_model_id', region=None, bucket_uri=None, timeout=None, force_overwrite=False, oci_auth=auth_mock)
+