Fix random windows CI failure

examples/test/test_bert_ptq_cpu_aml.py

@@ -6,7 +6,7 @@
 from pathlib import Path
 
 import pytest
-from utils import check_output, patch_config
+from utils import check_output, patch_config, set_azure_identity_logging
 
 
 @pytest.fixture(scope="module", autouse=True)
@@ -15,6 +15,7 @@ def setup():
     cur_dir = Path(__file__).resolve().parent.parent
     example_dir = cur_dir / "bert"
     os.chdir(example_dir)
+    set_azure_identity_logging()
     yield
     os.chdir(cur_dir)
 

examples/test/test_resnet_ptq_cpu.py

@@ -8,7 +8,7 @@
 import pytest
 from onnxruntime import __version__ as OrtVersion
 from packaging import version
-from utils import check_output, patch_config
+from utils import check_output, patch_config, set_azure_identity_logging
 
 from olive.common.utils import retry_func, run_subprocess
 
@@ -23,14 +23,15 @@ def setup():
     # prepare model and data
     # retry since it fails randomly
     retry_func(run_subprocess, kwargs={"cmd": "python prepare_model_data.py", "check": True})
+    set_azure_identity_logging()
 
     yield
     os.chdir(cur_dir)
 
 
 @pytest.mark.parametrize("search_algorithm", ["random"])
 @pytest.mark.parametrize("execution_order", ["pass-by-pass"])
-@pytest.mark.parametrize("system", ["local_system", "aml_system"])
+@pytest.mark.parametrize("system", ["aml_system", "local_system"])
 @pytest.mark.parametrize("olive_json", ["resnet_ptq_cpu.json", "resnet_ptq_cpu_aml_dataset.json"])
 @pytest.mark.skipif(
     version.parse(OrtVersion) == version.parse("1.16.0"),

examples/test/utils.py

@@ -3,7 +3,9 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 import json
+import logging
 import os
+import sys
 
 # pylint: disable=broad-exception-raised
 
@@ -149,3 +151,13 @@
     with open(download_path, "wb") as my_blob:
         blob_data = blob.download_blob()
         blob_data.readinto(my_blob)
+
+
+def set_azure_identity_logging():
+    from azure.identity import DefaultAzureCredential  # noqa: F401
+
+    identity_logger = logging.getLogger("azure.identity")
+    identity_logger.setLevel(logging.DEBUG)
+    if not (identity_logger.handlers and isinstance(identity_logger.handlers[0], logging.StreamHandler)):
+        handler = logging.StreamHandler(stream=sys.stdout)
+        identity_logger.addHandler(handler)

olive/azureml/azureml_client.py

@@ -131,7 +131,6 @@ def _get_credentials(self):
             credential = DefaultAzureCredential(**default_auth_params)
             # Check if given credential can get token successfully.
             credential.get_token("https://management.azure.com/.default")
-            logger.debug("Using DefaultAzureCredential")
         except Exception:
             logger.warning("Using InteractiveBrowserCredential since of default credential errors", exc_info=True)
             # Fall back to InteractiveBrowserCredential in case DefaultAzureCredential not work