Adds dvclive tracker

docs/source/usage_guides/tracking.md

@@ -20,14 +20,15 @@ There are a large number of experiment tracking API's available, however getting
 
 ## Integrated Trackers
 
-Currently `Accelerate` supports six trackers out-of-the-box:
+Currently `Accelerate` supports seven trackers out-of-the-box:
 
 - TensorBoard
 - WandB
 - CometML
 - Aim
 - MLFlow
 - ClearML
+- DVCLive
 
 To use any of them, pass in the selected type(s) to the `log_with` parameter in [`Accelerate`]:
 ```python

examples/by_feature/deepspeed_with_config_support.py

@@ -220,7 +220,7 @@ def parse_args():
         default="all",
         help=(
             'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
-            ' `"wandb"` and `"comet_ml"`. Use `"all"` (default) to report to all integrations.'
+            ' `"wandb"`, `"comet_ml"`, and `"dvclive"`. Use `"all"` (default) to report to all integrations.'
             "Only applicable when `--with_tracking` is passed."
         ),
     )

examples/by_feature/megatron_lm_gpt_pretraining.py

@@ -216,7 +216,7 @@ def parse_args():
         default="all",
         help=(
             'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
-            ' `"wandb"` and `"comet_ml"`. Use `"all"` (default) to report to all integrations.'
+            ' `"wandb"`, `"comet_ml"`, and `"dvclive"`. Use `"all"` (default) to report to all integrations.'
             "Only applicable when `--with_tracking` is passed."
         ),
     )

setup.py

@@ -25,7 +25,7 @@
 extras["testing"] = extras["test_prod"] + extras["test_dev"]
 extras["rich"] = ["rich"]
 
-extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard"]
+extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard", "dvclive"]
 extras["dev"] = extras["quality"] + extras["testing"] + extras["rich"]
 
 extras["sagemaker"] = [

src/accelerate/test_utils/testing.py

@@ -35,6 +35,7 @@
     is_comet_ml_available,
     is_datasets_available,
     is_deepspeed_available,
+    is_dvclive_available,
     is_mps_available,
     is_pandas_available,
     is_tensorboard_available,
@@ -231,6 +232,13 @@ def require_clearml(test_case):
     return unittest.skipUnless(is_clearml_available(), "test requires clearml")(test_case)
 
 
+def require_dvclive(test_case):
+    """
+    Decorator marking a test that requires dvclive installed. These tests are skipped when dvclive isn't installed
+    """
+    return unittest.skipUnless(is_dvclive_available(), "test requires dvclive")(test_case)
+
+
 def require_pandas(test_case):
     """
     Decorator marking a test that requires pandas installed. These tests are skipped when pandas isn't installed

src/accelerate/tracking.py

@@ -30,6 +30,7 @@
     is_aim_available,
     is_clearml_available,
     is_comet_ml_available,
+    is_dvclive_available,
     is_mlflow_available,
     is_tensorboard_available,
     is_wandb_available,
@@ -57,6 +58,9 @@
 if is_clearml_available():
     _available_trackers.append(LoggerType.CLEARML)
 
+if is_dvclive_available():
+    _available_trackers.append(LoggerType.DVCLIVE)
+
 logger = get_logger(__name__)
 
 
@@ -837,13 +841,87 @@ def _get_title_series(name):
         return name, "train"
 
 
+class DVCLiveTracker(GeneralTracker):
+    """
+    A `Tracker` class that supports `dvclive`. Should be initialized at the start of your script.
+
+    Args:
+        run_name (`str`, *optional*):
+            Ignored for dvclive. See `kwargs` instead.
+        kwargs:
+            Additional key word arguments passed along to [`dvclive.Live()`](https://dvc.org/doc/dvclive/live).
+
+    Example:
+
+    ```py
+    from accelerate import Accelerator
+
+    accelerator = Accelerator(log_with="dvclive")
+    accelerator.init_trackers(project_name="my_project", init_kwargs={"dvclive": {"dir": "my_directory"}})
+    ```
+    """
+
+    name = "dvclive"
+    requires_logging_directory = False
+
+    @on_main_process
+    def __init__(self, run_name: Optional[str] = None, live: Optional[Any] = None, **kwargs):
+        from dvclive import Live
+
+        super().__init__()
+        self.live = live if live is not None else Live(**kwargs)
+
+    @property
+    def tracker(self):
+        return self.live
+
+    @on_main_process
+    def store_init_configuration(self, values: dict):
+        """
+        Logs `values` as hyperparameters for the run. Should be run at the beginning of your experiment. Stores the
+        hyperparameters in a yaml file for future use.
+
+        Args:
+            values (Dictionary `str` to `bool`, `str`, `float`, `int`, or a List or Dict of those types):
+                Values to be stored as initial hyperparameters as key-value pairs. The values need to have type `bool`,
+                `str`, `float`, or `int`.
+        """
+        self.live.log_params(values)
+
+    @on_main_process
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
+        """
+        Logs `values` to the current run.
+
+        Args:
+            values (Dictionary `str` to `str`, `float`, or `int`):
+                Values to be logged as key-value pairs. The values need to have type `str`, `float`, or `int`.
+            step (`int`, *optional*):
+                The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to `dvclive.Live.log_metric()`.
+        """
+        if step is not None:
+            self.live.step = step
+        for k, v in values.items():
+            self.live.log_metric(k, v, **kwargs)
+
+    @on_main_process
+    def finish(self):
+        """
+        Closes `dvclive.Live()`.
+        """
+        self.live.end()
+
+
 LOGGER_TYPE_TO_CLASS = {
     "aim": AimTracker,
     "comet_ml": CometMLTracker,
     "mlflow": MLflowTracker,
     "tensorboard": TensorBoardTracker,
     "wandb": WandBTracker,
     "clearml": ClearMLTracker,
+    "dvclive": DVCLiveTracker,
 }
 
 
@@ -866,6 +944,7 @@ def filter_trackers(
             - `"wandb"`
             - `"comet_ml"`
             - `"mlflow"`
+            - `"dvclive"`
             If `"all"` is selected, will pick up all available trackers in the environment and initialize them. Can
             also accept implementations of `GeneralTracker` for custom trackers, and can be combined with `"all"`.
         logging_dir (`str`, `os.PathLike`, *optional*):

src/accelerate/utils/__init__.py

@@ -52,6 +52,7 @@
     is_cuda_available,
     is_datasets_available,
     is_deepspeed_available,
+    is_dvclive_available,
     is_fp8_available,
     is_ipex_available,
     is_megatron_lm_available,

src/accelerate/utils/dataclasses.py

@@ -340,6 +340,7 @@ class LoggerType(BaseEnum):
         - **TENSORBOARD** -- TensorBoard as an experiment tracker
         - **WANDB** -- wandb as an experiment tracker
         - **COMETML** -- comet_ml as an experiment tracker
+        - **DVCLIVE** -- dvclive as an experiment tracker
     """
 
     ALL = "all"
@@ -349,6 +350,7 @@ class LoggerType(BaseEnum):
     COMETML = "comet_ml"
     MLFLOW = "mlflow"
     CLEARML = "clearml"
+    DVCLIVE = "dvclive"
 
 
 class PrecisionType(BaseEnum):

src/accelerate/utils/imports.py

@@ -297,3 +297,7 @@ def is_xpu_available(check_device=False):
         except RuntimeError:
             return False
     return hasattr(torch, "xpu") and torch.xpu.is_available()
+
+
+def is_dvclive_available():
+    return _is_package_available("dvclive")

tests/test_examples.py

@@ -205,7 +205,7 @@ def test_multi_process_metrics(self):
         run_command(self._launch_args + testargs)
 
     @require_trackers
-    @mock.patch.dict(os.environ, {"WANDB_MODE": "offline"})
+    @mock.patch.dict(os.environ, {"WANDB_MODE": "offline", "DVCLIVE_TEST": "true"})
     def test_tracking(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             testargs = f"""

tests/test_tracking.py

@@ -35,13 +35,19 @@
     TempDirTestCase,
     require_clearml,
     require_comet_ml,
+    require_dvclive,
     require_pandas,
     require_tensorboard,
     require_wandb,
     skip,
 )
 from accelerate.tracking import CometMLTracker, GeneralTracker
-from accelerate.utils import ProjectConfiguration, is_comet_ml_available, is_tensorboard_available
+from accelerate.utils import (
+    ProjectConfiguration,
+    is_comet_ml_available,
+    is_dvclive_available,
+    is_tensorboard_available,
+)
 
 
 if is_comet_ml_available():
@@ -52,6 +58,11 @@
 
     import tensorboard.compat.proto.event_pb2 as event_pb2
 
+if is_dvclive_available():
+    from dvclive.plots.metric import Metric
+    from dvclive.serialize import load_yaml
+    from dvclive.utils import parse_metrics
+
 logger = logging.getLogger(__name__)
 
 
@@ -473,3 +484,41 @@ def test_log(self):
                     "some_string": "",
                 }
                 self.assertDictEqual(data, truth)
+
+
+@require_dvclive
+@mock.patch("dvclive.live.get_dvc_repo", return_value=None)
+class DVCLiveTrackingTest(unittest.TestCase):
+    def test_init_trackers(self, mock_repo):
+        project_name = "test_project_with_config"
+        with tempfile.TemporaryDirectory() as dirpath:
+            accelerator = Accelerator(log_with="dvclive")
+            config = {
+                "num_iterations": 12,
+                "learning_rate": 1e-2,
+                "some_boolean": False,
+                "some_string": "some_value",
+            }
+            init_kwargs = {"dvclive": {"dir": dirpath, "save_dvc_exp": False, "dvcyaml": None}}
+            accelerator.init_trackers(project_name, config, init_kwargs)
+            accelerator.end_training()
+            live = accelerator.trackers[0].live
+            params = load_yaml(live.params_file)
+            assert params == config
+
+    def test_log(self, mock_repo):
+        project_name = "test_project_with_log"
+        with tempfile.TemporaryDirectory() as dirpath:
+            accelerator = Accelerator(log_with="dvclive", project_dir=dirpath)
+            init_kwargs = {"dvclive": {"dir": dirpath, "save_dvc_exp": False, "dvcyaml": None}}
+            accelerator.init_trackers(project_name, init_kwargs=init_kwargs)
+            values = {"total_loss": 0.1, "iteration": 1, "my_text": "some_value"}
+            accelerator.log(values, step=0)
+            accelerator.end_training()
+            live = accelerator.trackers[0].live
+            logs, latest = parse_metrics(live)
+            assert latest == values
+            scalars = os.path.join(live.plots_dir, Metric.subfolder)
+            assert os.path.join(scalars, "total_loss.tsv") in logs
+            assert os.path.join(scalars, "iteration.tsv") in logs
+            assert os.path.join(scalars, "my_text.tsv") in logs