test: enable internal pipeline job e2e tests (Azure#27214)

* test: enable internal pipeline job e2e tests

* refactor: apply new cancel util to more tests

* refactor: update cancel logic for import job tests

Author: elliotzh

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

@@ -1,4 +1,3 @@
-import time
 import uuid
 from contextlib import contextmanager
 from pathlib import Path
@@ -10,11 +9,11 @@
 from azure.ai.ml import MLClient, load_batch_deployment, load_batch_endpoint, load_environment, load_model
 from azure.ai.ml._utils._arm_id_utils import AMLVersionedArmId
 from azure.ai.ml.constants._common import AssetTypes
-from azure.ai.ml.entities import BatchDeployment, BatchEndpoint, Job
+from azure.ai.ml.entities import BatchDeployment, BatchEndpoint
 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
+from azure.ai.ml.operations._run_history_constants import JobStatus
+
+from test_utilities.utils import wait_until_done
 
 
 @contextmanager
@@ -137,19 +136,6 @@ def test_batch_deployment_dependency_label_resolution(self, client: MLClient, ra
 
     @pytest.mark.e2etest
     def test_batch_job_download(self, client: MLClient, tmp_path: Path) -> str:
-        def wait_until_done(job: Job, timeout: int = None) -> None:
-            poll_start_time = time.time()
-            while job.status not in RunHistoryConstants.TERMINAL_STATUSES:
-                time.sleep(_wait_before_polling(time.time() - poll_start_time))
-                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
-                    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
-
         endpoint = load_batch_endpoint(
             "./tests/test_configs/endpoints/batch/batch_endpoint_mlflow_new.yaml",
             params_override=[{"name": "batch-ept-" + uuid.uuid4().hex[:15]}],
@@ -174,7 +160,7 @@ def wait_until_done(job: Job, timeout: int = None) -> None:
             # Adding a timeout value of 5 minutes to avoid future scenarios where wait
             # long periods for the test to finish. Instead if the job takes longer
             # than 5 minutes to finish we will make the test fail
-            job_status = wait_until_done(batchjob, 300)
+            job_status = wait_until_done(job=batchjob, client=client, timeout=300)
 
             # Check if the job was canceled due to a timeout.
             # If the job's status is CANCELED, we will make the test fail

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

@@ -1,10 +1,9 @@
-import time
 from pathlib import Path
 from time import sleep
 from typing import Callable
 
 from azure.ai.ml.entities import AmlTokenConfiguration
-from devtools_testutils import AzureRecordedTestCase, is_live, set_bodiless_matcher
+from devtools_testutils import AzureRecordedTestCase, is_live
 import jwt
 import pytest
 
@@ -18,19 +17,17 @@
 from azure.ai.ml.entities._job.distribution import MpiDistribution
 from azure.ai.ml.entities._job.job import Job
 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
 
+from test_utilities.utils import wait_until_done
+
 # 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
 # here and in the script.
 TEST_PARAMS = {"a_param": "1", "another_param": "2"}
 
-
-@pytest.mark.fixture(autouse=True)
-def bodiless_matching(test_proxy):
-    set_bodiless_matcher()
+# previous bodiless_matcher fixture doesn't take effect because of typo, please add it in method level if needed
 
 
 @pytest.mark.timeout(600)
@@ -345,19 +342,13 @@ def get_job_list():
     def test_command_job_download(self, tmp_path: Path, randstr: Callable[[], str], client: MLClient) -> None:
         client: MLClient = client
 
-        def wait_until_done(job: Job) -> None:
-            poll_start_time = time.time()
-            while job.status not in RunHistoryConstants.TERMINAL_STATUSES:
-                time.sleep(_wait_before_polling(time.time() - poll_start_time))
-                job = client.jobs.get(job.name)
-
         job = client.jobs.create_or_update(
             load_job(
                 source="./tests/test_configs/command_job/command_job_quick_with_output.yml",
                 params_override=[{"name": randstr("name")}],
             )
         )
-        wait_until_done(job)
+        wait_until_done(client=client, job=job)
 
         client.jobs.download(name=job.name, download_path=tmp_path, all=True)
 
@@ -374,20 +365,14 @@ def wait_until_done(job: Job) -> None:
     def test_command_job_local_run_download(self, tmp_path: Path, randstr: Callable[[], str], client: MLClient) -> None:
         client: MLClient = client
 
-        def wait_until_done(job: Job) -> None:
-            poll_start_time = time.time()
-            while job.status not in RunHistoryConstants.TERMINAL_STATUSES:
-                time.sleep(_wait_before_polling(time.time() - poll_start_time))
-                job = client.jobs.get(job.name)
-
         job = client.jobs.create_or_update(
             load_job(
                 source="./tests/test_configs/command_job/command_job_quick_with_output.yml",
                 params_override=[{"name": randstr("name")}, {"compute": LOCAL_COMPUTE_TARGET}],
             )
         )
 
-        wait_until_done(job)
+        wait_until_done(client=client, job=job)
 
         client.jobs.download(name=job.name, download_path=tmp_path, all=True)
 

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

@@ -9,7 +9,7 @@
 import pytest
 
 from azure.ai.ml import MLClient, MpiDistribution, load_component, load_environment
-from azure.ai.ml._restclient.v2022_05_01.models import ComponentContainerData, ListViewType
+from azure.ai.ml._restclient.v2022_05_01.models import ListViewType
 from azure.ai.ml._utils._arm_id_utils import is_ARM_id_for_resource
 from azure.ai.ml.constants._common import (
     ANONYMOUS_COMPONENT_NAME,
@@ -22,8 +22,8 @@
 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 test_utilities.utils import assert_job_cancel
 from .._util import _COMPONENT_TIMEOUT_SECOND
 from ..unittests.test_component_schema import load_component_entity_from_rest_json
 
@@ -84,10 +84,7 @@ def create_tensorflow_distribution(has_strs: bool = False):
 
     return create_tensorflow_distribution
 
-
-@pytest.mark.fixture(autouse=True)
-def bodiless_matching(test_proxy):
-    set_bodiless_matcher()
+# previous bodiless_matcher fixture doesn't take effect because of typo, please add it in method level if needed
 
 
 def assert_component_basic_workflow(
@@ -839,13 +836,7 @@ def test_create_pipeline_component_from_job(self, client: MLClient, randstr: Cal
             "./tests/test_configs/dsl_pipeline/pipeline_with_pipeline_component/pipeline.yml",
             params_override=params_override,
         )
-        job = client.jobs.create_or_update(pipeline_job)
-        try:
-            cancel_poller = client.jobs.begin_cancel(job.name)
-            assert isinstance(cancel_poller, LROPoller)
-            assert cancel_poller.result() is None
-        except Exception:
-            pass
+        job = assert_job_cancel(pipeline_job, client)
         name = randstr("component_name_1")
         component = PipelineComponent(name=name, source_job_id=job.id)
         rest_component = client.components.create_or_update(component)

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

@@ -30,11 +30,11 @@
     add_remove_header_sanitizer,
     is_live,
     set_custom_default_matcher,
-    test_proxy,
+    set_bodiless_matcher,
 )
 from devtools_testutils.fake_credentials import FakeTokenCredential
 from devtools_testutils.helpers import is_live_and_not_recording
-from devtools_testutils.proxy_fixtures import VariableRecorder, variable_recorder
+from devtools_testutils.proxy_fixtures import VariableRecorder
 from pytest_mock import MockFixture
 
 from test_utilities.constants import Test_Registry_Name, Test_Resource_Group, Test_Subscription, Test_Workspace_Name
@@ -555,3 +555,22 @@ def enable_internal_components():
     with environment_variable_overwrite(AZUREML_INTERNAL_COMPONENTS_ENV_VAR, "True"):
         # need to call _try_init_internal_components manually as environment variable is set after _internal is imported
         try_enable_internal_components()
+
+
+@pytest.fixture()
+def bodiless_matching(test_proxy):
+    set_bodiless_matcher()
+
+
+def pytest_configure(config):
+    # register customized pytest markers
+    for marker, description in [
+        ("e2etest", "marks tests as end to end tests, which involve requests to the server"),
+        ("unittest", "marks tests as unit tests, which do not involve requests to the server"),
+        ("pipeline_test", "marks tests as pipeline tests, which will create pipeline jobs during testing"),
+        ("automl_test", "marks tests as automl tests, which will create automl jobs during testing"),
+        ("production_experience_test", "marks tests as production experience tests"),
+    ]:
+        config.addinivalue_line("markers", f"{marker}: {description}")
+
+    config.addinivalue_line("markers", f"{marker}: {description}")

sdk/ml/azure-ai-ml/tests/dataset/e2etests/test_data.py

@@ -11,12 +11,9 @@
 from azure.core.paging import ItemPaged
 
 
-from devtools_testutils import AzureRecordedTestCase, set_bodiless_matcher
+from devtools_testutils import AzureRecordedTestCase
 
-
-@pytest.mark.fixture(autouse=True)
-def bodiless_matching(test_proxy):
-    set_bodiless_matcher()
+# previous bodiless_matcher fixture doesn't take effect because of typo, please add it in method level if needed
 
 
 @pytest.mark.e2etest

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

@@ -29,11 +29,9 @@
 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.exceptions import HttpResponseError
-from azure.core.polling import LROPoller
 from devtools_testutils import AzureRecordedTestCase
 from pipeline_job.e2etests.test_pipeline_job import assert_job_input_output_types
-from test_utilities.utils import _PYTEST_TIMEOUT_METHOD, omit_with_wildcard
+from test_utilities.utils import _PYTEST_TIMEOUT_METHOD, omit_with_wildcard, assert_job_cancel
 
 from .._util import _DSL_TIMEOUT_SECOND
 
@@ -57,18 +55,6 @@
 ]
 
 
-def assert_job_cancel(pipeline, client: MLClient, experiment_name=None):
-    job = client.jobs.create_or_update(pipeline, experiment_name=experiment_name)
-    try:
-        cancel_poller = client.jobs.begin_cancel(job.name)
-        assert isinstance(cancel_poller, LROPoller)
-        # skip wait for cancel result to reduce test run duration.
-        # assert cancel_poller.result() is None
-    except HttpResponseError:
-        pass
-    return job
-
-
 @pytest.mark.usefixtures(
     "enable_environment_id_arm_expansion",
     "enable_pipeline_private_preview_features",

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

@@ -1,30 +1,17 @@
 import pytest
 from pathlib import Path
-from test_utilities.utils import _PYTEST_TIMEOUT_METHOD
+from test_utilities.utils import _PYTEST_TIMEOUT_METHOD, assert_job_cancel
 
 from azure.ai.ml import MLClient, load_component, Input, load_model
 from azure.ai.ml.dsl import pipeline
 from azure.ai.ml.constants import AssetTypes
 from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
-from azure.core.polling import LROPoller
 
 from .._util import _DSL_TIMEOUT_SECOND
 
 from devtools_testutils import AzureRecordedTestCase
 
 
-def assert_job_cancel(pipeline, client: MLClient, experiment_name=None):
-    job = client.jobs.create_or_update(pipeline, experiment_name=experiment_name)
-    try:
-        cancel_poller = client.jobs.begin_cancel(job.name)
-        assert isinstance(cancel_poller, LROPoller)
-        # skip wait for cancel result to reduce test run duration.
-        # assert cancel_poller.result() is None
-    except HttpResponseError:
-        pass
-    return job
-
-
 @pytest.mark.usefixtures("enable_pipeline_private_preview_features", "recorded_test")
 @pytest.mark.timeout(timeout=_DSL_TIMEOUT_SECOND, method=_PYTEST_TIMEOUT_METHOD)
 @pytest.mark.e2etest