apply PluginVariant and make_plugin_registry to classes (#3346)

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

deepmd/backend/backend.py

@@ -16,8 +16,8 @@
 )
 
 from deepmd.utils.plugin import (
-    Plugin,
     PluginVariant,
+    make_plugin_registry,
 )
 
 if TYPE_CHECKING:
@@ -33,7 +33,7 @@
     )
 
 
-class Backend(PluginVariant):
+class Backend(PluginVariant, make_plugin_registry("backend")):
     r"""General backend class.
 
     Examples
@@ -44,24 +44,6 @@ class Backend(PluginVariant):
     ...     pass
     """
 
-    __plugins = Plugin()
-
-    @staticmethod
-    def register(key: str) -> Callable[[object], object]:
-        """Register a backend plugin.
-
-        Parameters
-        ----------
-        key : str
-            the key of a backend
-
-        Returns
-        -------
-        Callable[[object], object]
-            the decorator to register backend
-        """
-        return Backend.__plugins.register(key.lower())
-
     @staticmethod
     def get_backend(key: str) -> Type["Backend"]:
         """Get the backend by key.
@@ -76,12 +58,7 @@ def get_backend(key: str) -> Type["Backend"]:
         Backend
             the backend
         """
-        try:
-            backend = Backend.__plugins.get_plugin(key.lower())
-        except KeyError:
-            raise KeyError(f"Backend {key} is not registered.")
-        assert isinstance(backend, type)
-        return backend
+        return Backend.get_class_by_type(key)
 
     @staticmethod
     def get_backends() -> Dict[str, Type["Backend"]]:
@@ -92,7 +69,7 @@ def get_backends() -> Dict[str, Type["Backend"]]:
         list
             all the registered backends
         """
-        return Backend.__plugins.plugins
+        return Backend.get_plugins()
 
     @staticmethod
     def get_backends_by_feature(
@@ -112,7 +89,7 @@ def get_backends_by_feature(
         """
         return {
             key: backend
-            for key, backend in Backend.__plugins.plugins.items()
+            for key, backend in Backend.get_backends().items()
             if backend.features & feature
         }
 

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -4,10 +4,8 @@
     abstractmethod,
 )
 from typing import (
-    Callable,
     List,
     Optional,
-    Type,
 )
 
 from deepmd.common import (
@@ -17,7 +15,8 @@
     DPPath,
 )
 from deepmd.utils.plugin import (
-    Plugin,
+    PluginVariant,
+    make_plugin_registry,
 )
 
 
@@ -37,45 +36,14 @@ def make_base_descriptor(
 
     """
 
-    class BD(ABC):
+    class BD(ABC, PluginVariant, make_plugin_registry("descriptor")):
         """Base descriptor provides the interfaces of descriptor."""
 
-        __plugins = Plugin()
-
-        @staticmethod
-        def register(key: str) -> Callable:
-            """Register a descriptor plugin.
-
-            Parameters
-            ----------
-            key : str
-                the key of a descriptor
-
-            Returns
-            -------
-            Descriptor
-                the registered descriptor
-
-            Examples
-            --------
-            >>> @Descriptor.register("some_descrpt")
-                class SomeDescript(Descriptor):
-                    pass
-            """
-            return BD.__plugins.register(key)
-
         def __new__(cls, *args, **kwargs):
             if cls is BD:
                 cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__))
             return super().__new__(cls)
 
-        @classmethod
-        def get_class_by_type(cls, descrpt_type: str) -> Type["BD"]:
-            if descrpt_type in BD.__plugins.plugins:
-                return BD.__plugins.plugins[descrpt_type]
-            else:
-                raise RuntimeError("Unknown descriptor type: " + descrpt_type)
-
         @abstractmethod
         def get_rcut(self) -> float:
             """Returns the cut-off radius."""

deepmd/dpmodel/fitting/make_base_fitting.py

@@ -4,10 +4,8 @@
     abstractmethod,
 )
 from typing import (
-    Callable,
     Dict,
     Optional,
-    Type,
 )
 
 from deepmd.common import (
@@ -17,7 +15,8 @@
     FittingOutputDef,
 )
 from deepmd.utils.plugin import (
-    Plugin,
+    PluginVariant,
+    make_plugin_registry,
 )
 
 
@@ -37,45 +36,14 @@ def make_base_fitting(
 
     """
 
-    class BF(ABC):
+    class BF(ABC, PluginVariant, make_plugin_registry("fitting")):
         """Base fitting provides the interfaces of fitting net."""
 
-        __plugins = Plugin()
-
-        @staticmethod
-        def register(key: str) -> Callable[[object], object]:
-            """Register a descriptor plugin.
-
-            Parameters
-            ----------
-            key : str
-                the key of a descriptor
-
-            Returns
-            -------
-            callable[[object], object]
-                the registered descriptor
-
-            Examples
-            --------
-            >>> @Fitting.register("some_fitting")
-                class SomeFitting(Fitting):
-                    pass
-            """
-            return BF.__plugins.register(key)
-
         def __new__(cls, *args, **kwargs):
             if cls is BF:
                 cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__))
             return super().__new__(cls)
 
