Repository: microsoft/hi-ml

Issue:
Enable using the submitting user's identity when accessing data
Most of our data stores have the access key stored with them, meaning that all data can be accessed by anyone who has access to the AML workspace.

This [docu page](https://learn.microsoft.com/en-us/azure/machine-learning/v1/how-to-identity-based-data-access) lists all the steps to create a datastore without stored credentials, and also says how to create a dataset.

To submit a job that relies on user identity, use the equivalent of 
```
job = command(
     code="<my-local-code-location>",
     command="python <my-script>.py --input_data ${{inputs.input_data}}",
     inputs=my_job_inputs,
     environment="AzureML-sklearn-0.24-ubuntu18.04-py37-cpu:9",
     compute="<my-compute-cluster-name>",
     identity= UserIdentity()
)
```

Can we make our code flexible enough to check if a datastore is authenticated (then use the stored credentials), and otherwise user identity?



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
@@ -20,8 +20,8 @@
 from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
 
 from azure.ai.ml import MLClient, Input, Output, command
-from azure.ai.ml.constants import AssetTypes, InputOutputModes
-from azure.ai.ml.entities import Data, Job, Command, Sweep
+from azure.ai.ml.constants import InputOutputModes
+from azure.ai.ml.entities import Data, Job, Command, Sweep, UserIdentityConfiguration
 from azure.ai.ml.entities import Environment as EnvironmentV2
 from azure.ai.ml.entities._job.distribution import MpiDistribution, PyTorchDistribution
 
@@ -41,8 +41,7 @@
                                 is_run_and_child_runs_completed, is_running_in_azure_ml, register_environment,
                                 run_duration_string_to_seconds, to_azure_friendly_string, RUN_CONTEXT, get_workspace,
                                 PathOrString, DEFAULT_ENVIRONMENT_VARIABLES, get_ml_client,
-                                create_python_environment_v2, register_environment_v2, V2_INPUT_DATASET_PATTERN,
-                                V2_OUTPUT_DATASET_PATTERN, wait_for_job_completion)
+                                create_python_environment_v2, register_environment_v2, wait_for_job_completion)
 from health_azure.datasets import (DatasetConfig, StrOrDatasetConfig, setup_local_datasets,
                                    _input_dataset_key, _output_dataset_key, _replace_string_datasets,
                                    _get_or_create_v2_data_asset)
@@ -70,6 +69,9 @@
 SAMPLING_ALGORITHM_ARG = "sampling_algorithm"
 GOAL_ARG = "goal"
 
+V2_INPUT_PATTERN = "INPUT_"
+V2_OUTPUT_PATTERN = "OUTPUT_"
+
 
 @dataclass
 class AzureRunInfo:
@@ -386,36 +388,6 @@ def create_script_run(
         arguments=script_params)
 
 
