Repository: pyg-team/pytorch_geometric
Pull request #8297 (2023-10-31 08:41:13)

Do not silently override configurations in `ConfigStore`; make duplicate node logic more robust

codecov[bot]: We're currently processing your upload.  This comment will be updated when the results are available.

Patch:
diff --git a/torch_geometric/config_store.py b/torch_geometric/config_store.py
index ea70c8290..6d15a809c 100644
--- a/torch_geometric/config_store.py
+++ b/torch_geometric/config_store.py
@@ -20,26 +20,34 @@ except ImportError:
 
 try:
     from hydra.core.config_store import ConfigStore
+    WITH_HYDRA = True
 
     def get_node(cls: Union[str, Any]) -> Optional[Any]:
         if (not isinstance(cls, str)
                 and cls.__module__ in {'builtins', 'typing'}):
             return None
 
-        def _recursive_get_node(repo: Dict[str, Any]) -> Optional[Any]:
+        def _get_candidates(repo: Dict[str, Any]) -> List[Any]:
+            outs: List[Any] = []
             for key, value in repo.items():
                 if isinstance(value, dict):
-                    out = _recursive_get_node(value)
-                    if out is not None:
-                        return out
-                elif isinstance(cls, str) and key == f'{cls}.yaml':
-                    return value.node
+                    outs.extend(_get_candidates(value))
                 elif getattr(value.node._metadata, 'object_type', None) == cls:
-                    return value.node
+                    outs.append(value.node)
                 elif getattr(value.node._metadata, 'orig_type', None) == cls:
-                    return value.node
+                    outs.append(value.node)
+                elif isinstance(cls, str) and key == f'{cls}.yaml':
+                    outs.append(value.node)
+
+            return outs
+
+        candidates = _get_candidates(get_config_store().repo)
+
+        if len(candidates) > 1:
+            raise ValueError(f"Found multiple entries in the configuration "
+                             f"store for the same node '{candidates[0].name}'")
 
-        return _recursive_get_node(get_config_store().repo)
+        return candidates[0] if len(candidates) == 1 else None
 
     def dataclass_from_class(cls: Union[str, Any]) -> Optional[Any]:
         r"""Returns the :obj:`dataclass` of a class registered in the global
@@ -55,6 +63,8 @@ try:
 
 except ImportError:
 
+    WITH_HYDRA = False
+
     class Singleton(type):
         _instances: Dict[type, Any] = {}
 
@@ -84,31 +94,48 @@ except ImportError:
         def instance(cls, *args, **kwargs) -> 'ConfigStore':
             return cls(*args, **kwargs)
 
-        def store(self, name: str, node: Any, group: Optional[str] = None):
+        def store(
+            self,
+            name: str,
+            node: Any,
+            group: Optional[str] = None,
+            orig_type: Optional[Any] = None,
+        ):
             cur = self.repo
             if group is not None:
                 cur = cur[group]
-            cur[name] = ConfigNode(name, node, group)
+            if name in cur:
+                raise KeyError(f"Configuration '{name}' already registered. "
+                               f"Please store it under a different group.")
+            metadata = Metadata(orig_type=orig_type)
+            cur[name] = ConfigNode(name, node, group, metadata)
 
     def get_node(cls: Union[str, Any]) -> Optional[ConfigNode]:
         if (not isinstance(cls, str)
                 and cls.__module__ in {'builtins', 'typing'}):
             return None
 
-        def _recursive_get_node(repo: Dict[str, Any]) -> Optional[ConfigNode]:
+        def _get_candidates(repo: Dict[str, Any]) -> List[ConfigNode]:
+            outs: List[ConfigNode] = []
             for key, value in repo.items():
                 if isinstance(value, dict):
-                    out = _recursive_get_node(value)
-                    if out is not None:
-                        return out
-                elif isinstance(cls, str) and key == cls:
-                    return value
+                    outs.extend(_get_candidates(value))
                 elif value.node == cls:
-                    return value
+                    outs.append(value)
                 elif value._metadata.orig_type == cls:
-                    return value
+                    outs.append(value)
+                elif isinstance(cls, str) and key == cls:
+                    outs.append(value)
+
+            return outs
 
-        return _recursive_get_node(get_config_store().repo)
+        candidates = _get_candidates(get_config_store().repo)
+
+        if len(candidates) > 1:
+            raise ValueError(f"Found multiple entries in the configuration "
+                             f"store for the same node '{candidates[0].name}'")
+
+        return candidates[0] if len(candidates) == 1 else None
 
     def dataclass_from_class(cls: Union[str, Any]) -> Optional[Any]:
         r"""Returns the :obj:`dataclass` of a class registered in the global
