Update trackio integration to use Buckets and "freeze" Space after training

src/transformers/integrations/integration_utils.py

@@ -23,6 +23,7 @@
 import numbers
 import os
 import re
+import secrets
 import shutil
 import sys
 import tempfile
@@ -934,33 +935,43 @@ def on_predict(self, args, state, control, metrics, **kwargs):
 class TrackioCallback(TrainerCallback):
     """
     A [`TrainerCallback`] that logs metrics to Trackio.
-
-    It records training metrics, model (including PEFT) configuration.
-
-    **Requires**:
-    ```bash
-    pip install trackio
-    ```
     """
 
     SPACE_URL = "https://huggingface.co/spaces/{space_id}"
+    MIN_TRACKIO_VERSION_FOR_FREEZE = "0.21.1"
+    BADGE_MARKDOWN = (
+        '<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>'
+    )
+
+    @staticmethod
+    def _space_repo_name_from_project(project: str) -> str:
+        """Build a Hub-safe name for a static Space from the Trackio project and a short random suffix."""
+        s = project.strip().lower().replace("/", "-")
+        s = re.sub(r"[^a-z0-9._-]", "-", s)
+        s = re.sub(r"-+", "-", s).strip("-_.")
+        if not s:
+            s = "trackio-project"
+        base = s[:81].rstrip("-")
+        suffix = secrets.token_hex(3)
+        return f"{base}-static-{suffix}"
 
     def __init__(self):
-        has_trackio = is_trackio_available()
-        if not has_trackio:
+        if not is_trackio_available():
             raise RuntimeError("TrackioCallback requires trackio to be installed. Run `pip install trackio`.")
-        if has_trackio:
-            import trackio
+        import trackio
 
-            self._trackio = trackio
+        self._trackio = trackio
         self._initialized = False
+        self._space_id: str | None = None
+        self._static_space_id: str | None = None
 
     def setup(self, args, state, model, **kwargs):
         """
         Setup the optional Trackio integration.
 
-        To customize the setup you can also set the arguments `project`, `trackio_space_id` and `hub_private_repo` in
-        [`TrainingArguments`]. Please refer to the docstring of for more details.
+        To customize the setup you can also set `project`, `trackio_space_id`, `trackio_bucket_id`,
+        `trackio_static_space_id`, and `hub_private_repo` in [`TrainingArguments`].
         """
         if state.is_world_process_zero:
             combined_dict = {**args.to_dict()}
@@ -977,8 +988,11 @@ def setup(self, args, state, model, **kwargs):
                 space_id=args.trackio_space_id,
                 resume="allow",
                 private=args.hub_private_repo,
+                bucket_id=args.trackio_bucket_id,
             )
-
+            # The Trackio space_id may have been set by an environment variable, or set explicitly in the training arguments
+            # but without the full space_id. This ensures that self._space_id is set to the full space_id.
+            self._space_id = self._trackio.context_vars.current_space_id.get()
             # Add config parameters (run may have been created manually)
             self._trackio.config.update(combined_dict, allow_val_change=True)
 
@@ -989,13 +1003,78 @@ def setup(self, args, state, model, **kwargs):
                 logger.info("Could not log the number of model parameters in Trackio due to an AttributeError.")
         self._initialized = True
 
