From 3b93b2cc5256f8f98fd98168f3062b663b6f8758 Mon Sep 17 00:00:00 2001
From: Ilya Matiach <ilmat@microsoft.com>
Date: Wed, 27 Nov 2024 10:34:21 -0500
Subject: [PATCH] Fix RAI image builds

---
 .../environments/responsibleai-tabular/asset.yaml  |  2 +-
 .../responsibleai-tabular/tests/requirements.txt   |  6 +++---
 .../tests/responsibleai_sample_test.py             | 12 ++++++++++--
 .../tests/test_conda_env.yaml                      | 14 ++++++++++++++
 .../environments/responsibleai-text/asset.yaml     |  4 ++--
 .../tests/responsibleai_text_sample_test.py        | 12 ++++++++++--
 .../responsibleai-text/tests/test_conda_env.yaml   | 14 ++++++++++++++
 .../environments/responsibleai-vision/asset.yaml   |  2 +-
 .../tests/responsibleai_vision_sample_test.py      | 12 ++++++++++--
 .../responsibleai-vision/tests/test_conda_env.yaml | 14 ++++++++++++++
 10 files changed, 79 insertions(+), 13 deletions(-)
 create mode 100644 assets/responsibleai/environments/responsibleai-tabular/tests/test_conda_env.yaml
 create mode 100644 assets/responsibleai/environments/responsibleai-text/tests/test_conda_env.yaml
 create mode 100644 assets/responsibleai/environments/responsibleai-vision/tests/test_conda_env.yaml

diff --git a/assets/responsibleai/environments/responsibleai-tabular/asset.yaml b/assets/responsibleai/environments/responsibleai-tabular/asset.yaml
index c00c4b91fe..4e269e64cb 100644
--- a/assets/responsibleai/environments/responsibleai-tabular/asset.yaml
+++ b/assets/responsibleai/environments/responsibleai-tabular/asset.yaml
@@ -6,6 +6,6 @@ extra_config: environment.yaml
 test:
   pytest:
     enabled: true 
-    pip_requirements: tests/requirements.txt
+    conda_environment: tests/test_conda_env.yaml
     tests_dir: tests
 categories: ["Responsible AI"]
\ No newline at end of file
diff --git a/assets/responsibleai/environments/responsibleai-tabular/tests/requirements.txt b/assets/responsibleai/environments/responsibleai-tabular/tests/requirements.txt
index ee8748efec..c12e7a9183 100644
--- a/assets/responsibleai/environments/responsibleai-tabular/tests/requirements.txt
+++ b/assets/responsibleai/environments/responsibleai-tabular/tests/requirements.txt
@@ -1,4 +1,4 @@
-azure-ai-ml~=1.0.0
-azure.identity==1.10.0
-azureml-core~=1.53.0
+azure-ai-ml==1.22.4
+azure.identity==1.19.0
+azureml-core>=1.58.0
 azureml-mlflow
diff --git a/assets/responsibleai/environments/responsibleai-tabular/tests/responsibleai_sample_test.py b/assets/responsibleai/environments/responsibleai-tabular/tests/responsibleai_sample_test.py
index 2950015eae..d6ca3aed86 100644
--- a/assets/responsibleai/environments/responsibleai-tabular/tests/responsibleai_sample_test.py
+++ b/assets/responsibleai/environments/responsibleai-tabular/tests/responsibleai_sample_test.py
@@ -11,7 +11,7 @@
 from azure.ai.ml.entities import Environment, BuildContext
 from azure.ai.ml import automl
 from azure.ai.ml.constants import AssetTypes
-from azure.identity import AzureCliCredential
+from azure.identity import AzureCliCredential, DefaultAzureCredential
 
 
 BUILD_CONTEXT = Path("../context")
@@ -53,8 +53,16 @@ def test_responsibleai():
     resource_group = os.environ.get("resource_group")
     workspace_name = os.environ.get("workspace")
 
+    try:
+        credential = AzureCliCredential()
+        credential.get_token("https://management.azure.com/.default")
+    except Exception as ex:
+        print(f"Unable to authenticate via Azure CLI:\n{ex}")
+        credential = DefaultAzureCredential()
+        credential.get_token("https://management.azure.com/.default")
+
     ml_client = MLClient(
-        AzureCliCredential(), subscription_id, resource_group, workspace_name
+        credential, subscription_id, resource_group, workspace_name
     )
 
     env_name = "responsibleai"
diff --git a/assets/responsibleai/environments/responsibleai-tabular/tests/test_conda_env.yaml b/assets/responsibleai/environments/responsibleai-tabular/tests/test_conda_env.yaml
new file mode 100644
index 0000000000..3552e3165e
--- /dev/null
+++ b/assets/responsibleai/environments/responsibleai-tabular/tests/test_conda_env.yaml
@@ -0,0 +1,14 @@
+channels:
+  - conda-forge
+  - defaults
+  - anaconda
+dependencies:
+  - python=3.12.0
+  - pip
+  - pip:
+    - -r requirements.txt
+    ## Test requirements
+    - pytest
+    - tqdm
+    - azure-core
+name: rai_test
diff --git a/assets/responsibleai/environments/responsibleai-text/asset.yaml b/assets/responsibleai/environments/responsibleai-text/asset.yaml
index 0ae9db86d7..d7c5fda6fc 100644
--- a/assets/responsibleai/environments/responsibleai-text/asset.yaml
+++ b/assets/responsibleai/environments/responsibleai-text/asset.yaml
@@ -5,7 +5,7 @@ spec: spec.yaml
 extra_config: environment.yaml
 test:
   pytest:
