Interoperability with model catalog (#260)

Author: KshitizLohia

ads/feature_store/dataset.py

@@ -498,6 +498,7 @@ def with_model_details(self, model_details: ModelDetails) -> "Dataset":
                 "The argument `model_details` has to be of type `ModelDetails`"
                 "but is of type: `{}`".format(type(model_details))
             )
+
         return self.set_spec(self.CONST_MODEL_DETAILS, model_details.to_dict())
 
     def add_models(self, model_details: ModelDetails) -> "Dataset":
@@ -516,9 +517,20 @@ def add_models(self, model_details: ModelDetails) -> "Dataset":
         if existing_model_details and existing_model_details.items:
             items = existing_model_details["items"]
             for item in items:
-                model_details.items.append(item)
+                if item not in model_details.items:
+                    model_details.items.append(item)
         self.with_model_details(model_details)
-        return self.update()
+        try:
+            return self.update()
+        except Exception as ex:
+            logger.error(
+                f"Dataset update Failed with : {type(ex)} with error message: {ex}"
+            )
+            if existing_model_details:
+                self.with_model_details(ModelDetails().with_items(existing_model_details["items"]))
+            else:
+                self.with_model_details(ModelDetails().with_items([]))
+                return self
 
     def remove_models(self, model_details: ModelDetails) -> "Dataset":
         """remove model details from the dataset, remove from the existing dataset model id list

ads/feature_store/docs/source/dataset.rst

@@ -306,7 +306,7 @@ Use the ``show()`` method on the ``Dataset`` instance to visualize the lineage o
 
 The ``show()`` method takes the following optional parameter:
 
-  - ``rankdir: (str, optional)``. Defaults to ``LR``. The allowed values are ``TB`` or ``LR``. This parameter is applicable only for ``graph`` mode and it renders the direction of the graph as either top to bottom (TB) or left to right (LR).
+- ``rankdir: (str, optional)``. Defaults to ``LR``. The allowed values are ``TB`` or ``LR``. This parameter is applicable only for ``graph`` mode and it renders the direction of the graph as either top to bottom (TB) or left to right (LR).
 
 
 .. code-block:: python3
@@ -317,3 +317,16 @@ Below is an example of the output.
 
 .. figure:: figures/dataset_lineage.png
   :width: 400
+
+
+Add Model Details
+=================
+
+You can call the ``add_models()`` method of the Dataset instance to add model ids to dataset.
+The ``.add_models()`` method takes the following parameter:
+
+- ``model_details: ModelDetails``.  ModelDetails takes ``items: List[str]`` as parameter and model ids to be passed as items.
+
+.. code-block:: python3
+
+  dataset.add_models(ModelDetails().with_items([<ocid1.datasciencemodel..<unique_id>]))

ads/feature_store/docs/source/quickstart.rst

@@ -55,8 +55,9 @@ Background reading to understand the concepts of Feature Store and OCI Data Scie
     compartment_id = "ocid1.compartment.<unique_id>"
     metastore_id = "ocid1.datacatalogmetastore.oc1.iad.<unique_id>"
     api_gateway_endpoint = "https://**.{region}.oci.customer-oci.com/20230101"
+    os.environ["OCI_FS_SERVICE_ENDPOINT"] = api_gateway_endpoint
 
-    ads.set_auth(auth="user_principal", client_kwargs={"service_endpoint": api_gateway_endpoint})
+    ads.set_auth(auth="api_key")
 
     # step1: Create feature store
     feature_store_resource = (

ads/feature_store/mixin/oci_feature_store.py

@@ -6,13 +6,19 @@
 
 from ads.common.oci_mixin import OCIModelMixin
 import oci.feature_store
+import os
 
 
 class OCIFeatureStoreMixin(OCIModelMixin):
     @classmethod
     def init_client(
         cls, **kwargs
     ) -> oci.feature_store.feature_store_client.FeatureStoreClient:
+        # TODO: Getting the endpoint from authorizer
+        fs_service_endpoint = os.environ.get("OCI_FS_SERVICE_ENDPOINT")
+        if fs_service_endpoint:
+            kwargs = {"service_endpoint": fs_service_endpoint}
+
         client = cls._init_client(
             client=oci.feature_store.feature_store_client.FeatureStoreClient, **kwargs
         )

ads/feature_store/validation_output.py

@@ -23,10 +23,10 @@ def to_pandas(self) -> pd.DataFrame:
             The validation output information as a pandas DataFrame.
         """
         if self.content:
-            validation_output_json = (
-                json.loads(self.content)
-            )
-            profile_result = pd.json_normalize(validation_output_json.get("results")).transpose()
+            validation_output_json = json.loads(self.content)
+            profile_result = pd.json_normalize(
+                validation_output_json.get("results")
+            ).transpose()
             return profile_result
 
     def to_summary(self) -> pd.DataFrame:
