[ML] Integrate MonitorSchedule with ScheduleOperations (Azure#30049)

Author: needuv

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/monitoring/input_data.py

@@ -8,12 +8,12 @@
 
 from azure.ai.ml.constants._monitoring import MonitorDatasetContext
 from azure.ai.ml._schema.core.schema import PatchedSchemaMeta
-from azure.ai.ml._schema.core.fields import NestedField, StringTransformedEnum
-from azure.ai.ml._schema.job.input_output_entry import MLTableInputSchema
+from azure.ai.ml._schema.core.fields import NestedField, StringTransformedEnum, UnionField
+from azure.ai.ml._schema.job.input_output_entry import MLTableInputSchema, DataInputSchema
 
 
 class MonitorInputDataSchema(metaclass=PatchedSchemaMeta):
-    input_dataset = NestedField(MLTableInputSchema)
+    input_dataset = UnionField(union_fields=[NestedField(DataInputSchema), NestedField(MLTableInputSchema)])
     dataset_context = StringTransformedEnum(allowed_values=[o.value for o in MonitorDatasetContext])
     target_column_name = fields.Str()
     pre_processing_component = fields.Str()

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_arm_id_utils.py

@@ -22,6 +22,7 @@
     SINGULARITY_FULL_NAME_REGEX_FORMAT,
     SINGULARITY_ID_REGEX_FORMAT,
     SINGULARITY_SHORT_NAME_REGEX_FORMAT,
+    NAMED_RESOURCE_ID_FORMAT_WITH_PARENT,
 )
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
 
