Repository: microsoft/hi-ml

Issue:
Bug: Retrieving a V2 data asset reads version 1 by default
See https://github.com/microsoft/hi-ml/blob/main/hi-ml-azure/src/health_azure/himl.py#L709
Code should read the latest version unless specified otherwise



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
@@ -140,18 +140,21 @@ def _get_latest_v2_asset_version(ml_client: MLClient, data_asset_name: str) -> s
     :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:
+    highest_version = float('-inf')
 
+    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
+            continue
 
         if data_asset_version > highest_version:
             highest_version = data_asset_version
 
+    if highest_version == float('-inf'):
+        raise ResourceNotFoundError(f"No data asset found with the provided name: {data_asset_name}")
+
     return str(highest_version)
 
 
@@ -174,7 +177,7 @@ def _retrieve_v2_data_asset(
         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"
+    assert aml_data.version == version, f"Data asset version ({aml_data.version}) does not match version ({version})"
     return aml_data
 
 
@@ -212,7 +215,7 @@ def _create_v2_data_asset(
         version=version,
     )
 
-    ml_client.data.create_or_update(azureml_data_asset)
+    azureml_data_asset = ml_client.data.create_or_update(azureml_data_asset)
     return azureml_data_asset
 
 
@@ -236,7 +239,7 @@ def _get_or_create_v2_data_asset(
     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."
+            "Attempting to create a new data asset with specified name and version."
         )
         azureml_data_asset = _create_v2_data_asset(ml_client, datastore_name, data_asset_name, version)
 
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
@@ -44,7 +44,8 @@
                                 create_python_environment_v2, register_environment_v2, V2_INPUT_DATASET_PATTERN,
                                 V2_OUTPUT_DATASET_PATTERN, wait_for_job_completion)
 from health_azure.datasets import (DatasetConfig, StrOrDatasetConfig, setup_local_datasets,
-                                   _input_dataset_key, _output_dataset_key, _replace_string_datasets)
+                                   _input_dataset_key, _output_dataset_key, _replace_string_datasets,
+                                   _get_or_create_v2_data_asset)
 
 
 logger = logging.getLogger('health_azure')
@@ -704,8 +705,17 @@ def create_v2_inputs(ml_client: MLClient, input_datasets: List[DatasetConfig]) -
     inputs: Dict[str, Input] = {}
     for i, input_dataset in enumerate(input_datasets):
         input_name = f"INPUT_{i}"
-        version = input_dataset.version or 1
-        data_asset: Data = ml_client.data.get(input_dataset.name, version=str(version))
+        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:


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
@@ -461,7 +461,7 @@ def test_create_v2_data_asset(asset_name: str, version: Optional[str]) -> None:
 
     else:
         assert isinstance(data_asset, Data)
-        assert data_asset.path == f"azureml://datastores/{TEST_DATASTORE_NAME}/paths/{TEST_DATA_ASSET_NAME}/"
+        assert f"{TEST_DATASTORE_NAME}/paths/{TEST_DATA_ASSET_NAME}/" in str(data_asset.path)
         assert data_asset.type == "uri_folder"
         assert data_asset.name == TEST_DATASET_NAME
         if version:
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
@@ -28,6 +28,7 @@
 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
 from azureml._restclient.constants import RunStatus
 from azureml.core import ComputeTarget, Environment, RunConfiguration, ScriptRunConfig, Workspace
 from azureml.data.azure_storage_datastore import AzureBlobDatastore
@@ -1553,7 +1554,8 @@ def test_submit_to_azure_if_needed_with_hyperdrive(mock_sys_args: MagicMock,
 
 
 @pytest.mark.fast
-def test_create_v2_inputs() -> None:
+@pytest.mark.parametrize("already_exists", [True, False])
+def test_create_v2_inputs(already_exists: bool) -> None:
     mock_ml_client = MagicMock()
     mock_data_name = "some_arbitrary_name"
     # These values are copied from an actual Data item
@@ -1563,16 +1565,27 @@ def test_create_v2_inputs() -> None:
     mock_data_path = "azureml://subscriptions/123/resourcegroups/myrg/workspaces/myws/datastores/ds/paths/foldername/**"
     # This is normally "uri_folder", but we want to test if that value is passed through unchanged
     mock_data_type = "some_arbitrary_type"
-    mock_ml_client.data.get.return_value = Data(
+    mock_asset = Data(
         name=mock_data_name,
         version=mock_data_version,
         id=mock_data_id,
         path=mock_data_path,
         type=mock_data_type,
     )
+    if already_exists:
+        mock_ml_client.data.get.return_value = mock_asset
+    else:
+        def mock_v2_data_get(*args: Any, **kwargs: Any) -> None:
+            raise ResourceNotFoundError("dummy error")
+        mock_ml_client.data.get.side_effect = mock_v2_data_get
+        mock_ml_client.data.create_or_update.return_value = mock_asset
 
     for use_mounting in [True, False]:
-        mock_input_dataconfigs = [DatasetConfig(name="dummy_dataset", use_mounting=use_mounting)]
+        mock_input_dataconfigs = [DatasetConfig(
+            name="dummy_dataset",
+            use_mounting=use_mounting,
+            version=int(mock_data_version)
+        )]
         inputs = himl.create_v2_inputs(mock_ml_client, mock_input_dataconfigs)
         assert isinstance(inputs, Dict)
         assert len(inputs) == len(mock_input_dataconfigs)
@@ -1590,12 +1603,13 @@ def test_create_v2_inputs() -> None:
 @pytest.mark.fast
 @pytest.mark.parametrize("missing", [None, ""])
 def test_create_v2_inputs_fails(missing: Any) -> None:
-    mock_ml_client = MagicMock()
-    # For this mock, we can't use the Data class because the constructor always fills in a non-empty path
-    mock_ml_client.data.get.return_value = MagicMock(path=missing)
-    mock_input_dataconfigs = [DatasetConfig(name="dummy_dataset")]
-    with pytest.raises(ValueError, match="has no path"):
-        himl.create_v2_inputs(mock_ml_client, mock_input_dataconfigs)
+
+    with patch("health_azure.himl._get_or_create_v2_data_asset") as mock_get_or_create_v2_data_asset:
+        mock_get_or_create_v2_data_asset.return_value = MagicMock(path=missing)
+        mock_ml_client = MagicMock()
+        mock_input_dataconfigs = [DatasetConfig(name="dummy_dataset")]
+        with pytest.raises(ValueError, match="has no path"):
+            himl.create_v2_inputs(mock_ml_client, mock_input_dataconfigs)
 
 
 @pytest.mark.fast