Skip to content

Add register_tracker_class to register custom trackers by name#4060

Open
javierdejesusda wants to merge 1 commit into
huggingface:mainfrom
javierdejesusda:register-custom-tracker-class
Open

Add register_tracker_class to register custom trackers by name#4060
javierdejesusda wants to merge 1 commit into
huggingface:mainfrom
javierdejesusda:register-custom-tracker-class

Conversation

@javierdejesusda
Copy link
Copy Markdown

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 both as a module-level register_tracker_class and as Accelerator.register_tracker_class.

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

Who can review?

@SunMarc @BenjaminBossan

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 both a module-level register_tracker_class and a thin Accelerator.register_tracker_class static method. 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature Request] Allows registering custom trackers to internal tracker type registry

1 participant