Repository: mlcommons/medperf (python)

Issue:
[BUG] MLCube assets are not updated after content or url change
**Describe the bug**
MLCubes are constructed from multiple assets, all of which are downloaded the first time such MLCube is requested by the CLI. After subsequent requests to the MLCube, cached assets are preferred. This is okay as long as the contents of the remote MLCube don't change. But, if assets are updated, contents are modified or the hash changes, while the MLCube metadata gets updated, assets are not redownloaded, or rechecked with the new hash. This leads to people running outdated MLCubes without knowledge, and can also lead to broken MLCube registration updates, because tests on the user side will pass due to local cache being valid instead of remote contents being revised.

**To Reproduce**
Steps to reproduce the behavior:
1. Submit an MLCube and run it
2. Change the URL of one of the assets (MLCube URL for example) to something nonsensical through the admin dashboard
3. Run the MLCube again. It runs as usual
4. Change the hash of one of the assets (parameters hash for example) to something nonsensical through the admin dashboard
5. Run the MLCube again. It still runs like normal
6. Remove mepderf's cache of that MLCube
7. Run the MLCube again. It now fails.

**Expected behavior**
When the asset url or hash changes, the local cache should be invalidated or removed before proceeding with the download.

**Additional context**
This raises a question. We can detect asset change by hash change or url change. What happens when the contents of the asset change, but not the url? We cannot detect content changes without downloading the asset again. Someone could submit an MLCube with, for example, a url for the latest mlcube.yaml in a repository. In that case, contents may change, but the link is the same. Users that have downloaded that MLCube previously will not have any issues later on because local cache is used for that asset, but new users of that MLCube will encounter issues because of hash mismatch. How can we proceed here? Should we prefer redownloading lightweight assets like yaml files? What about heavy assets like additional files?


Pull request:
Fix caching bug in resources.download
Details of the problem can be found in the issue this PR fixes.
Closes #549 

Patch:
diff --git a/cli/medperf/comms/entity_resources/resources.py b/cli/medperf/comms/entity_resources/resources.py
index c826ace..09dc7c0 100644
--- a/cli/medperf/comms/entity_resources/resources.py
+++ b/cli/medperf/comms/entity_resources/resources.py
@@ -14,49 +14,85 @@ will not be re-downloaded.
 
 import shutil
 import os
+import logging
+import yaml
 import medperf.config as config
-from medperf.utils import generate_tmp_path, get_cube_image_name, remove_path, untar
+from medperf.utils import (
+    generate_tmp_path,
+    get_cube_image_name,
+    remove_path,
+    untar,
+    get_file_hash,
+)
 from .utils import download_resource
 
 
-def get_cube(url: str, cube_path: str, expected_hash: str = None) -> str:
-    """Downloads and writes an mlcube.yaml file. If the hash is provided,
+def _should_get_regular_file(output_path, expected_hash):
+    if os.path.exists(output_path) and expected_hash:
+        calculated_hash = get_file_hash(output_path)
+        logging.debug(
+            f"{output_path}: Expected {expected_hash}, found {calculated_hash}."
+        )
+        if expected_hash == calculated_hash:
+            logging.debug(f"{output_path} exists and is up to date")
+            return False
+        logging.debug(f"{output_path} exists but is out of date")
+    return True
+
+
+def _should_get_cube_additional(
+    additional_files_folder, expected_tarball_hash, mlcube_local_cache_metadata
+):
+    # Read the hash of the tarball file whose extracted contents may already exist
+    additional_files_cached_hash = None
+    if os.path.exists(mlcube_local_cache_metadata):
+        with open(mlcube_local_cache_metadata) as f:
+            contents = yaml.safe_load(f)
+        additional_files_cached_hash = contents.get(
+            "additional_files_cached_hash", None
+        )
+
+    # Return if the additional files already exist and seem to originate from a valid tarball
+    if os.path.exists(additional_files_folder) and expected_tarball_hash:
+        if additional_files_cached_hash == expected_tarball_hash:
+            logging.debug(f"{additional_files_folder} exists and is up to date")
+            return False
+        logging.debug(f"{additional_files_folder} exists but is out of date")
+    return True
+
+
+def _get_regular_file(url: str, output_path: str, expected_hash: str = None) -> str:
+    """Downloads and writes a regular file. If the hash is provided,
     the file's integrity will be checked upon download.
+    Used for parameters.yaml and mlcube.yaml
 
     Args:
         url (str): URL where the mlcube.yaml file can be downloaded.