@@ -328,6 +329,20 @@ def is_ARM_id_for_resource(name: Any, resource_type: str = ".*", sub_workspace_r
     return False
 
 
+def is_ARM_id_for_parented_resource(name: str, parent_resource_type: str, child_resource_type: str) -> bool:
+    resource_regex = NAMED_RESOURCE_ID_FORMAT_WITH_PARENT.format(
+        ".*",
+        ".*",
+        AZUREML_RESOURCE_PROVIDER,
+        ".*",
+        parent_resource_type,
+        ".*",
+        child_resource_type,
+        "*",
+    )
+    return re.match(resource_regex, name, re.IGNORECASE) is not None
+
+
 def is_registry_id_for_resource(name: Any) -> bool:
     if isinstance(name, str) and re.match(REGISTRY_URI_REGEX_FORMAT, name, re.IGNORECASE):
         return True

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_common.py

@@ -26,6 +26,7 @@
 AZUREML_RESOURCE_PROVIDER = "Microsoft.MachineLearningServices"
 RESOURCE_ID_FORMAT = "/subscriptions/{}/resourceGroups/{}/providers/{}/workspaces/{}"
 NAMED_RESOURCE_ID_FORMAT = "/subscriptions/{}/resourceGroups/{}/providers/{}/workspaces/{}/{}/{}"
+NAMED_RESOURCE_ID_FORMAT_WITH_PARENT = "/subscriptions/{}/resourceGroups/{}/providers/{}/workspaces/{}/{}/{}/{}/{}"
 LEVEL_ONE_NAMED_RESOURCE_ID_FORMAT = "/subscriptions/{}/resourceGroups/{}/providers/{}/{}/{}"
 VERSIONED_RESOURCE_ID_FORMAT = "/subscriptions/{}/resourceGroups/{}/providers/{}/workspaces/{}/{}/{}/versions/{}"
 LABELLED_RESOURCE_ID_FORMAT = "/subscriptions/{}/resourceGroups/{}/providers/{}/workspaces/{}/{}/{}/labels/{}"

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_monitoring/input_data.py

@@ -21,7 +21,7 @@ def __init__(
         target_column_name: str = None,
         pre_processing_component: str = None,
     ):
-        self.input_dataset = input_dataset
+        self.input_dataset = input_dataset if isinstance(input_dataset, Input) else Input(**input_dataset)
         self.dataset_context = dataset_context
         self.target_column_name = target_column_name
         self.pre_processing_component = pre_processing_component

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_schedule_operations.py

@@ -14,14 +14,24 @@
 
 from azure.ai.ml._telemetry import ActivityType, monitor_with_activity, monitor_with_telemetry_mixin
 from azure.ai.ml._utils._logger_utils import OpsLogger
-from azure.ai.ml.entities import Job, Schedule
+from azure.ai.ml.entities import Job, JobSchedule, Schedule
+from azure.ai.ml.entities._monitoring.schedule import MonitorSchedule
 from azure.core.credentials import TokenCredential
 from azure.core.polling import LROPoller
 from azure.core.tracing.decorator import distributed_trace
 
 from .._restclient.v2022_10_01.models import ScheduleListViewType
+from .._utils._arm_id_utils import is_ARM_id_for_parented_resource
+from .._utils.utils import snake_to_camel
 from .._utils._azureml_polling import AzureMLPolling
-from ..constants._common import AzureMLResourceType, LROConfigurations
+from ..constants._common import (
+    ARM_ID_PREFIX,
+    AzureMLResourceType,
+    LROConfigurations,
+    NAMED_RESOURCE_ID_FORMAT_WITH_PARENT,
+    AZUREML_RESOURCE_PROVIDER,
+)
+from ..constants._monitoring import MonitorSignalType
 from . import JobOperations
 from ._job_ops_helper import stream_logs_until_completion
 from ._operation_orchestrator import OperationOrchestrator
@@ -171,10 +181,14 @@ def begin_create_or_update(
         :rtype: Union[LROPoller, ~azure.ai.ml.entities.Schedule]
         """
 
-        schedule._validate(raise_error=True)
-        if isinstance(schedule.create_job, Job):
-            # Create all dependent resources for job inside schedule
-            self._job_operations._resolve_arm_id_or_upload_dependencies(schedule.create_job)
+        if isinstance(schedule, JobSchedule):
+            schedule._validate(raise_error=True)
+            if isinstance(schedule.create_job, Job):
+                # Create all dependent resources for job inside schedule
+                self._job_operations._resolve_arm_id_or_upload_dependencies(schedule.create_job)
+        elif isinstance(schedule, MonitorSchedule):
+            # resolve ARM id for target, compute, and input datasets for each signal
+            self._resolve_monitor_schedule_arm_id(schedule)
         # Create schedule
         schedule_data = schedule._to_rest_object()
         poller = self.service_client.begin_create_or_update(
@@ -221,3 +235,64 @@ def begin_disable(
         schedule = self.get(name=name)
         schedule._is_enabled = False
         return self.begin_create_or_update(schedule)
+
+    def _resolve_monitor_schedule_arm_id(self, schedule: MonitorSchedule) -> None:
+        # resolve compute ID
+        schedule.create_monitor.compute = self._orchestrators.get_asset_arm_id(
+            schedule.create_monitor.compute, AzureMLResourceType.COMPUTE, register_asset=False
+        )
+
+        # resolve target ARM ID
+        target = schedule.create_monitor.monitoring_target
+        if target and target.endpoint_deployment_id:
+            target.endpoint_deployment_id = (
+                target.endpoint_deployment_id[len(ARM_ID_PREFIX) :]
+                if target.endpoint_deployment_id.startswith(ARM_ID_PREFIX)
+                else target.endpoint_deployment_id
+            )
+
+            # if it is an ARM ID, don't process it
+            if not is_ARM_id_for_parented_resource(
+                target.endpoint_deployment_id,
+                snake_to_camel(AzureMLResourceType.ONLINE_ENDPOINT),
+                AzureMLResourceType.DEPLOYMENT,
+            ):
+                endpoint_name, deployment_name = target.endpoint_deployment_id.split(":")
+                target.endpoint_deployment_id = NAMED_RESOURCE_ID_FORMAT_WITH_PARENT.format(
+                    self._subscription_id,
+                    self._resource_group_name,
+                    AZUREML_RESOURCE_PROVIDER,
+                    self._workspace_name,
+                    snake_to_camel(AzureMLResourceType.ONLINE_ENDPOINT),
+                    endpoint_name,
+                    AzureMLResourceType.DEPLOYMENT,
+                    deployment_name,
+                )
+
+        elif target and target.model_id:
+            target.model_id = self._orchestrators.get_asset_arm_id(
+                target.model_id,
+                AzureMLResourceType.MODEL,
+                register_asset=False,
+            )
+
+        # resolve input paths and preprocessing component ids
+        for signal in schedule.create_monitor.monitoring_signals.values():
+            if signal.type == MonitorSignalType.CUSTOM:
+                for input_value in signal.input_datasets.values():
+                    self._job_operations._resolve_job_input(input_value.input_dataset, schedule._base_path)
+                    input_value.pre_processing_component = self._orchestrators.get_asset_arm_id(
+                        asset=input_value.pre_processing_component, azureml_type=AzureMLResourceType.COMPONENT
+                    )
+            else:
+                self._job_operations._resolve_job_inputs(
+                    [signal.target_dataset.dataset.input_dataset, signal.baseline_dataset.input_dataset],
+                    schedule._base_path,
+                )
+                signal.target_dataset.dataset.pre_processing_component = self._orchestrators.get_asset_arm_id(
+                    asset=signal.target_dataset.dataset.pre_processing_component,
+                    azureml_type=AzureMLResourceType.COMPONENT,
+                )
+                signal.baseline_dataset.pre_processing_component = self._orchestrators.get_asset_arm_id(
+                    asset=signal.baseline_dataset.pre_processing_component, azureml_type=AzureMLResourceType.COMPONENT
+                )

sdk/ml/azure-ai-ml/tests/monitoring/e2etests/test_monitor_schedule.py

@@ -0,0 +1,54 @@
+from typing import Callable
+
+from devtools_testutils import AzureRecordedTestCase
+import pytest
+
+from azure.ai.ml import MLClient
+from azure.ai.ml.constants._common import AzureMLResourceType
+from azure.ai.ml.entities._load_functions import load_schedule
+from azure.ai.ml.entities._monitoring.schedule import MonitorSchedule
+from azure.ai.ml._utils._arm_id_utils import is_ARM_id_for_resource
+
+
+@pytest.mark.timeout(600)
+@pytest.mark.usefixtures("recorded_test")
+@pytest.mark.core_sdk_test
+class TestMonitorSchedule(AzureRecordedTestCase):
+    def test_data_drift_schedule_create(
+        self, client: MLClient, data_with_2_versions: str, randstr: Callable[[str], str]
+    ):
+        test_path = "tests/test_configs/monitoring/yaml_configs/data_drift.yaml"
+
+        schedule_name = randstr("schedule_name")
+
+        params_override = [
+            {"name": schedule_name},
+            {
+                "create_monitor.monitoring_signals.testSignal.target_dataset.dataset.input_dataset.path": f"azureml:{data_with_2_versions}:1"
+            },
+            {"create_monitor.monitoring_signals.testSignal.target_dataset.dataset.input_dataset.type": "uri_folder"},
+            {
+                "create_monitor.monitoring_signals.testSignal.baseline_dataset.input_dataset.path": f"azureml:{data_with_2_versions}:2"
+            },
+            {"create_monitor.monitoring_signals.testSignal.baseline_dataset.input_dataset.type": "uri_folder"},
+        ]
+
+        schedule = load_schedule(test_path, params_override=params_override)
+        # not testing monitoring target expansion right now
+        schedule.create_monitor.monitoring_target = None
+        # bug in service when deserializing lookback_period is not supported yet
+        schedule.create_monitor.monitoring_signals["testSignal"].target_dataset.lookback_period = None
+
+        created_schedule = client.schedules.begin_create_or_update(schedule).result()
+
+        # test ARM id resolution
+        assert isinstance(created_schedule, MonitorSchedule)
+        assert is_ARM_id_for_resource(created_schedule.create_monitor.compute, AzureMLResourceType.COMPUTE)
+
+        data_drift_signal = created_schedule.create_monitor.monitoring_signals["testSignal"]
+        assert data_drift_signal.target_dataset
+        assert is_ARM_id_for_resource(
+            data_drift_signal.target_dataset.dataset.input_dataset.path, AzureMLResourceType.DATA
+        )
+        assert data_drift_signal.baseline_dataset
+        assert is_ARM_id_for_resource(data_drift_signal.baseline_dataset.input_dataset.path, AzureMLResourceType.DATA)