Repository: microsoft/hi-ml

Issue:
Defaulting to Identity-Based data access breaks cross-region compute jobs
It seems that if identity-based access is not available, the v2 API does not fall back to using account keys / SAS tokens. As such, when the submission fails to cross-region compute because of identity-based access not working, then the whole job fails instead of defaulting to account keys / SAS tokens.

A flag should be added to make the identity-based access an opt-in feature, instead defaulting to account keys / SAS tokens still.


Discussion:

Resolving patch:
diff --git a/hi-ml-azure/src/health_azure/himl.py b/hi-ml-azure/src/health_azure/himl.py
--- a/hi-ml-azure/src/health_azure/himl.py
+++ b/hi-ml-azure/src/health_azure/himl.py
@@ -425,6 +425,7 @@ def submit_run_v2(workspace: Optional[Workspace],
                   tags: Optional[Dict[str, str]] = None,
                   docker_shm_size: str = "",
                   wait_for_completion: bool = False,
+                  identity_based_auth: bool = False,
                   workspace_config_path: Optional[PathOrString] = None,
                   ml_client: Optional[MLClient] = None,
                   hyperparam_args: Optional[Dict[str, Any]] = None,
@@ -516,7 +517,7 @@ def create_command_job(cmd: str) -> Command:
             display_name=display_name,
             instance_count=num_nodes,
             distribution=distribution,
-            identity=UserIdentityConfiguration(),
+            identity=UserIdentityConfiguration() if identity_based_auth else None,
         )
 
     if hyperparam_args:
@@ -748,6 +749,7 @@ def submit_to_azure_if_needed(  # type: ignore
         hyperdrive_config: Optional[HyperDriveConfig] = None,
         hyperparam_args: Optional[Dict[str, Any]] = None,
         strictly_aml_v1: bool = False,
+        identity_based_auth: bool = False,
         pytorch_processes_per_node_v2: Optional[int] = None,
         display_name: Optional[str] = None,
 ) -> AzureRunInfo:  # pragma: no cover
@@ -963,6 +965,7 @@ def submit_to_azure_if_needed(  # type: ignore
                                 display_name=display_name,
                                 docker_shm_size=docker_shm_size,
                                 wait_for_completion=wait_for_completion,
+                                identity_based_auth=identity_based_auth,
                                 hyperparam_args=hyperparam_args,
                                 num_nodes=num_nodes,
                                 pytorch_processes_per_node=pytorch_processes_per_node_v2,
diff --git a/hi-ml/src/health_ml/experiment_config.py b/hi-ml/src/health_ml/experiment_config.py
--- a/hi-ml/src/health_ml/experiment_config.py
+++ b/hi-ml/src/health_ml/experiment_config.py
@@ -46,6 +46,11 @@ class ExperimentConfig(param.Parameterized):
                                                          "should only be used when debugging issues")
     strictly_aml_v1: bool = param.Boolean(default=False, doc="If True, use AzureML v1 SDK. If False (default), use "
                                                              "the v2 of the SDK")
