feat(tracking): register custom tracker classes by name

Add register_tracker_class to register a custom GeneralTracker subclass into the LOGGER_TYPE_TO_CLASS registry so it can be selected by string name in Accelerator(log_with=...), the same way as the built-in trackers.

Expose a module-level register_tracker_class in accelerate.tracking. Validate that the class subclasses GeneralTracker and defines name and requires_logging_directory, and warn when an existing name is overwritten. Extend filter_trackers so registered custom names resolve through the registry, including when combined with 'all', with the same logging-directory guard as the built-in trackers. Add tests and documentation for registering and using a custom tracker by name.

Fixes #2734

Author: javierdejesusda

docs/source/package_reference/tracking.md

@@ -19,6 +19,10 @@ rendered properly in your Markdown viewer.
 
 [[autodoc]] tracking.GeneralTracker
 
+## register_tracker_class
+
+[[autodoc]] tracking.register_tracker_class
+
 ## TensorBoardTracker
 
 [[autodoc]] tracking.TensorBoardTracker

docs/source/usage_guides/tracking.md

@@ -170,6 +170,25 @@ tracker = MyCustomTracker("some_run_name")
 accelerator = Accelerator(log_with=[tracker, "all"])
 ```
 
+### Registering a custom tracker by name
+
+Instead of passing an instance, you can register a custom tracker class so it can be selected by its `name`, just like the built-in trackers. This is handy when the trackers to use are read from a config file or the command line:
+
+```python
+from accelerate import Accelerator
+from accelerate.tracking import GeneralTracker, register_tracker_class
+
+
+class MyCustomTracker(GeneralTracker):
+    name = "my_tracker"
+    requires_logging_directory = False
+    # ... implement the rest of the `GeneralTracker` interface
+
+
+register_tracker_class(MyCustomTracker)
+accelerator = Accelerator(log_with="my_tracker")
+```
+
 ## Accessing the internal tracker 
 
 If some custom interactions with a tracker might be wanted directly, you can quickly access one using the 

src/accelerate/tracking.py

@@ -1259,6 +1259,56 @@ def finish(self):
 }
 
 
+def register_tracker_class(tracker_class: type[GeneralTracker]):
+    """
+    Registers a custom [`GeneralTracker`] subclass so it can be referenced by its `name` in the `log_with` argument of
+    [`Accelerator`], the same way as the built-in trackers.
+
+    The tracker must be registered before instantiating the [`Accelerator`] that uses it.
+
+    Args:
+        tracker_class (subclass of [`GeneralTracker`]):
+            The tracker class to register. It must subclass [`GeneralTracker`] and define a non-empty `name` class
+            attribute and a `requires_logging_directory` class attribute. The class is instantiated with the
+            `project_name` passed to [`Accelerator.init_trackers`] as the first positional argument. When
+            `requires_logging_directory` is `True`, the logging directory is passed as the second positional argument
+            (matching the built-in tracker convention). Tracker-specific keyword arguments from `init_kwargs` are
+            forwarded as `**kwargs`.
+
+    Example:
+
+    ```python
+    from accelerate import Accelerator
+    from accelerate.tracking import GeneralTracker, register_tracker_class
+
+
+    class MyTracker(GeneralTracker):
+        name = "my_tracker"
+        requires_logging_directory = False
+        # ... implement the rest of the `GeneralTracker` interface
+
+
+    register_tracker_class(MyTracker)
+    accelerator = Accelerator(log_with="my_tracker")
+    ```
+    """
+    if not (isinstance(tracker_class, type) and issubclass(tracker_class, GeneralTracker)):
+        raise ValueError(f"`tracker_class` must be a subclass of `GeneralTracker`, but got {tracker_class}.")
+    if not getattr(tracker_class, "name", None):
+        raise ValueError("The tracker class to register must define a non-empty `name` attribute.")
+    if not hasattr(tracker_class, "requires_logging_directory"):
+        raise ValueError(
+            "The tracker class to register must define a `requires_logging_directory` class attribute. "
+            "Set it to `True` if the tracker needs a logging directory, or `False` otherwise."
+        )
+    if tracker_class.name in LOGGER_TYPE_TO_CLASS:
+        logger.warning(
+            f"A tracker with the name '{tracker_class.name}' is already registered and will be overwritten by "
+            f"{tracker_class}. This is especially significant when shadowing a built-in tracker."
+        )
+    LOGGER_TYPE_TO_CLASS[tracker_class.name] = tracker_class
+
+
 def filter_trackers(
     log_with: list[Union[str, LoggerType, GeneralTracker]],
     logging_dir: Optional[Union[str, os.PathLike]] = None,
@@ -1293,14 +1343,28 @@ def filter_trackers(
         if not isinstance(log_with, (list, tuple)):
             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()
+            loggers = [o for o in log_with if issubclass(type(o), GeneralTracker)]
+            # Also keep custom trackers registered via `register_tracker_class` that are named alongside "all".
+            # Apply the same logging-dir guard as the explicit-name branch for consistency.
+            for custom_name in log_with:
+                if custom_name not in LoggerType and custom_name in LOGGER_TYPE_TO_CLASS:
+                    tracker_init = LOGGER_TYPE_TO_CLASS[str(custom_name)]
+                    if tracker_init.requires_logging_directory and logging_dir is None:
+                        raise ValueError(f"Logging with `{custom_name}` requires a `logging_dir` to be passed in.")
+                    if custom_name not in loggers:
+                        loggers.append(custom_name)
+            loggers += get_available_trackers()
         else:
             for log_type in log_with:
-                if log_type not in LoggerType and not issubclass(type(log_type), GeneralTracker):
+                if (
+                    log_type not in LoggerType
+                    and not issubclass(type(log_type), GeneralTracker)
+                    and str(log_type) not in LOGGER_TYPE_TO_CLASS
+                ):
                     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 +1377,12 @@ def filter_trackers(
                             loggers.append(log_type)
                         else:
                             logger.debug(f"Tried adding logger {log_type}, but package is unavailable in the system.")
+                else:
+                    # Custom tracker class registered via `register_tracker_class`
+                    if log_type not in loggers:
+                        tracker_init = LOGGER_TYPE_TO_CLASS[str(log_type)]
+                        if tracker_init.requires_logging_directory and logging_dir is None:
+                            raise ValueError(f"Logging with `{log_type}` requires a `logging_dir` to be passed in.")
+                        loggers.append(log_type)
 
     return loggers

tests/test_tracking.py

@@ -50,6 +50,7 @@
     skip,
 )
 from accelerate.tracking import (
+    LOGGER_TYPE_TO_CLASS,
     AimTracker,
     ClearMLTracker,
     CometMLTracker,
@@ -60,6 +61,8 @@
     TensorBoardTracker,
     TrackioTracker,
     WandBTracker,
+    filter_trackers,
+    register_tracker_class,
 )
 from accelerate.utils import (
     ProjectConfiguration,
@@ -731,6 +734,192 @@ def test_log(self):
                 assert data == truth
 
 
+class MyRegisteredTracker(GeneralTracker):
+    "Minimal tracker following the built-in `(run_name, **kwargs)` convention, used to test class registration."
+
+    name = "my_registered_tracker"
+    requires_logging_directory = False
+
+    def __init__(self, run_name: str, **kwargs):
+        super().__init__()
+        self.run_name = run_name
+        self.config = {}
+        self.values = {}
+
+    def start(self):
+        pass
+
+    @property
+    def tracker(self):
+        return self
+
+    def store_init_configuration(self, values: dict):
+        self.config.update(values)
+
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
+        self.values.update(values)
+
+    def finish(self):
+        pass
+
+
+class RegisterTrackerClassTest(unittest.TestCase):
+    def setUp(self):
+        self._registry_snapshot = dict(LOGGER_TYPE_TO_CLASS)
+        super().setUp()
+
+    def tearDown(self):
+        LOGGER_TYPE_TO_CLASS.clear()
+        LOGGER_TYPE_TO_CLASS.update(self._registry_snapshot)
+        super().tearDown()
+
+    def test_register_and_use_by_name(self):
+        register_tracker_class(MyRegisteredTracker)
+        assert MyRegisteredTracker.name in LOGGER_TYPE_TO_CLASS
+
+        accelerator = Accelerator(log_with=MyRegisteredTracker.name)
+        config = {"num_iterations": 12, "learning_rate": 1e-2}
+        accelerator.init_trackers("some_project", config)
+        accelerator.log({"total_loss": 0.1}, step=0)
+        tracker = accelerator.get_tracker(MyRegisteredTracker.name)
+        accelerator.end_training()
+
+        assert isinstance(tracker, MyRegisteredTracker)
+        assert tracker.config == config
+        assert tracker.values == {"total_loss": 0.1}
+
+    def test_registered_tracker_included_with_all(self):
+        register_tracker_class(MyRegisteredTracker)
+        loggers = filter_trackers([MyRegisteredTracker.name, "all"])
+        assert MyRegisteredTracker.name in loggers
+
+    def test_register_non_tracker_raises(self):
+        with self.assertRaises(ValueError):
+            register_tracker_class(dict)
+
+    def test_register_without_name_raises(self):
+        class NamelessTracker(GeneralTracker):
+            requires_logging_directory = False
+
+        with self.assertRaises(ValueError):
+            register_tracker_class(NamelessTracker)
+
+    def test_register_missing_requires_logging_directory_raises(self):
+        class NoAttrTracker(GeneralTracker):
+            name = "no_attr_tracker"
+            # intentionally missing requires_logging_directory
+
+        with self.assertRaises(ValueError):
+            register_tracker_class(NoAttrTracker)
+
+    def test_register_overwrites_existing_warns(self):
+        class FirstTracker(GeneralTracker):
+            name = "collision_tracker"
+            requires_logging_directory = False
+
+            @property
+            def tracker(self):
+                return self
+
+            def store_init_configuration(self, values: dict):
+                pass
+
+            def log(self, values: dict, step=None, **kwargs):
+                pass
+
+            def finish(self):
+                pass
+
+        class SecondTracker(GeneralTracker):
+            name = "collision_tracker"
+            requires_logging_directory = False
+
+            @property
+            def tracker(self):
+                return self
+
+            def store_init_configuration(self, values: dict):
+                pass
+
+            def log(self, values: dict, step=None, **kwargs):
+                pass
+
+            def finish(self):
+                pass
+
+        register_tracker_class(FirstTracker)
+        # Registering a second class with the same name must warn
+        with mock.patch("accelerate.tracking.logger") as mock_logger:
+            register_tracker_class(SecondTracker)
+            mock_logger.warning.assert_called_once()
+            warning_msg = mock_logger.warning.call_args[0][0]
+            assert "collision_tracker" in warning_msg
+
+    def test_all_branch_requires_logging_directory(self):
+        class DirRequiringTracker(GeneralTracker):
+            name = "dir_requiring_custom"
+            requires_logging_directory = True
+
+            @property
+            def tracker(self):
+                return self
+
+            def store_init_configuration(self, values: dict):
+                pass
+
+            def log(self, values: dict, step=None, **kwargs):
+                pass
+
+            def finish(self):
+                pass
+
+        register_tracker_class(DirRequiringTracker)
+        with self.assertRaises(ValueError):
+            # No logging_dir provided - should raise because custom tracker requires one
+            filter_trackers([DirRequiringTracker.name, "all"], logging_dir=None)
+
+    def test_requires_logging_directory_custom_tracker(self):
+        class DirTracker(GeneralTracker):
+            name = "dir_tracker_test"
+            requires_logging_directory = True
+
+            def __init__(self, run_name: str, logging_dir: str, **kwargs):
+                super().__init__()
+                self.run_name = run_name
+                self.logging_dir = logging_dir
+                self.config = {}
+                self.values = {}
+
+            @property
+            def tracker(self):
+                return self
+
+            def store_init_configuration(self, values: dict):
+                self.config.update(values)
+
+            def log(self, values: dict, step=None, **kwargs):
+                self.values.update(values)
+
+            def finish(self):
+                pass
+
+        register_tracker_class(DirTracker)
+
+        # Without project_dir, Accelerator should raise ValueError when trying to use it
+        with self.assertRaises(ValueError):
+            Accelerator(log_with=DirTracker.name)
+
+        # With a project_dir it should succeed and produce the right tracker type
+        with tempfile.TemporaryDirectory() as tmpdir:
+            accelerator = Accelerator(log_with=DirTracker.name, project_dir=tmpdir)
+            accelerator.init_trackers("test_dir_project")
+            tracker = accelerator.get_tracker(DirTracker.name)
+            assert isinstance(tracker, DirTracker)
+            assert tracker.run_name == "test_dir_project"
+            assert tracker.logging_dir == tmpdir
+            accelerator.end_training()
+
+
 @require_dvclive
 @mock.patch("dvclive.live.get_dvc_repo", return_value=None)
 class DVCLiveTrackingTest(unittest.TestCase):