[ML] Azure standards related fixes for Job Operations (Azure#26371)

* [ML] Azure standards related fixes for Job Operations

* Revert few changes picked from yet to be merged pr

* Updated changelog, removed None[]

* updates

* fix recordings for cancel job

Co-authored-by: Douglas Xiao <xiake@microsoft.com>
Co-authored-by: Neehar Duvvuri <neduvvur@microsoft.com>

Author: 3 people

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

@@ -29,13 +29,15 @@
  - Enable using @dsl.pipeline without brackets when no additional parameters.
  - Expose Azure subscription Id and resource group name from MLClient objects.
  - Added Idle Shutdown support for Compute Instances, allowing instances to shutdown after a set period of inactivity.
+ - JobOperations.cancel() returns a LROPoller.
 
 ### Breaking Changes
  - Change (begin_)create_or_update typehints to use generics.
  - Remove invalid option from create_or_update typehints.
  - Change error returned by (begin_)create_or_update invalid input to TypeError.
  - Rename set_image_model APIs for all vision tasks to set_training_parameters
  - JobOperations.download defaults to "." instead of Path.cwd()
+ - JobOperations.cancel() is renamed to JobOperations.begin_cancel() and it returns LROPoller
  - Workspace.list_keys renamed to Workspace.get_keys.
 
 ### Bugs Fixed

sdk/ml/azure-ai-ml/azure/ai/ml/entities/__init__.py

@@ -84,6 +84,7 @@
 from ._schedule.schedule import JobSchedule
 from ._schedule.trigger import CronTrigger, RecurrencePattern, RecurrenceTrigger
 from ._system_data import SystemData
+from ._validation import ValidationResult
 from ._workspace.connections.workspace_connection import WorkspaceConnection
 from ._workspace.customer_managed_key import CustomerManagedKey
 from ._workspace.identity import ManagedServiceIdentity
@@ -198,6 +199,7 @@
     "AmlComputeNodeInfo",
     "SystemCreatedAcrAccount",
     "SystemCreatedStorageAccount",
+    "ValidationResult",
     "RegistryRegionArmDetails",
     "Registry",
     "SynapseSparkCompute",

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

@@ -91,6 +91,8 @@
 from azure.ai.ml.sweep import SweepJob
 from azure.core.credentials import TokenCredential
 from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
+from azure.core.polling import LROPoller
+from azure.core.tracing.decorator import distributed_trace
 
 from .._utils._experimental import experimental
 from ..constants._component import ComponentSource
@@ -209,11 +211,12 @@ def _api_url(self):
             self._api_base_url = self._get_workspace_url(url_key=API_URL_KEY)
         return self._api_base_url
 
+    @distributed_trace
     @monitor_with_activity(logger, "Job.List", ActivityType.PUBLICAPI)
     def list(
         self,
-        parent_job_name: str = None,
         *,
+        parent_job_name: str = None,
         list_view_type: ListViewType = ListViewType.ACTIVE_ONLY,
         **kwargs,
     ) -> Iterable[Job]:
@@ -251,6 +254,7 @@ def _handle_rest_errors(self, job_object):
         except JobParsingError:
             pass
 
+    @distributed_trace
     @monitor_with_telemetry_mixin(logger, "Job.Get", ActivityType.PUBLICAPI)
     def get(self, name: str) -> Job:
         """Get a job resource.
@@ -294,13 +298,16 @@ def _show_services(self, name: str, node_index: int):
             k: ServiceInstance._from_rest_object(v, node_index) for k, v in service_instances_dict.instances.items()
         }
 
+    @distributed_trace
     @monitor_with_activity(logger, "Job.Cancel", ActivityType.PUBLICAPI)
-    def cancel(self, name: str) -> None:
+    def begin_cancel(self, name: str) -> LROPoller[None]:
         """Cancel job resource.
 
         :param str name: Name of the job.
         :return: None, or the result of cls(response)
         :rtype: None
+        :return: A poller to track the operation status.
+        :rtype: ~azure.core.polling.LROPoller[None]
         :raise: ResourceNotFoundError if can't find a job matching provided name.
         """
         return self._operation_2022_06_preview.begin_cancel(
@@ -344,6 +351,7 @@ def try_get_compute_arm_id(self, compute: Union[Compute, str]):
                 raise ResourceNotFoundError(response=response)
         return None
 
+    @distributed_trace
     @experimental
     @monitor_with_telemetry_mixin(logger, "Job.Validate", ActivityType.PUBLICAPI)
     def validate(self, job: Job, *, raise_on_failure: bool = False, **kwargs) -> ValidationResult:
@@ -410,6 +418,7 @@ def _validate(
         validation_result.resolve_location_for_diagnostics(job._source_path)
         return validation_result.try_raise(raise_error=raise_on_failure, error_target=ErrorTarget.PIPELINE)
 
+    @distributed_trace
     @monitor_with_telemetry_mixin(logger, "Job.CreateOrUpdate", ActivityType.PUBLICAPI)
     def create_or_update(
         self,
@@ -539,6 +548,7 @@ def _archive_or_restore(self, name: str, is_archived: bool):
             body=job_object,
         )
 
+    @distributed_trace
     @monitor_with_telemetry_mixin(logger, "Job.Archive", ActivityType.PUBLICAPI)
     def archive(self, name: str) -> None:
         """Archive a job or restore an archived job.
@@ -550,6 +560,7 @@ def archive(self, name: str) -> None:
 
         self._archive_or_restore(name=name, is_archived=True)
 
+    @distributed_trace
     @monitor_with_telemetry_mixin(logger, "Job.Restore", ActivityType.PUBLICAPI)
     def restore(self, name: str) -> None:
         """Archive a job or restore an archived job.
@@ -561,6 +572,7 @@ def restore(self, name: str) -> None:
 
         self._archive_or_restore(name=name, is_archived=False)
 
+    @distributed_trace
     @monitor_with_activity(logger, "Job.Stream", ActivityType.PUBLICAPI)
     def stream(self, name: str) -> None:
         """Stream logs of a job.
@@ -577,6 +589,7 @@ def stream(self, name: str) -> None:
             self._runs_operations, job_object, self._datastore_operations, requests_pipeline=self._requests_pipeline
         )
 