-def _generate_input_dataset_command(input_datasets_v2: Dict[str, Input]) -> str:
-    """
-    Generate command line arguments to pass AML v2 data assets into a script
-
-    :param input_datasets_v2: A dictionary of Input objects that have been passed into the AML command
-    :return: A string representing the input datasets that the script should expect
-    """
-    input_cmd = ""
-    for i, (input_data_name, input_dataset_v2) in enumerate(input_datasets_v2.items()):
-        input_name = f"INPUT_{i}"
-        input_str = "${{inputs." + f"{input_name}" + "}}"
-        input_cmd += f" --{input_name}={input_str}"
-    return input_cmd
-
-
-def _generate_output_dataset_command(output_datasets_v2: Dict[str, Output]) -> str:
-    """
-    Generate command line arguments to pass AML v2 outputs into a script
-
-    :param output_datasets_v2: A dictionary of Output objects that have been passed into the AML command
-    :return: A string representing the output values that the script should expect
-    """
-    output_cmd = ""
-    for i, (output_data_name, output_dataset_v2) in enumerate(output_datasets_v2.items()):
-        output_name = f"OUTPUT_{i}"
-        output_str = "${{outputs." + f"{output_name}" + "}}"
-        output_cmd += f" --{output_name}={output_str}"
-    return output_cmd
-
-
 def effective_experiment_name(experiment_name: Optional[str],
                               entry_script: Optional[PathOrString] = None) -> str:
     """Choose the experiment name to use for the run. If provided in the environment variable HIML_EXPERIMENT_NAME,
@@ -512,16 +484,6 @@ def submit_run_v2(workspace: Optional[Workspace],
     script_params = script_params or []
     cmd = " ".join(["python", str(entry_script), *script_params])
 
-    if input_datasets_v2:
-        cmd += _generate_input_dataset_command(input_datasets_v2)
-    else:
-        input_datasets_v2 = {}
-
-    if output_datasets_v2:
-        cmd += _generate_output_dataset_command(output_datasets_v2)
-    else:
-        output_datasets_v2 = {}
-
     job_to_submit: Union[Command, Sweep]
 
     # number of nodes and processes per node cannot be less than one
@@ -553,11 +515,15 @@ def create_command_job(cmd: str) -> Command:
             display_name=display_name,
             instance_count=num_nodes,
             distribution=distribution,
+            identity=UserIdentityConfiguration(),
         )
 
     if hyperparam_args:
         param_sampling = hyperparam_args[PARAM_SAMPLING_ARG]
 
+        if input_datasets_v2 is None:
+            input_datasets_v2 = {}
+
         for sample_param, choices in param_sampling.items():
             input_datasets_v2[sample_param] = choices.values[0]
             cmd += f" --{sample_param}=" + "${{inputs." + sample_param + "}}"
@@ -685,6 +651,33 @@ def _str_to_path(s: Optional[PathOrString]) -> Optional[Path]:
     return s
 
 
+def get_data_asset_from_config(ml_client: MLClient, dataset_config: DatasetConfig) -> Data:
+    """Given a list of dataset configs, generates and returns a list of data assets.
+
+    :param ml_client: An MLClient object.
+    :param dataset_list: The list of datasets to create data assets for.
+    :raises ValueError: Raised if a data asset has no path.
+    :return: A list of data assets.
+    """
+
+    version = dataset_config.version
+    logging.info(
+        f"Trying to access data asset {dataset_config.name} version {version}, datastore {dataset_config.datastore}"
+    )
+
+    # if version is None, this function gets the latest version
+    data_asset: Data = _get_or_create_v2_data_asset(
+        ml_client,
+        dataset_config.datastore,
+        dataset_config.name,
+        version=str(version) if version else None,
+    )
+    if not data_asset.path:
+        raise ValueError(f"Data asset {data_asset.id} has no path.")
+
+    return data_asset
+
+
 def create_v2_inputs(ml_client: MLClient, input_datasets: List[DatasetConfig]) -> Dict[str, Input]:
     """
     Create a dictionary of Azure ML v2 Input objects, required for passing input data in to an AML job
@@ -693,59 +686,35 @@ def create_v2_inputs(ml_client: MLClient, input_datasets: List[DatasetConfig]) -
     :param input_datasets: A list of DatasetConfigs to convert to Inputs.
     :return: A dictionary in the format "input_name": Input.
     """
-    inputs: Dict[str, Input] = {}
-    for i, input_dataset in enumerate(input_datasets):
-        input_name = f"INPUT_{i}"
-        version = input_dataset.version
-        # if version is None, this function gets the latest version
-        logging.info(
-            f"Trying to access data asset {input_dataset.name} version {version}, datastore {input_dataset.datastore}"
-        )
-        data_asset: Data = _get_or_create_v2_data_asset(
-            ml_client,
-            input_dataset.datastore,
-            input_dataset.name,
-            version=str(version) if version else None,
-        )
-        if not data_asset.path:
-            raise ValueError(f"Data asset {data_asset.id} has no path.")
-        # Some mismatches with the documentation here:
-        # data_path = data_asset.id: This works in some jobs, but in other gives a
-        # DataAccessError(InvalidInput { message: "invalid uri format", source: None }))
-        # Unclear what the difference is between failing and successful jobs
-        # Alternative: data_path = f"azureml:{data_asset.name}:{version}"
-        # This does not work at all, neither with v1 nor v2 data assets. In both cases, we get
-        # InvalidInput { message: "invalid uri format"
-        data_path = data_asset.path
-
-        inputs[input_name] = Input(
-            # Data assets can be of type "uri_folder", "uri_file", "mltable", all of which are value types in Input
+    input_assets = [get_data_asset_from_config(ml_client, input_dataset) for input_dataset in input_datasets]
+    # Data assets can be of type "uri_folder", "uri_file", "mltable", all of which are value types in Input
+    return {
+        f"{V2_INPUT_PATTERN}{i}": Input(  # type: ignore
             type=data_asset.type,  # type: ignore
-            path=data_path,
-            mode=InputOutputModes.MOUNT if input_dataset.use_mounting else InputOutputModes.DOWNLOAD
-        )
-    return inputs
+            path=data_asset.path,
+            mode=InputOutputModes.MOUNT if input_datasets[i].use_mounting else InputOutputModes.DOWNLOAD
+        ) for i, data_asset in enumerate(input_assets)
+    }
 
 
-def create_v2_outputs(output_datasets: List[DatasetConfig]) -> Dict[str, Output]:
+def create_v2_outputs(ml_client: MLClient, output_datasets: List[DatasetConfig]) -> Dict[str, Output]:
     """
     Create a dictionary of Azure ML v2 Output objects, required for passing output data in to an AML job
 
