Remove registration for TrainingStrategy

Author: crowbat

src/coreai_opt/palettization/spec/spec.py

@@ -64,9 +64,11 @@ class PalettizationSpec(CompressionSpec):
         training_strategy_config: Which training-time behavior this weight's
             fake-palettize module uses, and that strategy's settings.
             ``DefaultTrainingConfig()`` is post-training, one-shot k-means
-            (today's KMeansPalettizer behavior). Additional strategies can be
-            registered via ``TrainingStrategy.register()`` /
-            ``TrainingStrategyConfig.register()``. Default: DefaultTrainingConfig().
+            (today's KMeansPalettizer behavior). Additional strategies are added
+            by subclassing ``TrainingStrategyConfig`` (registered via
+            ``TrainingStrategyConfig.register()``) and pointing its
+            ``_strategy_cls`` at a ``TrainingStrategy`` subclass. Default:
+            DefaultTrainingConfig().
 
     Example:
         >>> # Basic 4-bit palettization

src/coreai_opt/palettization/spec/training_strategy.py

@@ -7,22 +7,19 @@
 
 from __future__ import annotations
 
-from abc import abstractmethod
-from typing import TYPE_CHECKING, Any
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import torch
 from pydantic import BaseModel, ConfigDict, model_serializer
 
-from coreai_opt._utils.registry_utils import (
-    ClassRegistryMixin as _ClassRegistryMixin,
-    ConfigRegistryMixin as _ConfigRegistryMixin,
-)
+from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin
 
 if TYPE_CHECKING:
     from coreai_opt.palettization.kmeans.kmeans_fake_palettize import _KMeansFakePalettize
 
 
-class TrainingStrategy(_ClassRegistryMixin):
+class TrainingStrategy(ABC):
     """Contract for a fake-palettize module's training-time forward pass."""
 
     @abstractmethod
@@ -38,7 +35,6 @@ def train_forward(self, module: _KMeansFakePalettize, weight: torch.Tensor) -> t
         raise NotImplementedError
 
 
-@TrainingStrategy.register("default")
 class _DefaultTrainingStrategy(TrainingStrategy):
     """Post-training, one-shot k-means — today's KMeansPalettizer behavior.
 
@@ -61,13 +57,16 @@ def train_forward(self, module: _KMeansFakePalettize, weight: torch.Tensor) -> t
 class TrainingStrategyConfig(BaseModel, _ConfigRegistryMixin):
     """Base class for a fake-palettize module's training-strategy settings.
 
-    Each subclass is registered under the same key as its paired
-    ``TrainingStrategy`` behavior class (e.g. ``"default"``); ``build_strategy()``
-    resolves and constructs that pair from this config's own fields.
+    Each subclass points ``_strategy_cls`` at its paired ``TrainingStrategy``
+    behavior class; ``build_strategy()`` constructs that strategy from this
+    config's own fields.
     """
 
     model_config = ConfigDict(frozen=True, extra="forbid")
 
+    # Each subclass points this at its paired TrainingStrategy behavior class.
+    _strategy_cls: ClassVar[type[TrainingStrategy]]
+
     @model_serializer
     def _serialize_model(self) -> dict[str, Any]:
         """Custom serializer that includes the registry type."""
@@ -91,15 +90,12 @@ def _serialize_model(self) -> dict[str, Any]:
 
     def build_strategy(self) -> TrainingStrategy:
         """Construct this config's paired ``TrainingStrategy`` behavior instance."""
-        for key, registered_class in TrainingStrategyConfig.REGISTRY.items():
-            if registered_class is type(self):
-                kwargs = {k: v for k, v in self.model_dump().items() if k != "type"}
-                return TrainingStrategy.resolve(key)(**kwargs)
-        raise ValueError(
-            f"{type(self).__name__} is not registered in TrainingStrategyConfig's registry."
-        )
+        kwargs = {k: v for k, v in self.model_dump().items() if k != "type"}
+        return self._strategy_cls(**kwargs)
 
 
 @TrainingStrategyConfig.register("default")
 class DefaultTrainingConfig(TrainingStrategyConfig):
     """Settings for the default, post-training one-shot k-means strategy. No fields."""
+
+    _strategy_cls = _DefaultTrainingStrategy

tests/palettization/test_kmeans_fake_palettize.py

@@ -1611,18 +1611,19 @@ def test_lut_consistent_with_quantized_lut_and_scale(self):
 
 
 class TestTrainingStrategy:
-    """The training-strategy registry and its paired config objects.
+    """The training-strategy config registry and its paired behavior classes.
 
     Scope is the generic OSS machinery (default strategy only); concrete
-    strategies registered elsewhere are tested with those strategies.
+    strategies defined elsewhere are tested with those strategies.
     """
 
-    def test_resolve_default(self):
-        assert TrainingStrategy.resolve("default") is _DefaultTrainingStrategy
+    def test_default_config_points_at_default_strategy(self):
+        assert DefaultTrainingConfig._strategy_cls is _DefaultTrainingStrategy
+        assert issubclass(DefaultTrainingConfig._strategy_cls, TrainingStrategy)
 
-    def test_resolve_unregistered_raises(self):
-        with pytest.raises(ValueError):
-            TrainingStrategy.resolve("nonexistent")
+    def test_build_from_dict_unregistered_raises(self):
+        with pytest.raises(KeyError):
+            TrainingStrategyConfig.maybe_build_from_dict({"type": "nonexistent"})
 
     def test_default_config_builds_default_strategy(self):
         assert isinstance(DefaultTrainingConfig().build_strategy(), _DefaultTrainingStrategy)