Repository: microsoft/hi-ml

Issue:
Submitting repeated training run from a VM creates multiple versions of the same data asset
This is probably a behaviour only with SDK v2 - why?


Discussion:

Resolving patch:
diff --git a/hi-ml-azure/src/health_azure/datasets.py b/hi-ml-azure/src/health_azure/datasets.py
--- a/hi-ml-azure/src/health_azure/datasets.py
+++ b/hi-ml-azure/src/health_azure/datasets.py
@@ -130,63 +130,117 @@ def _get_or_create_v1_dataset(datastore_name: str, dataset_name: str, workspace:
     return azureml_dataset
 
 
-def _retrieve_v2_dataset(dataset_name: str, ml_client: MLClient) -> Data:
+def _get_latest_v2_asset_version(ml_client: MLClient, data_asset_name: str) -> str:
+    """
+    Retrieve the version of a v2 data asset that has the highest value (version numbers are assumed to be integers).
+
+    :param ml_client: An Azure MLClient object for interacting with Azure resources.
+    :param data_asset_name: The name of the data asset to look for.
+    :raises ResourceNotFoundError: If no data asset can be found with a matching name.
+    :return: The latest version of the data asset if found, else None.
+    """
+    data_assets = ml_client.data.list(name=data_asset_name)
+    highest_version = 0
+    for data_asset in data_assets:
+
+        try:
+            data_asset_version = int(data_asset.version)
+        except ValueError as val_er:
+            logging.warning(f"Failed to convert data asset version to int: {val_er}")
+            data_asset_version = -1
+
+        if data_asset_version > highest_version:
+            highest_version = data_asset_version
+
+    return str(highest_version)
+
+
+def _retrieve_v2_data_asset(
+    ml_client: MLClient,
+    data_asset_name: str,
+    version: Optional[str] = None,
+) -> Data:
     """
     Attempt to retrieve a v2 Data Asset using a provided Azure ML Workspace connection. If
     no Data asset can be found with a matching name, the underlying code will raise an Exception
 
-    :param dataset_name: The name of the dataset to look for.
     :param ml_client: An Azure MLClient object for interacting with Azure resources.
-    :return: An Azure Data asset representing the dataset if found, otherwise an Exception will be raised.
+    :param data_asset_name: The name of the data asset to look for.
+    :param version: version of the data asset to retrieve, latest version retrieved if None.
+    :return: An Azure Data asset representing the data asset if found.
     """
-    aml_data = ml_client.data.get(name=dataset_name)
+
+    if version is None:
+        version = _get_latest_v2_asset_version(ml_client, data_asset_name)
+
+    aml_data = ml_client.data.get(name=data_asset_name, version=version)
+    assert aml_data.version == version, "Version of retrieved data asset does not match requested version"
     return aml_data
 
 
-def _create_v2_dataset(datastore_name: str, dataset_name: str, ml_client: MLClient) -> Data:
+def _create_v2_data_asset(
+    ml_client: MLClient,
+    datastore_name: str,
+    data_asset_name: str,
+    version: Optional[str] = None
+) -> Data:
     """
     Create or update a v2 Data asset in the specified Datastore
 
-    :param datastore_name: The name of the datastore in which to create or update the Data asset.
-    :param dataset_name: The name of the dataset to be created.
     :param ml_client: An Azure MLClient object for interacting with Azure resources.
-    :raises ValueError: If no datastore name is provided to define where to create the data
+    :param datastore_name: The name of the datastore in which to create or update the Data asset.
+    :param data_asset_name: The name of the data asset to be created.
+    :param version: The version of the data asset to be created.
+    :raises ValueError: If no datastore name is provided to define where to create the data.
     :return: The created or updated Data asset.
     """
-    if not dataset_name:
-        raise ValueError(f"Cannot create data asset without a valid dataset name (received {dataset_name})")
+    if not data_asset_name:
+        raise ValueError("Cannot create data asset with empty name.")
 
     if not datastore_name:
         default_datastore = ml_client.datastores.get_default()
         datastore_name = default_datastore.name
 
-    logging.info(f"Creating a new Data asset from data in folder '{dataset_name}' in the datastore '{datastore_name}'")
+    logging.info(
+        f"Creating a new Data asset from data in folder '{data_asset_name}' in the datastore '{datastore_name}'"
+    )
     azureml_data_asset = Data(
-        path=f"azureml://datastores/{datastore_name}/paths/{dataset_name}/",
+        path=f"azureml://datastores/{datastore_name}/paths/{data_asset_name}/",
         type=AssetTypes.URI_FOLDER,
         description="<description>",
-        name=dataset_name,
-        version=None
+        name=data_asset_name,
+        version=version,
     )
+
     ml_client.data.create_or_update(azureml_data_asset)
     return azureml_data_asset
 
 
-def _get_or_create_v2_dataset(datastore_name: str, dataset_name: str, ml_client: MLClient) -> Data:
+def _get_or_create_v2_data_asset(
+    ml_client: MLClient,
+    datastore_name: str,
+    data_asset_name: str,
+    version: Optional[str] = None
+) -> Data:
     """
-    Attempt to retrieve a v2 Dataset object and return that, otherwise attempt to create and register
-    a v2 Dataset and return that.
+    Attempt to retrieve a v2 data asset object and return that, otherwise attempt to create and register
+    a v2 data asset and return that.
 
-    :param datastore_name: The name of the Datastore to either retrieve or create and register the Data asset in.
-    :param dataset_name: The name of the Data asset to be retrieved or registered.
     :param ml_client: An Azure MLClient object for interacting with Azure resources.
-    :return: An Azure Data asset object with the provided dataset name, in the provided datastore
+    :param datastore_name: The name of the Datastore to either retrieve or create and register the Data asset in.
+    :param data_asset_name: The name of the Data asset to be retrieved or registered.
+    :return: An Azure Data asset object with the provided asset name, in the provided datastore
     """
     try:
-        azureml_dataset = _retrieve_v2_dataset(dataset_name, ml_client)
-    except Exception:
-        azureml_dataset = _create_v2_dataset(datastore_name, dataset_name, ml_client)
-    return azureml_dataset
+        azureml_data_asset = _retrieve_v2_data_asset(ml_client, data_asset_name, version)
+    except ResourceNotFoundError:  # catch the exception and create the dataset, raise all other types of exceptions
+        logging.info(
+            f"Data asset {data_asset_name} not found in datastore {datastore_name}. Version specified: {version}."
+            "Attempting to create a new data asseet with specified name and version."
+        )
+        azureml_data_asset = _create_v2_data_asset(ml_client, datastore_name, data_asset_name, version)
+
+    return azureml_data_asset
 
 
 def get_or_create_dataset(datastore_name: str,
@@ -221,7 +275,7 @@ def get_or_create_dataset(datastore_name: str,
     else:
         try:
             ml_client = get_ml_client(ml_client=ml_client)
-            aml_dataset = _get_or_create_v2_dataset(datastore_name, dataset_name, ml_client)
+            aml_dataset = _get_or_create_v2_data_asset(ml_client, datastore_name, dataset_name)
         except HttpResponseError as e:
             if "Cannot create v2 Data Version in v1 Data Container" in e.message:
                 logging.info("This appears to be a v1 Data Container. Reverting to API v1 to create this Dataset")
@@ -311,11 +365,13 @@ def to_input_dataset_local(self,
         if workspace is None:
             raise ValueError(f"Unable to make dataset '{self.name} available for a local run because no AzureML "
                              "workspace has been provided. Provide a workspace, or set a folder for local execution.")
-        azureml_dataset = get_or_create_dataset(workspace=workspace,
-                                                ml_client=ml_client,
-                                                dataset_name=self.name,
-                                                datastore_name=self.datastore,
-                                                strictly_aml_v1=strictly_aml_v1)
+        azureml_dataset = get_or_create_dataset(
+            datastore_name=self.datastore,
+            dataset_name=self.name,
+            workspace=workspace,
+            strictly_aml_v1=strictly_aml_v1,
+            ml_client=ml_client,
+        )
         if isinstance(azureml_dataset, FileDataset):
             target_path = self.target_folder or Path(tempfile.mkdtemp())
             use_mounting = self.use_mounting if self.use_mounting is not None else False
@@ -350,11 +406,13 @@ def to_input_dataset(self,
         :param ml_client: An Azure MLClient object for interacting with Azure resources.
         """
         status = f"In AzureML, dataset {self.name} (index {dataset_index}) will be "
-        azureml_dataset = get_or_create_dataset(workspace=workspace,
-                                                ml_client=ml_client,
-                                                dataset_name=self.name,
-                                                datastore_name=self.datastore,
-                                                strictly_aml_v1=strictly_aml_v1)
+        azureml_dataset = get_or_create_dataset(
+            datastore_name=self.datastore,
+            dataset_name=self.name,
+            workspace=workspace,
+            strictly_aml_v1=strictly_aml_v1,
+            ml_client=ml_client,
+        )
         # If running on windows then self.target_folder may be a WindowsPath, make sure it is
         # in posix format for Azure.
         use_mounting = False if self.use_mounting is None else self.use_mounting
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
@@ -10,12 +10,13 @@
 from pathlib import Path
 from typing import Dict, Generator, Optional
 
+from azure.ai.ml import MLClient
 from azureml.core import Run
-
 from health_azure.utils import (ENV_EXPERIMENT_NAME, WORKSPACE_CONFIG_JSON, UnitTestWorkspaceWrapper,
                                 to_azure_friendly_string)
 from health_azure import create_aml_run_object
 from health_azure.himl import effective_experiment_name
+from health_azure.utils import get_ml_client, get_workspace
 
 DEFAULT_DATASTORE = "himldatasets"
 FALLBACK_SINGLE_RUN = "refs_pull_545_merge:refs_pull_545_merge_1626538212_d2b07afd"
@@ -83,3 +84,13 @@ def create_unittest_run_object(snapshot_directory: Optional[Path] = None) -> Run
     return create_aml_run_object(experiment_name=effective_experiment_name("himl-tests"),
                                  workspace=DEFAULT_WORKSPACE.workspace,
                                  snapshot_directory=snapshot_directory)
+
+
+def get_test_ml_client() -> MLClient:
+    """Generates an MLClient object for use in tests.
+
+    :return: MLClient object
+    """
+
+    workspace = get_workspace()
+    return get_ml_client(aml_workspace=workspace)
diff --git a/hi-ml-cpath/src/SSL/configs/HistoSimCLRContainer.py b/hi-ml-cpath/src/SSL/configs/HistoSimCLRContainer.py
--- a/hi-ml-cpath/src/SSL/configs/HistoSimCLRContainer.py
+++ b/hi-ml-cpath/src/SSL/configs/HistoSimCLRContainer.py
@@ -50,11 +50,11 @@ def _get_transforms(self, augmentation_config: Optional[CfgNode],
         return train_transforms, val_transforms
 
     def get_preprocessing_transforms(self) -> List[Callable]:
-        return([Lambda(lambda x: x)])
+        return ([Lambda(lambda x: x)])
 
     def get_augmentations(self) -> List[Callable]:
         # SimClr augmentations
-        return([RandomResizedCrop(size=224),
+        return ([RandomResizedCrop(size=224),
                 RandomHorizontalFlip(p=0.5),
                 RandomApply([ColorJitter(brightness=0.8, contrast=0.8, saturation=0.8, hue=0.2)], 0.8),
                 RandomGrayscale(p=0.2),


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
@@ -6,27 +6,35 @@
 Test the data input and output functionality
 """
 from pathlib import Path
-from unittest.mock import create_autospec, DEFAULT, MagicMock, patch
-from health_azure.utils import PathOrString, get_ml_client
-from typing import List, Union, Optional
+from typing import List, Optional, Union
+from unittest.mock import DEFAULT, MagicMock, create_autospec, patch
 
 import pytest
 from azure.ai.ml import MLClient
 from azure.ai.ml.entities import Data
 from azure.ai.ml.operations import DatastoreOperations
-from azure.core.exceptions import HttpResponseError
+from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
 from azureml._restclient.exceptions import ServiceException
 from azureml.core import Dataset, Workspace
 from azureml.data import FileDataset, OutputFileDatasetConfig
 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 health_azure.datasets import (
+    DatasetConfig, _create_v1_dataset, _create_v2_data_asset, _get_or_create_v1_dataset, _get_or_create_v2_data_asset,
+    _input_dataset_key, _output_dataset_key, _replace_string_datasets, _retrieve_v1_dataset, _retrieve_v2_data_asset,
+    create_dataset_configs, get_datastore, get_or_create_dataset, _get_latest_v2_asset_version,
+)
+from health_azure.utils import PathOrString, get_ml_client
 
-from health_azure.datasets import (DatasetConfig, _input_dataset_key, _output_dataset_key,
-                                   _replace_string_datasets, get_datastore, get_or_create_dataset,
-                                   create_dataset_configs, _get_or_create_v1_dataset, _get_or_create_v2_dataset,
-                                   _retrieve_v1_dataset, _create_v1_dataset, _retrieve_v2_dataset, _create_v2_dataset)
-from testazure.utils_testazure import DEFAULT_DATASTORE, DEFAULT_WORKSPACE
+
+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:
@@ -35,9 +43,9 @@ def test_datasetconfig_init() -> None:
     assert "name of the dataset must be a non-empty string" in str(ex)
 
 
-def test_get_datastore() -> None:
+def test_get_datastore_v1() -> None:
     """
-    Test retrieving a datastore from the AML workspace.
+    Test retrieving a datastore from the AML workspace using the v1 SDK.
     """
     # Retrieving a datastore that does not exist should fail
     does_not_exist = "does_not_exist"
@@ -66,26 +74,31 @@ def test_get_datastore() -> None:
     assert isinstance(single_store, AzureBlobDatastore)
     assert single_store.name == name
 
+
+def test_get_datastore_v2() -> None:
+    """
+    Test retrieving a datastore from the AML workspace using the v2 SDK.
+    """
     # Now test retrieving a v2 datastore by name
-    mock_v2_dataset_name = "dummy_v2_datastore"
+    mock_v2_datastore_name = "dummy_v2_datastore"
     mock_returned_datastore = MagicMock()
-    mock_returned_datastore.name = mock_v2_dataset_name
+    mock_returned_datastore.name = mock_v2_datastore_name
     mock_workspace = MagicMock()
     mock_workspace.datastores = create_autospec(DatastoreOperations)
     mock_workspace.datastores.get.return_value = mock_returned_datastore
-    v2_datastore = get_datastore(mock_workspace, datastore_name=mock_v2_dataset_name)
-    assert v2_datastore.name == mock_v2_dataset_name
+    v2_datastore = get_datastore(mock_workspace, datastore_name=mock_v2_datastore_name)
+    assert v2_datastore.name == mock_v2_datastore_name
 
     # Test retrieving a default v2 datastore
     mock_workspace.datastores.list.return_value = [mock_returned_datastore]
     v2_default_datastore = get_datastore(mock_workspace, datastore_name="")
-    assert v2_default_datastore.name == mock_v2_dataset_name
+    assert v2_default_datastore.name == mock_v2_datastore_name
 
     # Mock case where list is empty but get_default returns a value
     mock_workspace.datastores.list.return_value = []
     mock_workspace.datastores.get_default.return_value = mock_returned_datastore
     v2_default_datastore = get_datastore(mock_workspace, datastore_name="")
-    assert v2_default_datastore.name == mock_v2_dataset_name
+    assert v2_default_datastore.name == mock_v2_datastore_name
 
     # If datastores has an unknown format, an exception should be raised
     mock_workspace = MagicMock()
@@ -196,7 +209,9 @@ def test_datasets_from_string() -> None:
 
 
 def test_get_or_create_dataset() -> None:
-    def _mock_retrieve_or_create_v2_dataset_fails(datastore_name: str, dataset_name: str, ml_client: MLClient) -> None:
+    def _mock_retrieve_or_create_v2_dataset_fails(
+        datastore_name: str, dataset_name: str, ml_client: MLClient
+    ) -> None:
         raise HttpResponseError("Cannot create v2 Data Version in v1 Data Container")
 
     data_asset_name = "himl_tiny_data_asset"
@@ -215,7 +230,7 @@ def _mock_retrieve_or_create_v2_dataset_fails(datastore_name: str, dataset_name:
     mock_v1_dataset = "v1_dataset"
     with patch.multiple("health_azure.datasets",
                         _get_or_create_v1_dataset=DEFAULT,
-                        _get_or_create_v2_dataset=DEFAULT) as mocks:
+                        _get_or_create_v2_data_asset=DEFAULT) as mocks:
         mocks["_get_or_create_v1_dataset"].return_value = mock_v1_dataset
         dataset = get_or_create_dataset(workspace=workspace,
                                         ml_client=ml_client,
@@ -223,30 +238,30 @@ def _mock_retrieve_or_create_v2_dataset_fails(datastore_name: str, dataset_name:
                                         dataset_name=data_asset_name,
                                         strictly_aml_v1=True)
         mocks["_get_or_create_v1_dataset"].assert_called_once()
-        mocks["_get_or_create_v2_dataset"].assert_not_called()
+        mocks["_get_or_create_v2_data_asset"].assert_not_called()
         assert dataset == mock_v1_dataset
 
         # Now pass strictly_aml_v1 as False
         mock_v2_dataset = "v2_dataset"
-        mocks["_get_or_create_v2_dataset"].return_value = mock_v2_dataset
+        mocks["_get_or_create_v2_data_asset"].return_value = mock_v2_dataset
         dataset = get_or_create_dataset(workspace=workspace,
                                         ml_client=ml_client,
                                         datastore_name="himldatasetsv2",
                                         dataset_name=data_asset_name,
                                         strictly_aml_v1=False)
         mocks["_get_or_create_v1_dataset"].assert_called_once()
-        mocks["_get_or_create_v2_dataset"].assert_called_once()
+        mocks["_get_or_create_v2_data_asset"].assert_called_once()
         assert dataset == mock_v2_dataset
 
         # if  trying to get or create a v2 dataset fails, should revert back to _get_or_create_v1_dataset
-        mocks["_get_or_create_v2_dataset"].side_effect = _mock_retrieve_or_create_v2_dataset_fails
+        mocks["_get_or_create_v2_data_asset"].side_effect = _mock_retrieve_or_create_v2_dataset_fails
         dataset = get_or_create_dataset(workspace=workspace,
                                         ml_client=ml_client,
                                         datastore_name="himldatasetsv2",
                                         dataset_name=data_asset_name,
                                         strictly_aml_v1=False)
         assert mocks["_get_or_create_v1_dataset"].call_count == 2
-        assert mocks["_get_or_create_v2_dataset"].call_count == 2
+        assert mocks["_get_or_create_v2_data_asset"].call_count == 2
         assert dataset == mock_v1_dataset
 
 
@@ -271,25 +286,29 @@ def _mock_error_from_retrieve_v1_dataset(dataset_name: str, workspace: Workspace
         mocks["_create_v1_dataset"].assert_called_once()
 
 
-def test_get_or_create_v2_dataset() -> None:
-    def _mock_error_from_retrieve_v2_dataset(dataset_name: str, workspace: Workspace) -> None:
-        raise Exception("Error Message")
+def test_get_or_create_v2_data_asset() -> None:
+    def _mock_error_from_retrieve_v2_data_asset(
+        ml_client: MLClient,
+        data_asset_name: str,
+        version: Optional[str] = None,
+    ) -> None:
+        raise ResourceNotFoundError("Error Message")
 
     ml_client = MagicMock()
     datastore = "dummy_datastore"
     dataset_name = "foo"
 
     with patch.multiple("health_azure.datasets",
-                        _retrieve_v2_dataset=DEFAULT,
-                        _create_v2_dataset=DEFAULT) as mocks:
-        _get_or_create_v2_dataset(datastore, dataset_name, ml_client)
-        mocks["_retrieve_v2_dataset"].assert_called_once()
-        mocks["_create_v2_dataset"].assert_not_called()
+                        _retrieve_v2_data_asset=DEFAULT,
+                        _create_v2_data_asset=DEFAULT) as mocks:
+        _get_or_create_v2_data_asset(ml_client, datastore, dataset_name)
+        mocks["_retrieve_v2_data_asset"].assert_called_once()
+        mocks["_create_v2_data_asset"].assert_not_called()
 
-        mocks["_retrieve_v2_dataset"].side_effect = _mock_error_from_retrieve_v2_dataset
-        _get_or_create_v2_dataset(datastore, dataset_name, ml_client)
-        assert mocks["_retrieve_v2_dataset"].call_count == 2
-        mocks["_create_v2_dataset"].assert_called_once()
+        mocks["_retrieve_v2_data_asset"].side_effect = _mock_error_from_retrieve_v2_data_asset
+        _get_or_create_v2_data_asset(ml_client, datastore, dataset_name)
+        assert mocks["_retrieve_v2_data_asset"].call_count == 2
+        mocks["_create_v2_data_asset"].assert_called_once()
 
 
 def test_retrieve_v1_dataset() -> None:
@@ -350,25 +369,103 @@ def test_create_v1_dataset() -> None:
         pass
 
 
-def test_retrieve_v2_dataset() -> None:
-    dataset_name = "dummydataset"
-    mock_ml_client = MagicMock()
-    mock_retrieved_dataset = "dummy_data"
-    mock_ml_client.data.get.return_value = mock_retrieved_dataset
-    data_asset = _retrieve_v2_dataset(dataset_name, mock_ml_client)
-    assert data_asset == mock_retrieved_dataset
-
-
-def test_create_v2_dataset() -> None:
-    dataset_name = "dummydataset"
-    datastore = "dummydataset"
-    mock_ml_client = MagicMock()
-    # mock_ml_client.data.create_or_update.return_value = None
-    data_asset = _create_v2_dataset(datastore, dataset_name, mock_ml_client)
-    assert isinstance(data_asset, Data)
-    assert data_asset.path == "azureml://datastores/dummydataset/paths/dummydataset/"
-    assert data_asset.type == "uri_folder"
-    assert data_asset.name == dataset_name
+@pytest.mark.parametrize(["asset_name", "asset_version"],
+                         [(TEST_DATA_ASSET_NAME, None),
+                          (TEST_DATA_ASSET_NAME, "1"),
+                          (TEST_INVALID_DATA_ASSET_NAME, None),
+                          (TEST_INVALID_DATA_ASSET_NAME, "1")])
+def test_retrieve_v2_data_asset(asset_name: str, asset_version: Optional[str]) -> None:
+    with patch("health_azure.datasets._get_latest_v2_asset_version") as mock_get_v2_asset_version:
+        mock_get_v2_asset_version.side_effect = _get_latest_v2_asset_version
+        try:
+            data_asset = _retrieve_v2_data_asset(
+                ml_client=TEST_ML_CLIENT,
+                data_asset_name=asset_name,
+                version=asset_version,
+            )
+        except ResourceNotFoundError as ex:
+            if asset_name == TEST_INVALID_DATA_ASSET_NAME:
+                if asset_version is None:
+                    expected_error_message = f"{TEST_INVALID_DATA_ASSET_NAME} container was not found."
+                else:
+                    expected_error_message = \
+                        f"{TEST_INVALID_DATA_ASSET_NAME}:{asset_version} (dataContainerName:version) not found."
+
+                assert expected_error_message in str(ex)
+            else:
+                pytest.fail(f"Unexpected error: {ex}")
+        else:
+            assert isinstance(data_asset, Data)
+            if asset_version is not None:
+                assert data_asset.version == asset_version
+            else:
+                mock_get_v2_asset_version.assert_called_once()
+            data_asset_path = data_asset.path
+            assert isinstance(data_asset_path, str)  # makes pyright happy
+            assert f"{TEST_DATASTORE_NAME}/paths/{TEST_DATA_ASSET_NAME}" in data_asset_path
+
+
+def test_retrieve_v2_data_asset_invalid_version() -> None:
+    invalid_asset_version = str(int(_get_latest_v2_asset_version(TEST_ML_CLIENT, TEST_DATA_ASSET_NAME)) + 1)
+    with pytest.raises(ResourceNotFoundError) as ex:
+        _retrieve_v2_data_asset(
+            ml_client=TEST_ML_CLIENT,
+            data_asset_name=TEST_DATA_ASSET_NAME,
+            version=invalid_asset_version,
+        )
+        assert f"{TEST_DATA_ASSET_NAME}:{invalid_asset_version} (dataContainerName:version) not found." in str(ex)
+
+
+def test_retrieving_v2_data_asset_does_not_increment() -> None:
+    """Test if calling the get_or_create_data_asset on an existing asset does not increment the version number.
+    """
+
+    with patch("health_azure.datasets._create_v2_data_asset") as mock_create_v2_data_asset:
+        asset_version_before_get_or_create = _get_latest_v2_asset_version(TEST_ML_CLIENT, TEST_DATA_ASSET_NAME)
+        get_or_create_dataset(
+            TEST_DATASTORE_NAME,
+            TEST_DATA_ASSET_NAME,
+            DEFAULT_WORKSPACE,
+            strictly_aml_v1=False,
+            ml_client=TEST_ML_CLIENT,
+        )
+        asset_version_after_get_or_create = _get_latest_v2_asset_version(TEST_ML_CLIENT, TEST_DATA_ASSET_NAME)
+
+        mock_create_v2_data_asset.assert_not_called()
+        assert asset_version_before_get_or_create == asset_version_after_get_or_create
+
+
+@pytest.mark.parametrize(["asset_name", "version"],
+                         [(TEST_DATA_ASSET_NAME, None),
+                          (TEST_DATA_ASSET_NAME, "1"),
+                          ("", 1)])
+def test_create_v2_data_asset(asset_name: str, version: Optional[str]) -> None:
+    try:
+        data_asset = _create_v2_data_asset(
+            ml_client=TEST_ML_CLIENT,
+            datastore_name=TEST_DATASTORE_NAME,
+            data_asset_name=asset_name,
+            version=version,
+        )
+    except HttpResponseError as ex:
+        if version is not None:
+            assert "A data version with this name and version already exists" in str(ex)
+        else:
+            pytest.fail(f"Unexpected error: {ex}")
+
+    except ValueError as ex:
+        if asset_name == "":
+            assert "Cannot create data asset with empty name." in str(ex)
+        else:
+            pytest.fail(f"Unexpected error: {ex}")
+
+    else:
+        assert isinstance(data_asset, Data)
+        assert data_asset.path == f"azureml://datastores/{TEST_DATASTORE_NAME}/paths/{TEST_DATA_ASSET_NAME}/"
+        assert data_asset.type == "uri_folder"
+        assert data_asset.name == TEST_DATASET_NAME
+        if version:
+            assert data_asset.version == version
 
 
 def test_dataset_keys() -> 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
@@ -1279,12 +1279,11 @@ class TestOutputDataset:
 
 
 @pytest.mark.parametrize(["run_target", "local_folder", "strictly_aml_v1"],
-                         [(RunTarget.LOCAL, False, False),
-                          (RunTarget.LOCAL, True, False),
+                         [(RunTarget.LOCAL, True, False),
                           (RunTarget.AZUREML, False, True),
-                          # Test with AML SDK v2 fails, logged as https://github.com/microsoft/hi-ml/issues/763
-                          # (RunTarget.AZUREML, False, False)
                           ])
+# Test with AML SDK v2 fails, logged as https://github.com/microsoft/hi-ml/issues/763
+# (RunTarget.AZUREML, False, False)
 def test_invoking_hello_world_datasets(run_target: RunTarget,
                                        local_folder: bool,
                                        strictly_aml_v1: bool,