+    :ml_client: An MLClient object.
     :param output_datasets: A list of DatasetConfigs to convert to Outputs.
     :return: A dictionary in the format "output_name": Output.
     """
-    outputs = {}
-    for i, output_dataset in enumerate(output_datasets):
-        output_name = f"OUTPUT_{i}"
-        v1_datastore_path = f"azureml://datastores/{output_dataset.datastore}/paths/{output_dataset.name}"
-        # Note that there are alternative formats that the output path can take, such as:
-        # v2_data_asset_path = f"azureml:{output_dataset.name}@latest"
-        outputs[output_name] = Output(  # type: ignore
-            type=AssetTypes.URI_FOLDER,
-            path=v1_datastore_path,
-            mode=InputOutputModes.DIRECT,
-        )
-    return outputs
+
+    output_assets = [get_data_asset_from_config(ml_client, output_dataset) for output_dataset in output_datasets]
+    return {
+        # Data assets can be of type "uri_folder", "uri_file", "mltable", all of which are value types in Input
+        f"{V2_OUTPUT_PATTERN}{i}": Output(  # type: ignore
+            type=data_asset.type,  # type: ignore
+            path=data_asset.path,
+            mode=InputOutputModes.MOUNT,  # hard-coded to mount for now, as this is the only mode that doesn't break
+        ) for i, data_asset in enumerate(output_assets)
+    }
 
 
 def submit_to_azure_if_needed(  # type: ignore
@@ -978,7 +947,7 @@ def submit_to_azure_if_needed(  # type: ignore
             ml_client = get_ml_client(ml_client=ml_client, aml_workspace=workspace)
             registered_env = register_environment_v2(environment, ml_client)
             input_datasets_v2 = create_v2_inputs(ml_client, cleaned_input_datasets)
-            output_datasets_v2 = create_v2_outputs(cleaned_output_datasets)
+            output_datasets_v2 = create_v2_outputs(ml_client, cleaned_output_datasets)
 
             job = submit_run_v2(workspace=workspace,
                                 input_datasets_v2=input_datasets_v2,
@@ -997,6 +966,7 @@ def submit_to_azure_if_needed(  # type: ignore
                                 num_nodes=num_nodes,
                                 pytorch_processes_per_node=pytorch_processes_per_node_v2,
                                 )
+
             if after_submission is not None:
                 after_submission(job, ml_client)  # type: ignore
 
@@ -1108,21 +1078,23 @@ def _get_dataset_names_from_string(sys_arg: str, pattern: str) -> Path:
     return dataset_path
 
 
-def _extract_v2_inputs_outputs_from_args() -> Tuple[List[Path], List[Path]]:
-    """
-    Extract all command line arguments of the format INPUT_i=path_to_input or OUTPUT_i=path_to_output (where i is any
-    integer) and return a list of the Paths for each.
+def _extract_v2_inputs_outputs_from_env_vars() -> Tuple[List[Path], List[Path]]:
+    """Provides paths to the input and output datasets for v2 jobs by extracting them from the environment variables.
 
     :return: A list of Input paths and a list of Output paths
     """
     returned_input_datasets: List[Path] = []
     returned_output_datasets: List[Path] = []
 
-    for sys_arg in sys.argv:
-        if re.match(V2_INPUT_DATASET_PATTERN, sys_arg):
-            returned_input_datasets += [_get_dataset_names_from_string(sys_arg, V2_INPUT_DATASET_PATTERN)]
-        if re.match(V2_OUTPUT_DATASET_PATTERN, sys_arg):
-            returned_output_datasets += [_get_dataset_names_from_string(sys_arg, V2_OUTPUT_DATASET_PATTERN)]
+    input_pattern_string = "AZURE_ML_INPUT_" + V2_INPUT_PATTERN + r"\d+"
+    output_pattern_string = r"AZURE_ML_OUTPUT_" + V2_OUTPUT_PATTERN + r"\d+"
+
+    for env_var in os.environ:  # input and output env vars set by V2 SDK on job submission
+        if re.match(input_pattern_string, env_var):
+            returned_input_datasets.append(Path(os.environ[env_var]))
+        elif re.match(output_pattern_string, env_var):
+            returned_output_datasets.append(Path(os.environ[env_var]))
+
     return returned_input_datasets, returned_output_datasets
 
 
