Add AzureML in integrations via dedicated callback

src/transformers/integrations.py

@@ -53,14 +53,20 @@
     except ImportError:
         _has_tensorboard = False
 
+try:
+    from azureml.core.run import Run  # noqa: F401
+
+    _has_azureml = True
+except ImportError:
+    _has_azureml = False
+
 try:
     import mlflow  # noqa: F401
 
     _has_mlflow = True
 except ImportError:
     _has_mlflow = False
 
-
 # No transformer imports above this point
 
 from .file_utils import is_torch_tpu_available
@@ -93,6 +99,10 @@ def is_ray_available():
     return _has_ray
 
 
+def is_azureml_available():
+    return _has_azureml
+
+
 def is_mlflow_available():
     return _has_mlflow
 
@@ -420,6 +430,22 @@ def on_log(self, args, state, control, model=None, logs=None, **kwargs):
                 experiment._log_metrics(logs, step=state.global_step, epoch=state.epoch, framework="transformers")
 
 
+class AzureMLCallback(TrainerCallback):
+    def __init__(self, azureml_run=None):
+        assert _has_azureml, "AzureMLCallback requires azureml to be installed. Run `pip install azureml-sdk`."
+        self.azureml_run = azureml_run
+
+    def on_init_end(self, args, state, control, **kwargs):
+        if self.azureml_run is None and state.is_world_process_zero:
+            self.azureml_run = Run.get_context()
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if self.azureml_run:
+            for k, v in logs.items():
+                if isinstance(v, (int, float)):
+                    self.azureml_run.log(k, v, description=k)
+
+
 class MLflowCallback(TrainerCallback):
     """
     A :class:`~transformers.TrainerCallback` that sends the logs to `MLflow <https://www.mlflow.org/>`__.

src/transformers/trainer.py

@@ -40,6 +40,7 @@
 from .integrations import (
     default_hp_search_backend,
     hp_params,
+    is_azureml_available,
     is_comet_available,
     is_mlflow_available,
     is_optuna_available,
@@ -151,6 +152,11 @@
 if is_ray_available():
     from ray import tune
 
+if is_azureml_available():
+    from .integrations import AzureMLCallback
+
+    DEFAULT_CALLBACKS.append(AzureMLCallback)
+
 logger = logging.get_logger(__name__)