feat: Deprecate ai-models.json in favour of anemoi.json (#247)

## Description
The default metadata name was `ai-models.json`, this PR deprecates this
in favour of `anemoi.json` with a warning logged to the console.
Logging was done as this should be a seemless change for users with old
checkpoints working fine.

Author: HCookie

src/anemoi/utils/checkpoints.py

@@ -1,4 +1,4 @@
-# (C) Copyright 2024 Anemoi contributors.
+# (C) Copyright 2024- Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
@@ -19,15 +19,19 @@
 import zipfile
 from collections.abc import Callable
 from tempfile import TemporaryDirectory
+from typing import Literal
+from typing import overload
 
 import numpy as np
 import tqdm
 
 LOG = logging.getLogger(__name__)
 
-DEFAULT_NAME = "ai-models.json"
+DEFAULT_NAME = "anemoi.json"
 DEFAULT_FOLDER = "anemoi-metadata"
 
+DEPRECATED_NAME = "ai-models.json"
+
 
 def has_metadata(path: str, *, name: str = DEFAULT_NAME) -> bool:
     """Check if a checkpoint file has a metadata file.
@@ -45,14 +49,11 @@ def has_metadata(path: str, *, name: str = DEFAULT_NAME) -> bool:
         True if the metadata file is found
     """
     with zipfile.ZipFile(path, "r") as f:
-        for b in f.namelist():
-            if os.path.basename(b) == name:
-                return True
-    return False
+        return any(os.path.basename(b) == name for b in f.namelist())
 
 
-def metadata_root(path: str, *, name: str = DEFAULT_NAME) -> str:
-    """Get the root directory of the metadata file.
+def get_metadata_path(path: str, *, name: str = DEFAULT_NAME) -> str:
+    """Get the full path of the metadata file in the checkpoint.
 
     Parameters
     ----------
@@ -64,21 +65,50 @@ def metadata_root(path: str, *, name: str = DEFAULT_NAME) -> str:
     Returns
     -------
     str
-        The root directory of the metadata file
+        The full path of the metadata file in the zip archive
 
     Raises
     ------
-    ValueError
+    FileNotFoundError
         If the metadata file is not found
+    ValueError
+        If multiple metadata files are found
     """
     with zipfile.ZipFile(path, "r") as f:
-        for b in f.namelist():
-            if os.path.basename(b) == name:
-                return os.path.dirname(b)
-    raise ValueError(f"Could not find '{name}' in {path}.")
+        metadata_file = list(filter(lambda b: os.path.basename(b) == name, f.namelist()))
+        if len(metadata_file) == 0:
+            raise FileNotFoundError(f"Could not find '{name}' in {path}.")
+        if len(metadata_file) > 1:
+            raise ValueError(f"Found two or more '{name}' in {path}.")
+        return metadata_file[0]
+
+
+def _support_metadata_name_deprecation(path: str, name: str) -> str:
+    """Support deprecated metadata name, automatically switching if needed and logging a warning."""
+    if name == DEFAULT_NAME and not has_metadata(path, name=DEFAULT_NAME):
+        if has_metadata(path, name=DEPRECATED_NAME):
+            LOG.warning(
+                "The metadata file '%s' is deprecated. New versions of checkpoints will write to '%s' instead.",
+                DEPRECATED_NAME,
+                DEFAULT_NAME,
+            )
+            name = DEPRECATED_NAME
+    return name
+
+
+# TODO: Refactor this function to reduce complexity
+@overload
+def load_metadata(path: str, *, supporting_arrays: Literal[False] = False, name: str = DEFAULT_NAME) -> dict:  # type: ignore[reportOverlappingOverload]
+    ...
+
+
+@overload
+def load_metadata(
+    path: str, *, supporting_arrays: Literal[True] = True, name: str = DEFAULT_NAME
+) -> tuple[dict, dict]: ...
 
 
-def load_metadata(path: str, *, supporting_arrays: bool = False, name: str = DEFAULT_NAME) -> dict:
+def load_metadata(path: str, *, supporting_arrays: bool = False, name: str = DEFAULT_NAME) -> dict | tuple[dict, dict]:
     """Load metadata from a checkpoint file.
 
     Parameters
@@ -102,24 +132,15 @@ def load_metadata(path: str, *, supporting_arrays: bool = False, name: str = DEF
     ValueError
         If the metadata file is not found
     """
-    with zipfile.ZipFile(path, "r") as f:
-        metadata = None
-        for b in f.namelist():
-            if os.path.basename(b) == name:
-                if metadata is not None:
-                    raise ValueError(f"Found two or more '{name}' in {path}.")
-                metadata = b
-
-    if metadata is not None:
-        with zipfile.ZipFile(path, "r") as f:
-            metadata = json.load(f.open(metadata, "r"))
-            if supporting_arrays:
-                arrays = load_supporting_arrays(f, metadata.get("supporting_arrays_paths", {}))
-                return metadata, arrays
+    name = _support_metadata_name_deprecation(path, name)
+    metadata = get_metadata_path(path, name=name)
 
-            return metadata
-    else:
-        raise ValueError(f"Could not find '{name}' in {path}.")
+    with zipfile.ZipFile(path, "r") as f:
+        metadata = json.load(f.open(metadata, "r"))
+        if supporting_arrays:
+            arrays = load_supporting_arrays(f, metadata.get("supporting_arrays_paths", {}))
+            return metadata, arrays
+        return metadata
 
 
 def load_supporting_arrays(zipf: zipfile.ZipFile, entries: dict) -> dict:
@@ -190,7 +211,12 @@ def _write_array_to_bytes(array: dict | np.ndarray, name: str, entry: dict, zipf
 
 
 def save_metadata(
-    path: str, metadata: dict, *, supporting_arrays: dict = None, name: str = DEFAULT_NAME, folder: str = DEFAULT_FOLDER
+    path: str,
+    metadata: dict,
+    *,
+    supporting_arrays: dict | None = None,
+    name: str = DEFAULT_NAME,
+    folder: str = DEFAULT_FOLDER,
 ) -> None:
     """Save metadata to a checkpoint file.
 
@@ -200,7 +226,7 @@ def save_metadata(
         The path to the checkpoint file
     metadata : dict
         A JSON serializable object
-    supporting_arrays : dict, optional
+    supporting_arrays : dict | None, optional
         A dictionary of supporting NumPy arrays
     name : str, optional
         The name of the metadata file in the zip archive
@@ -257,20 +283,14 @@ def _edit_metadata(path: str, name: str, callback: Callable, supporting_arrays:
     """
     new_path = f"{path}.anemoi-edit-{time.time()}-{os.getpid()}.tmp"
 
-    with zipfile.ZipFile(path, "r") as source_zip:
-        file_list = source_zip.namelist()
+    target_file = get_metadata_path(path, name=name)
+    if target_file is None:
+        raise FileNotFoundError(f"Could not find '{name}' in {path}")
 
-        # Find the target file and its directory
-        target_file = None
-        directory = None
-        for file_path in file_list:
-            if os.path.basename(file_path) == name:
-                target_file = file_path
-                directory = os.path.dirname(file_path)
-                break
+    directory = os.path.dirname(target_file)
 
-        if target_file is None:
-            raise ValueError(f"Could not find '{name}' in {path}")
+    with zipfile.ZipFile(path, "r") as source_zip:
+        file_list = source_zip.namelist()
 
         # Calculate total files for progress bar
         total_files = len(file_list)
@@ -313,7 +333,9 @@ def _edit_metadata(path: str, name: str, callback: Callable, supporting_arrays:
     LOG.info("Updated metadata in %s", path)
 
 
-def replace_metadata(path: str, metadata: dict, supporting_arrays: dict = None, *, name: str = DEFAULT_NAME) -> None:
+def replace_metadata(
+    path: str, metadata: dict, supporting_arrays: dict | None = None, *, name: str = DEFAULT_NAME
+) -> None:
     """Replace metadata in a checkpoint file.
 
     Parameters
@@ -337,6 +359,7 @@ def callback(full):
         with open(full, "w") as f:
             json.dump(metadata, f)
 
+    name = _support_metadata_name_deprecation(path, name)
     return _edit_metadata(path, name, callback, supporting_arrays)
 
 
@@ -350,6 +373,7 @@ def remove_metadata(path: str, *, name: str = DEFAULT_NAME) -> None:
     name : str, optional
         The name of the metadata file in the zip archive
     """
+    name = _support_metadata_name_deprecation(path, name)
 
     def callback(full):
         os.remove(full)

tests/test_checkpoints.py

@@ -6,6 +6,8 @@
 import numpy as np
 import pytest
 
+from anemoi.utils.checkpoints import DEFAULT_NAME
+from anemoi.utils.checkpoints import DEPRECATED_NAME
 from anemoi.utils.checkpoints import _edit_metadata
 from anemoi.utils.checkpoints import has_metadata
 from anemoi.utils.checkpoints import load_metadata
@@ -149,7 +151,7 @@ def test_edit_metadata_file_not_found(self, sample_checkpoint):
         def dummy_callback(file_path):
             pass
 
-        with pytest.raises(ValueError, match="Could not find 'nonexistent.json'"):
+        with pytest.raises(FileNotFoundError, match="Could not find 'nonexistent.json'"):
             _edit_metadata(sample_checkpoint, "nonexistent.json", dummy_callback)
 
     def test_edit_metadata_callback_exception_handling(self, sample_checkpoint):
@@ -244,7 +246,7 @@ def test_metadata_no_arrays(self, sample_checkpoint):
         save_metadata(sample_checkpoint, metadata, supporting_arrays=None, name="metadata.json")
 
         # Load and verify
-        loaded_metadata = load_metadata(sample_checkpoint, supporting_arrays=None, name="metadata.json")
+        loaded_metadata = load_metadata(sample_checkpoint, supporting_arrays=False, name="metadata.json")
         assert loaded_metadata["test"] is True
 
         # Edit with _edit_metadata
@@ -256,3 +258,9 @@ def update_callback(file_path):
                 json.dump(data, f)
 
         _edit_metadata(sample_checkpoint, "metadata.json", update_callback)
+
+    def test_automatic_deprecation_handling(self, sample_checkpoint: str):
+        """Test that deprecated metadata name is handled automatically."""
+        save_metadata(sample_checkpoint, {"version": "1.0", "test": "Deprecation"}, name=DEPRECATED_NAME)
+        metadata = load_metadata(sample_checkpoint, name=DEFAULT_NAME)  # Should auto-switch and log warning
+        assert metadata["test"] == "Deprecation"