-        cube_path (str): Cube location.
+        output_path (str): Path to store the downloaded file at.
         expected_hash (str, optional): expected hash of the downloaded file
 
     Returns:
-        output_path (str): location where the mlcube.yaml file is stored locally.
+        output_path (str): location where the regular file is stored locally.
         hash_value (str): The hash of the downloaded file
     """
-    output_path = os.path.join(cube_path, config.cube_filename)
-    if os.path.exists(output_path):
+    if not _should_get_regular_file(output_path, expected_hash):
         return output_path, expected_hash
+    if os.path.exists(output_path):
+        remove_path(output_path)
     hash_value = download_resource(url, output_path, expected_hash)
     return output_path, hash_value
 
 
-def get_cube_params(url: str, cube_path: str, expected_hash: str = None) -> str:
-    """Downloads and writes a cube parameters file. If the hash is provided,
-    the file's integrity will be checked upon download.
+def get_cube(url: str, cube_path: str, expected_hash: str = None):
+    """Downloads and writes a cube mlcube.yaml file"""
+    output_path = os.path.join(cube_path, config.cube_filename)
+    return _get_regular_file(url, output_path, expected_hash)
 
-    Args:
-        url (str): URL where the parameters.yaml file can be downloaded.
-        cube_path (str): Cube location.
-        expected_hash (str, optional): expected hash of the downloaded file
 
-    Returns:
-        output_path (str): location where the parameters file is stored locally.
-        hash_value (str): The hash of the downloaded file
-    """
+def get_cube_params(url: str, cube_path: str, expected_hash: str = None):
+    """Downloads and writes a cube parameters.yaml file"""
     output_path = os.path.join(cube_path, config.workspace_path, config.params_filename)
-    if os.path.exists(output_path):
-        return output_path, expected_hash
-    hash_value = download_resource(url, output_path, expected_hash)
-    return output_path, hash_value
+    return _get_regular_file(url, output_path, expected_hash)
 
 
 def get_cube_image(url: str, cube_path: str, hash_value: str = None) -> str:
@@ -118,12 +154,14 @@ def get_cube_additional(
         tarball_hash (str): The hash of the downloaded tarball file
     """
     additional_files_folder = os.path.join(cube_path, config.additional_path)
-
-    if os.path.exists(additional_files_folder):
+    mlcube_cache_file = os.path.join(cube_path, config.mlcube_cache_file)
+    if not _should_get_cube_additional(
+        additional_files_folder, expected_tarball_hash, mlcube_cache_file
+    ):
         return expected_tarball_hash
 
-    # make sure files are uncompressed while in tmp storage, to avoid any clutter
-    # objects if uncompression fails for some reason.
+    # Download the additional files. Make sure files are extracted in tmp storage
+    # to avoid any clutter objects if uncompression fails for some reason.
     tmp_output_folder = generate_tmp_path()
     output_tarball_path = os.path.join(tmp_output_folder, config.tarball_filename)
     tarball_hash = download_resource(url, output_tarball_path, expected_tarball_hash)