-        @classmethod
-        def get_class_by_type(cls, fitting_type: str) -> Type["BF"]:
-            if fitting_type in BF.__plugins.plugins:
-                return BF.__plugins.plugins[fitting_type]
-            else:
-                raise RuntimeError("Unknown fitting type: " + fitting_type)
-
         @abstractmethod
         def output_def(self) -> FittingOutputDef:
             """Returns the output def of the fitting net."""

deepmd/dpmodel/model/base_model.py

@@ -11,12 +11,13 @@
 )
 
 from deepmd.utils.plugin import (
+    PluginVariant,
     make_plugin_registry,
 )
 
 
 def make_base_model() -> Type[object]:
-    class BaseBaseModel(ABC, make_plugin_registry("model")):
+    class BaseBaseModel(ABC, PluginVariant, make_plugin_registry("model")):
         """Base class for final exported model that will be directly used for inference.
 
         The class defines some abstractmethods that will be directly called by the

deepmd/pt/model/descriptor/descriptor.py

@@ -5,7 +5,6 @@
     abstractmethod,
 )
 from typing import (
-    Callable,
     Dict,
     List,
     Optional,
@@ -22,60 +21,34 @@
 from deepmd.pt.utils.env_mat_stat import (
     EnvMatStatSe,
 )
-from deepmd.pt.utils.plugin import (
-    Plugin,
-)
 from deepmd.utils.env_mat_stat import (
     StatItem,
 )
 from deepmd.utils.path import (
     DPPath,
 )
+from deepmd.utils.plugin import (
+    make_plugin_registry,
+)
 
 log = logging.getLogger(__name__)
 
 
-class DescriptorBlock(torch.nn.Module, ABC):
+class DescriptorBlock(torch.nn.Module, ABC, make_plugin_registry("DescriptorBlock")):
     """The building block of descriptor.
     Given the input descriptor, provide with the atomic coordinates,
     atomic types and neighbor list, calculate the new descriptor.
     """
 
-    __plugins = Plugin()
     local_cluster = False
 
-    @staticmethod
-    def register(key: str) -> Callable:
-        """Register a DescriptorBlock plugin.
-
-        Parameters
-        ----------
-        key : str
-            the key of a DescriptorBlock
-
-        Returns
-        -------
-        DescriptorBlock
-            the registered DescriptorBlock
-
-        Examples
-        --------
-        >>> @DescriptorBlock.register("some_descrpt")
-            class SomeDescript(DescriptorBlock):
-                pass
-        """
-        return DescriptorBlock.__plugins.register(key)
-
     def __new__(cls, *args, **kwargs):
         if cls is DescriptorBlock:
             try:
                 descrpt_type = kwargs["type"]
             except KeyError:
                 raise KeyError("the type of DescriptorBlock should be set by `type`")
-            if descrpt_type in DescriptorBlock.__plugins.plugins:
-                cls = DescriptorBlock.__plugins.plugins[descrpt_type]
-            else:
-                raise RuntimeError("Unknown DescriptorBlock type: " + descrpt_type)
+            cls = cls.get_class_by_type(descrpt_type)
         return super().__new__(cls)
 
     @abstractmethod

deepmd/tf/descriptor/descriptor.py

@@ -4,7 +4,6 @@
 )
 from typing import (
     Any,
-    Callable,
     Dict,
     List,
     Optional,
@@ -21,12 +20,14 @@
     tf,
 )
 from deepmd.tf.utils import (
-    Plugin,
     PluginVariant,
 )
+from deepmd.utils.plugin import (
+    make_plugin_registry,
+)
 
 
-class Descriptor(PluginVariant):
+class Descriptor(PluginVariant, make_plugin_registry("descriptor")):
     r"""The abstract class for descriptors. All specific descriptors should
     be based on this class.
 
@@ -45,37 +46,6 @@ class Descriptor(PluginVariant):
     that can be called by other classes.
     """
 
-    __plugins = Plugin()
-
-    @staticmethod
-    def register(key: str) -> Callable:
-        """Register a descriptor plugin.
-
-        Parameters
-        ----------
-        key : str
-            the key of a descriptor
-
-        Returns
-        -------
-        Descriptor
-            the registered descriptor
-
-        Examples
-        --------
-        >>> @Descriptor.register("some_descrpt")
-            class SomeDescript(Descriptor):
-                pass
-        """
-        return Descriptor.__plugins.register(key)
-
-    @classmethod
-    def get_class_by_type(cls, descrpt_type: str):
-        if descrpt_type in Descriptor.__plugins.plugins:
-            return Descriptor.__plugins.plugins[descrpt_type]
-        else:
-            raise RuntimeError("Unknown descriptor type: " + descrpt_type)
-
     def __new__(cls, *args, **kwargs):
         if cls is Descriptor:
             cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__))