Do not report feature data if deployment has it disabled. (#1662)

Author: nickolai-dr

custom_model_runner/datarobot_drum/drum/language_predictors/base_language_predictor.py

@@ -104,6 +104,16 @@ def __init__(
         self._mlops = None
         self._schema_validator = None
         self._prompt_column_name = DEFAULT_PROMPT_COLUMN_NAME
+        self._deployment = None
+
+        self._tracking_settings = {
+            "target_drift": {"enabled": True},
+            "feature_drift": {"enabled": True},
+        }
+        self._data_collection = {"enabled": True}
+
+        self._settings_refresh_time = time.monotonic()
+        self._settings_refresh_interval = 60  # sec
 
     def configure(self, params):
         """
@@ -154,6 +164,7 @@ def _init_mlops(self):
         if to_bool(self._params.get("allow_dr_api_access")):
             try:
                 self._deployment = dr.Deployment.get(deployment_id)
+                self._refresh_tracking_settings()
             except Exception as e:
                 logger.warning(f"Failed to get deployment info: {e}", exc_info=True)
 
@@ -172,6 +183,14 @@ def _init_mlops(self):
 
         self._mlops.init()
 
+    def _refresh_tracking_settings(self):
+        deployment_id = self._params.get("deployment_id", None)
+        if to_bool(self._params.get("allow_dr_api_access")) and deployment_id is not None:
+            self._deployment = dr.Deployment.get(deployment_id)
+            self._tracking_settings = self._deployment.get_drift_tracking_settings()
+            self._data_collection = self._deployment.get_predictions_data_collection_settings()
+            self._settings_refresh_time = time.monotonic()
+
     def _configure_mlops(self):
         # If monitor_settings were provided (e.g. for testing) use them, otherwise we will
         # use the API spooler as the default config.
@@ -318,6 +337,16 @@ def _mlops_report_chat_prediction(
         except DRCommonException:
             logger.exception("Failed to report deployment stats")
 
+        if self._deployment is not None:
+            if time.monotonic() - self._settings_refresh_time > self._settings_refresh_interval:
+                self._refresh_tracking_settings()
+
+        is_drift = self._tracking_settings["feature_drift"]["enabled"]
+        is_collection = self._data_collection["enabled"]
+
+        if not (is_drift or is_collection):
+            return
+
         prompt_content = completion_create_params["messages"][-1]["content"]
         if isinstance(prompt_content, str):
             latest_message = completion_create_params["messages"][-1]["content"]

tests/unit/datarobot_drum/drum/language_predictors/test_base_language_predictor.py

@@ -123,7 +123,7 @@ def _language_predictor_with_mlops_params_dr_api_access(self):
 
     @pytest.fixture
     def mock_dr_client(self):
-        with patch.object(dr, "Client") as _:
+        with patch.object(dr, "Client") as m:
             yield
 
     def test_mlops_init(self, language_predictor_with_mlops, mock_mlops):
@@ -271,7 +271,10 @@ def test_association_id(self, language_predictor_with_mlops, mock_mlops):
             mock_chat.assert_called_once_with(ANY, association_id)
             hasattr(completion, "datarobot_association_id")
 
-    def test_prompt_column_name(self, chat_python_model_adapter, mock_mlops, mock_dr_client):
+    @pytest.mark.parametrize("row_storage_enabled", [False, True])
+    def test_prompt_column_name(
+        self, chat_python_model_adapter, mock_mlops, mock_dr_client, row_storage_enabled
+    ):
         language_predictor = TestLanguagePredictor()
         language_predictor_with_mlops_params = (
             self._language_predictor_with_mlops_params_dr_api_access()
@@ -282,6 +285,13 @@ def test_prompt_column_name(self, chat_python_model_adapter, mock_mlops, mock_dr
             deployment_instance.return_value.get_champion_model_package.return_value = Mock()
             mock_deployment.get.return_value = deployment_instance
 
+            deployment_instance.get_drift_tracking_settings.return_value = {
+                "target_drift": {"enabled": False},
+                "feature_drift": {"enabled": False},
+            }
+            deployment_instance.get_predictions_data_collection_settings.return_value = {
+                "enabled": row_storage_enabled
+            }
             language_predictor.configure(language_predictor_with_mlops_params)
 
         def chat_hook(completion_request):
@@ -298,16 +308,19 @@ def chat_hook(completion_request):
             }
         )
 
-        mock_mlops.report_predictions_data.assert_called_once_with(
-            ANY,
-            ["How are you"],
-            association_ids=ANY,
-        )
-        # Compare features dataframe separately as this doesn't play nice with assert_called
-        assert (
-            mock_mlops.report_predictions_data.call_args.args[0]["newPromptName"].values[0]
-            == "Hello!"
-        )
+        if row_storage_enabled:
+            mock_mlops.report_predictions_data.assert_called_once_with(
+                ANY,
+                ["How are you"],
+                association_ids=ANY,
+            )
+            # Compare features dataframe separately as this doesn't play nice with assert_called
+            assert (
+                mock_mlops.report_predictions_data.call_args.args[0]["newPromptName"].values[0]
+                == "Hello!"
+            )
+        else:
+            mock_mlops.report_predictions_data.assert_not_called()
 
     @pytest.mark.parametrize("stream", [False, True])
     def test_failing_hook_with_mlops(self, language_predictor_with_mlops, mock_mlops, stream):