+    def _point_model_card_at_space(self, model) -> None:
+        """
+        Adds a link from a model's model card to the Trackio Space (preferring the `self._static_space_id` if not None).
+        Also adds two tags to the model to indicate that it was trained with Trackio:
+        - "trackio"
+        - "trackio:<space_url>" (this is used to embed the Space in the model's "Training Metrics" tab)
+        """
+        if self._space_id is None and self._static_space_id is None:
+            return
+
+        prev = getattr(model, "model_tags", None) or []
+        old_tag = next((t for t in prev if isinstance(t, str) and t.startswith("trackio:https://")), None)
+        if old_tag:
+            old_url = old_tag.split(":", 1)[1]
+            old_comment = self.BADGE_MARKDOWN.format(space_url=old_url)
+            cmt = modelcard.AUTOGENERATED_TRAINER_COMMENT or ""
+            if old_comment in cmt:
+                modelcard.AUTOGENERATED_TRAINER_COMMENT = cmt.replace(old_comment, "")
+
+        preferred_space_id = self._static_space_id or self._space_id
+        preferred_url = self.SPACE_URL.format(space_id=preferred_space_id)
+        new_tag = f"trackio:{preferred_url}"
+        new_comment = self.BADGE_MARKDOWN.format(space_url=preferred_url)
+
+        kept = [t for t in prev if not (isinstance(t, str) and t.startswith("trackio:https://"))]
+        for extra in ("trackio", new_tag):
+            if extra not in kept:
+                kept.append(extra)
+        model.model_tags = kept
+        c = modelcard.AUTOGENERATED_TRAINER_COMMENT or ""
+        modelcard.AUTOGENERATED_TRAINER_COMMENT = new_comment + c
+
+    def _freeze_space(self, args: TrainingArguments, model):
+        """
+        Freezes the Gradio Space after training is complete, if `trackio_static_space_id` is set to a `str` or `None`.
+        """
+        if args.trackio_static_space_id is False:
+            return
+        if packaging.version.parse(self._trackio.__version__) < packaging.version.parse(
+            self.MIN_TRACKIO_VERSION_FOR_FREEZE
+        ):
+            logger.warning(
+                "An older version of Trackio is installed; the post-training static snapshot Space (`trackio.freeze`) "
+                "will not be created. Upgrade with "
+                f"`pip install trackio>={self.MIN_TRACKIO_VERSION_FOR_FREEZE}` to enable it, or set "
+                "`trackio_static_space_id=False` to silence this warning."
+            )
+            return
+        static_space_id = (
+            self._static_space_id or args.trackio_static_space_id or self._space_repo_name_from_project(args.project)
+        )
+        self._static_space_id = self._trackio.freeze(
+            space_id=self._space_id,
+            project=args.project,
+            new_space_id=static_space_id,
+            private=args.hub_private_repo,
+        )
+        self._point_model_card_at_space(model)
+
     def on_train_begin(self, args, state, control, model=None, **kwargs):
         if not self._initialized:
             self.setup(args, state, model, **kwargs)
 
     def on_train_end(self, args: TrainingArguments, state, control, model=None, processing_class=None, **kwargs):
-        if state.is_world_process_zero and self._initialized:
-            self._trackio.finish()
+        if not state.is_world_process_zero or not self._initialized:
+            return
+        self._trackio.finish()
+        if self._space_id:
+            try:
+                self._freeze_space(args, model)
+            except Exception as e:
+                logger.warning(f"Trackio could not freeze the Gradio Space after training: {e}")
 
     def on_log(self, args, state, control, model=None, logs=None, **kwargs):
         single_value_scalars = [
@@ -1030,33 +1109,22 @@ def on_push_begin(self, args, state, control, model, **kwargs):
             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."
-            )
+        if self._space_id or args.trackio_static_space_id is False:
+            # If there's a Gradio space, it will be frozen after training is complete, so we don't need to sync it here.
+            # If a user has explicitly set trackio_static_space_id to False, we also don't sync their logs.
             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>'
+        static_space_id = (
+            self._static_space_id or args.trackio_static_space_id or self._space_repo_name_from_project(args.project)
         )
-        if badge_markdown not in modelcard.AUTOGENERATED_TRAINER_COMMENT:
-            modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n{badge_markdown}"
-
-        trackio_tags = ["trackio", f"trackio:{space_url}"]
-        if getattr(model, "model_tags", None) is not None:
-            if "trackio" not in model.model_tags:
-                model.model_tags.extend(trackio_tags)
-        else:
-            model.model_tags = trackio_tags
+        self._static_space_id = self._trackio.sync(
+            project=current_project,
+            sdk="static",
+            space_id=static_space_id,
+            private=args.hub_private_repo,
+            bucket_id=args.trackio_bucket_id,
+            force=True,
+        )
+        self._point_model_card_at_space(model)
 
 
 class CometCallback(TrainerCallback):

src/transformers/training_args.py

@@ -21,7 +21,7 @@
 from datetime import timedelta
 from enum import Enum
 from functools import cached_property
