From 037e466b1056b49f29e731c4ebe15889c227a63c Mon Sep 17 00:00:00 2001
From: Kevin Canwen Xu <canwenxu@126.com>
Date: Wed, 23 Jun 2021 14:53:09 +0800
Subject: [PATCH] Add CodeCarbon Integration (#12304)

* Add optional dependency

* Add CodeCarbon integration

* Add CodeCarbon integration

* Add CodeCarbon integration

* typo
---
 setup.py                                      |  3 ++
 src/transformers/dependency_versions_table.py |  1 +
 src/transformers/integrations.py              | 35 +++++++++++++++++++
 3 files changed, 39 insertions(+)

diff --git a/setup.py b/setup.py
index 5c4f075fc71..a9bf9c631bf 100644
--- a/setup.py
+++ b/setup.py
@@ -87,6 +87,7 @@ if stale_egg_info.exists():
 _deps = [
     "Pillow",
     "black==21.4b0",
+    "codecarbon==1.2.0",
     "cookiecutter==1.7.2",
     "dataclasses",
     "datasets",
@@ -252,6 +253,7 @@ extras["serving"] = deps_list("pydantic", "uvicorn", "fastapi", "starlette")
 extras["speech"] = deps_list("soundfile", "torchaudio")
 extras["vision"] = deps_list("Pillow")
 extras["timm"] = deps_list("timm")
+extras["codecarbon"] = deps_list("codecarbon")
 
 extras["sentencepiece"] = deps_list("sentencepiece", "protobuf")
 extras["testing"] = (
@@ -274,6 +276,7 @@ extras["all"] = (
     + extras["vision"]
     + extras["integrations"]
     + extras["timm"]
+    + extras["codecarbon"]
 )
 
 extras["docs_specific"] = deps_list(
diff --git a/src/transformers/dependency_versions_table.py b/src/transformers/dependency_versions_table.py
index 69718eaddf5..47da3c05cb6 100644
--- a/src/transformers/dependency_versions_table.py
+++ b/src/transformers/dependency_versions_table.py
@@ -4,6 +4,7 @@
 deps = {
     "Pillow": "Pillow",
     "black": "black==21.4b0",
+    "codecarbon": "codecarbon==1.2.0",
     "cookiecutter": "cookiecutter==1.7.2",
     "dataclasses": "dataclasses",
     "datasets": "datasets",
diff --git a/src/transformers/integrations.py b/src/transformers/integrations.py
index 9750d6883e7..552ab22e0af 100644
--- a/src/transformers/integrations.py
+++ b/src/transformers/integrations.py
@@ -100,6 +100,10 @@ def is_neptune_available():
     return importlib.util.find_spec("neptune") is not None
 
 
+def is_codecarbon_available():
+    return importlib.util.find_spec("codecarbon") is not None
+
+
 def hp_params(trial):
     if is_optuna_available():
         import optuna
@@ -259,6 +263,8 @@ def get_available_reporting_integrations():
         integrations.append("tensorboard")
     if is_wandb_available():
         integrations.append("wandb")
+    if is_codecarbon_available():
+        integrations.append("codecarbon")
     return integrations
 
 
@@ -718,6 +724,34 @@ class NeptuneCallback(TrainerCallback):
             pass
 
 
+class CodeCarbonCallback(TrainerCallback):
+    """
+    A :class:`~transformers.TrainerCallback` that tracks the CO2 emission of training.
+    """
+
+    def __init__(self):
+        assert (
+            is_codecarbon_available()
+        ), "CodeCarbonCallback requires `codecarbon` to be installed. Run `pip install codecarbon`."
+        import codecarbon
+
+        self._codecarbon = codecarbon
+        self.tracker = None
+
+    def on_init_end(self, args, state, control, **kwargs):
+        if self.tracker is None and state.is_local_process_zero:
+            # CodeCarbon will automatically handle environment variables for configuration
+            self.tracker = self._codecarbon.EmissionsTracker(output_dir=args.output_dir)
+
+    def on_train_begin(self, args, state, control, model=None, **kwargs):
+        if self.tracker and state.is_local_process_zero:
+            self.tracker.start()
+
+    def on_train_end(self, args, state, control, **kwargs):
+        if self.tracker and state.is_local_process_zero:
+            self.tracker.stop()
+
+
 INTEGRATION_TO_CALLBACK = {
     "azure_ml": AzureMLCallback,
     "comet_ml": CometCallback,
@@ -725,6 +759,7 @@ INTEGRATION_TO_CALLBACK = {
     "neptune": NeptuneCallback,
     "tensorboard": TensorBoardCallback,
     "wandb": WandbCallback,
+    "codecarbon": CodeCarbonCallback,
 }