Add register_tracker_class() for custom tracker registration by name

src/accelerate/__init__.py

@@ -28,6 +28,7 @@
 from .launchers import debug_launcher, notebook_launcher
 from .parallelism_config import ParallelismConfig
 from .state import PartialState
+from .tracking import register_tracker_class
 from .utils import (
     AutocastKwargs,
     DataLoaderConfiguration,

src/accelerate/accelerator.py

@@ -42,7 +42,7 @@
 from .parallelism_config import ParallelismConfig
 from .scheduler import AcceleratedScheduler
 from .state import AcceleratorState, GradientState, PartialState
-from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers
+from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers, register_tracker_class
 from .utils import (
     MODEL_NAME,
     SAFE_WEIGHTS_INDEX_NAME,
@@ -3319,6 +3319,37 @@ def init_trackers(self, project_name: str, config: dict | None = None, init_kwar
             for tracker in self.trackers:
                 tracker.store_init_configuration(config)
 
+    @staticmethod
+    def register_tracker_class(tracker_cls):
+        """
+        Register a custom tracker class so it can be used by name in `log_with` and `init_trackers`, just like
+        built-in trackers.
+
+        The class must be a subclass of [`~tracking.GeneralTracker`] and define a `name` class attribute.
+
+        Args:
+            tracker_cls (`type`):
+                A subclass of `GeneralTracker` with a `name` attribute.
+
+        Example:
+
+        ```python
+        >>> from accelerate import Accelerator
+        >>> from accelerate.tracking import GeneralTracker
+
+        >>> class MyTracker(GeneralTracker):
+        ...     name = "my_tracker"
+        ...     requires_logging_directory = False
+        ...     @property
+        ...     def tracker(self):
+        ...         return None
+
+        >>> Accelerator.register_tracker_class(MyTracker)
+        >>> accelerator = Accelerator(log_with="my_tracker")
+        ```
+        """
+        register_tracker_class(tracker_cls)
+
     def get_tracker(self, name: str, unwrap: bool = False):
         """
         Returns a `tracker` from `self.trackers` based on `name` on the main process only.

src/accelerate/tracking.py

@@ -1259,6 +1259,48 @@ def finish(self):
 }
 
 
+def register_tracker_class(tracker_cls):
+    """
+    Register a custom tracker class so it can be used by name in `Accelerator(log_with=...)` and
+    `Accelerator.init_trackers(...)`, just like built-in trackers.
+
+    The class must be a subclass of [`~tracking.GeneralTracker`] and define a `name` class attribute. Its `__init__`
+    must accept `run_name` as the first positional argument (the project name passed to `init_trackers`), plus any
+    additional keyword arguments supplied via `init_kwargs`.
+
+    Registration is global: once registered, the tracker is available to all `Accelerator` instances.
+
+    Args:
+        tracker_cls (`type`):
+            A subclass of `GeneralTracker` with a `name` attribute.
+
+    Example:
+
+    ```python
+    >>> from accelerate.tracking import GeneralTracker, register_tracker_class
+
+    >>> class MyTracker(GeneralTracker):
+    ...     name = "my_tracker"
+    ...     requires_logging_directory = False
+    ...     @property
+    ...     def tracker(self):
+    ...         return None
+
+    >>> register_tracker_class(MyTracker)
+    ```
+    """
+    if not (isinstance(tracker_cls, type) and issubclass(tracker_cls, GeneralTracker)):
+        raise TypeError(f"`tracker_cls` must be a subclass of `GeneralTracker`, got {tracker_cls}.")
+    name = getattr(tracker_cls, "name", None)
+    if not name or not isinstance(name, str):
+        raise ValueError(
+            f"The tracker class must define a `name` class attribute as a non-empty string, got {name!r}."
+        )
+    if name in LOGGER_TYPE_TO_CLASS:
+        logger.warning(f"Overwriting existing tracker class for '{name}'.")
+    LOGGER_TYPE_TO_CLASS[name] = tracker_cls
+
+
 def filter_trackers(
     log_with: list[Union[str, LoggerType, GeneralTracker]],
     logging_dir: Optional[Union[str, os.PathLike]] = None,
@@ -1285,6 +1327,7 @@ def filter_trackers(
             - `"swanlab"`
             If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can
             also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
+            Custom tracker classes registered via `register_tracker_class` can also be referenced by their string name.
         logging_dir (`str`, `os.PathLike`, *optional*):
             A path to a directory for storing logs of locally-compatible loggers.
     """
@@ -1294,13 +1337,16 @@ def filter_trackers(
             log_with = [log_with]
         if "all" in log_with or LoggerType.ALL in log_with:
             loggers = [o for o in log_with if issubclass(type(o), GeneralTracker)] + get_available_trackers()
+            # Include custom-registered trackers not in the built-in LoggerType enum
+            existing = {str(t) for t in loggers}
+            for name in LOGGER_TYPE_TO_CLASS:
+                if name not in existing and name not in LoggerType:
+                    loggers.append(name)
         else:
             for log_type in log_with:
-                if log_type not in LoggerType and not issubclass(type(log_type), GeneralTracker):
-                    raise ValueError(f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}")
                 if issubclass(type(log_type), GeneralTracker):
                     loggers.append(log_type)
-                else:
+                elif log_type in LoggerType:
                     log_type = LoggerType(log_type)
                     if log_type not in loggers:
                         if log_type in get_available_trackers():
@@ -1313,5 +1359,16 @@ def filter_trackers(
                             loggers.append(log_type)
                         else:
                             logger.debug(f"Tried adding logger {log_type}, but package is unavailable in the system.")
+                elif isinstance(log_type, str) and log_type in LOGGER_TYPE_TO_CLASS:
+                    if log_type not in loggers:
+                        tracker_init = LOGGER_TYPE_TO_CLASS[log_type]
+                        if getattr(tracker_init, "requires_logging_directory", False) and logging_dir is None:
+                            raise ValueError(f"Logging with `{log_type}` requires a `logging_dir` to be passed in.")
+                        loggers.append(log_type)
+                else:
+                    raise ValueError(
+                        f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}"
+                        " or register a custom tracker with `register_tracker_class()`."
+                    )
 
     return loggers

tests/test_tracking.py

@@ -60,6 +60,7 @@
     TensorBoardTracker,
     TrackioTracker,
     WandBTracker,
+    register_tracker_class,
 )
 from accelerate.utils import (
     ProjectConfiguration,
@@ -731,6 +732,144 @@ def test_log(self):
                 assert data == truth
 
 
+class RegisterTrackerClassTestCase(unittest.TestCase):
+    def setUp(self):
+        from accelerate.tracking import LOGGER_TYPE_TO_CLASS
+
+        self._original_registry = dict(LOGGER_TYPE_TO_CLASS)
+
+    def tearDown(self):
+        from accelerate.tracking import LOGGER_TYPE_TO_CLASS
+
+        LOGGER_TYPE_TO_CLASS.clear()
+        LOGGER_TYPE_TO_CLASS.update(self._original_registry)
+
+    def _make_simple_tracker(self, name):
+        """Helper to create a minimal tracker class with the given name."""
+
+        class _Tracker(GeneralTracker):
+            requires_logging_directory = False
+
+            def __init__(self, run_name=None, **kwargs):
+                super().__init__()
+                self.run_name = run_name
+                self._config = {}
+                self._logs = []
+
+            @property
+            def tracker(self):
+                return self
+
+            def store_init_configuration(self, values: dict):
+                self._config = values
+
+            def log(self, values: dict, step=None, **kwargs):
+                self._logs.append(values)
+
+        _Tracker.name = name
+        return _Tracker
+
+    def test_register_and_use_by_name(self):
+        """Test that a registered custom tracker can be referenced by its string name."""
+        tracker_cls = self._make_simple_tracker("name_registered_tracker")
+        register_tracker_class(tracker_cls)
+        accelerator = Accelerator(log_with="name_registered_tracker")
+        config = {"learning_rate": 1e-2, "num_iterations": 12}
+        accelerator.init_trackers("test_project", config)
+        accelerator.log({"loss": 0.5}, step=0)
+        accelerator.end_training()
+        tracker = accelerator.get_tracker("name_registered_tracker")
+        assert tracker._config == config
+        assert tracker._logs == [{"loss": 0.5}]
+
+    def test_register_rejects_non_subclass(self):
+        """Test that register_tracker_class rejects objects that aren't GeneralTracker subclasses."""
+        with self.assertRaises(TypeError):
+            register_tracker_class("not a class")
+
+    def test_register_rejects_missing_name(self):
+        """Test that register_tracker_class rejects tracker classes without a name attribute."""
+
+        class NoNameTracker(GeneralTracker):
+            requires_logging_directory = False
+
+            @property
+            def tracker(self):
+                return None
+
+        with self.assertRaises(ValueError):
+            register_tracker_class(NoNameTracker)
+
+    def test_register_via_accelerator_static_method(self):
+        """Test that Accelerator.register_tracker_class works as a static method."""
+        tracker_cls = self._make_simple_tracker("static_method_tracker")
+        Accelerator.register_tracker_class(tracker_cls)
+        from accelerate.tracking import LOGGER_TYPE_TO_CLASS
+
+        assert "static_method_tracker" in LOGGER_TYPE_TO_CLASS
+
+    def test_register_overwrites_existing(self):
+        """Test that re-registering a tracker with the same name overwrites the previous one."""
+        first_cls = self._make_simple_tracker("overwrite_tracker")
+        second_cls = self._make_simple_tracker("overwrite_tracker")
+        register_tracker_class(first_cls)
+        register_tracker_class(second_cls)
+        from accelerate.tracking import LOGGER_TYPE_TO_CLASS
+
+        assert LOGGER_TYPE_TO_CLASS["overwrite_tracker"] is second_cls
+
+    def test_requires_logging_directory(self):
+        """Test that a custom tracker with requires_logging_directory=True raises without logging_dir."""
+
+        class DirRequiredTracker(GeneralTracker):
+            name = "dir_required_tracker"
+            requires_logging_directory = True
+
+            def __init__(self, run_name=None, logging_dir=None, **kwargs):
+                super().__init__()
+
+            @property
+            def tracker(self):
+                return None
+
+        register_tracker_class(DirRequiredTracker)
+        with self.assertRaises(ValueError, msg="should require logging_dir"):
+            Accelerator(log_with="dir_required_tracker")
+
+    def test_init_kwargs_passed_to_custom_tracker(self):
+        """Test that init_kwargs are forwarded to the custom tracker constructor."""
+
+        class KwargsTracker(GeneralTracker):
+            name = "kwargs_tracker"
+            requires_logging_directory = False
+
+            def __init__(self, run_name=None, **kwargs):
+                super().__init__()
+                self.run_name = run_name
+                self.extra = kwargs
+
+            @property
+            def tracker(self):
+                return self
+
+        register_tracker_class(KwargsTracker)
+        accelerator = Accelerator(log_with="kwargs_tracker")
+        accelerator.init_trackers("test_project", init_kwargs={"kwargs_tracker": {"custom_key": 42}})
+        tracker = accelerator.get_tracker("kwargs_tracker")
+        assert tracker.extra == {"custom_key": 42}
+        accelerator.end_training()
+
+    def test_filter_trackers_all_includes_custom(self):
+        """Test that log_with='all' includes custom-registered trackers."""
+        from accelerate.tracking import filter_trackers
+
+        tracker_cls = self._make_simple_tracker("all_test_tracker")
+        register_tracker_class(tracker_cls)
+        result = filter_trackers(["all"])
+        result_names = [str(t) for t in result]
+        assert "all_test_tracker" in result_names
+
+
 @require_dvclive
 @mock.patch("dvclive.live.get_dvc_repo", return_value=None)
 class DVCLiveTrackingTest(unittest.TestCase):