@@ -131,8 +169,17 @@ def get_cube_additional(
     untar(output_tarball_path)
     parent_folder = os.path.dirname(os.path.normpath(additional_files_folder))
     os.makedirs(parent_folder, exist_ok=True)
+    if os.path.exists(additional_files_folder):
+        # handle the possibility of having clutter uncompressed files
+        remove_path(additional_files_folder)
     os.rename(tmp_output_folder, additional_files_folder)
 
+    # Store the downloaded tarball hash to be used later for verifying that the
+    # local cache is up to date
+    with open(mlcube_cache_file, "w") as f:  # assumes parent folder already exists
+        contents = {"additional_files_cached_hash": tarball_hash}
+        yaml.dump(contents, f)
+
     return tarball_hash
 
 
diff --git a/cli/medperf/comms/entity_resources/utils.py b/cli/medperf/comms/entity_resources/utils.py
index c4c74af..bccffc1 100644
--- a/cli/medperf/comms/entity_resources/utils.py
+++ b/cli/medperf/comms/entity_resources/utils.py
@@ -1,8 +1,9 @@
 import os
+import logging
 from typing import Optional
-from medperf.utils import generate_tmp_path, get_file_hash, verify_hash
+from medperf.utils import generate_tmp_path, get_file_hash
 from .sources import supported_sources
-from medperf.exceptions import InvalidArgumentError
+from medperf.exceptions import InvalidArgumentError, InvalidEntityError
 
 
 def __parse_resource(resource: str):
@@ -51,22 +52,6 @@ def tmp_download_resource(resource):
     return tmp_output_path
 
 
-def verify_or_get_hash(tmp_output_path: str, expected_hash: str) -> str:
-    """Checks if the downloaded file matches the passed expected hash
-    if provided. The function returns the calculated hash.
-
-    Args:
-        tmp_output_path (str): path to the asset that will be checked
-        expected_hash (str): expected hash of the asset
-
-    Returns:
-        str: Calculated hash of the asset. Only returns if the hashes match
-    """
-    calculated_hash = get_file_hash(tmp_output_path)
-    verify_hash(calculated_hash, expected_hash)
-    return calculated_hash
-
-
 def to_permanent_path(tmp_output_path, output_path):
     """Writes a file from the temporary storage to the desired output path."""
     output_folder = os.path.dirname(os.path.abspath(output_path))
@@ -95,7 +80,11 @@ def download_resource(
     """
     tmp_output_path = tmp_download_resource(resource)
 
-    calculated_hash = verify_or_get_hash(tmp_output_path, expected_hash)
+    calculated_hash = get_file_hash(tmp_output_path)
+
+    if expected_hash and calculated_hash != expected_hash:
+        logging.debug(f"{resource}: Expected {expected_hash}, found {calculated_hash}.")
+        raise InvalidEntityError(f"Hash mismatch: {resource}")
 
     to_permanent_path(tmp_output_path, output_path)
 
diff --git a/cli/medperf/config.py b/cli/medperf/config.py
index 0391e04..892649b 100644
--- a/cli/medperf/config.py
+++ b/cli/medperf/config.py
@@ -140,6 +140,7 @@ cube_metadata_filename = "mlcube-meta.yaml"
 log_file = "medperf.log"
 tarball_filename = "tmp.tar.gz"
 demo_dset_paths_file = "paths.yaml"
+mlcube_cache_file = ".cache_metadata.yaml"
 
 # MLCube assets conventions
 cube_filename = "mlcube.yaml"
@@ -185,7 +186,7 @@ inline_parameters = [
     "platform",
     "gpus",
     "cleanup",
-    "container_loglevel"
+    "container_loglevel",
 ]
 configurable_parameters = inline_parameters + [
     "server",
diff --git a/cli/medperf/entities/cube.py b/cli/medperf/entities/cube.py
index 6e6f656..60f6b76 100644
--- a/cli/medperf/entities/cube.py
+++ b/cli/medperf/entities/cube.py
@@ -9,7 +9,6 @@ from medperf.utils import (
     combine_proc_sp_text,
     log_storage,
     remove_path,
-    verify_hash,
     generate_tmp_path,
     spawn_and_kill,
 )
@@ -257,7 +256,10 @@ class Cube(Entity, Uploadable, MedperfSchema, DeployableSchema):
             mlcube_details = yaml.safe_load(f)
         remove_path(tmp_out_yaml)
         local_hash = mlcube_details["hash"]
-        verify_hash(local_hash, self.image_hash)
+        if self.image_hash and local_hash != self.image_hash:
+            raise InvalidEntityError(
+                f"Hash mismatch. Expected {self.image_hash}, found {local_hash}."
+            )
         self.image_hash = local_hash
 
     def _get_image_from_registry(self):
diff --git a/cli/medperf/utils.py b/cli/medperf/utils.py
index c054d64..8068c62 100644
--- a/cli/medperf/utils.py
+++ b/cli/medperf/utils.py
@@ -21,7 +21,7 @@ from colorama import Fore, Style
 from pexpect.exceptions import TIMEOUT
 from git import Repo, GitCommandError
 import medperf.config as config
-from medperf.exceptions import ExecutionError, MedperfException, InvalidEntityError
+from medperf.exceptions import ExecutionError, MedperfException
 
 
 def get_file_hash(path: str) -> str:
@@ -402,22 +402,6 @@ def get_cube_image_name(cube_path: str) -> str:
         raise MedperfException(msg)
 
 
-def verify_hash(obtained_hash: str, expected_hash: str):
-    """Checks hash exact match, and throws an error if not a match
-
-    Args:
-        obtained_hash (str): local hash computed from asset
-        expected_hash (str): expected hash obtained externally
-
-    Raises:
-        InvalidEntityError: Thrown if hashes don't match
-    """
-    if expected_hash and expected_hash != obtained_hash:
-        raise InvalidEntityError(
-            f"Hash mismatch. Expected {expected_hash}, found {obtained_hash}."
-        )
-
-
 def filter_latest_associations(associations, entity_key):
     """Given a list of entity-benchmark associations, this function
     retrieves a list containing the latest association of each


Test patch:
diff --git a/cli/medperf/tests/comms/entity_resources/test_resources.py b/cli/medperf/tests/comms/entity_resources/test_resources.py
index 776c737..437e518 100644
--- a/cli/medperf/tests/comms/entity_resources/test_resources.py
+++ b/cli/medperf/tests/comms/entity_resources/test_resources.py
@@ -69,15 +69,112 @@ class TestGetAdditionalFiles:
     def class_setup(self, mocker):
         mocker.patch.object(resources, "untar")
 
-    def test_get_additional_files_does_not_download_if_folder_exists(self, mocker, fs):
+    def test_get_additional_files_does_not_download_if_folder_exists_and_hash_valid(
+        self, mocker, fs
+    ):
         # Arrange
         cube_path = "cube/1"
         additional_files_folder = os.path.join(cube_path, config.additional_path)
         fs.create_dir(additional_files_folder)
         spy = mocker.spy(resources, "download_resource")
+        exp_hash = resources.get_cube_additional(url, cube_path)
 
         # Act
+        resources.get_cube_additional(url, cube_path, exp_hash)
+
+        # Assert
+        spy.assert_called_once()  # second time shouldn't download
+
+    def test_get_additional_files_will_download_if_folder_exists_and_hash_outdated(
+        self, mocker, fs
+    ):
+        # Arrange
+        cube_path = "cube/1"
+        additional_files_folder = os.path.join(cube_path, config.additional_path)
+        fs.create_dir(additional_files_folder)
+        spy = mocker.spy(resources, "download_resource")
         resources.get_cube_additional(url, cube_path)
 
+        # Act
+        resources.get_cube_additional(url, cube_path, "incorrect hash")
+
         # Assert
-        spy.assert_not_called()
+        assert spy.call_count == 2
+
+    def test_get_additional_files_will_download_if_folder_exists_and_hash_valid_but_no_cached_hash(
+        self, mocker, fs
+    ):  # a test for existing installation before this feature
+        # Arrange
+        cube_path = "cube/1"
+        additional_files_folder = os.path.join(cube_path, config.additional_path)
+        fs.create_dir(additional_files_folder)
+        spy = mocker.spy(resources, "download_resource")
+        exp_hash = resources.get_cube_additional(url, cube_path)
+        hash_cache_file = os.path.join(cube_path, config.mlcube_cache_file)
+        os.remove(hash_cache_file)
+
+        # Act
+        resources.get_cube_additional(url, cube_path, exp_hash)
+
+        # Assert
+        assert spy.call_count == 2
+
+
+class TestGetCube:
+    def test_get_cube_does_not_download_if_folder_exists_and_hash_valid(
+        self, mocker, fs
+    ):
+        # Arrange
+        cube_path = "cube/1"
+        spy = mocker.spy(resources, "download_resource")
+        _, exp_hash = resources.get_cube(url, cube_path)
+
+        # Act
+        resources.get_cube(url, cube_path, exp_hash)
+
+        # Assert
+        spy.assert_called_once()  # second time shouldn't download
+
+    def test_get_cube_does_will_download_if_folder_exists_and_hash_outdated(
+        self, mocker, fs
+    ):
+        # Arrange
+        cube_path = "cube/1"
+        spy = mocker.spy(resources, "download_resource")
+        resources.get_cube(url, cube_path)
+
+        # Act
+        resources.get_cube(url, cube_path, "incorrect hash")
+
+        # Assert
+        assert spy.call_count == 2
+
+
+class TestGetCubeParams:
+    def test_get_cube_params_does_not_download_if_folder_exists_and_hash_valid(
+        self, mocker, fs
+    ):
+        # Arrange
+        cube_path = "cube/1"
+        spy = mocker.spy(resources, "download_resource")
+        _, exp_hash = resources.get_cube_params(url, cube_path)
+
+        # Act
+        resources.get_cube_params(url, cube_path, exp_hash)
+
+        # Assert
+        spy.assert_called_once()  # second time shouldn't download
+
+    def test_get_cube_params_does_will_download_if_folder_exists_and_hash_outdated(
+        self, mocker, fs
+    ):
+        # Arrange
+        cube_path = "cube/1"
+        spy = mocker.spy(resources, "download_resource")
+        resources.get_cube_params(url, cube_path)
+
+        # Act
+        resources.get_cube_params(url, cube_path, "incorrect hash")
+
+        # Assert
+        assert spy.call_count == 2