Merge branch 'main' into feature/odsc-65115

Author: codeloop

.github/workflows/run-forecast-unit-tests.yml

@@ -56,6 +56,6 @@ jobs:
           $CONDA/bin/conda init
           source /home/runner/.bashrc
           pip install -r test-requirements-operators.txt
-          pip install "oracle-automlx[forecasting]>=24.4.0"
+          pip install "oracle-automlx[forecasting]>=24.4.1"
           pip install pandas>=2.2.0
           python -m pytest -v -p no:warnings --durations=5 tests/operators/forecast

ads/aqua/common/enums.py

@@ -52,6 +52,9 @@ class InferenceContainerTypeFamily(str, metaclass=ExtendedEnumMeta):
     AQUA_VLLM_CONTAINER_FAMILY = "odsc-vllm-serving"
     AQUA_TGI_CONTAINER_FAMILY = "odsc-tgi-serving"
     AQUA_LLAMA_CPP_CONTAINER_FAMILY = "odsc-llama-cpp-serving"
+
+
+class CustomInferenceContainerTypeFamily(str, metaclass=ExtendedEnumMeta):
     AQUA_TEI_CONTAINER_FAMILY = "odsc-tei-serving"
 
 

ads/aqua/common/utils.py

@@ -11,6 +11,7 @@
 import random
 import re
 import shlex
+import shutil
 import subprocess
 from datetime import datetime, timedelta
 from functools import wraps
@@ -21,6 +22,8 @@
 import fsspec
 import oci
 from cachetools import TTLCache, cached