@@ -268,6 +295,15 @@ def get_config_store() -> ConfigStore:
     return ConfigStore.instance()
 
 
+def clear_config_store() -> ConfigStore:
+    r"""Clears the global configuration store."""
+    config_store = get_config_store()
+    for key in list(config_store.repo.keys()):
+        if key != 'hydra' and not key.endswith('.yaml'):
+            del config_store.repo[key]
+    return config_store
+
+
 def register(
     cls: Optional[Any] = None,
     data_cls: Optional[Any] = None,
@@ -281,21 +317,32 @@ def register(
             return a decorator. (default: :obj:`None`)
         data_cls (Any, optional): The data class to register. If set to
             :obj:`None`, will dynamically create the data class according to
-            :class:`~torch_geometric.graphgym.config_store.to_dataclass`.
+            :class:`~torch_geometric.config_store.to_dataclass`.
             (default: :obj:`None`)
         group (str, optional): The group in the global configuration store.
             (default: :obj:`None`)
         **kwargs (optional): Additional arguments of
-            :class:`~torch_geometric.graphgym.config_store.to_dataclass`.
+            :class:`~torch_geometric.config_store.to_dataclass`.
     """
     if cls is not None:
         name = cls.__name__
 
+        if get_node(cls):
+            raise ValueError(f"The class '{name}' is already registered in "
+                             "the global configuration store")
+
         if data_cls is None:
             data_cls = to_dataclass(cls, **kwargs)
-
-        get_config_store().store(name, data_cls, group)
-        get_node(name)._metadata.orig_type = cls
+        elif get_node(data_cls):
+            raise ValueError(
+                f"The data class '{data_cls.__name__}' is already registered "
+                f"in the global configuration store")
+
+        if WITH_HYDRA:
+            get_config_store().store(name, data_cls, group)
+            get_node(name)._metadata.orig_type = cls
+        else:
+            get_config_store().store(name, data_cls, group, cls)
 
         return data_cls
 


Test patch:
diff --git a/test/test_config_store.py b/test/test_config_store.py
index e74933f85..718207667 100644
--- a/test/test_config_store.py
+++ b/test/test_config_store.py
@@ -2,6 +2,7 @@ from typing import Any, Optional
 
 from torch_geometric.config_store import (
     class_from_dataclass,
+    clear_config_store,
     dataclass_from_class,
     fill_config_store,
     get_config_store,
@@ -9,6 +10,11 @@ from torch_geometric.config_store import (
     to_dataclass,
 )
 from torch_geometric.testing import withPackage
+from torch_geometric.transforms import AddSelfLoops
+
+
+def teardown_function():
+    clear_config_store()
 
 
 def test_to_dataclass():
@@ -39,10 +45,8 @@ def test_to_dataclass():
 
 
 def test_register():
-    from torch_geometric.transforms import AddSelfLoops
-
-    register(AddSelfLoops, group='transforms')
-    assert 'transforms' in get_config_store().repo
+    register(AddSelfLoops, group='transform')
+    assert 'transform' in get_config_store().repo
 
     AddSelfLoopsConfig = dataclass_from_class('AddSelfLoops')
 
@@ -57,6 +61,18 @@ def test_register():
     assert ConfigCls == AddSelfLoopsConfig
 
 
+def test_fill_config_store():
+    fill_config_store()
+
+    assert {
+        'transform',
+        'dataset',
+        'model',
+        'optimizer',
+        'lr_scheduler',
+    }.issubset(get_config_store().repo.keys())
+
+
 @withPackage('hydra')
 def test_hydra_config_store():
     import hydra