@@ -1136,7 +1108,7 @@ def _generate_v2_azure_datasets(cleaned_input_datasets: List[DatasetConfig],
     :param cleaned_output_datasets: The list of output dataset configs
     :return: The AzureRunInfo containing the AzureML input and output dataset lists etc.
     """
-    returned_input_datasets, returned_output_datasets = _extract_v2_inputs_outputs_from_args()
+    returned_input_datasets, returned_output_datasets = _extract_v2_inputs_outputs_from_env_vars()
 
     return AzureRunInfo(
         input_datasets=returned_input_datasets,  # type: ignore
diff --git a/hi-ml-azure/testazure/testazure/utils_testazure.py b/hi-ml-azure/testazure/testazure/utils_testazure.py
--- a/hi-ml-azure/testazure/testazure/utils_testazure.py
+++ b/hi-ml-azure/testazure/testazure/utils_testazure.py
@@ -19,6 +19,16 @@
 from health_azure.utils import get_ml_client, get_workspace
 
 DEFAULT_DATASTORE = "himldatasets"
+
+TEST_DATASET_NAME = "test_dataset"
+TEST_DATA_ASSET_NAME = "test_dataset"
+TEST_INVALID_DATA_ASSET_NAME = "non_existent_dataset"
+TEST_DATASTORE_NAME = "test_datastore"
+
+USER_IDENTITY_TEST_DATASTORE = "test_identity_based_datastore"
+USER_IDENTITY_TEST_ASSET = "test_identity_based_data_asset"
+USER_IDENTITY_TEST_ASSET_OUTPUT = "test_identity_based_data_asset_output"
+USER_IDENTITY_TEST_FILE = "test_identity_based_file.txt"
 FALLBACK_SINGLE_RUN = "refs_pull_545_merge:refs_pull_545_merge_1626538212_d2b07afd"
 
 # List of root folders to add to .amlignore


Test patch:
diff --git a/hi-ml-azure/testazure/testazure/test_datasets.py b/hi-ml-azure/testazure/testazure/test_datasets.py
--- a/hi-ml-azure/testazure/testazure/test_datasets.py
+++ b/hi-ml-azure/testazure/testazure/test_datasets.py
@@ -20,7 +20,10 @@
 from azureml.data.azure_storage_datastore import AzureBlobDatastore
 from azureml.data.dataset_consumption_config import DatasetConsumptionConfig
 from azureml.exceptions._azureml_exception import UserErrorException
-from testazure.utils_testazure import DEFAULT_DATASTORE, DEFAULT_WORKSPACE, get_test_ml_client
+from testazure.utils_testazure import (
+    DEFAULT_DATASTORE, DEFAULT_WORKSPACE, TEST_DATASET_NAME, TEST_DATA_ASSET_NAME, TEST_INVALID_DATA_ASSET_NAME,
+    TEST_DATASTORE_NAME, get_test_ml_client,
+)
 
 from health_azure.datasets import (
     DatasetConfig, _create_v1_dataset, _create_v2_data_asset, _get_or_create_v1_dataset, _get_or_create_v2_data_asset,
@@ -31,10 +34,6 @@
 
 
 TEST_ML_CLIENT = get_test_ml_client()
-TEST_DATASET_NAME = "test_dataset"
-TEST_DATA_ASSET_NAME = "test_dataset"
-TEST_INVALID_DATA_ASSET_NAME = "non_existent_dataset"
-TEST_DATASTORE_NAME = "test_datastore"
 
 
 def test_datasetconfig_init() -> None:
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
+from azure.ai.ml.entities import Data, Job, UserIdentityConfiguration
 from azure.ai.ml.entities._job.distribution import MpiDistribution, PyTorchDistribution
 from azure.ai.ml.sweep import Choice
 from azure.core.exceptions import ResourceNotFoundError
@@ -58,10 +58,17 @@
     get_workspace,
     is_running_in_azure_ml,
     get_driver_log_file_text,
+    get_ml_client,
 )
 from testazure.test_data.make_tests import render_environment_yaml, render_test_script
 from testazure.utils_testazure import (
     DEFAULT_DATASTORE,
+    USER_IDENTITY_TEST_DATASTORE,
+    USER_IDENTITY_TEST_ASSET,
+    USER_IDENTITY_TEST_ASSET_OUTPUT,
+    USER_IDENTITY_TEST_FILE,
+    TEST_DATA_ASSET_NAME,
+    TEST_DATASTORE_NAME,
     change_working_directory,
     current_test_name,
     get_shared_config_json,
@@ -73,6 +80,8 @@
 GITHUB_SHIBBOLETH = "GITHUB_RUN_ID"  # https://docs.github.com/en/actions/reference/environment-variables
 AZUREML_FLAG = himl.AZUREML_FLAG
 
+TEST_ML_CLIENT = get_ml_client()
+
 logger = logging.getLogger('test.health_azure')
 logger.setLevel(logging.DEBUG)
 
@@ -743,13 +752,21 @@ def _mock_sweep(*args: Any, **kwargs: Any) -> MagicMock:
     dummy_input_data_name = "my_input_dataset"
     dummy_input_path = "path_to_my_input_data"
     dummy_inputs = {
-        dummy_input_data_name: Input(type=AssetTypes.URI_FOLDER, path=dummy_input_path, mode=InputOutputModes.MOUNT)
+        dummy_input_data_name: Input(  # type: ignore
+            type=AssetTypes.URI_FOLDER,
+            path=dummy_input_path,
+            mode=InputOutputModes.MOUNT
+        )
     }
 
     dummy_output_data_name = "my_output_dataset"
     dummy_output_path = "path_to_my_output_data"
     dummy_outputs = {
-        dummy_output_data_name: Output(type=AssetTypes.URI_FOLDER, path=dummy_output_path, mode=InputOutputModes.DIRECT)
+        dummy_output_data_name: Output(  # type: ignore
+            type=AssetTypes.URI_FOLDER,
+            path=dummy_output_path,
+            mode=InputOutputModes.MOUNT
+        )
     }
 
     dummy_root_directory = tmp_path
@@ -784,11 +801,8 @@ def _mock_sweep(*args: Any, **kwargs: Any) -> MagicMock:
             )
 
             expected_arg_str = " ".join(dummy_script_params)
-            expected_inputs_str = "--INPUT_0=${{inputs.INPUT_0}}"
-            expected_outputs_str = "--OUTPUT_0=${{outputs.OUTPUT_0}}"
             relative_entry_script = dummy_entry_script.relative_to(dummy_root_directory)
-            expected_command = f"python {relative_entry_script} {expected_arg_str} {expected_inputs_str} "\
-                f"{expected_outputs_str}"
+            expected_command = f"python {relative_entry_script} {expected_arg_str}"
 
             mock_command.assert_called_once_with(
                 code=str(dummy_root_directory),
@@ -802,7 +816,8 @@ def _mock_sweep(*args: Any, **kwargs: Any) -> MagicMock:
                 shm_size=dummy_docker_shm_size,
                 display_name=dummy_display_name,
                 distribution=MpiDistribution(process_count_per_instance=1),
-                instance_count=1
+                instance_count=1,
+                identity=UserIdentityConfiguration(),
             )
 
             # job with hyperparameter sampling:
@@ -863,7 +878,8 @@ def _mock_sweep(*args: Any, **kwargs: Any) -> MagicMock:
                 shm_size=dummy_docker_shm_size,
                 display_name=dummy_display_name,
                 distribution=MpiDistribution(process_count_per_instance=1),
-                instance_count=1
+                instance_count=1,
+                identity=UserIdentityConfiguration(),
             )
 
             mock_command.assert_any_call(**param_sampling)
@@ -918,6 +934,51 @@ def spawn_and_monitor_subprocess(process: str, args: List[str],
     return p.wait(), stdout_lines
 
 
+def validate_v1_job_outputs(captured: str, run_path: Path, extra_options: Dict[str, Any]) -> str:
+    """Assert that a v1 job/run has completed successfully.
+
+    :param captured: The captured output from the test script.
+    :param run_path: File path to where the run was executed.
+    :raises ValueError: Raises a value error if missing expected output files.
+    :return: Log text stored in the v1 log driver file.
+    """
+    assert EXPECTED_QUEUED in captured
+    with check_config_json(run_path, shared_config_json=get_shared_config_json()):
+        workspace = get_workspace(aml_workspace=None, workspace_config_path=run_path / WORKSPACE_CONFIG_JSON)
+
+    run = get_most_recent_run(
+        run_recovery_file=run_path / himl.RUN_RECOVERY_FILE,
+        workspace=workspace,
+    )
+    if run.status not in [RunStatus.FAILED, RunStatus.COMPLETED, RunStatus.CANCELED]:
+        run.wait_for_completion()
+    assert run.status == "Completed"
+    if "display_name" in extra_options:
+        assert run.display_name == extra_options["display_name"], "Display name has not been set"
+
+    # test error case mocking where no log file is present
+    log_text_undownloaded = get_driver_log_file_text(run=run, download_file=False)
+    assert log_text_undownloaded is None
+
+    # TODO: upgrade to walrus operator when upgrading python version to 3.8+
+    # if log_text := get_driver_log_file_text(run=run):
+    log_text = get_driver_log_file_text(run=run)
+
+    if log_text is None:
+        raise ValueError(
+            "The run does not contain any of the following log files: "
+            f"{[log_file_path for log_file_path in VALID_LOG_FILE_PATHS]}"
+        )
+
+    return log_text
+
+
+def validate_v2_job_outputs() -> str:
+    # TODO: implement this using run recovery files, issue open here:
+    # https://github.com/microsoft/hi-ml/issues/785
+    return "Not implemented yet"
+
+
 def render_and_run_test_script(path: Path,
                                run_target: RunTarget,
                                extra_options: Dict[str, Any],
@@ -1022,33 +1083,12 @@ def spawn() -> Tuple[int, List[str]]:
             assert EXPECTED_QUEUED not in captured
         return captured
     else:
-        assert EXPECTED_QUEUED in captured
-        with check_config_json(path, shared_config_json=get_shared_config_json()):
-            workspace = get_workspace(aml_workspace=None, workspace_config_path=path / WORKSPACE_CONFIG_JSON)
-
-        run = get_most_recent_run(run_recovery_file=path / himl.RUN_RECOVERY_FILE,
-                                  workspace=workspace)
-        if run.status not in ["Failed", "Completed", "Cancelled"]:
-            run.wait_for_completion()
-        assert run.status == "Completed"
-        if "display_name" in extra_options:
-            assert run.display_name == extra_options["display_name"], "Display name has not been set"
-
-        # test error case mocking where no log file is present
-        log_text_undownloaded = get_driver_log_file_text(run=run, download_file=False)
-        assert log_text_undownloaded is None
-
-        # TODO: upgrade to walrus operator when upgrading python version to 3.8+
-        # if log_text := get_driver_log_file_text(run=run):
-        log_text = get_driver_log_file_text(run=run)
-
-        if log_text is None:
-            raise ValueError(
-                "The run does not contain any of the following log files: "
-                f"{[log_file_path for log_file_path in VALID_LOG_FILE_PATHS]}"
-            )
 
-        return log_text
+        if "strictly_aml_v1" not in extra_options or extra_options["strictly_aml_v1"] == "True":
+            return validate_v1_job_outputs(captured, path, extra_options)
+
+        else:
+            return validate_v2_job_outputs()
 
 
 @pytest.mark.parametrize("run_target", [RunTarget.LOCAL, RunTarget.AZUREML])
@@ -1276,7 +1316,7 @@ class TestInputDataset:
     # Name of container for this dataset in blob storage.
     blob_name: str
     # Local folder for this dataset when running locally.
-    folder_name: Path
+    folder_name: Optional[Path] = None
     # Contents of test file.
     contents: str = ""
     # Local folder str
@@ -1288,7 +1328,7 @@ class TestOutputDataset:
     # Name of container for this dataset in blob storage.
     blob_name: str
     # Local folder for this dataset when running locally or when testing after running in Azure.
-    folder_name: Path
+    folder_name: Optional[Path] = None
 
 
 @pytest.mark.parametrize(["run_target", "local_folder", "strictly_aml_v1"],
@@ -1347,6 +1387,7 @@ def test_invoking_hello_world_datasets(run_target: RunTarget,
             assert downloaded == 1
 
             # Check that the input file is downloaded
+            assert input_dataset.folder_name is not None
             downloaded_dummy_txt_file = input_dataset.folder_name / input_dataset.blob_name / input_dataset.filename
             # Check it has expected contents
             assert input_dataset.contents == downloaded_dummy_txt_file.read_text()
@@ -1354,6 +1395,7 @@ def test_invoking_hello_world_datasets(run_target: RunTarget,
 
     if run_target == RunTarget.LOCAL:
         for output_dataset in output_datasets:
+            assert output_dataset.folder_name is not None
             output_blob_folder = output_dataset.folder_name / output_dataset.blob_name
             output_blob_folder.mkdir(parents=True)
     else:
@@ -1450,13 +1492,72 @@ def test_invoking_hello_world_datasets(run_target: RunTarget,
                     show_progress=True)
                 assert downloaded == 1
 
+            assert isinstance(output_dataset.folder_name, Path)
             output_dummy_txt_file = output_dataset.folder_name / output_dataset.blob_name / input_dataset.filename
             assert input_dataset.contents == output_dummy_txt_file.read_text()
 
-
 # endregion Elevate to AzureML unit tests
 
 
+def test_invoking_user_identity_datasets(tmp_path: Path) -> None:
+    output_test_file_name = f"test_output_{uuid4().hex}.txt"
+    extra_options: Dict[str, str] = {
+        'imports': """
+import shutil
+import sys
+import os
+""",
+        'default_datastore': f"'{USER_IDENTITY_TEST_DATASTORE}'",
+        'input_datasets': f"['{USER_IDENTITY_TEST_ASSET}']",
+        'output_datasets': f"['{USER_IDENTITY_TEST_ASSET_OUTPUT}']",
+        'strictly_aml_v1': str(False),
+        'body': f"""
+
+    input_folder = run_info.input_datasets[0]
+    output_folder = run_info.output_datasets[0]
+
+    print("input dir contents: " + str(os.listdir(input_folder)))
+
+    output_dir_contents = os.listdir(output_folder)
+    num_output_items = len(output_dir_contents)
+
+    print("output dir contents before copying: " + str(output_dir_contents))
+    print("Number of items in output dir before copying: " + str(num_output_items))
+
+    input_file = input_folder / "{USER_IDENTITY_TEST_FILE}"
+    output_file = output_folder / "{output_test_file_name}"
+
+    print('input file: ' + str(input_file) + ', output file: ' + str(output_file))
+
+    if os.path.exists(input_file):
+        print("Input file exists")
+
+    if not os.path.exists(output_file):
+        print("Output file does not exist (yet)")
+
+    print("Copying file...")
+    shutil.copy(input_file, output_file)
+
+    num_output_dir_items_after_copying = len(os.listdir(output_folder))
+
+    assert num_output_dir_items_after_copying == num_output_items + 1, "Copied file not present in output dir"
+
+    print("File successfully copied!")
+        """,
+    }
+
+    extra_args: List[str] = []
+
+    render_and_run_test_script(
+        tmp_path,
+        RunTarget.AZUREML,
+        extra_options,
+        extra_args,
+        expected_pass=True,
+        upload_package=False,
+    )
+
+
 @pytest.mark.fast
 # Azure ML expects run_config to be instance of ScriptRunConfig
 @patch("azureml.train.hyperdrive.runconfig.isinstance", return_value=True)
@@ -1565,6 +1666,32 @@ def test_submit_to_azure_if_needed_with_hyperdrive(mock_sys_args: MagicMock,
                         mock_hyperdrive_config.assert_called_once()
 
 
+def test_get_data_asset_from_config() -> None:
+    n_configs = 3
+    test_dataset_configs = [
+        DatasetConfig(
+            name=TEST_DATA_ASSET_NAME,
+            datastore=TEST_DATASTORE_NAME,
+        ) for _ in range(n_configs)
+    ]
+
+    test_assets = [
+        himl.get_data_asset_from_config(TEST_ML_CLIENT, test_dataset_config)
+        for test_dataset_config in test_dataset_configs
+    ]
+    assert len(test_assets) == n_configs
+    assert all([asset.name == TEST_DATA_ASSET_NAME for asset in test_assets])
+
+    test_version = 1
+    test_versioned_config = DatasetConfig(
+        name=TEST_DATA_ASSET_NAME,
+        datastore=TEST_DATASTORE_NAME,
+        version=test_version,
+    )
+    test_versioned_asset = himl.get_data_asset_from_config(TEST_ML_CLIENT, test_versioned_config)
+    assert test_versioned_asset.version == str(test_version)
+
+
 @pytest.mark.fast
 @pytest.mark.parametrize("already_exists", [True, False])
 def test_create_v2_inputs(already_exists: bool) -> None:
@@ -1626,18 +1753,20 @@ def test_create_v2_inputs_fails(missing: Any) -> None:
 
 @pytest.mark.fast
 def test_create_v2_outputs() -> None:
-    mock_datastore_name = "dummy_datastore"
-    mock_data_name = "dummy_dataset"
+    test_dataset_config = DatasetConfig(
+        name=TEST_DATA_ASSET_NAME,
+        datastore=TEST_DATASTORE_NAME,
+    )
+    outputs = himl.create_v2_outputs(TEST_ML_CLIENT, [test_dataset_config])
+
+    output_key = "OUTPUT_0"
+    assert output_key in outputs
+    assert len(outputs) == 1
 
-    mock_output_dataconfigs = [DatasetConfig(name=mock_data_name, datastore=mock_datastore_name)]
-    outputs = himl.create_v2_outputs(mock_output_dataconfigs)
-    assert isinstance(outputs, Dict)
-    assert len(outputs) == len(mock_output_dataconfigs)
-    output_entry = outputs["OUTPUT_0"]
-    assert isinstance(output_entry, Output)
-    assert output_entry.type == AssetTypes.URI_FOLDER
-    expected_path = f"azureml://datastores/{mock_datastore_name}/paths/{mock_data_name}"
-    assert expected_path in output_entry['path']
+    output = outputs[output_key]
+    assert output.path is not None
+    assert isinstance(output, Output)
+    assert output.mode == InputOutputModes.MOUNT
 
 
 def test_submit_to_azure_if_needed_v2() -> None:
@@ -1686,42 +1815,6 @@ def test_submit_to_azure_if_needed_v2() -> None:
                 assert return_value is None
 
 
-@pytest.mark.fast
-def test_generate_input_dataset_command() -> None:
-    input_datasets = {"INPUT_0": Input(), "INPUT_1": Input()}
-    input_data_cmd = himl._generate_input_dataset_command(input_datasets)
-    assert input_data_cmd == " --INPUT_0=${{inputs.INPUT_0}} --INPUT_1=${{inputs.INPUT_1}}"
-
-
-@pytest.mark.fast
-def test_generate_output_dataset_command() -> None:
-    output_datasets = {"OUTPUT_0": Output(), "OUTPUT_1": Output()}
-    output_data_cmd = himl._generate_output_dataset_command(output_datasets)
-    assert output_data_cmd == " --OUTPUT_0=${{outputs.OUTPUT_0}} --OUTPUT_1=${{outputs.OUTPUT_1}}"
-
-
-@pytest.mark.fast
-def test_extract_v2_inputs_outputs_from_args() -> None:
-    path_to_input_0 = "path_to_input_0"
-    path_to_output_0 = "path_to_output_0"
-    mock_args = [f"--INPUT_0={path_to_input_0}", "--INPUT_1=path_to_input_1", f"--OUTPUT_0={path_to_output_0}",
-                 "--a=foo", "--b=bar"]
-    with patch.object(sys, "argv", new=mock_args):
-        input_datasets, output_datasets = himl._extract_v2_inputs_outputs_from_args()
-        assert len(input_datasets) == 2
-        assert input_datasets[0] == Path(path_to_input_0)
-        assert len(output_datasets) == 1
-        assert output_datasets[0] == Path(path_to_output_0)
-
-    # similar args should be ignored
-    mock_args_similar = [f"--input_0={path_to_input_0}", "--input_1=path_to_input_1", f"--output_0={path_to_output_0}",
-                         "--a=foo", "--b=bar"]
-    with patch.object(sys, "argv", new=mock_args_similar):
-        input_datasets, output_datasets = himl._extract_v2_inputs_outputs_from_args()
-        assert len(input_datasets) == 0
-        assert len(output_datasets) == 0
-
-
 @pytest.mark.parametrize("wait_for_completion", [True, False])
 def test_submitting_script_with_sdk_v2(tmp_path: Path, wait_for_completion: bool) -> None:
     """
@@ -1959,3 +2052,29 @@ def test_submit_to_azure_v2_distributed() -> None:
                 assert call_kwargs.get("instance_count") == num_nodes
                 distribution = call_kwargs.get("distribution")
                 assert distribution == MpiDistribution(process_count_per_instance=1)
+
+
+@pytest.mark.fast
+def test_extract_v2_inputs_outputs_from_env_vars() -> None:
+    valid_mock_environment = {
+        "AZURE_ML_INPUT_INPUT_0": "input_0",
+        "AZURE_ML_INPUT_INPUT_1": "input_1",
+        "AZURE_ML_OUTPUT_OUTPUT_0": "output_0",
+        "AZURE_ML_OUTPUT_OUTPUT_1": "output_1",
+    }
+
+    with patch.dict(os.environ, valid_mock_environment):
+        input_datasets, output_datasets = himl._extract_v2_inputs_outputs_from_env_vars()
+        assert input_datasets == [Path("input_0"), Path("input_1")]
+        assert output_datasets == [Path("output_0"), Path("output_1")]
+
+    invalid_mock_environment = {
+        "ML_INPUT_INPUT_0": "input_0",
+        "AZURE_ML_INPUT_INPUT__1": "input_1",
+        "AZURE_ML_OUTPUT_OUTPUT_test": "output_0",
+        "bad_string": "output_1",
+    }
+
+    with patch.dict(os.environ, invalid_mock_environment):
+        input_datasets, output_datasets = himl._extract_v2_inputs_outputs_from_env_vars()
+        assert input_datasets == output_datasets == []