Add register_tracker_class to register custom trackers by name by javierdejesusda · Pull Request #4060 · huggingface/accelerate

javierdejesusda · 2026-05-30T18:37:57Z

What does this PR do?

Adds register_tracker_class, a small public API to register a custom GeneralTracker subclass into the internal LOGGER_TYPE_TO_CLASS registry, so it can be selected by string name in Accelerator(log_with=...), exactly like the built-in trackers.

This removes the boilerplate described in #2734, where users currently have to special-case custom trackers when building log_with:

trackers = [name if name != "custom" else CustomTracker() for name in args.tracker_names]
accelerator = Accelerator(log_with=trackers)

With this change:

from accelerate.tracking import register_tracker_class

register_tracker_class(MyTracker)  # MyTracker.name == "my_tracker"
accelerator = Accelerator(log_with="my_tracker")

The API matches what was agreed in the issue thread (name taken from the class attribute, per @muellerzr's suggestion). It is exposed as a module-level register_tracker_class in accelerate.tracking.

Details

Validates that the class subclasses GeneralTracker and defines name and requires_logging_directory.
Warns (does not raise) when a registration overwrites an existing name, so shadowing a built-in is intentional but visible.
Extends filter_trackers so registered names resolve through the registry, including when combined with "all", with the same logging-directory guard applied to the built-in trackers.
Adds tests covering registration, use-by-name end to end, the requires_logging_directory path, overwrite warnings, validation errors, and the "all" combination, plus usage-guide and API-reference docs.

Fixes #2734

Before submitting

Did you read the contributor guideline, Pull Request section?
Was this discussed/approved via a GitHub issue or the forum? Yes, [Feature Request] Allows registering custom trackers to internal tracker type registry #2734 (the API was approved by @muellerzr and @SunMarc in the issue thread).
Did you make sure to update the documentation with your changes?
Did you write any new necessary tests?

Who can review?

@SunMarc @BenjaminBossan

SunMarc

Left a couple of comments thanks !

SunMarc · 2026-06-09T15:00:27Z

+            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)


well here, log_with is either "all" or "LoggerType.ALL" no ?

log_with is the full list here, so it can be e.g. ["my_tracker", "all"] the loop picks up a tracker registered by name when it's listed alongside "all", since "all" on its own only pulls in the built-ins. Happy to drop it and keep just instance + "all" if you'd prefer the smaller surface.

yeah, i think it is better to drop it. it's a bit counter intuitive to pass ["my_tracker", "all"]

Done in 155194d: "all" is back to main's behavior, so a registered name passed alongside "all" is no longer picked up. Selecting by name still works on its own via log_with="my_tracker".

SunMarc · 2026-06-09T15:01:29Z

                assert data == truth


+class MyRegisteredTracker(GeneralTracker):


keep tests simple

Will trim these to the shared fixture plus the core cases, together with the filter_trackers decision above (some depend on it).

javierdejesusda · 2026-06-10T18:13:50Z

Extra commit d22cb63: the overwrite warning used logger.warning, but accelerate's logger raises RuntimeError before PartialState/Accelerator init, so shadowing or re-registering a name in the documented "register first" flow would crash instead of warn. Switched that call to warnings.warn (matches accelerator.py) and the test now asserts via assertWarns.

github-actions · 2026-07-05T15:31:29Z

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

javierdejesusda · 2026-07-06T17:20:41Z

Not stale, this is still ready for review. @SunMarc when you have a moment, could you take another look? Happy to adjust anything else if needed.

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 huggingface#2734

Custom trackers registered via register_tracker_class are selected by listing their name explicitly, not when combined with "all". The "all" branch of filter_trackers returns to its original behavior (instances plus available built-ins). Removes the test covering the dropped behavior and its now-unused filter_trackers import.

register_tracker_class is documented to run before Accelerator()/PartialState(), so the accelerate logger (a MultiProcessAdapter) raises RuntimeError when the name-collision warning fires (shadowing a built-in or re-registering a name). Use warnings.warn so the warning is emitted regardless of state, matching how tracker warnings are reported elsewhere. The test now exercises the real warning instead of mocking the logger.

SunMarc

Thanks !

HuggingFaceDocBuilderDev · 2026-07-07T15:23:13Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

SunMarc reviewed Jun 9, 2026

View reviewed changes

SunMarc mentioned this pull request Jun 9, 2026

Add register_tracker_class() for custom tracker registration by name #4061

Closed

4 of 5 tasks

javierdejesusda force-pushed the register-custom-tracker-class branch 3 times, most recently from 893e814 to 3b9c71b Compare June 9, 2026 16:10

javierdejesusda added 3 commits July 6, 2026 19:21

javierdejesusda force-pushed the register-custom-tracker-class branch from d22cb63 to 2edb1a6 Compare July 6, 2026 17:22

SunMarc approved these changes Jul 7, 2026

View reviewed changes

SunMarc merged commit 665444c into huggingface:main Jul 7, 2026
23 of 25 checks passed

		assert data == truth


		class MyRegisteredTracker(GeneralTracker):

Conversation

javierdejesusda commented May 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Details

Before submitting

Who can review?

Uh oh!

SunMarc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

SunMarc Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

javierdejesusda Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

SunMarc Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

javierdejesusda Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

SunMarc Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

javierdejesusda Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

javierdejesusda commented Jun 10, 2026

Uh oh!

github-actions Bot commented Jul 5, 2026

Uh oh!

javierdejesusda commented Jul 6, 2026

Uh oh!

SunMarc left a comment

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Jul 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

javierdejesusda commented May 30, 2026 •

edited

Loading