Repository: huggingface/transformers
Pull request: Add `.on_push_begin()` callback to Trainer and implement for `TrackioCallback`

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Changed files:
filename: src/transformers/integrations/integration_utils.py
patch: @@ -940,6 +940,8 @@ class TrackioCallback(TrainerCallback):
     ```
     """
 
+    SPACE_URL = "https://huggingface.co/spaces/{space_id}"
+
     def __init__(self):
         has_trackio = is_trackio_available()
         if not has_trackio:
@@ -1058,6 +1060,39 @@ def on_predict(self, args, state, control, metrics, **kwargs):
             metrics = rewrite_logs(metrics)
             self._trackio.log(metrics)
 
+    def on_push_begin(self, args, state, control, model, **kwargs):
+        if not state.is_world_process_zero or self._trackio is None:
+            return
+        if (current_project := self._trackio.context_vars.current_project.get()) is None:
+            return
+        trackio_version = packaging.version.parse(self._trackio.__version__)
+        if trackio_version < packaging.version.parse("0.13.0"):
+            warnings.warn(
+                "The version of `trackio` that is installed is <=0.13.0, so "
+                "the local Trackio project will not be pushed to Hugging Face. Run "
+                "`pip install --upgrade trackio` to fix this."
+            )
+            return
+
+        space_id = self._trackio.context_vars.current_space_id.get()
+        if space_id is None:
+            space_id = self._trackio.sync(current_project, force=True)
+        space_url = self.SPACE_URL.format(space_id=space_id)
+
+        badge_markdown = (
+            f'<a href="{space_url}" target="_blank"><img src="https://raw.githubusercontent.com/gradio-app/trackio/refs/heads/main/trackio/assets/badge.png" alt="Visualize in Trackio"'
+            ' title="Visualize in Trackio" style="height: 40px;"/></a>'
+        )
+        if badge_markdown not in modelcard.AUTOGENERATED_TRAINER_COMMENT:
+            modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n{badge_markdown}"
+
+        if getattr(model, "model_tags", None) is not None:
+            if "trackio" not in model.model_tags:
+                model.model_tags.append("trackio")
+                model.model_tags.append(f"trackio::{space_url}")
+        else:
+            model.model_tags = ["trackio", f"trackio:{space_url}"]
+
 
 class CometCallback(TrainerCallback):
     """

filename: src/transformers/trainer.py
patch: @@ -4827,6 +4827,7 @@ def _push_from_checkpoint(self, checkpoint_folder):
         if not self.args.hub_always_push and self.push_in_progress is not None and not self.push_in_progress.is_done():
             return
 
+        self.callback_handler.on_push_begin(self.args, self.state, self.control)
         output_dir = self.args.output_dir
         # To avoid a new synchronization of all model weights, we just copy the file from the checkpoint folder
         modeling_files = [CONFIG_NAME, GENERATION_CONFIG_NAME, WEIGHTS_NAME, SAFE_WEIGHTS_NAME]
@@ -4921,6 +4922,8 @@ def push_to_hub(
             The URL of the repository where the model was pushed if `blocking=False`, or a `Future` object tracking the
             progress of the commit if `blocking=True`.
         """
+        self.callback_handler.on_push_begin(self.args, self.state, self.control)
+
         model_name = kwargs.pop("model_name", None)
         if model_name is None and self.args.should_save:
             if self.args.hub_model_id is None:

filename: src/transformers/trainer_callback.py
patch: @@ -420,6 +420,11 @@ def on_prediction_step(self, args: TrainingArguments, state: TrainerState, contr
         Event called after a prediction step.
         """
 
+    def on_push_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        """
+        Event called before pushing the model to the hub, at the beginning of Trainer.push_to_hub and Trainer._push_from_checkpoint.
+        """
+
 
 class CallbackHandler(TrainerCallback):
     """Internal class that just calls the list of callbacks in order."""
@@ -532,6 +537,9 @@ def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerC
     def on_prediction_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
         return self.call_event("on_prediction_step", args, state, control)
 
+    def on_push_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        return self.call_event("on_push_begin", args, state, control, **kwargs)
+
     def call_event(self, event, args, state, control, **kwargs):
         for callback in self.callbacks:
             result = getattr(callback, event)(

filename: tests/trainer/test_trainer_callback.py
patch: @@ -102,6 +102,9 @@ def on_log(self, args, state, control, **kwargs):
     def on_prediction_step(self, args, state, control, **kwargs):
         self.events.append("on_prediction_step")
 
+    def on_push_begin(self, args, state, control, **kwargs):
+        self.events.append("on_push_begin")
+
 
 @require_torch
 class TrainerCallbackTest(unittest.TestCase):
@@ -443,3 +446,14 @@ def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control, **
         trainer = self.get_trainer(max_steps=2, save_strategy="epoch", callbacks=[OnEndCallback])
         trainer.train()
         assert times_saved == 1
+
+    def test_on_push_begin(self):
+        trainer = self.get_trainer(callbacks=[MyTestTrainerCallback], max_steps=1)
+        trainer.train()
+        callback = [cb for cb in trainer.callback_handler.callbacks if isinstance(cb, MyTestTrainerCallback)][0]
+        initial_event_count = len(callback.events)
+
+        trainer.callback_handler.on_push_begin(trainer.args, trainer.state, trainer.control)
+        assert "on_push_begin" in callback.events
+        assert callback.events.count("on_push_begin") == 1
+        assert len(callback.events) == initial_event_count + 1