Upgrade pipeline job and component to 24-01 api

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

sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py

@@ -24,9 +24,6 @@
 from azure.ai.ml._restclient.v2020_09_01_dataplanepreview import (
     AzureMachineLearningWorkspaces as ServiceClient092020DataplanePreview,
 )
-from azure.ai.ml._restclient.workspace_dataplane import (
-    AzureMachineLearningWorkspaces as ServiceClientWorkspaceDataplane,
-)
 from azure.ai.ml._restclient.v2022_02_01_preview import AzureMachineLearningWorkspaces as ServiceClient022022Preview
 from azure.ai.ml._restclient.v2022_05_01 import AzureMachineLearningWorkspaces as ServiceClient052022
 from azure.ai.ml._restclient.v2022_10_01 import AzureMachineLearningWorkspaces as ServiceClient102022
@@ -39,6 +36,10 @@
 
 # Same object, but was renamed starting in v2023_08_01_preview
 from azure.ai.ml._restclient.v2023_10_01 import AzureMachineLearningServices as ServiceClient102023
+from azure.ai.ml._restclient.v2024_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient0124Preview
+from azure.ai.ml._restclient.workspace_dataplane import (
+    AzureMachineLearningWorkspaces as ServiceClientWorkspaceDataplane,
+)
 from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationsContainer, OperationScope
 from azure.ai.ml._telemetry.logging_handler import get_appinsights_log_handler
 from azure.ai.ml._user_agent import USER_AGENT
@@ -344,6 +345,13 @@ def __init__(
             **kwargs,
         )
 
+        self._service_client_01_2024_preview = ServiceClient0124Preview(
+            credential=self._credential,
+            subscription_id=self._operation_scope._subscription_id,
+            base_url=base_url,
+            **kwargs,
+        )
+
         # A general purpose, user-configurable pipeline for making
         # http requests
         self._requests_pipeline = HttpPipeline(**kwargs)
