🐛fix(storage): materialize hub-loaded constants and keep disk memmaps (#345)

Author: casenave

CHANGELOG.md

@@ -25,6 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixes
 
+- (storage/common/reader) clarify constants loading policy: local disk keeps numeric constants as `np.memmap` for efficiency, while Hub metadata loading materializes them to in-memory arrays to avoid temporary-directory lifetime issues.
 - (sample) get_global in the scalar case now returns the scalar with the original type.
 - (datasets) fix missing location use in get_field_names, and improve corresponding tests.
 - (cgns_helpers) update_features_for_CGNS_compatibility: fix behavior where input variable was modified by the function.

docs/source/tutorials/storage.md

@@ -23,6 +23,10 @@ Notes and requirements:
 - Parallel mode demonstrates sharding and num_proc controls; writer_batch_size / num_workers are used when uploading to the hub.
 - The tutorial includes utilities to inspect CGNS trees and to save single Plaid samples to disk for visualization (e.g., Paraview).
 
+Constants loading policy (metadata):
+- Loading metadata from **local disk** keeps numeric constants as file-backed `np.memmap` for memory efficiency on large datasets.
+- Loading metadata from the **Hub** materializes numeric constants into in-memory arrays before returning them, to avoid lifetime issues with temporary download folders.
+
 Use these examples as templates to:
 - Adapt generators to your raw data format,
 - Choose the backend that fits your workflow (hf_datasets for hub integration, cgns for native CGNS interchange, zarr for efficient chunked numeric storage),

src/plaid/storage/cgns/reader.py

@@ -132,6 +132,9 @@ def sample_generator(
             sample = Sample(
                 path=Path(temp_folder) / "data" / f"{split}" / f"sample_{idx:09d}"
             )
+            # Sample data are eagerly loaded in memory during initialization;
+            # clear the transient on-disk path before leaving the temp dir.
+            sample.path = None
         yield sample
 
 

src/plaid/storage/common/reader.py

@@ -25,6 +25,22 @@
 
 logger = logging.getLogger(__name__)
 
+
+def _materialize_memmaps(
+    flat_cst: dict[str, dict[str, Any]],
+) -> dict[str, dict[str, Any]]:
+    """Return constants with numeric memmaps materialized into in-memory arrays.
+
+    This is used for metadata loaded from ephemeral locations (e.g. temporary
+    download directories), where memmap-backed file lifetimes are not guaranteed.
+    """
+    for split, split_cst in flat_cst.items():
+        for key, value in split_cst.items():
+            if isinstance(value, np.memmap):
+                flat_cst[split][key] = np.asarray(value).copy()
+    return flat_cst
+
+
 # ------------------------------------------------------
 # Load from disk
 # ------------------------------------------------------
@@ -107,8 +123,8 @@ def load_constants_from_disk(path):
     Returns:
         tuple:
             flat_cst (dict[str, dict[str, Any]]): Mapping split -> {constant_name: numpy array | None}.
-                - Numeric constants are returned as numpy arrays with the dtype and shape specified
-                  in the schema.
+                - Numeric constants are returned as ``np.memmap`` arrays backed by
+                  ``data.mmap`` in the dataset directory.
                 - String constants are returned as 1-element numpy arrays of Python str decoded using ASCII.
                 - If layout entry for a key is None, the value is returned as None.
             constant_schema (dict[str, dict[str, Any]]): Mapping split -> loaded constant schema (from YAML).
@@ -181,7 +197,8 @@ def load_metadata_from_disk(
 
     Returns:
         tuple[dict[str, Any], dict[str, Any], dict[str, Any], dict[str, Any]]:
-            - flat_cst: constant features dictionary
+            - flat_cst: constant features dictionary (numeric constants kept as
+              file-backed ``np.memmap``)
             - variable_schema: variable schema dictionary
             - constant_schema: constant schema dictionary
             - cgns_types: CGNS types dictionary
@@ -261,7 +278,8 @@ def load_metadata_from_hub(
 
     Returns:
         tuple[dict[str, Any], dict[str, Any], dict[str, Any], dict[str, Any]]:
-            - flat_cst: constant features dictionary
+            - flat_cst: constant features dictionary (numeric constants are
+              materialized to in-memory arrays)
             - variable_schema: variable schema dictionary
             - constant_schema: constant schema dictionary
             - cgns_types: CGNS types dictionary
@@ -275,6 +293,9 @@ def load_metadata_from_hub(
             local_dir=temp_folder,
         )
         flat_cst, constant_schema = load_constants_from_disk(temp_folder)
+        # Hub metadata is downloaded under a temporary directory: materialize
+        # memmaps so returned constants remain valid after temp cleanup.
+        flat_cst = _materialize_memmaps(flat_cst)
 
     # variable_schema
     yaml_path = hf_hub_download(

tests/storage/test_storage.py

@@ -7,14 +7,17 @@
 
 # %% Imports
 
+import json
+import shutil
 from copy import deepcopy
 from functools import partial
 from pathlib import Path
 from typing import Callable
 
+import numpy as np
 import pytest
+import yaml
 
-# from plaid.bridges import huggingface_bridge
 from plaid.containers.dataset import Dataset
 from plaid.containers.sample import Sample
 from plaid.problem_definition import ProblemDefinition
@@ -28,6 +31,110 @@
 )
 
 
+def test_load_metadata_from_hub_materializes_memmaps(tmp_path, monkeypatch):
+    """Hub metadata loader must return arrays independent from temp files."""
+    from plaid.storage.common import reader as common_reader
+
+    repo_root = tmp_path / "fake_hub_repo"
+    constants_dir = repo_root / "constants" / "train"
+    constants_dir.mkdir(parents=True)
+
+    data = np.arange(6, dtype=np.float32).reshape(2, 3)
+    with open(constants_dir / "data.mmap", "wb") as f:
+        f.write(data.tobytes(order="C"))
+
+    with open(constants_dir / "layout.json", "w", encoding="utf-8") as f:
+        json.dump(
+            {
+                "Global/cst_numeric": {
+                    "offset": 0,
+                    "shape": list(data.shape),
+                    "dtype": str(data.dtype),
+                }
+            },
+            f,
+        )
+
+    with open(constants_dir / "constant_schema.yaml", "w", encoding="utf-8") as f:
+        yaml.safe_dump({"Global/cst_numeric": {"dtype": str(data.dtype), "ndim": 2}}, f)
+
+    with open(repo_root / "variable_schema.yaml", "w", encoding="utf-8") as f:
+        yaml.safe_dump({"Global/var": {"dtype": "float32", "ndim": 1}}, f)
+
+    with open(repo_root / "cgns_types.yaml", "w", encoding="utf-8") as f:
+        yaml.safe_dump({"Global": "DataArray_t"}, f)
+
+    def _fake_snapshot_download(**kwargs):
+        local_dir = Path(kwargs["local_dir"])
+        shutil.copytree(
+            repo_root / "constants", local_dir / "constants", dirs_exist_ok=True
+        )
+        return str(local_dir)
+
+    def _fake_hf_hub_download(**kwargs):
+        return str(repo_root / kwargs["filename"])
+
+    monkeypatch.setattr(common_reader, "snapshot_download", _fake_snapshot_download)
+    monkeypatch.setattr(common_reader, "hf_hub_download", _fake_hf_hub_download)
+
+    flat_cst, variable_schema, constant_schema, cgns_types = (
+        common_reader.load_metadata_from_hub("dummy/repo")
+    )
+
+    loaded = flat_cst["train"]["Global/cst_numeric"]
+    assert isinstance(loaded, np.ndarray)
+    assert not isinstance(loaded, np.memmap)
+    assert np.array_equal(loaded, data)
+    assert variable_schema["Global/var"]["dtype"] == "float32"
+    assert "Global/cst_numeric" in constant_schema["train"]
+    assert cgns_types["Global"] == "DataArray_t"
+
+
+def test_load_metadata_from_disk_keeps_memmaps(tmp_path):
+    """Local metadata loader keeps memmap-backed numeric constants."""
+    from plaid.storage.common import reader as common_reader
+
+    dataset_root = tmp_path / "dataset"
+    constants_dir = dataset_root / "constants" / "train"
+    constants_dir.mkdir(parents=True)
+
+    data = np.arange(6, dtype=np.float32).reshape(2, 3)
+    with open(constants_dir / "data.mmap", "wb") as f:
+        f.write(data.tobytes(order="C"))
+
+    with open(constants_dir / "layout.json", "w", encoding="utf-8") as f:
+        json.dump(
+            {
+                "Global/cst_numeric": {
+                    "offset": 0,
+                    "shape": list(data.shape),
+                    "dtype": str(data.dtype),
+                }
+            },
+            f,
+        )
+
+    with open(constants_dir / "constant_schema.yaml", "w", encoding="utf-8") as f:
+        yaml.safe_dump({"Global/cst_numeric": {"dtype": str(data.dtype), "ndim": 2}}, f)
+
+    with open(dataset_root / "variable_schema.yaml", "w", encoding="utf-8") as f:
+        yaml.safe_dump({"Global/var": {"dtype": "float32", "ndim": 1}}, f)
+
+    with open(dataset_root / "cgns_types.yaml", "w", encoding="utf-8") as f:
+        yaml.safe_dump({"Global": "DataArray_t"}, f)
+
+    flat_cst, variable_schema, constant_schema, cgns_types = (
+        common_reader.load_metadata_from_disk(dataset_root)
+    )
+
+    loaded = flat_cst["train"]["Global/cst_numeric"]
+    assert isinstance(loaded, np.memmap)
+    assert np.array_equal(np.asarray(loaded), data)
+    assert variable_schema["Global/var"]["dtype"] == "float32"
+    assert "Global/cst_numeric" in constant_schema["train"]
+    assert cgns_types["Global"] == "DataArray_t"
+
+
 @pytest.fixture()
 def current_directory():
     return Path(__file__).absolute().parent