feat: added various training loggers using confit.Draft

aphp · Feb 18, 2025 · 0bf9310 · 0bf9310
1 parent 30ec72a
commit 0bf9310
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 45 deletions.
diff --git a/edsnlp/training/loggers.py b/edsnlp/training/loggers.py
@@ -365,31 +365,6 @@ def WandBLogger(
     return accelerate.tracking.WandBTracker(project_name, **kwargs)
 
 
-@edsnlp.registry.loggers.register("clearml", auto_draft_in_config=True)
-def ClearMLLogger(
-    project_name: str,
-    **kwargs,
-) -> "accelerate.tracking.ClearMLTracker":  # pragma: no cover
-    """
-    Logger for
-    [ClearML](https://clear.ml/docs/latest/docs/getting_started/ds/ds_first_steps/).
-    This logger is also available via the loggers registry as `clearml`.
-
-    Parameters
-    ----------
-    project_name: str
-        Name of the experiment. Environment variables `CLEARML_PROJECT` and
-        `CLEARML_TASK` have priority over this argument.
-    kwargs: Dict
-        Additional keyword arguments to pass to the ClearML Task object.
-
-    Returns
-    -------
-    accelerate.tracking.ClearMLTracker
-    """
-    return accelerate.tracking.ClearMLTracker(project_name, **kwargs)
-
-
 @edsnlp.registry.loggers.register("mlflow", auto_draft_in_config=True)
 def MLflowLogger(
     project_name: str,
@@ -471,24 +446,63 @@ def CometMLLogger(
     return accelerate.tracking.CometMLTracker(project_name, **kwargs)
 
 
-@edsnlp.registry.loggers.register("dvclive", auto_draft_in_config=True)
-def DVCLiveLogger(
-    live: Any = None,
-    **kwargs,
-) -> "accelerate.tracking.DVCLiveTracker":
-    """
-    Logger for [DVC Live](https://dvc.org/doc/dvclive).
-    This logger is also available via the loggers registry as `dvclive`.
+try:
+    from accelerate.tracking import ClearMLTracker as _ClearMLTracker
 
-    Parameters
-    ----------
-    live: dvclive.Live
-        An instance of `dvclive.Live` to use for logging.
-    kwargs: Dict
-        Additional keyword arguments to pass to the `dvclive.Live` constructor.
+    @edsnlp.registry.loggers.register("clearml", auto_draft_in_config=True)
+    def ClearMLLogger(
+        project_name: str,
+        **kwargs,
+    ) -> "accelerate.tracking.ClearMLTracker":  # pragma: no cover
+        """
+        Logger for
+        [ClearML](https://clear.ml/docs/latest/docs/getting_started/ds/ds_first_steps/).
+        This logger is also available via the loggers registry as `clearml`.
 
-    Returns
-    -------
-    accelerate.tracking.DVCLiveTracker
-    """
-    return accelerate.tracking.DVCLiveTracker(None, live=live, **kwargs)
+        Parameters
+        ----------
+        project_name: str
+            Name of the experiment. Environment variables `CLEARML_PROJECT` and
+            `CLEARML_TASK` have priority over this argument.
+        kwargs: Dict
+            Additional keyword arguments to pass to the ClearML Task object.
+
+        Returns
+        -------
+        accelerate.tracking.ClearMLTracker
+        """
+        return _ClearMLTracker(project_name, **kwargs)
+except ImportError:  # pragma: no cover
+
+    def ClearMLLogger(*args, **kwargs):
+        raise ImportError("ClearMLLogger is not available.")
+
+
+try:
+    from accelerate.tracking import DVCLiveTracker as _DVCLiveTracker
+
+    @edsnlp.registry.loggers.register("dvclive", auto_draft_in_config=True)
+    def DVCLiveLogger(
+        live: Any = None,
+        **kwargs,
+    ) -> "accelerate.tracking.DVCLiveTracker":  # pragma: no cover
+        """
+        Logger for [DVC Live](https://dvc.org/doc/dvclive).
+        This logger is also available via the loggers registry as `dvclive`.
+
+        Parameters
+        ----------
+        live: dvclive.Live
+            An instance of `dvclive.Live` to use for logging.
+        kwargs: Dict
+            Additional keyword arguments to pass to the `dvclive.Live` constructor.
+
+        Returns
+        -------
+        accelerate.tracking.DVCLiveTracker
+        """
+        return _DVCLiveTracker(None, live=live, **kwargs)
+except ImportError:  # pragma: no cover
+
+    def DVCLiveLogger(*args, **kwargs):
+        raise ImportError("DVCLiveLogger is not available.")
diff --git a/pyproject.toml b/pyproject.toml
@@ -296,7 +296,7 @@ where = ["."]
 "csv"         = "edsnlp.training.loggers:CSVLogger"
 "json"        = "edsnlp.training.loggers:JSONLogger"
 "rich"        = "edsnlp.training.loggers:RichLogger"
-"tensorboard" = "edsnlp.training.loggers:TensorboardLogger"
+"tensorboard" = "edsnlp.training.loggers:TensorBoardLogger"
 "aim"         = "edsnlp.training.loggers:AimLogger"
 "wandb"       = "edsnlp.training.loggers:WandBLogger"
 "clearml"     = "edsnlp.training.loggers:ClearMLLogger"