Schedule operations rewire for feature and custom (Azure#31903)

LazaUK · Aug 31, 2023 · b292fe4 · b292fe4
1 parent 24e443f
commit b292fe4
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_schedule_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_schedule_operations.py
@@ -340,10 +340,22 @@ def _resolve_monitor_schedule_arm_id(  # pylint:disable=too-many-branches,too-ma
             if signal.type == MonitorSignalType.CUSTOM:
                 for input_value in signal.input_literals.values():
                     self._job_operations._resolve_job_input(input_value, schedule._base_path)
-                    input_value.pre_processing_component = self._orchestrators.get_asset_arm_id(
-                        asset=input_value.pre_processing_component, azureml_type=AzureMLResourceType.COMPONENT
+                for prod_data in signal.input_datasets.values():
+                    prod_data.pre_processing_component = self._orchestrators.get_asset_arm_id(
+                        asset=prod_data.pre_processing_component, azureml_type=AzureMLResourceType.COMPONENT
                     )
                 continue
+            if signal.type == MonitorSignalType.FEATURE_ATTRIBUTION_DRIFT:
+                for prod_data in signal.production_data:
+                    self._job_operations._resolve_job_input(prod_data.input_data, schedule._base_path)
+                    prod_data.pre_processing_component = self._orchestrators.get_asset_arm_id(
+                        asset=prod_data.pre_processing_component, azureml_type=AzureMLResourceType.COMPONENT
+                    )
+                self._job_operations._resolve_job_input(signal.reference_data.input_data, schedule._base_path)
+                signal.reference_data.pre_processing_component = self._orchestrators.get_asset_arm_id(
+                    asset=signal.reference_data.pre_processing_component, azureml_type=AzureMLResourceType.COMPONENT
+                )
+                continue
             error_messages = []
             if not signal.production_data or not signal.reference_data:
                 # if there is no target dataset, we check the type of signal