+    @distributed_trace
     @monitor_with_activity(logger, "Job.Download", ActivityType.PUBLICAPI)
     def download(
         self,

sdk/ml/azure-ai-ml/tests/batch_services/e2etests/test_batch_deployment.py

@@ -14,6 +14,7 @@
 from azure.ai.ml.entities._inputs_outputs import Input
 from azure.ai.ml.operations._job_ops_helper import _wait_before_polling
 from azure.ai.ml.operations._run_history_constants import JobStatus, RunHistoryConstants
+from azure.core.polling import LROPoller
 
 
 @contextmanager
@@ -142,7 +143,9 @@ def wait_until_done(job: Job, timeout: int = None) -> None:
                 job = client.jobs.get(job.name)
                 if timeout is not None and time.time() - poll_start_time > timeout:
                     # if timeout is passed in, execute job cancel if timeout and directly return CANCELED status
-                    client.jobs.cancel(job.name)
+                    cancel_poller = client.jobs.begin_cancel(job.name)
+                    assert isinstance(cancel_poller, LROPoller)
+                    assert cancel_poller.result() is None
                     return JobStatus.CANCELED
             return job.status
 

sdk/ml/azure-ai-ml/tests/command_job/e2etests/test_command_job.py

@@ -19,6 +19,7 @@
 from azure.ai.ml.exceptions import ValidationException
 from azure.ai.ml.operations._job_ops_helper import _wait_before_polling
 from azure.ai.ml.operations._run_history_constants import JobStatus, RunHistoryConstants
+from azure.core.polling import LROPoller
 
 # These params are logged in ..\test_configs\python\simple_train.py. test_command_job_with_params asserts these parameters are
 # logged in the training script, so any changes to parameter logging in simple_train.py must preserve this logging or change it both
@@ -250,7 +251,9 @@ def test_command_job_cancel(self, randstr: Callable[[], str], client: MLClient)
         )
         command_job_resource = client.jobs.create_or_update(job=job)
         assert command_job_resource.name == job_name