@@ -39,9 +39,7 @@ def to_summary(self) -> pd.DataFrame:
             The validation output summary information as a pandas DataFrame.
         """
         if self.content:
-            validation_output_json = (
-                json.loads(self.content)
-            )
+            validation_output_json = json.loads(self.content)
             profile_result = pd.json_normalize(validation_output_json).transpose()
             summary_df = profile_result.drop("results")
             return summary_df

ads/model/generic_model.py

@@ -34,6 +34,7 @@
 from ads.evaluations import EvaluatorMixin
 from ads.feature_engineering import ADSImage
 from ads.feature_engineering.schema import Schema
+from ads.feature_store.model_details import ModelDetails
 from ads.model.artifact import ModelArtifact
 from ads.model.common.utils import (
     _extract_locals,
@@ -65,7 +66,7 @@
     Framework,
     ModelCustomMetadata,
     ModelProvenanceMetadata,
-    ModelTaxonomyMetadata,
+    ModelTaxonomyMetadata, MetadataCustomCategory,
 )
 from ads.model.model_metadata_mixin import MetadataMixin
 from ads.model.model_properties import ModelProperties
@@ -1825,6 +1826,7 @@ def save(
         remove_existing_artifact: Optional[bool] = True,
         model_version_set: Optional[Union[str, ModelVersionSet]] = None,
         version_label: Optional[str] = None,
+        featurestore_dataset=None,
         **kwargs,
     ) -> str:
         """Saves model artifacts to the model catalog.
@@ -1856,6 +1858,8 @@ def save(
             The model version set OCID, or model version set name, or `ModelVersionSet` instance.
         version_label: (str, optional). Defaults to None.
             The model version lebel.
+        featurestore_dataset: (Dataset, optional).
+            The feature store dataset
         kwargs:
             project_id: (str, optional).
                 Project OCID. If not specified, the value will be taken either
@@ -1937,6 +1941,15 @@ def save(
         # variables in case of saving model in context of model version set.
         model_version_set_id = _extract_model_version_set_id(model_version_set)
 
+        if featurestore_dataset:
+            dataset_details = {
+                "dataset-id": featurestore_dataset.id,
+                "dataset-name": featurestore_dataset.name
+            }
+            self.metadata_custom.add("featurestore.dataset", value=str(dataset_details),
+                                     category=MetadataCustomCategory.TRAINING_AND_VALIDATION_DATASETS,
+                                     description="feature store dataset", replace=True)
+
         self.dsc_model = (
             self.dsc_model.with_compartment_id(self.properties.compartment_id)
             .with_project_id(self.properties.project_id)
@@ -1965,6 +1978,10 @@ def save(
             .with_infrastructure(ModelDeploymentInfrastructure())
             .with_runtime(ModelDeploymentContainerRuntime())
         )
+        # Add the model id to the feature store dataset
+        if featurestore_dataset:
+            model_details = ModelDetails().with_items([self.model_id])
+            featurestore_dataset.add_models(model_details)
 
         return self.model_id
 

tests/integration/feature_store/test_base.py

@@ -20,11 +20,8 @@
 from ads.feature_store.statistics_config import StatisticsConfig
 
 
-client_kwargs = dict(
-    retry_strategy=oci.retry.NoneRetryStrategy,
-    service_endpoint=os.getenv("service_endpoint"),
-)
-ads.set_auth(client_kwargs=client_kwargs)
+ads.set_auth()
+os.environ["OCI_FS_SERVICE_ENDPOINT"] = os.getenv("service_endpoint")
 
 try:
     from ads.feature_store.feature_store import FeatureStore
@@ -37,7 +34,9 @@
 
 class FeatureStoreTestCase:
     # networks compartment in feature store
-    TIME_NOW = str.format("{}_{}",datetime.utcnow().strftime("%Y_%m_%d_%H_%M_%S"),int(random()*1000))
+    TIME_NOW = str.format(
+        "{}_{}", datetime.utcnow().strftime("%Y_%m_%d_%H_%M_%S"), int(random() * 1000)
+    )
     TENANCY_ID = "ocid1.tenancy.oc1..aaaaaaaa462hfhplpx652b32ix62xrdijppq2c7okwcqjlgrbknhgtj2kofa"
     COMPARTMENT_ID = "ocid1.tenancy.oc1..aaaaaaaa462hfhplpx652b32ix62xrdijppq2c7okwcqjlgrbknhgtj2kofa"
     METASTORE_ID = "ocid1.datacatalogmetastore.oc1.iad.amaaaaaabiudgxyap7tizm4gscwz7amu7dixz7ml3mtesqzzwwg3urvvdgua"

tests/integration/feature_store/test_dataset_validations.py

@@ -91,9 +91,8 @@ def test_dataset_model_details(self):
         assert dataset.oci_dataset.id
 
         dataset.materialise()
-        updated_dataset = dataset.add_models(ModelDetails().with_items(["model_ocid"]))
-        updated_dataset.show()
-        assert updated_dataset.model_details is not None
+        dataset.add_models(ModelDetails().with_items(["model_ocid_invalid"]))
+        assert len(dataset.model_details.get("items")) == 0
         self.clean_up_dataset(dataset)
         self.clean_up_feature_group(fg)
         self.clean_up_entity(entity)