-    enabled: true 
-    pip_requirements: tests/requirements.txt
+    enabled: true
+    conda_environment: tests/test_conda_env.yaml
     tests_dir: tests
 categories: ["Responsible AI"]
\ No newline at end of file
diff --git a/assets/responsibleai/environments/responsibleai-text/tests/responsibleai_text_sample_test.py b/assets/responsibleai/environments/responsibleai-text/tests/responsibleai_text_sample_test.py
index 0750307f76..855be2710c 100644
--- a/assets/responsibleai/environments/responsibleai-text/tests/responsibleai_text_sample_test.py
+++ b/assets/responsibleai/environments/responsibleai-text/tests/responsibleai_text_sample_test.py
@@ -9,7 +9,7 @@
 from azure.ai.ml import command
 from azure.ai.ml.operations._run_history_constants import JobStatus
 from azure.ai.ml.entities import Environment, BuildContext
-from azure.identity import AzureCliCredential
+from azure.identity import AzureCliCredential, DefaultAzureCredential
 
 BUILD_CONTEXT = Path("../context")
 JOB_SOURCE_CODE = "src"
@@ -25,8 +25,16 @@ def test_responsibleai_text():
     resource_group = os.environ.get("resource_group")
     workspace_name = os.environ.get("workspace")
 
+    try:
+        credential = AzureCliCredential()
+        credential.get_token("https://management.azure.com/.default")
+    except Exception as ex:
+        print(f"Unable to authenticate via Azure CLI:\n{ex}")
+        credential = DefaultAzureCredential()
+        credential.get_token("https://management.azure.com/.default")
+
     ml_client = MLClient(
-        AzureCliCredential(), subscription_id, resource_group, workspace_name
+        credential, subscription_id, resource_group, workspace_name
     )
 
     env_name = "responsibleai-text"
diff --git a/assets/responsibleai/environments/responsibleai-text/tests/test_conda_env.yaml b/assets/responsibleai/environments/responsibleai-text/tests/test_conda_env.yaml
new file mode 100644
index 0000000000..3552e3165e
--- /dev/null
+++ b/assets/responsibleai/environments/responsibleai-text/tests/test_conda_env.yaml
@@ -0,0 +1,14 @@
+channels:
+  - conda-forge
+  - defaults
+  - anaconda
+dependencies:
+  - python=3.12.0
+  - pip
+  - pip:
+    - -r requirements.txt
+    ## Test requirements
+    - pytest
+    - tqdm
+    - azure-core
+name: rai_test
diff --git a/assets/responsibleai/environments/responsibleai-vision/asset.yaml b/assets/responsibleai/environments/responsibleai-vision/asset.yaml
index 9782ec7a65..c7ef0778bd 100644
--- a/assets/responsibleai/environments/responsibleai-vision/asset.yaml
+++ b/assets/responsibleai/environments/responsibleai-vision/asset.yaml
@@ -6,6 +6,6 @@ extra_config: environment.yaml
 test:
   pytest:
     enabled: true 
-    pip_requirements: tests/requirements.txt
+    conda_environment: tests/test_conda_env.yaml
     tests_dir: tests
 categories: ["Responsible AI"]
\ No newline at end of file
diff --git a/assets/responsibleai/environments/responsibleai-vision/tests/responsibleai_vision_sample_test.py b/assets/responsibleai/environments/responsibleai-vision/tests/responsibleai_vision_sample_test.py
index e299d37c08..3c6124fe3e 100644
--- a/assets/responsibleai/environments/responsibleai-vision/tests/responsibleai_vision_sample_test.py
+++ b/assets/responsibleai/environments/responsibleai-vision/tests/responsibleai_vision_sample_test.py
@@ -8,7 +8,7 @@
 from azure.ai.ml import MLClient
 from azure.ai.ml import command
 from azure.ai.ml.entities import Environment, BuildContext
-from azure.identity import AzureCliCredential
+from azure.identity import AzureCliCredential, DefaultAzureCredential
 
 BUILD_CONTEXT = Path("../context")
 JOB_SOURCE_CODE = "src"
@@ -24,8 +24,16 @@ def test_responsibleai_vision():
     resource_group = os.environ.get("resource_group")
     workspace_name = os.environ.get("workspace")
 
+    try:
+        credential = AzureCliCredential()
+        credential.get_token("https://management.azure.com/.default")
+    except Exception as ex:
+        print(f"Unable to authenticate via Azure CLI:\n{ex}")
+        credential = DefaultAzureCredential()
+        credential.get_token("https://management.azure.com/.default")
+
     ml_client = MLClient(
-        AzureCliCredential(), subscription_id, resource_group, workspace_name
+        credential, subscription_id, resource_group, workspace_name
     )
 
     env_name = "responsibleai-vision"
diff --git a/assets/responsibleai/environments/responsibleai-vision/tests/test_conda_env.yaml b/assets/responsibleai/environments/responsibleai-vision/tests/test_conda_env.yaml
new file mode 100644
index 0000000000..3552e3165e
--- /dev/null
+++ b/assets/responsibleai/environments/responsibleai-vision/tests/test_conda_env.yaml
@@ -0,0 +1,14 @@
+channels:
+  - conda-forge
+  - defaults
+  - anaconda
+dependencies:
+  - python=3.12.0
+  - pip
+  - pip:
+    - -r requirements.txt
+    ## Test requirements
+    - pytest
+    - tqdm
+    - azure-core
+name: rai_test