[v1] refactor registry plugin structure and params (#10641)

src/llamafactory/v1/utils/plugin.py

@@ -12,96 +12,98 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Lightweight plugin routing and shared parameter parsing helpers."""
 
-from collections import defaultdict
-from collections.abc import Callable
-from typing import Any
+from __future__ import annotations
+
+from dataclasses import fields, is_dataclass
+from typing import Any, TypeVar
 
 from . import logging
 
 
 logger = logging.get_logger(__name__)
+ParamsT = TypeVar("ParamsT")
+
+
+def ensure_methods_implemented(cls: type) -> None:
+    """Raise when a static method-group implementation is incomplete."""
+    required: set[str] = set()
+    for base in cls.__mro__[1:]:
+        required |= getattr(base, "__abstractmethods__", frozenset())
+
+    missing = sorted(name for name in required if getattr(getattr(cls, name, None), "__isabstractmethod__", False))
+    if missing:
+        raise TypeError(f"{cls.__name__} does not implement all required methods: {missing}")
 
 
 class BasePlugin:
-    """Base class for plugins.
+    """Route a plugin name to one function or static method-group class.
 
-    A plugin is a callable object that can be registered and called by name.
-
-    Example usage:
-    ```python
-    class PrintPlugin(BasePlugin):
-        def again(self):  # optional
-            self["again"]()
-
-
-    @PrintPlugin("hello").register()
-    def print_hello():
-        print("Hello world!")
-
-
-    @PrintPlugin("hello").register("again")
-    def print_hello_again():
-        print("Hello world! Again.")
-
-
-    PrintPlugin("hello")()
-    PrintPlugin("hello").again()
-    ```
+    Every plugin family subclass owns an isolated registry. Parameter schemas
+    deliberately do not live here; each plugin entrypoint parses its own config.
     """
 
-    _registry: dict[str, dict[str, Callable]] = defaultdict(dict)
+    _registry: dict[str, Any] = {}
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        cls._registry = {}
 
     def __init__(self, name: str | None = None) -> None:
-        """Initialize the plugin with a name."""
         self.name = name
 
-    def register(self, method_name: str = "__call__") -> Callable:
-        """Decorator to register a function as a plugin."""
+    def register(self):
+        """Register one implementation object under this plugin name."""
         if self.name is None:
             raise ValueError("Plugin name should be specified.")
 
-        if method_name in self._registry[self.name]:
-            logger.warning_rank0_once(f"Method {method_name} of plugin {self.name} is already registered.")
+        cls = type(self)
+        if self.name in cls._registry:
+            logger.warning_rank0_once(f"Plugin {self.name!r} is already registered under {cls.__name__}.")
 
-        def decorator(func: Callable) -> Callable:
-            self._registry[self.name][method_name] = func
-            return func
+        def decorator(obj: Any) -> Any:
+            cls._registry[self.name] = obj
+            return obj
 
         return decorator
 
+    @classmethod
+    def parse_params(cls, config: Any, params_cls: type[ParamsT]) -> ParamsT:
+        """Strictly convert config to the params dataclass used by one plugin entrypoint."""
+        if not is_dataclass(params_cls):
+            raise TypeError(f"{cls.__name__} params must be a dataclass type, got {params_cls!r}.")
+        if isinstance(config, params_cls):
+            return config
+        if config is None:
+            values = {}
+        elif isinstance(config, dict):
+            values = dict(config)
+        else:
+            raise TypeError(
+                f"{cls.__name__} config must be a mapping or {params_cls.__name__}, got {type(config).__name__}."
+            )
+
+        known = {item.name for item in fields(params_cls)}
+        unknown = set(values) - known
+        if unknown:
+            raise ValueError(
+                f"Unknown params for {cls.__name__}.{params_cls.__name__}: {sorted(unknown)}. "
+                f"Expected: {sorted(known)}"
+            )
+
+        return params_cls(**values)
+
+    def _resolve(self) -> Any:
+        cls = type(self)
+        if self.name is None:
+            raise ValueError(f"{cls.__name__} must be constructed with a name.")
+        if self.name not in cls._registry:
+            raise ValueError(f"Plugin {self.name!r} is not registered under {cls.__name__}.")
+        return cls._registry[self.name]
+
     def __call__(self, *args, **kwargs) -> Any:
-        """Call the registered function with the given arguments."""
-        return self["__call__"](*args, **kwargs)
+        return self._resolve()(*args, **kwargs)
 
-    def __getattr__(self, method_name: str) -> Callable:
-        """Get the registered function with the given name."""
-        return self[method_name]
-
-    def __getitem__(self, method_name: str) -> Callable:
-        """Get the registered function with the given name."""
-        if method_name not in self._registry[self.name]:
-            raise ValueError(f"Method {method_name} of plugin {self.name} is not registered.")
-
-        return self._registry[self.name][method_name]
-
-
-if __name__ == "__main__":
-    """
-    python -m llamafactory.v1.utils.plugin
-    """
-
-    class PrintPlugin(BasePlugin):
-        def again(self):  # optional
-            self["again"]()
-
-    @PrintPlugin("hello").register()
-    def print_hello():
-        print("Hello world!")
-
-    @PrintPlugin("hello").register("again")
-    def print_hello_again():
-        print("Hello world! Again.")
-
-    PrintPlugin("hello")()
-    PrintPlugin("hello").again()
+    def __getattr__(self, attr: str) -> Any:
+        return getattr(self._resolve(), attr)

src/llamafactory/v1/utils/types.py

@@ -78,19 +78,6 @@ class DatasetInfo(TypedDict, total=False):
     """Is streaming dataset, default to False."""
 
 
-class DistributedConfig(TypedDict, total=False):
-    mp_replicate_size: NotRequired[int]
-    """Model parallel replicate size, default to 1."""
-    mp_shard_size: NotRequired[int]
-    """Model parallel shard size, default to world_size // mp_replicate_size."""
-    dp_size: NotRequired[int]
-    """Data parallel size, default to world_size // cp_size."""
-    cp_size: NotRequired[int]
-    """Context parallel size, default to 1."""
-    timeout: NotRequired[int]
-    """Timeout for distributed communication, default to 600."""
-
-
 class Content(TypedDict):
     type: Literal["text", "reasoning", "tool_call", "image_url", "video_url", "audio_url"]
     """Type of the content."""