diff --git a/src/transformers/integrations/integration_utils.py b/src/transformers/integrations/integration_utils.py index 6bd21560b849..2656b7169c62 100755 --- a/src/transformers/integrations/integration_utils.py +++ b/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 = ( + 'Visualize in Trackio' + ) + + @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:" (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'Visualize in Trackio' + 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): diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py index d4d4d49931cb..e66af264b89a 100644 --- a/src/transformers/training_args.py +++ b/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(