@@ -554,7 +562,7 @@ def __init__(
         self._components = ComponentOperations(
             self._operation_scope,
             self._operation_config,
-            self._service_client_10_2021_dataplanepreview if registry_name else self._service_client_10_2022,
+            self._service_client_10_2021_dataplanepreview if registry_name else self._service_client_01_2024_preview,
             self._operation_container,
             self._preflight,
             **ops_kwargs,
@@ -568,14 +576,14 @@ def __init__(
             self._credential,
             _service_client_kwargs=kwargs,
             requests_pipeline=self._requests_pipeline,
-            service_client_08_2023_preview=self._service_client_08_2023_preview,
+            service_client_01_2024_preview=self._service_client_01_2024_preview,
             **ops_kwargs,
         )
         self._operation_container.add(AzureMLResourceType.JOB, self._jobs)
         self._schedules = ScheduleOperations(
             self._operation_scope,
             self._operation_config,
-            self._service_client_06_2023_preview,
+            self._service_client_01_2024_preview,
             self._operation_container,
             self._credential,
             _service_client_kwargs=kwargs,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/component.py

@@ -10,7 +10,7 @@
 from marshmallow import INCLUDE
 from typing_extensions import Literal
 
-from ..._restclient.v2023_08_01_preview.models import (
+from ..._restclient.v2024_01_01_preview.models import (
     ComponentContainer,
     ComponentContainerProperties,
     ComponentVersion,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job.py

@@ -16,7 +16,7 @@
 from azure.ai.ml._restclient.runhistory.models import Run
 from azure.ai.ml._restclient.v2023_04_01_preview.models import JobBase, JobService
 from azure.ai.ml._restclient.v2023_04_01_preview.models import JobType as RestJobType
-from azure.ai.ml._restclient.v2023_08_01_preview.models import JobBase as JobBase_2308
+from azure.ai.ml._restclient.v2024_01_01_preview.models import JobBase as JobBase_2401
 from azure.ai.ml._utils._html_utils import make_link, to_html
 from azure.ai.ml._utils.utils import dump_yaml_to_file
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY, CommonYamlFields
@@ -280,7 +280,7 @@ def _load(
 
     @classmethod
     def _from_rest_object(  # pylint: disable=too-many-return-statements
-        cls, obj: Union[JobBase, JobBase_2308, Run]
+        cls, obj: Union[JobBase, JobBase_2401, Run]
     ) -> "Job":  # pylint: disable=too-many-return-statements
         from azure.ai.ml.entities import PipelineJob
         from azure.ai.ml.entities._builders.command import Command

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/pipeline/pipeline_job.py

@@ -10,8 +10,8 @@
 from pathlib import Path
 from typing import Dict, List, Optional, Union
 
-from azure.ai.ml._restclient.v2023_08_01_preview.models import JobBase
-from azure.ai.ml._restclient.v2023_08_01_preview.models import PipelineJob as RestPipelineJob
+from azure.ai.ml._restclient.v2024_01_01_preview.models import JobBase
+from azure.ai.ml._restclient.v2024_01_01_preview.models import PipelineJob as RestPipelineJob
 from azure.ai.ml._schema import PathAwareSchema
 from azure.ai.ml._schema.pipeline.pipeline_job import PipelineJobSchema
 from azure.ai.ml._utils._arm_id_utils import get_resource_name_from_arm_id_safe

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

@@ -14,8 +14,8 @@
 from azure.ai.ml._restclient.v2021_10_01_dataplanepreview import (
     AzureMachineLearningWorkspaces as ServiceClient102021Dataplane,
 )
-from azure.ai.ml._restclient.v2022_10_01 import AzureMachineLearningWorkspaces as ServiceClient102022
-from azure.ai.ml._restclient.v2022_10_01.models import ComponentVersion, ListViewType
+from azure.ai.ml._restclient.v2024_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient012024
+from azure.ai.ml._restclient.v2024_01_01_preview.models import ComponentVersion, ListViewType
 from azure.ai.ml._scope_dependent_operations import (
     OperationConfig,
     OperationsContainer,
@@ -91,7 +91,7 @@ def __init__(
         self,
         operation_scope: OperationScope,
         operation_config: OperationConfig,
-        service_client: Union[ServiceClient102022, ServiceClient102021Dataplane],
+        service_client: Union[ServiceClient012024, ServiceClient102021Dataplane],
         all_operations: OperationsContainer,
         preflight_operation: Optional[DeploymentsOperations] = None,
         **kwargs: Dict,

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

@@ -29,8 +29,8 @@
 from azure.ai.ml._restclient.runhistory.models import Run
 from azure.ai.ml._restclient.v2023_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient022023Preview
 from azure.ai.ml._restclient.v2023_04_01_preview.models import JobBase, ListViewType, UserIdentity
-from azure.ai.ml._restclient.v2023_08_01_preview.models import JobBase as JobBase_2308
 from azure.ai.ml._restclient.v2023_08_01_preview.models import JobType as RestJobType
+from azure.ai.ml._restclient.v2024_01_01_preview.models import JobBase as JobBase_2401
 from azure.ai.ml._scope_dependent_operations import (
     OperationConfig,
     OperationsContainer,
@@ -175,7 +175,7 @@ def __init__(
             self._all_operations, self._operation_scope, self._operation_config
         )  # pylint: disable=line-too-long
 
-        self.service_client_08_2023_preview = kwargs.pop("service_client_08_2023_preview", None)
+        self.service_client_01_2024_preview = kwargs.pop("service_client_01_2024_preview", None)
         self._kwargs = kwargs
 
         self._requests_pipeline: HttpPipeline = kwargs.pop("requests_pipeline")
@@ -692,7 +692,7 @@ def create_or_update(
             # request for submitting to ES. Once we request to ES and start the run, we
             # need to put the same body to MFE to append user tags etc.
             if rest_job_resource.properties.job_type == RestJobType.PIPELINE:
-                job_object = self._get_job_2308(rest_job_resource.name)
+                job_object = self._get_job_2401(rest_job_resource.name)
             else:
                 job_object = self._get_job(rest_job_resource.name)
             if result.properties.tags is not None:
@@ -717,10 +717,10 @@ def _create_or_update_with_different_version_api(  # pylint: disable=name-too-lo
         service_client_operation = self._operation_2023_02_preview
         # Upgrade api from 2023-04-01-preview to 2023-08-01 for pipeline job
         if rest_job_resource.properties.job_type == RestJobType.PIPELINE:
-            service_client_operation = self.service_client_08_2023_preview.jobs
+            service_client_operation = self.service_client_01_2024_preview.jobs
 
         if rest_job_resource.properties.job_type == RestJobType.SWEEP:
-            service_client_operation = self.service_client_08_2023_preview.jobs
+            service_client_operation = self.service_client_01_2024_preview.jobs
 
         result = service_client_operation.create_or_update(
             id=rest_job_resource.name,
@@ -735,7 +735,7 @@ def _create_or_update_with_different_version_api(  # pylint: disable=name-too-lo
     def _archive_or_restore(self, name: str, is_archived: bool) -> None:
         job_object = self._get_job(name)
         if job_object.properties.job_type == RestJobType.PIPELINE:
-            job_object = self._get_job_2308(name)
+            job_object = self._get_job_2401(name)
         if _is_pipeline_child_job(job_object):
             raise PipelineChildJobError(job_id=job_object.id)
         job_object.properties.is_archived = is_archived
@@ -1026,10 +1026,10 @@ def _get_job(self, name: str) -> JobBase:
             **self._kwargs,
         )
 
-    # Upgrade api from 2023-04-01-preview to 2023-08-01 for pipeline job
-    # We can remove this function once `_get_job` function has also been upgraded to 2023-08-01 api
-    def _get_job_2308(self, name: str) -> JobBase_2308:
-        service_client_operation = self.service_client_08_2023_preview.jobs
+    # Upgrade api from 2023-04-01-preview to 2024-01-01-preview for pipeline job
+    # We can remove this function once `_get_job` function has also been upgraded to the same version with pipeline
+    def _get_job_2401(self, name: str) -> JobBase_2401:
+        service_client_operation = self.service_client_01_2024_preview.jobs
         return service_client_operation.get(
             id=name,
             resource_group_name=self._operation_scope.resource_group_name,

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

@@ -2,9 +2,10 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 # pylint: disable=protected-access
+from datetime import datetime, timezone
 from typing import Any, Iterable, List, Tuple, cast
 
-from azure.ai.ml._restclient.v2023_06_01_preview import AzureMachineLearningWorkspaces as ServiceClient062023Preview
+from azure.ai.ml._restclient.v2024_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient010124Preview
 from azure.ai.ml._scope_dependent_operations import (
     OperationConfig,
     OperationsContainer,
@@ -24,6 +25,7 @@
 from azure.core.tracing.decorator import distributed_trace
 
 from .._restclient.v2022_10_01.models import ScheduleListViewType
+from .._restclient.v2024_01_01_preview.models import TriggerOnceRequest
 from .._utils._arm_id_utils import AMLNamedArmId, AMLVersionedArmId, is_ARM_id_for_parented_resource
 from .._utils._azureml_polling import AzureMLPolling
 from .._utils.utils import snake_to_camel
@@ -65,14 +67,14 @@ def __init__(
         self,
         operation_scope: OperationScope,
         operation_config: OperationConfig,
-        service_client_06_2023_preview: ServiceClient062023Preview,
+        service_client_01_2024_preview: ServiceClient010124Preview,
         all_operations: OperationsContainer,
         credential: TokenCredential,
         **kwargs: Any,
     ):
         super(ScheduleOperations, self).__init__(operation_scope, operation_config)
         ops_logger.update_info(kwargs)
-        self.service_client = service_client_06_2023_preview.schedules
+        self.service_client = service_client_01_2024_preview.schedules
         self._all_operations = all_operations
         self._stream_logs_until_completion = stream_logs_until_completion
         # Dataplane service clients are lazily created as they are needed
@@ -240,7 +242,6 @@ def begin_create_or_update(
             self._resolve_monitor_schedule_arm_id(schedule)
         # Create schedule
         schedule_data = schedule._to_rest_object()
-        print(schedule_data.properties.tags)
         poller = self.service_client.begin_create_or_update(
             resource_group_name=self._operation_scope.resource_group_name,
             workspace_name=self._workspace_name,
@@ -289,6 +290,28 @@ def begin_disable(
         schedule._is_enabled = False
         return self.begin_create_or_update(schedule, **kwargs)
 
+    @distributed_trace
+    @monitor_with_activity(logger, "Schedule.Trigger", ActivityType.PUBLICAPI)
+    def trigger(
+        self,
+        name: str,
+        **kwargs: Any,
+    ) -> LROPoller[Schedule]:
+        """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
+        """
+        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()),
+            **kwargs,
+        )
+
     def _resolve_monitor_schedule_arm_id(  # pylint:disable=too-many-branches,too-many-statements
         self, schedule: MonitorSchedule
     ) -> None: