Adds a FlyteCallback

src/transformers/integrations.py

@@ -30,7 +30,7 @@
 import numpy as np
 
 from . import __version__ as version
-from .utils import flatten_dict, is_datasets_available, is_torch_available, logging
+from .utils import flatten_dict, is_datasets_available, is_pandas_available, is_torch_available, logging
 from .utils.versions import importlib_metadata
 
 
@@ -146,6 +146,10 @@ def is_codecarbon_available():
     return importlib.util.find_spec("codecarbon") is not None
 
 
+def is_flytekit_available():
+    return importlib.util.find_spec("flytekit") is not None
+
+
 def hp_params(trial):
     if is_optuna_available():
         import optuna
@@ -1537,6 +1541,63 @@ def on_save(self, args, state, control, **kwargs):
             self._clearml_task.update_output_model(artifact_path, iteration=state.global_step, auto_delete_file=False)
 
 
+class FlyteCallback(TrainerCallback):
+    """A [`TrainerCallback`] that sends the logs to [Flyte](https://flyte.org/).
+    NOTE: This callback only works within a Flyte task.
+
+    Args:
+        save_log_history (bool, optional, defaults to `True`):
+            Determines whether or not to save the training logs to the task's Flyte Deck.
+
+        sync_checkpoints (bool, optional, defaults to `True`):
+            When set to True, checkpoints are synced to Flyte and can be used to resume training in the case of an
+            interruption.
+
+    Example:
+        ```python
+        # Note: This example skips over some setup steps for brevity.
+        from flytekit import current_context, task
+
+
+        @task
+        def train_hf_transformer():
+            cp = current_context().checkpoint
+            trainer = Trainer(..., callbacks=[FlyteCallback()])
+            output = trainer.train(resume_from_checkpoint=cp.restore())
+        ```
+    """
+
+    def __init__(self, save_log_history: bool = True, sync_checkpoints: bool = True):
+        super().__init__()
+        if not is_flytekit_available():
+            raise ImportError("FlyteCallback requires flytekit to be installed. Run `pip install flytekit`.")
+
+        from flytekit import current_context
+
+        self.cp = current_context().checkpoint
+        self.save_log_history = save_log_history
+        self.sync_checkpoints = sync_checkpoints
+
+    def on_save(self, args, state, control, **kwargs):
+        if self.sync_checkpoints and state.is_world_process_zero:
+            ckpt_dir = f"checkpoint-{state.global_step}"
+            artifact_path = os.path.join(args.output_dir, ckpt_dir)
+
+            logger.info(f"Saving checkpoint in {ckpt_dir} to Flyte. This may take time.")
+            self.cp.save(artifact_path)
+
+    def on_train_end(self, args, state, control, **kwargs):
+        if self.save_log_history:
+            import pandas as pd
+            from flytekit import Deck
+
+            if is_pandas_available():
+                Deck("Log History", pd.DataFrame(state.log_history).to_html())
+            else:
+                logger.warning("Install pandas for optimal Flyte log formatting.")
+                Deck("Log History", str(state.log_history))
+
+
 INTEGRATION_TO_CALLBACK = {
     "azure_ml": AzureMLCallback,
     "comet_ml": CometCallback,
@@ -1547,6 +1608,7 @@ def on_save(self, args, state, control, **kwargs):
     "codecarbon": CodeCarbonCallback,
     "clearml": ClearMLCallback,
     "dagshub": DagsHubCallback,
+    "flyte": FlyteCallback,
 }