Add trigger api to operation

Signed-off-by: Brynn Yin <[email protected]>

sdk/ml/azure-ai-ml/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Features Added
 - Remove `experimental` tag for  `ml_client.jobs.validate`.
+- Support `ml_client.schedules.trigger(name='my_schedule')` function to trigger a schedule once.
 
 ### Bugs Fixed
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_schedule/schedule.py

@@ -10,12 +10,13 @@
 
 from typing_extensions import Literal
 
-from azure.ai.ml._restclient.v2023_06_01_preview.models import JobBase as RestJobBase
-from azure.ai.ml._restclient.v2023_06_01_preview.models import JobScheduleAction
-from azure.ai.ml._restclient.v2023_06_01_preview.models import PipelineJob as RestPipelineJob
-from azure.ai.ml._restclient.v2023_06_01_preview.models import Schedule as RestSchedule
-from azure.ai.ml._restclient.v2023_06_01_preview.models import ScheduleActionType as RestScheduleActionType
-from azure.ai.ml._restclient.v2023_06_01_preview.models import ScheduleProperties
+from azure.ai.ml._restclient.v2024_01_01_preview.models import JobBase as RestJobBase
+from azure.ai.ml._restclient.v2024_01_01_preview.models import JobScheduleAction
+from azure.ai.ml._restclient.v2024_01_01_preview.models import PipelineJob as RestPipelineJob
+from azure.ai.ml._restclient.v2024_01_01_preview.models import Schedule as RestSchedule
+from azure.ai.ml._restclient.v2024_01_01_preview.models import ScheduleActionType as RestScheduleActionType
+from azure.ai.ml._restclient.v2024_01_01_preview.models import ScheduleProperties
+from azure.ai.ml._restclient.v2024_01_01_preview.models import TriggerRunSubmissionDto as RestTriggerRunSubmissionDto
 from azure.ai.ml._schema.schedule.schedule import JobScheduleSchema
 from azure.ai.ml._utils.utils import camel_to_snake, dump_yaml_to_file, is_private_preview_enabled
 from azure.ai.ml.constants import JobType
@@ -464,3 +465,16 @@ def _get_telemetry_values(self, *args, **kwargs) -> Dict[Literal["trigger_type"]
         :rtype: Dict[Literal["trigger_type"], str]
         """
         return {"trigger_type": type(self.trigger).__name__}
+
+
+class TriggerRunSubmissionDto:
+    def __init__(self, schedule_action_type, submission_id, **kwargs):  # pylint: disable=unused-argument
+        self.schedule_action_type = schedule_action_type
+        self.submission_id = submission_id
+
+    @classmethod
+    def _from_rest_object(cls, obj: RestTriggerRunSubmissionDto) -> "TriggerRunSubmissionDto":
+        return cls(
+            schedule_action_type=obj.schedule_action_type,
+            submission_id=obj.submission_id,
+        )

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

@@ -46,6 +46,7 @@
     MonitorDatasetContext,
     MonitorSignalType,
 )
+from ..entities._schedule.schedule import TriggerRunSubmissionDto
 from . import DataOperations, JobOperations, OnlineDeploymentOperations
 from ._job_ops_helper import stream_logs_until_completion
 from ._operation_orchestrator import OperationOrchestrator
@@ -296,19 +297,20 @@ def trigger(
         self,
         name: str,
         **kwargs: Any,
-    ) -> LROPoller[Schedule]:
+    ) -> TriggerRunSubmissionDto:
         """Trigger a schedule once.
 
         :param name: Schedule name.
         :type name: str
         :return: TriggerRunSubmissionDto, or the result of cls(response)
-        :rtype: ~azure.mgmt.machinelearningservices.models.TriggerRunSubmissionDto
+        :rtype: ~azure.ai.ml.entities._schedule.schedule.TriggerRunSubmissionDto
         """
         return self.service_client.trigger(
             name=name,
             resource_group_name=self._operation_scope.resource_group_name,
             workspace_name=self._workspace_name,
             body=TriggerOnceRequest(schedule_time=datetime.now(timezone.utc).isoformat()),
+            cls=lambda _, obj, __: TriggerRunSubmissionDto._from_rest_object(obj),
             **kwargs,
         )
 

sdk/ml/azure-ai-ml/tests/schedule/e2etests/test_schedule.py

@@ -41,6 +41,10 @@ def test_schedule_lifetime(self, client: MLClient, randstr: Callable[[], str]):
         assert rest_schedule._is_enabled is True
         job: PipelineJob = rest_schedule.create_job
         assert isinstance(job.identity, AmlTokenConfiguration)
+        # trigger once
+        result = client.schedules.trigger(schedule.name)
+        assert result.submission_id is not None
+        assert client.jobs.get(name=result.submission_id) is not None
         # disable
         rest_schedule = client.schedules.begin_disable(schedule.name).result(timeout=LROConfigurations.POLLING_TIMEOUT)
         assert rest_schedule._is_enabled is False