+from huggingface_hub.constants import HF_HUB_CACHE
+from huggingface_hub.file_download import repo_folder_name
 from huggingface_hub.hf_api import HfApi, ModelInfo
 from huggingface_hub.utils import (
     GatedRepoError,
@@ -821,6 +824,48 @@ def upload_folder(
     return f"oci://{os_details.bucket}@{os_details.namespace}" + "/" + object_path
 
 
+def cleanup_local_hf_model_artifact(
+    model_name: str,
+    local_dir: str = None,
+):
+    """
+    Helper function that deletes local artifacts downloaded from Hugging Face to free up disk space.
+    Parameters
+    ----------
+    model_name (str): Name of the huggingface model
+    local_dir (str): Local directory where the object is downloaded
+
+    """
+    if local_dir and os.path.exists(local_dir):
+        model_dir = os.path.join(local_dir, model_name)
+        model_dir = (
+            os.path.dirname(model_dir)
+            if "/" in model_name or os.sep in model_name
+            else model_dir
+        )
+        shutil.rmtree(model_dir, ignore_errors=True)
+        if os.path.exists(model_dir):
+            logger.debug(
+                f"Could not delete local model artifact directory: {model_dir}"
+            )
+        else:
+            logger.debug(f"Deleted local model artifact directory: {model_dir}.")
+
+    hf_local_path = os.path.join(
+        HF_HUB_CACHE, repo_folder_name(repo_id=model_name, repo_type="model")
+    )
+    shutil.rmtree(hf_local_path, ignore_errors=True)
+
+    if os.path.exists(hf_local_path):
+        logger.debug(
+            f"Could not clear the local Hugging Face cache directory {hf_local_path} for the model {model_name}."
+        )
+    else:
+        logger.debug(
+            f"Cleared contents of local Hugging Face cache directory {hf_local_path} for the model {model_name}."
+        )
+
+
 def is_service_managed_container(container):
     return container and container.startswith(SERVICE_MANAGED_CONTAINER_URI_SCHEME)
 

ads/aqua/extension/model_handler.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# Copyright (c) 2024 Oracle and/or its affiliates.
+# Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 from typing import Optional
@@ -8,6 +8,9 @@
 from tornado.web import HTTPError
 
 from ads.aqua.common.decorator import handle_exceptions
+from ads.aqua.common.enums import (
+    CustomInferenceContainerTypeFamily,
+)
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
     get_hf_model_info,
@@ -128,6 +131,10 @@ def post(self, *args, **kwargs):  # noqa: ARG002
         download_from_hf = (
             str(input_data.get("download_from_hf", "false")).lower() == "true"
         )
+        local_dir = input_data.get("local_dir")
+        cleanup_model_cache = (
+            str(input_data.get("cleanup_model_cache", "true")).lower() == "true"
+        )
         inference_container_uri = input_data.get("inference_container_uri")
         allow_patterns = input_data.get("allow_patterns")
         ignore_patterns = input_data.get("ignore_patterns")
@@ -139,6 +146,8 @@ def post(self, *args, **kwargs):  # noqa: ARG002
                 model=model,
                 os_path=os_path,
                 download_from_hf=download_from_hf,
+                local_dir=local_dir,
+                cleanup_model_cache=cleanup_model_cache,
                 inference_container=inference_container,
                 finetuning_container=finetuning_container,
                 compartment_id=compartment_id,
@@ -163,7 +172,9 @@ def put(self, id):
             raise HTTPError(400, Errors.NO_INPUT_DATA)
 
         inference_container = input_data.get("inference_container")
+        inference_container_uri = input_data.get("inference_container_uri")
         inference_containers = AquaModelApp.list_valid_inference_containers()
+        inference_containers.extend(CustomInferenceContainerTypeFamily.values())
         if (
             inference_container is not None
             and inference_container not in inference_containers
@@ -176,7 +187,13 @@ def put(self, id):
         task = input_data.get("task")
         app = AquaModelApp()
         self.finish(
-            app.edit_registered_model(id, inference_container, enable_finetuning, task)
+            app.edit_registered_model(
+                id,
+                inference_container,
+                inference_container_uri,
+                enable_finetuning,
+                task,
+            )
         )
         app.clear_model_details_cache(model_id=id)
 

ads/aqua/model/entities.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# Copyright (c) 2024 Oracle and/or its affiliates.
+# Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 """
@@ -283,6 +283,7 @@ class ImportModelDetails(CLIBuilderMixin):
     os_path: str
     download_from_hf: Optional[bool] = True
     local_dir: Optional[str] = None
+    cleanup_model_cache: Optional[bool] = True
     inference_container: Optional[str] = None
     finetuning_container: Optional[str] = None
     compartment_id: Optional[str] = None

ads/aqua/model/model.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# Copyright (c) 2024 Oracle and/or its affiliates.
+# Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 import os
 import pathlib
@@ -15,6 +15,7 @@
 from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID, logger
 from ads.aqua.app import AquaApp
 from ads.aqua.common.enums import (
+    CustomInferenceContainerTypeFamily,
     FineTuningContainerTypeFamily,
     InferenceContainerTypeFamily,
     Tags,
@@ -23,6 +24,7 @@
 from ads.aqua.common.utils import (
     LifecycleStatus,
     _build_resource_identifier,
+    cleanup_local_hf_model_artifact,
     copy_model_config,
     create_word_icon,
     generate_tei_cmd_var,
@@ -376,8 +378,10 @@ def delete_model(self, model_id):
                 f"Failed to delete model:{model_id}. Only registered models or finetuned model can be deleted."
             )
 
-    @telemetry(entry_point="plugin=model&action=delete", name="aqua")
-    def edit_registered_model(self, id, inference_container, enable_finetuning, task):
+    @telemetry(entry_point="plugin=model&action=edit", name="aqua")
+    def edit_registered_model(
+        self, id, inference_container, inference_container_uri, enable_finetuning, task
+    ):
         """Edits the default config of unverified registered model.
 
         Parameters
@@ -386,6 +390,8 @@ def edit_registered_model(self, id, inference_container, enable_finetuning, task
             The model OCID.
         inference_container: str.
             The inference container family name
+        inference_container_uri: str
+            The inference container uri for embedding models
         enable_finetuning: str
             Flag to enable or disable finetuning over the model. Defaults to None
         task:
@@ -401,19 +407,44 @@ def edit_registered_model(self, id, inference_container, enable_finetuning, task
         if ds_model.freeform_tags.get(Tags.BASE_MODEL_CUSTOM, None):
             if ds_model.freeform_tags.get(Tags.AQUA_SERVICE_MODEL_TAG, None):
                 raise AquaRuntimeError(
-                    f"Failed to edit model:{id}. Only registered unverified models can be edited."
+                    "Only registered unverified models can be edited."
                 )
             else:
                 custom_metadata_list = ds_model.custom_metadata_list
                 freeform_tags = ds_model.freeform_tags
                 if inference_container:
-                    custom_metadata_list.add(
-                        key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER,
-                        value=inference_container,
-                        category=MetadataCustomCategory.OTHER,
-                        description="Deployment container mapping for SMC",
-                        replace=True,
-                    )
+                    if (
+                        inference_container in CustomInferenceContainerTypeFamily
+                        and inference_container_uri is None
+                    ):
+                        raise AquaRuntimeError(
+                            "Inference container URI must be provided."
+                        )
+                    else:
+                        custom_metadata_list.add(
+                            key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER,
+                            value=inference_container,
+                            category=MetadataCustomCategory.OTHER,
+                            description="Deployment container mapping for SMC",
+                            replace=True,
+                        )
+                if inference_container_uri:
+                    if (
+                        inference_container in CustomInferenceContainerTypeFamily
+                        or inference_container is None
+                    ):
+                        custom_metadata_list.add(
+                            key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER_URI,
+                            value=inference_container_uri,
+                            category=MetadataCustomCategory.OTHER,
+                            description=f"Inference container URI for {ds_model.display_name}",
+                            replace=True,
+                        )
+                    else:
+                        raise AquaRuntimeError(
+                            f"Inference container URI can be edited only with container values: {CustomInferenceContainerTypeFamily.values()}"
+                        )
+
                 if enable_finetuning is not None:
                     if enable_finetuning.lower() == "true":
                         custom_metadata_list.add(
@@ -448,9 +479,7 @@ def edit_registered_model(self, id, inference_container, enable_finetuning, task
                 )
                 AquaApp().update_model(id, update_model_details)
         else:
-            raise AquaRuntimeError(
-                f"Failed to edit model:{id}. Only registered unverified models can be edited."
-            )
+            raise AquaRuntimeError("Only registered unverified models can be edited.")
 
     def _fetch_metric_from_metadata(
         self,
@@ -869,8 +898,7 @@ def _create_model_catalog_entry(
             # only add cmd vars if inference container is not an SMC
             if (
                 inference_container not in smc_container_set
-                and inference_container
-                == InferenceContainerTypeFamily.AQUA_TEI_CONTAINER_FAMILY
+                and inference_container in CustomInferenceContainerTypeFamily.values()
             ):
                 cmd_vars = generate_tei_cmd_var(os_path)
                 metadata.add(
@@ -1322,20 +1350,20 @@ def _download_model_from_hf(
         Returns
         -------
         model_artifact_path (str): Location where the model artifacts are downloaded.
-
         """
         # Download the model from hub
-        if not local_dir:
-            local_dir = os.path.join(os.path.expanduser("~"), "cached-model")
-        local_dir = os.path.join(local_dir, model_name)
-        os.makedirs(local_dir, exist_ok=True)
-        snapshot_download(
+        if local_dir:
+            local_dir = os.path.join(local_dir, model_name)
+            os.makedirs(local_dir, exist_ok=True)
+
+        # if local_dir is not set, the return value points to the cached data folder
+        local_dir = snapshot_download(
             repo_id=model_name,
             local_dir=local_dir,
             allow_patterns=allow_patterns,
             ignore_patterns=ignore_patterns,
         )
-        # Upload to object storage and skip .cache/huggingface/ folder
+        # Upload to object storage
         model_artifact_path = upload_folder(
             os_path=os_path,
             local_dir=local_dir,
@@ -1365,6 +1393,8 @@ def register(
                 ignore_patterns (list): Model files matching any of the patterns are not downloaded.
                     Example: ["*.json"] will ignore all .json files. ["folder/*"] will ignore all files under `folder`.
                     Patterns are Standard Wildcards (globbing patterns) and rules can be found here: https://docs.python.org/3/library/fnmatch.html
+                cleanup_model_cache (bool): Deletes downloaded files from local machine after model is successfully
+                registered. Set to True by default.
 
         Returns:
             AquaModel:
@@ -1474,6 +1504,14 @@ def register(
             detail=validation_result.telemetry_model_name,
         )
 
+        if (
+            import_model_details.download_from_hf
+            and import_model_details.cleanup_model_cache
+        ):
+            cleanup_local_hf_model_artifact(
+                model_name=model_name, local_dir=import_model_details.local_dir
+            )
+
         return AquaModel(**aqua_model_attributes)
 
     def _if_show(self, model: DataScienceModel) -> bool:

ads/opctl/config/merger.py

@@ -1,35 +1,33 @@
 #!/usr/bin/env python
-# -*- coding: utf-8; -*-
 
-# Copyright (c) 2022, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2022, 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 os
 from string import Template
 from typing import Dict
-import json
 
 import yaml
 
 from ads.common.auth import AuthType, ResourcePrincipal
 from ads.opctl import logger
 from ads.opctl.config.base import ConfigProcessor
-from ads.opctl.config.utils import read_from_ini, _DefaultNoneDict
-from ads.opctl.utils import is_in_notebook_session, get_service_pack_prefix
+from ads.opctl.config.utils import _DefaultNoneDict, read_from_ini
 from ads.opctl.constants import (
-    DEFAULT_PROFILE,
-    DEFAULT_OCI_CONFIG_FILE,
-    DEFAULT_CONDA_PACK_FOLDER,
-    DEFAULT_ADS_CONFIG_FOLDER,
-    ADS_JOBS_CONFIG_FILE_NAME,
     ADS_CONFIG_FILE_NAME,
-    ADS_ML_PIPELINE_CONFIG_FILE_NAME,
     ADS_DATAFLOW_CONFIG_FILE_NAME,
+    ADS_JOBS_CONFIG_FILE_NAME,
     ADS_LOCAL_BACKEND_CONFIG_FILE_NAME,
+    ADS_ML_PIPELINE_CONFIG_FILE_NAME,
     ADS_MODEL_DEPLOYMENT_CONFIG_FILE_NAME,
-    DEFAULT_NOTEBOOK_SESSION_CONDA_DIR,
     BACKEND_NAME,
+    DEFAULT_ADS_CONFIG_FOLDER,
+    DEFAULT_CONDA_PACK_FOLDER,
+    DEFAULT_NOTEBOOK_SESSION_CONDA_DIR,
+    DEFAULT_OCI_CONFIG_FILE,
+    DEFAULT_PROFILE,
 )
+from ads.opctl.utils import get_service_pack_prefix, is_in_notebook_session
 
 
 class ConfigMerger(ConfigProcessor):
@@ -41,8 +39,9 @@ class ConfigMerger(ConfigProcessor):
     """
 
     def process(self, **kwargs) -> None:
-        config_string = Template(json.dumps(self.config)).safe_substitute(os.environ)
-        self.config = json.loads(config_string)
+        for key, value in self.config.items():
+            if isinstance(value, str):  # Substitute only if the value is a string
+                self.config[key] = Template(value).safe_substitute(os.environ)
 
         if "runtime" not in self.config:
             self.config["runtime"] = {}