Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
luowyang opened this issue May 2, 2024 · 2 comments
Labels
enhancement New feature or request feature request Request for a new feature to be added to Accelerate

Comments

@luowyang
Copy link
Contributor

luowyang commented May 2, 2024

Currently accelerate maintains an internal known tracker registry LOGGER_TYPE_TO_CLASS, which is used for tracker initializing, filtering, among other things. However, this makes working with custom trackers annoying in user code. One such example is when initializing custom trackers, instead of something like

accelerator = Accelerator(log_with=args.tracker_names)

one has to write:

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

This is annoying in that users have to write many boilerplate code to manage the init args of both custom trackers and built-in trackers in one place. If the users can register their custom trackers into LOGGER_TYPE_TO_CLASS, it will be possible to initialize the custom trackers in the same (and uniform) way as the built-in trackers by:

init_kwargs = dict(
    wandb=dict(
        # init args for wandb
    ),
    custom_tracker=dict(
        # init args for custom tracker named `custom_tracker`
    ),
)
self.accelerator.init_trackers(self.args.project, config, init_kwargs)

The API for registering custom trackers could be something like:

accelerator.register_custom_tracker(name, CustomTracker)

I have examined the source code, and there seems no obstacle in doing so, as long as the GeneralTracker protocol is obeyed by the custom trackers. To make sure of it, it may be beneficial to mark GeneralTracker as an abstract base class, though.

@luowyang luowyang changed the title [Feature Request] Allows register custom trackers to internal tracker type registry [Feature Request] Allows registering custom trackers to internal tracker type registry May 2, 2024
@SunMarc SunMarc added feature request Request for a new feature to be added to Accelerate enhancement New feature or request labels May 17, 2024
@SunMarc
Copy link
Member

SunMarc commented May 17, 2024

Hi @luowyang, thanks for the detailed request ! What do you think about having an API that register custom tracker @muellerzr ?

@muellerzr
Copy link
Collaborator

Sure. I'd be open to this. However we can make name a class attribute (they already are) so the API is just:

accelerator.register_tracker_class(CustomTracker)

How does that sound :)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request feature request Request for a new feature to be added to Accelerate
Projects
None yet
Development

No branches or pull requests

3 participants