+    identity_based_auth: bool = param.Boolean(
+        default=False,
+        doc="If True, use identity based authentication to access blob storage data via datastores. If False (default),"
+            " then use SAS tokens / account keys if available. Only supported when using v2 sdk."
+    )
     workspace_config_path: Optional[Path] = \
         param.ClassSelector(class_=Path, default=None, allow_None=True,
                             doc="The path to the AzureML workspace configuration file. If not specified, the "
diff --git a/hi-ml/src/health_ml/runner.py b/hi-ml/src/health_ml/runner.py
--- a/hi-ml/src/health_ml/runner.py
+++ b/hi-ml/src/health_ml/runner.py
@@ -257,6 +257,7 @@ def submit_to_azureml_if_needed(self) -> AzureRunInfo:
                 display_name=self.lightning_container.tag,
                 tags=self.additional_run_tags(script_params),
                 strictly_aml_v1=self.experiment_config.strictly_aml_v1,
+                identity_based_auth=self.experiment_config.identity_based_auth,
             )
         else:
             azure_run_info = submit_to_azure_if_needed(


Test patch:
diff --git a/hi-ml-azure/testazure/testazure/test_data/make_tests.py b/hi-ml-azure/testazure/testazure/test_data/make_tests.py
--- a/hi-ml-azure/testazure/testazure/test_data/make_tests.py
+++ b/hi-ml-azure/testazure/testazure/test_data/make_tests.py
@@ -106,6 +106,7 @@ def render_test_script(entry_script_path: Path, extra_options: Dict[str, str],
     default_options['input_datasets'] = 'None'
     default_options['output_datasets'] = 'None'
     default_options['wait_for_completion'] = 'True'
+    default_options['identity_based_auth'] = 'False'
     default_options['wait_for_completion_show_output'] = 'True'
     default_options['args'] = ''
     default_options['body'] = ''
diff --git a/hi-ml-azure/testazure/testazure/test_data/simple/hello_world_template.txt b/hi-ml-azure/testazure/testazure/test_data/simple/hello_world_template.txt
--- a/hi-ml-azure/testazure/testazure/test_data/simple/hello_world_template.txt
+++ b/hi-ml-azure/testazure/testazure/test_data/simple/hello_world_template.txt
@@ -50,6 +50,7 @@ def main() -> None:
         input_datasets={{ input_datasets }},
         output_datasets={{ output_datasets }},
         wait_for_completion={{ wait_for_completion }},
+        identity_based_auth={{ identity_based_auth }},
         wait_for_completion_show_output={{ wait_for_completion_show_output }},
         tags={{ tags }},
         display_name='{{ display_name }}',
diff --git a/hi-ml-azure/testazure/testazure/test_himl.py b/hi-ml-azure/testazure/testazure/test_himl.py
--- a/hi-ml-azure/testazure/testazure/test_himl.py
+++ b/hi-ml-azure/testazure/testazure/test_himl.py
@@ -25,7 +25,7 @@
 from _pytest.capture import CaptureFixture
 from azure.ai.ml import Input, Output, MLClient
 from azure.ai.ml.constants import AssetTypes, InputOutputModes
-from azure.ai.ml.entities import Data, Job, UserIdentityConfiguration
+from azure.ai.ml.entities import Data, Job
 from azure.ai.ml.entities._job.distribution import MpiDistribution, PyTorchDistribution
 from azure.ai.ml.sweep import Choice
 from azure.core.exceptions import ResourceNotFoundError
@@ -817,7 +817,7 @@ def _mock_sweep(*args: Any, **kwargs: Any) -> MagicMock:
                 display_name=dummy_display_name,
                 distribution=MpiDistribution(process_count_per_instance=1),
                 instance_count=1,
-                identity=UserIdentityConfiguration(),
+                identity=None,
             )
 
             # job with hyperparameter sampling:
@@ -879,7 +879,7 @@ def _mock_sweep(*args: Any, **kwargs: Any) -> MagicMock:
                 display_name=dummy_display_name,
                 distribution=MpiDistribution(process_count_per_instance=1),
                 instance_count=1,
-                identity=UserIdentityConfiguration(),
+                identity=None,
             )
 
             mock_command.assert_any_call(**param_sampling)
@@ -1167,7 +1167,7 @@ def test_invoking_hello_world_using_azureml_flag(tmp_path: Path) -> None:
     """
     Test that invoking hello_world.py with the --azureml flag will submit to AzureML and not run locally.
 
-    :param tmp_path: PyTest test fixture for temporary path.
+    :param tmp_path: Pytest fixture for temporary path.
     """
 
     message_guid = uuid4().hex
@@ -1511,6 +1511,7 @@ def test_invoking_user_identity_datasets(tmp_path: Path) -> None:
         'input_datasets': f"['{USER_IDENTITY_TEST_ASSET}']",
         'output_datasets': f"['{USER_IDENTITY_TEST_ASSET_OUTPUT}']",
         'strictly_aml_v1': str(False),
+        'identity_based_auth': str(True),
         'body': f"""
 
     input_folder = run_info.input_datasets[0]