Copy verified configs for imported models (#858)

Author: VipulMascarenhas

ads/aqua/common/utils.py

@@ -33,7 +33,7 @@
 from ads.common.extended_enum import ExtendedEnumMeta
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.common.oci_resource import SEARCH_TYPE, OCIResource
-from ads.common.utils import get_console_link, upload_to_os
+from ads.common.utils import get_console_link, upload_to_os, copy_file
 from ads.config import AQUA_SERVICE_MODELS_BUCKET, CONDA_BUCKET_NS, TENANCY_OCID
 from ads.model import DataScienceModel, ModelVersionSet
 
@@ -227,12 +227,10 @@ def is_valid_ocid(ocid: str) -> bool:
     bool:
         Whether the given ocid is valid.
     """
-    # TODO: revisit pattern
-    pattern = (
-        r"^ocid1\.([a-z0-9_]+)\.([a-z0-9]+)\.([a-z0-9-]*)(\.[^.]+)?\.([a-z0-9_]+)$"
-    )
-    match = re.match(pattern, ocid)
-    return True
+
+    if not ocid:
+        return False
+    return ocid.lower().startswith("ocid")
 
 
 def get_resource_type(ocid: str) -> str:
@@ -557,7 +555,7 @@ def fetch_service_compartment() -> Union[str, None]:
             config_file_name=CONTAINER_INDEX,
         )
     except Exception as e:
-        logger.error(
+        logger.debug(
             f"Config file {config_file_name}/{CONTAINER_INDEX} to fetch service compartment OCID could not be found. "
             f"\n{str(e)}."
         )
@@ -824,3 +822,51 @@ def get_combined_params(params1: str = None, params2: str = None) -> str:
     ]
 
     return " ".join(combined_params)
+
+
+def copy_model_config(artifact_path: str, os_path: str, auth: dict = None):
+    """Copies the aqua model config folder from the artifact path to the user provided object storage path.
+    The config folder is overwritten if the files already exist at the destination path.
+
+    Parameters
+    ----------
+    artifact_path:
+        Path of the aqua model where config folder is available.
+    os_path:
+        User provided path where config folder will be copied.
+    auth: (Dict, optional). Defaults to None.
+        The default authentication is set using `ads.set_auth` API. If you need to override the
+        default, use the `ads.common.auth.api_keys` or `ads.common.auth.resource_principal` to create appropriate
+        authentication signer and kwargs required to instantiate IdentityClient object.
+
+    Returns
+    -------
+    None
+        Nothing.
+    """
+
+    try:
+        source_dir = ObjectStorageDetails(
+            AQUA_SERVICE_MODELS_BUCKET,
+            CONDA_BUCKET_NS,
+            f"{os.path.dirname(artifact_path).rstrip('/')}/config",
+        ).path
+        dest_dir = f"{os_path.rstrip('/')}/config"
+
+        oss_details = ObjectStorageDetails.from_path(source_dir)
+        objects = oss_details.list_objects(fields="name").objects
+
+        for obj in objects:
+            source_path = ObjectStorageDetails(
+                AQUA_SERVICE_MODELS_BUCKET, CONDA_BUCKET_NS, obj.name
+            ).path
+            destination_path = os.path.join(dest_dir, os.path.basename(obj.name))
+            copy_file(
+                uri_src=source_path,
+                uri_dst=destination_path,
+                force_overwrite=True,
+                auth=auth,
+            )
+    except Exception as ex:
+        logger.debug(ex)
+        logger.debug(f"Failed to copy config folder from {artifact_path} to {os_path}.")

ads/aqua/constants.py

@@ -41,3 +41,20 @@
 SERVICE_MANAGED_CONTAINER_URI_SCHEME = "dsmc://"
 SUPPORTED_FILE_FORMATS = ["jsonl"]
 MODEL_BY_REFERENCE_OSS_PATH_KEY = "artifact_location"
+
+CONSOLE_LINK_RESOURCE_TYPE_MAPPING = dict(
+    datasciencemodel="models",
+    datasciencemodeldeployment="model-deployments",
+    datasciencemodeldeploymentdev="model-deployments",
+    datasciencemodeldeploymentint="model-deployments",
+    datasciencemodeldeploymentpre="model-deployments",
+    datasciencejob="jobs",
+    datasciencejobrun="job-runs",
+    datasciencejobrundev="job-runs",
+    datasciencejobrunint="job-runs",
+    datasciencejobrunpre="job-runs",
+    datasciencemodelversionset="model-version-sets",
+    datasciencemodelversionsetpre="model-version-sets",
+    datasciencemodelversionsetint="model-version-sets",
+    datasciencemodelversionsetdev="model-version-sets",
+)

ads/aqua/evaluation/evaluation.py

@@ -50,6 +50,7 @@
     JOB_INFRASTRUCTURE_TYPE_DEFAULT_NETWORKING,
     NB_SESSION_IDENTIFIER,
     UNKNOWN,
+    CONSOLE_LINK_RESOURCE_TYPE_MAPPING,
 )
 from ads.aqua.evaluation.constants import *
 from ads.aqua.evaluation.entities import *
@@ -1344,7 +1345,7 @@ def _build_resource_identifier(
     ) -> AquaResourceIdentifier:
         """Constructs AquaResourceIdentifier based on the given ocid and display name."""
         try:
-            resource_type = utils.CONSOLE_LINK_RESOURCE_TYPE_MAPPING.get(
+            resource_type = CONSOLE_LINK_RESOURCE_TYPE_MAPPING.get(
                 utils.get_resource_type(id)
             )
 

ads/aqua/model/model.py

@@ -21,6 +21,7 @@
     is_service_managed_container,
     read_file,
     upload_folder,
+    copy_model_config,
 )
 from ads.aqua.constants import (
     LICENSE_TXT,
@@ -660,11 +661,18 @@ def _create_model_catalog_entry(
 
         try:
             # If verified model already has a artifact json, use that.
-            metadata.get(MODEL_BY_REFERENCE_OSS_PATH_KEY)
+            artifact_path = metadata.get(MODEL_BY_REFERENCE_OSS_PATH_KEY).value
             logger.info(
                 f"Found model artifact in the service bucket. "
                 f"Using artifact from service bucket instead of {os_path}"
             )
+
+            # todo: implement generic copy_folder method
+            # copy model config from artifact path to user bucket
+            copy_model_config(
+                artifact_path=artifact_path, os_path=os_path, auth=self._auth
+            )
+
         except:
             # Add artifact from user bucket
             metadata.add(

tests/unitary/with_extras/aqua/test_model.py

@@ -521,12 +521,24 @@ def test_get_model_fine_tuned(
             "evaluation_container": "odsc-llm-evaluate",
         }
 
+    @pytest.mark.parametrize(
+        "artifact_location_set",
+        [
+            True,
+            False,
+        ],
+    )
+    @patch("ads.aqua.common.utils.copy_file")
+    @patch("ads.common.object_storage_details.ObjectStorageDetails.list_objects")
     @patch("huggingface_hub.snapshot_download")
     @patch("subprocess.check_call")
     def test_import_verified_model(
         self,
         mock_subprocess,
         mock_snapshot_download,
+        mock_list_objects,
+        mock_copy_file,
+        artifact_location_set,
     ):
         ObjectStorageDetails.is_bucket_versioned = MagicMock(return_value=True)
         ads.common.oci_datascience.OCIDataScienceMixin.init_client = MagicMock()
@@ -536,6 +548,16 @@ def test_import_verified_model(
         DataScienceModel.sync = MagicMock()
         OCIDataScienceModel.create = MagicMock()
 
+        # The name attribute cannot be mocked during creation of the mock object,
+        # hence attach it separately to the mocked objects.
+        artifact_path = "service_models/model-name/commit-id/artifact"
+        obj1 = MagicMock(etag="12345-1234-1234-1234-123456789", size=150)
+        obj1.name = f"{artifact_path}/config/deployment_config.json"
+        obj2 = MagicMock(etag="12345-1234-1234-1234-123456789", size=150)
+        obj2.name = f"{artifact_path}/config/ft_config.json"
+        objects = [obj1, obj2]
+        mock_list_objects.return_value = MagicMock(objects=objects)
+
         ds_model = DataScienceModel()
         os_path = "oci://aqua-bkt@aqua-ns/prefix/path"
         hf_model = "oracle/aqua-1t-mega-model"
@@ -562,6 +584,14 @@ def test_import_verified_model(
         custom_metadata_list.add(
             **{"key": "evaluation-container", "value": "odsc-llm-evaluate"}
         )
+        if not artifact_location_set:
+            custom_metadata_list.add(
+                **{
+                    "key": "artifact_location",
+                    "value": artifact_path,
+                    "description": "artifact location",
+                }
+            )
         ds_model.with_custom_metadata_list(custom_metadata_list)
         ds_model.set_spec(ds_model.CONST_MODEL_FILE_DESCRIPTION, {})
         ds_model.dsc_model = MagicMock(id="test_model_id")
@@ -579,6 +609,8 @@ def test_import_verified_model(
                 local_dir=f"{str(tmpdir)}/{hf_model}",
                 local_dir_use_symlinks=False,
             )
+            if not artifact_location_set:
+                mock_copy_file.assert_called()
             mock_subprocess.assert_called_with(
                 shlex.split(
                     f"oci os object bulk-upload --src-dir {str(tmpdir)}/{hf_model} --prefix prefix/path/{hf_model}/ -bn aqua-bkt -ns aqua-ns --auth api_key --profile DEFAULT"
@@ -613,7 +645,9 @@ def test_import_verified_model(
                 },
                 {
                     "key": "artifact_location",
-                    "value": f"{os_path}/{hf_model}/",
+                    "value": f"{os_path}/{hf_model}/"
+                    if artifact_location_set
+                    else artifact_path,
                     "description": "artifact location",
                     "category": "Other",
                 },