-        client.jobs.cancel(job_name)
+        cancel_poller = client.jobs.begin_cancel(job_name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
         command_job_resource_2 = client.jobs.get(job_name)
         assert command_job_resource_2.status in (JobStatus.CANCEL_REQUESTED, JobStatus.CANCELED)
 
@@ -295,7 +298,9 @@ def test_command_job_dependency_label_resolution(self, randstr: Callable[[], str
             ],
         )
         command_job_resource = client.jobs.create_or_update(job=job)
-        client.jobs.cancel(job_name)
+        cancel_poller = client.jobs.begin_cancel(job_name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
 
         # Check that environment resolves to latest version
         assert command_job_resource.environment == f"{environment_name}:{environment_versions[-1]}"

sdk/ml/azure-ai-ml/tests/component/e2etests/test_component.py

@@ -23,16 +23,13 @@
 from azure.ai.ml.entities._load_functions import load_code, load_job
 from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
 from azure.core.paging import ItemPaged
+from azure.core.polling import LROPoller
 
 from .._util import _COMPONENT_TIMEOUT_SECOND
 from ..unittests.test_component_schema import load_component_entity_from_rest_json
 
 
-from devtools_testutils import (
-    AzureRecordedTestCase,
-    is_live,
-    set_bodiless_matcher
-)
+from devtools_testutils import AzureRecordedTestCase, is_live, set_bodiless_matcher
 
 
 def create_component(
@@ -269,7 +266,7 @@ def test_spark_component(self, client: MLClient, randstr: Callable[[], str]) ->
             path="./tests/test_configs/dsl_pipeline/spark_job_in_pipeline/add_greeting_column_component.yml",
             expected_dict=expected_dict,
             omit_fields=["name", "creation_context", "id", "code", "environment"],
-            recorded_component_name="spark_component_name"
+            recorded_component_name="spark_component_name",
         )
 
     @pytest.mark.parametrize(
@@ -381,10 +378,7 @@ def test_component_update(self, client: MLClient, randstr: Callable[[str], str])
         assert component_resource.display_name == display_name
 
     @pytest.mark.disable_mock_code_hash
-    @pytest.mark.skipif(
-        condition=not is_live(),
-        reason="non-deterministic upload fails in playback on CI"
-    )
+    @pytest.mark.skipif(condition=not is_live(), reason="non-deterministic upload fails in playback on CI")
     def test_component_create_twice_same_code_arm_id(
         self, client: MLClient, randstr: Callable[[str], str], tmp_path: Path
     ) -> None:
@@ -409,10 +403,7 @@ def test_component_create_twice_same_code_arm_id(
         # the code arm id should be the same
         assert component_resource1.code == component_resource2.code
 
-    @pytest.mark.skipif(
-        condition=not is_live(),
-        reason="non-deterministic upload fails in playback on CI"
-    )
+    @pytest.mark.skipif(condition=not is_live(), reason="non-deterministic upload fails in playback on CI")
     def test_component_update_code(self, client: MLClient, randstr: Callable[[str], str], tmp_path: Path) -> None:
         component_name = randstr("component_name")
         path = "./tests/test_configs/components/basic_component_code_local_path.yml"
@@ -851,7 +842,9 @@ def test_create_pipeline_component_from_job(self, client: MLClient, randstr: Cal
         )
         job = client.jobs.create_or_update(pipeline_job)
         try:
-            client.jobs.cancel(job.name)
+            cancel_poller = client.jobs.begin_cancel(job.name)
+            assert isinstance(cancel_poller, LROPoller)
+            assert cancel_poller.result() is None
         except Exception:
             pass
         component = PipelineComponent(name=randstr(), source_job_id=job.id)

sdk/ml/azure-ai-ml/tests/conftest.py

@@ -59,7 +59,7 @@ def fake_datastore_key() -> str:
 
 @pytest.fixture(autouse=True)
 def add_sanitizers(test_proxy, fake_datastore_key):
-    add_remove_header_sanitizer(headers="x-azureml-token")
+    add_remove_header_sanitizer(headers="x-azureml-token,Log-URL")
     set_custom_default_matcher(excluded_headers="x-ms-meta-name,x-ms-meta-version")
     add_body_key_sanitizer(json_path="$.key", value=fake_datastore_key)
     add_body_key_sanitizer(json_path="$....key", value=fake_datastore_key)

sdk/ml/azure-ai-ml/tests/dsl/e2etests/test_dsl_pipeline.py

@@ -31,6 +31,7 @@
 from azure.ai.ml.entities import Data, PipelineJob
 from azure.ai.ml.exceptions import ValidationException
 from azure.ai.ml.parallel import ParallelJob, RunFunction, parallel_run_function
+from azure.core.polling import LROPoller
 
 from .._util import _DSL_TIMEOUT_SECOND
 
@@ -1510,7 +1511,9 @@ def parallel_in_pipeline(job_data_path, score_model):
         )
         # submit pipeline job
         pipeline_job = client.jobs.create_or_update(pipeline, experiment_name="parallel_in_pipeline")
-        client.jobs.cancel(pipeline_job.name)
+        cancel_poller = client.jobs.begin_cancel(pipeline_job.name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
         # check required fields in job dict
         job_dict = pipeline_job._to_dict()
         expected_keys = ["status", "properties", "tags", "creation_context"]
@@ -1541,7 +1544,9 @@ def parallel_in_pipeline(job_data_path):
         )
         # submit pipeline job
         pipeline_job = client.jobs.create_or_update(pipeline, experiment_name="parallel_in_pipeline")
-        client.jobs.cancel(pipeline_job.name)
+        cancel_poller = client.jobs.begin_cancel(pipeline_job.name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
         # check required fields in job dict
         job_dict = pipeline_job._to_dict()
         expected_keys = ["status", "properties", "tags", "creation_context"]
@@ -1675,6 +1680,7 @@ def parallel_in_pipeline(job_data_path):
         assert_job_input_output_types(pipeline_job)
         assert pipeline_job.settings.default_compute == "cpu-cluster"
 
+    @pytest.mark.skip("TODO: re-record since job is in terminal state before cancel")
     def test_parallel_job(self, randstr: Callable[[str], str], client: MLClient):
         environment = "AzureML-sklearn-0.24-ubuntu18.04-py37-cpu:5"
         inputs = {
@@ -1740,7 +1746,9 @@ def parallel_in_pipeline(job_data_path):
             pipeline,
             experiment_name="parallel_in_pipeline",
         )
-        client.jobs.cancel(pipeline_job.name)
+        cancel_poller = client.jobs.begin_cancel(pipeline_job.name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
         omit_fields = [
             "jobs.parallel_node.task.code",
             "jobs.parallel_node.task.environment",
@@ -1826,7 +1834,9 @@ def parallel_in_pipeline(job_data_path):
 
         # submit pipeline job
         pipeline_job = client.jobs.create_or_update(pipeline, experiment_name="parallel_in_pipeline")
-        client.jobs.cancel(pipeline_job.name)
+        cancel_poller = client.jobs.begin_cancel(pipeline_job.name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
 
         omit_fields = [
             "jobs.*.task.code",
@@ -2260,7 +2270,9 @@ def spark_pipeline_from_yaml(iris_data):
 
         # submit pipeline job
         pipeline_job = client.jobs.create_or_update(pipeline, experiment_name="spark_in_pipeline")
-        client.jobs.cancel(pipeline_job.name)
+        cancel_poller = client.jobs.begin_cancel(pipeline_job.name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
         # check required fields in job dict
         job_dict = pipeline_job._to_dict()
         expected_keys = ["status", "properties", "tags", "creation_context"]

sdk/ml/azure-ai-ml/tests/dsl/e2etests/test_dsl_pipeline_samples.py

@@ -16,6 +16,7 @@
 from azure.ai.ml.entities import Job, PipelineJob
 from azure.ai.ml.operations._run_history_constants import JobStatus
 from azure.core.exceptions import HttpResponseError
+from azure.core.polling import LROPoller
 
 from .._util import _DSL_TIMEOUT_SECOND
 
@@ -37,7 +38,9 @@ def job_cancel_after_submit(pipeline, client: MLClient):
     #  the status before confirming whether there is a problem with pipeline cancel.
     job = client.jobs.create_or_update(pipeline)
     try:
-        client.jobs.cancel(job.name)
+        cancel_poller = client.jobs.begin_cancel(job.name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
     except HttpResponseError:
         pass
 

sdk/ml/azure-ai-ml/tests/import_job/e2etests/test_import_job.py

@@ -18,6 +18,7 @@
 from azure.ai.ml.operations._job_ops_helper import _wait_before_polling
 from azure.ai.ml.operations._operation_orchestrator import OperationOrchestrator
 from azure.ai.ml.operations._run_history_constants import JobStatus, RunHistoryConstants
+from azure.core.polling import LROPoller
 
 
 from devtools_testutils import AzureRecordedTestCase
@@ -84,7 +85,9 @@ def validate_import_job_submit_cancel(self, job: ImportJob, client: MLClient) ->
 
         # Test cancel with submit to save test resource.
         # The job not supposed to succeed and usually failed quickly so status can be 'failed' as well
-        client.jobs.cancel(import_job.name)
+        cancel_poller = client.jobs.begin_cancel(import_job.name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
         import_job_3 = client.jobs.get(import_job.name)
         assert import_job_3.status in (JobStatus.CANCEL_REQUESTED, JobStatus.CANCELED, JobStatus.FAILED)
 
@@ -179,7 +182,9 @@ def validate_test_import_pipepine_submit_cancel(
                 == import_pipeline.jobs[import_step].outputs["output"]._data.path
             )
 
-        client.jobs.cancel(import_pipeline.name)
+        cancel_poller = client.jobs.begin_cancel(import_pipeline.name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
         import_pipeline_3 = client.jobs.get(import_pipeline.name)
         assert import_pipeline_3.status in (JobStatus.CANCEL_REQUESTED, JobStatus.CANCELED)