-from typing import Any
+from typing import Any, Literal
 
 from .debug_utils import DebugOption
 from .trainer_utils import (
@@ -391,12 +391,23 @@ class TrainingArguments:
             [swanlab](https://swanlab.cn) logging.
         project (`str`, *optional*, defaults to `"huggingface"`):
             The name of the project to use for logging. Currently, only used by Trackio.
-        trackio_space_id (`str` or `None`, *optional*, defaults to `"trackio"`):
-            The Hugging Face Space ID to deploy to when using Trackio. Should be a complete Space name like
-            `'username/reponame'` or `'orgname/reponame'`, or just `'reponame'` in which case the Space will be
-            created in the currently-logged-in Hugging Face user's namespace. If `None`, will log to a local directory.
-            Note that this Space will be public unless you set `hub_private_repo=True` or your organization's default
-            is to create private Spaces."
+        trackio_space_id (`str` or `None`, *optional*, defaults to `None`):
+            The Hugging Face Space ID to use for live Trackio logging with a Gradio-based Space. Should be a full
+            Space name like `'username/reponame'` or `'orgname/reponame'`, or just `'reponame'` (the Space is created in
+            the currently logged-in user's namespace). If `None`, metrics are logged only to a **local** directory (no
+            Space on the Hub). That Gradio Space has **read and write** access to the Trackio **Bucket**, which is what you want while
+            training is **in progress**—for example when **resuming** a partial run or **aggregating logs** from multiple
+            machines. The Space will be **public** unless you set `hub_private_repo=True` or your organization's default is to
+            create private Spaces.
+        trackio_bucket_id (`str` or `None`, *optional*, defaults to `None`):
+            Optional Hugging Face Bucket id for Trackio. If unset, Trackio derives one. Used together with a Gradio Space
+            (`trackio_space_id`) and when deploying a static Space (`trackio_static_space_id` is not `False`).
+        trackio_static_space_id (`str`, `False`, or `None`, *optional*, defaults to `None`):
+            The Hugging Face Space ID to use for static Space created after training is complete. Should be a full
+            Space name like `'username/reponame'` or `'orgname/reponame'`, or just `'reponame'` (the Space is created in
+            the currently logged-in user's namespace). If False, no static Space will be created. If None, and model is pushed to the Hub,
+            a static Space will be created with a default name and this will be linked from the model card. The Space will be public
+            unless you set `hub_private_repo=True` or your organization's default is to create private Spaces.
 
         > Evaluation
 
@@ -1045,13 +1056,29 @@ class TrainingArguments:
         metadata={"help": "The name of the project to use for logging. Currently, only used by Trackio."},
     )
     trackio_space_id: str | None = field(
-        default="trackio",
+        default=None,
         metadata={
-            "help": "The Hugging Face Space ID to deploy to when using Trackio. Should be a complete Space name like "
-            "'username/reponame' or 'orgname/reponame', or just 'reponame' in which case the Space will be created in "
-            "the currently-logged-in Hugging Face user's namespace. If `None`, will log to a local directory. Note "
-            "that this Space will be public unless you set `hub_private_repo=True` or your organization's "
-            "default is to create private Spaces."
+            "help": (
+                "Hugging Face Space id for live Gradio-based Trackio logging (read/write Bucket access). Use "
+                "'username/reponame', 'orgname/reponame', or 'reponame' (current user's namespace). None: log only "
+                "locally, no Space. Prefer trackio_static_space_id for stable post-training dashboard links. Public "
+                "unless hub_private_repo=True or org default."
+            )
+        },
+    )
+    trackio_bucket_id: str | None = field(
+        default=None,
+        metadata={"help": "Optional HF Bucket id when using a Trackio Space; if unset, Trackio picks a default."},
+    )
+    trackio_static_space_id: str | None | Literal[False] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Static read-only Trackio Space over the Bucket (stable model-card links). False: no static sync on Hub "
+                "push and no freeze after training. None/str: allow static Space; for local-only logging, Hub push runs "
+                "sync(static); after training, freeze runs only if trackio_space_id was set (Gradio Space). str sets "
+                "explicit static Space id. Public unless hub_private_repo=True or org default."
+            )
         },
     )
 
@@ -1179,8 +1206,8 @@ class TrainingArguments:
         metadata={
             "help": "Whether to make the repo private. If `None` (default), the repo will be public unless the "
             "organization's default is private. This value is ignored if the repo already exists. If reporting to "
-            "Trackio with deployment to Hugging Face Spaces enabled, the same logic determines whether the Space is "
-            "private."
+            "Trackio Spaces created or synced (including on Hub push when `trackio_space_id` is None) use the same "
+            "logic for whether the Space is private."
         },
     )
     hub_model_id: str | None = field(