Merge branch 'main' into ODSC-66852/watch_aqua_ft_job

Author: VipulMascarenhas

.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

@@ -838,21 +838,31 @@ def cleanup_local_hf_model_artifact(
     """
     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):
-            shutil.rmtree(model_dir)
-            logger.debug(f"Deleted local model artifact directory: {model_dir}")
-
-        if not os.listdir(local_dir):
-            shutil.rmtree(local_dir)
-            logger.debug(f"Deleted local directory {model_dir} as it is empty.")
+            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):
-        shutil.rmtree(hf_local_path)
         logger.debug(
-            f"Deleted local Hugging Face cache directory {hf_local_path} for the model {model_name} "
+            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}."
         )
 
 

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

@@ -283,7 +283,7 @@ class ImportModelDetails(CLIBuilderMixin):
     os_path: str
     download_from_hf: Optional[bool] = True
     local_dir: Optional[str] = None
-    delete_from_local: Optional[bool] = True
+    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

@@ -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,
@@ -377,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
@@ -387,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:
@@ -402,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(
@@ -449,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,
@@ -870,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(
@@ -1328,7 +1355,9 @@ def _download_model_from_hf(
         if local_dir:
             local_dir = os.path.join(local_dir, model_name)
             os.makedirs(local_dir, exist_ok=True)
-        snapshot_download(
+
+        # 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,
@@ -1364,7 +1393,7 @@ 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
-                delete_from_local (bool): Deletes downloaded files from local machine after model is successfully
+                cleanup_model_cache (bool): Deletes downloaded files from local machine after model is successfully
                 registered. Set to True by default.
 
         Returns:
@@ -1477,7 +1506,7 @@ def register(
 
         if (
             import_model_details.download_from_hf
-            and import_model_details.delete_from_local
+            and import_model_details.cleanup_model_cache
         ):
             cleanup_local_hf_model_artifact(
                 model_name=model_name, local_dir=import_model_details.local_dir

ads/model/__init__.py

@@ -1,29 +1,26 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*--
 
-# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2021, 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 ads.model.generic_model import GenericModel, ModelState
 from ads.model.datascience_model import DataScienceModel
-from ads.model.model_properties import ModelProperties
+from ads.model.deployment.model_deployer import ModelDeployer
+from ads.model.deployment.model_deployment import ModelDeployment
+from ads.model.deployment.model_deployment_properties import ModelDeploymentProperties
 from ads.model.framework.automl_model import AutoMLModel
+from ads.model.framework.embedding_onnx_model import EmbeddingONNXModel
+from ads.model.framework.huggingface_model import HuggingFacePipelineModel
 from ads.model.framework.lightgbm_model import LightGBMModel
 from ads.model.framework.pytorch_model import PyTorchModel
 from ads.model.framework.sklearn_model import SklearnModel
+from ads.model.framework.spark_model import SparkPipelineModel
 from ads.model.framework.tensorflow_model import TensorFlowModel
 from ads.model.framework.xgboost_model import XGBoostModel
-from ads.model.framework.spark_model import SparkPipelineModel
-from ads.model.framework.huggingface_model import HuggingFacePipelineModel
-
-from ads.model.deployment.model_deployer import ModelDeployer
-from ads.model.deployment.model_deployment import ModelDeployment
-from ads.model.deployment.model_deployment_properties import ModelDeploymentProperties
-
+from ads.model.generic_model import GenericModel, ModelState
+from ads.model.model_properties import ModelProperties
+from ads.model.model_version_set import ModelVersionSet, experiment
 from ads.model.serde.common import SERDE
 from ads.model.serde.model_input import ModelInputSerializer
-
-from ads.model.model_version_set import ModelVersionSet, experiment
 from ads.model.service.oci_datascience_model_version_set import (
     ModelVersionSetNotExists,
     ModelVersionSetNotSaved,
@@ -42,6 +39,7 @@
     "XGBoostModel",
     "SparkPipelineModel",
     "HuggingFacePipelineModel",
+    "EmbeddingONNXModel",
     "ModelDeployer",
     "ModelDeployment",
     "ModelDeploymentProperties",

ads/model/artifact.py

@@ -1,28 +1,28 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*--
 
-# Copyright (c) 2022, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2022, 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 fnmatch
 import importlib
 import os
-import sys
 import shutil
+import sys
 import tempfile
 import uuid
-import fsspec
+from datetime import datetime
 from typing import Dict, Optional, Tuple
+
+import fsspec
+from jinja2 import Environment, PackageLoader
+
+from ads import __version__
 from ads.common import auth as authutil
 from ads.common import logger, utils
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.config import CONDA_BUCKET_NAME, CONDA_BUCKET_NS
 from ads.model.runtime.env_info import EnvInfo, InferenceEnvInfo, TrainingEnvInfo
 from ads.model.runtime.runtime_info import RuntimeInfo
-from jinja2 import Environment, PackageLoader
-import warnings
-from ads import __version__
-from datetime import datetime
 
 MODEL_ARTIFACT_VERSION = "3.0"
 REQUIRED_ARTIFACT_FILES = ("runtime.yaml", "score.py")
@@ -378,6 +378,45 @@ def prepare_score_py(
         ) as f:
             f.write(scorefn_template.render(context))
 
+    def prepare_schema(self, schema_name: str):
+        """Copies schema to artifact directory.
+
+        Parameters
+        ----------
+        schema_name: str
+            The schema name
+
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        FileExistsError
+            If `schema_name` doesn't exist.
+        """
+        uri_src = os.path.join(
+            os.path.abspath(os.path.join(os.path.dirname(__file__), "..")),
+            "templates",
+            "schemas",
+            f"{schema_name}",
+        )
+
+        if not os.path.exists(uri_src):
+            raise FileExistsError(
+                f"{schema_name} does not exists. "
+                "Ensure the schema name is valid or specify a different one."
+            )
+
+        uri_dst = os.path.join(self.artifact_dir, os.path.basename(uri_src))
+
+        utils.copy_file(
+            uri_src=uri_src,
+            uri_dst=uri_dst,
+            force_overwrite=True,
+            auth=self.auth,
+        )
+
     def reload(self):
         """Syncs the `score.py` to reload the model and predict function.