Bugfix/fetch ds endpoint in mixin (#514)

Author: mingkang111

ads/common/oci_datascience.py

@@ -1,17 +1,28 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2021, 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
+
 import oci.data_science
-from ads.common.oci_mixin import OCIModelMixin
+
 from ads.common.decorator.utils import class_or_instance_method
+from ads.common.oci_mixin import OCIModelMixin
+
+ENV_VAR_OCI_ODSC_SERVICE_ENDPOINT = "OCI_ODSC_SERVICE_ENDPOINT"
 
 
 class OCIDataScienceMixin(OCIModelMixin):
     @class_or_instance_method
     def init_client(cls, **kwargs) -> oci.data_science.DataScienceClient:
+        client_kwargs = kwargs.get("client_kwargs", {})
+        if os.environ.get(ENV_VAR_OCI_ODSC_SERVICE_ENDPOINT):
+            client_kwargs.update(
+                dict(service_endpoint=os.environ.get(ENV_VAR_OCI_ODSC_SERVICE_ENDPOINT))
+            )
+            kwargs.update(client_kwargs)
         return cls._init_client(client=oci.data_science.DataScienceClient, **kwargs)
 
     @property

ads/common/oci_mixin.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
-# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2021, 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 """Contains Mixins for integrating OCI data models
@@ -13,18 +13,19 @@
 import re
 import traceback
 from datetime import date, datetime
-from typing import Callable, Optional, Union
 from enum import Enum
+from typing import Callable, Optional, Union
 
 import oci
 import yaml
+from dateutil import tz
+from dateutil.parser import parse
+from oci._vendor import six
+
 from ads.common import auth
 from ads.common.decorator.utils import class_or_instance_method
 from ads.common.utils import camel_to_snake
 from ads.config import COMPARTMENT_OCID
-from dateutil import tz
-from dateutil.parser import parse
-from oci._vendor import six
 
 logger = logging.getLogger(__name__)
 
@@ -273,7 +274,7 @@ def deserialize(cls, data, to_cls):
         else:
             return cls.__deserialize_model(data, to_cls)
 
-    @classmethod
+    @class_or_instance_method
     def __deserialize_model(cls, data, to_cls):
         """De-serializes list or dict to model."""
         if isinstance(data, to_cls):

tests/unitary/with_extras/model/test_generic_model.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
+# Copyright (c) 2021, 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 """Unit tests for model frameworks. Includes tests for:
@@ -1153,15 +1153,15 @@ def test_update_deployment_class_level(
     def test_update_deployment_instance_level_with_id(
         self, mock_client, mock_signer, mock_update
     ):
+        mock_signer.return_value = {}
         test_model_deployment_id = "xxxx.datasciencemodeldeployment.xxxx"
         md_props = ModelDeploymentProperties(model_id=test_model_deployment_id)
         md = ModelDeployment(properties=md_props)
 
         test_model = MagicMock(model_deployment=md, _summary_status=SummaryStatus())
         mock_update.return_value = test_model
 
-        generic_model = GenericModel(estimator=TestEstimator())
-        test_result = generic_model.update_deployment(
+        test_result = self.generic_model.update_deployment(
             model_deployment_id=test_model_deployment_id,
             properties=None,
             wait_for_completion=True,
@@ -1430,11 +1430,10 @@ def test_restart_deployment(
         test_model_deployment_id = "xxxx.datasciencemodeldeployment.xxxx"
         md_props = ModelDeploymentProperties(model_id=test_model_deployment_id)
         md = ModelDeployment(properties=md_props)
-        generic_model = GenericModel(estimator=TestEstimator())
-        generic_model.model_deployment = md
+        self.generic_model.model_deployment = md
         mock_deactivate.return_value = md
         mock_activate.return_value = md
-        generic_model.restart_deployment(max_wait_time=2000, poll_interval=50)
+        self.generic_model.restart_deployment(max_wait_time=2000, poll_interval=50)
         mock_deactivate.assert_called_with(max_wait_time=2000, poll_interval=50)
         mock_activate.assert_called_with(max_wait_time=2000, poll_interval=50)