utils: Add catch_and_warn. (#673)

dvc: Extract `ensure_dir_is_tracked` to `dvc`.

Co-authored-by: dberenbaum <[email protected]>

Author: daavoo

src/dvclive/dvc.py

@@ -150,3 +150,19 @@ def find_overlapping_stage(dvc_repo: "Repo", path: StrPath) -> Optional["Stage"]
             if str(out.fs_path) in abs_path:
                 return stage
     return None
+
+
+def ensure_dir_is_tracked(directory: str, dvc_repo: "Repo") -> None:
+    from pathspec import PathSpec
+
+    dir_spec = PathSpec.from_lines("gitwildmatch", [directory])
+    outs_spec = PathSpec.from_lines(
+        "gitwildmatch", [str(o) for o in dvc_repo.index.outs]
+    )
+    paths_to_track = [
+        f
+        for f in dvc_repo.scm.untracked_files()
+        if (dir_spec.match_file(f) and not outs_spec.match_file(f))
+    ]
+    if paths_to_track:
+        dvc_repo.scm.add(paths_to_track)

src/dvclive/live.py

@@ -5,13 +5,14 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Set, Union
 
+from dvc.exceptions import DvcException
 from dvc_studio_client.post_live_metrics import post_live_metrics
 from funcy import set_in
-from pathspec import PathSpec
 from ruamel.yaml.representer import RepresenterError
 
 from . import env
 from .dvc import (
+    ensure_dir_is_tracked,
     find_overlapping_stage,
     get_dvc_repo,
     get_random_exp_name,
@@ -31,6 +32,7 @@
 from .studio import get_dvc_studio_config, get_studio_updates
 from .utils import (
     StrPath,
+    catch_and_warn,
     clean_and_copy_into,
     env2bool,
     inside_notebook,
@@ -124,6 +126,7 @@ def _init_cleanup(self):
         if self.dvc_file and os.path.exists(self.dvc_file):
             os.remove(self.dvc_file)
 
+    @catch_and_warn(DvcException, logger)
     def _init_dvc(self):
         from dvc.scm import NoSCM
 
@@ -453,35 +456,31 @@ def log_artifact(
                         name,
                     )
 
+    @catch_and_warn(DvcException, logger)
     def cache(self, path):
-        try:
-            if self._inside_dvc_exp:
-                existing_stage = find_overlapping_stage(self._dvc_repo, path)
-
-                if existing_stage:
-                    if existing_stage.cmd:
-                        logger.info(
-                            f"Skipping `dvc add {path}` because it is already being"
-                            " tracked automatically as an output of `dvc exp run`."
-                        )
-                        return  # skip caching
-                    logger.warning(
-                        f"To track '{path}' automatically during `dvc exp run`:"
-                        f"\n1. Run `dvc remove {existing_stage.addressing}` "
-                        "to stop tracking it outside the pipeline."
-                        "\n2. Add it as an output of the pipeline stage."
-                    )
-                else:
-                    logger.warning(
-                        f"To track '{path}' automatically during `dvc exp run`, "
-                        "add it as an output of the pipeline stage."
-                    )
+        if self._inside_dvc_exp:
+            existing_stage = find_overlapping_stage(self._dvc_repo, path)
 
-            stage = self._dvc_repo.add(str(path))
+            if existing_stage:
+                if existing_stage.cmd:
+                    logger.info(
+                        f"Skipping `dvc add {path}` because it is already being"
+                        " tracked automatically as an output of `dvc exp run`."
+                    )
+                    return  # skip caching
+                logger.warning(
+                    f"To track '{path}' automatically during `dvc exp run`:"
+                    f"\n1. Run `dvc remove {existing_stage.addressing}` "
+                    "to stop tracking it outside the pipeline."
+                    "\n2. Add it as an output of the pipeline stage."
+                )
+            else:
+                logger.warning(
+                    f"To track '{path}' automatically during `dvc exp run`, "
+                    "add it as an output of the pipeline stage."
+                )
 
-        except Exception as e:  # noqa: BLE001
-            logger.warning(f"Failed to dvc add {path}: {e}")
-            return
+        stage = self._dvc_repo.add(str(path))
 
         dvc_file = stage[0].addressing
 
@@ -539,7 +538,10 @@ def end(self):
         if self._dvcyaml:
             self.make_dvcyaml()
 
-        self._ensure_paths_are_tracked_in_dvc_exp()
+        if self._inside_dvc_exp and self._dvc_repo:
+            catch_and_warn(DvcException, logger)(ensure_dir_is_tracked)(
+                self.dir, self._dvc_repo
+            )
 
         self.save_dvc_exp()
 
@@ -582,35 +584,12 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self._inside_with = False
         self.end()
 
-    def _ensure_paths_are_tracked_in_dvc_exp(self):
-        if self._inside_dvc_exp and self._dvc_repo:
-            dir_spec = PathSpec.from_lines("gitwildmatch", [self.dir])
-            outs_spec = PathSpec.from_lines(
-                "gitwildmatch", [str(o) for o in self._dvc_repo.index.outs]
-            )
-            try:
-                paths_to_track = [
-                    f
-                    for f in self._dvc_repo.scm.untracked_files()
-                    if (dir_spec.match_file(f) and not outs_spec.match_file(f))
-                ]
-                if paths_to_track:
-                    self._dvc_repo.scm.add(paths_to_track)
-            except Exception as e:  # noqa: BLE001
-                logger.warning(f"Failed to git add paths:\n{e}")
-
+    @catch_and_warn(DvcException, logger, mark_dvclive_only_ended)
     def save_dvc_exp(self):
         if self._save_dvc_exp:
-            from dvc.exceptions import DvcException
-
-            try:
-                self._experiment_rev = self._dvc_repo.experiments.save(
-                    name=self._exp_name,
-                    include_untracked=self._include_untracked,
-                    force=True,
-                    message=self._exp_message,
-                )
-            except DvcException as e:
-                logger.warning(f"Failed to save experiment:\n{e}")
-            finally:
-                mark_dvclive_only_ended()
+            self._experiment_rev = self._dvc_repo.experiments.save(
+                name=self._exp_name,
+                include_untracked=self._include_untracked,
+                force=True,
+                message=self._exp_message,
+            )

src/dvclive/utils.py

@@ -155,3 +155,19 @@ def isinstance_without_import(val, module, name):
         if (cls.__module__, cls.__name__) == (module, name):
             return True
     return False
+
+
+def catch_and_warn(exception, logger, on_finally=None):
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except exception as e:
+                logger.warning(f"Error in {func.__name__}: {e}")
+            finally:
+                if on_finally is not None:
+                    on_finally()
+
+        return wrapper
+
+    return decorator

tests/test_dvc.py

@@ -1,6 +1,7 @@
 import os
 
 import pytest
+from dvc.exceptions import DvcException
 from dvc.repo import Repo
 from dvc.scm import NoSCM
 from PIL import Image
@@ -279,7 +280,7 @@ def test_errors_on_git_add_are_catched(tmp_dir, mocked_dvc_repo, monkeypatch):
     monkeypatch.setenv(DVC_EXP_BASELINE_REV, "foo")
     monkeypatch.setenv(DVC_EXP_NAME, "bar")
     mocked_dvc_repo.scm.untracked_files.return_value = ["dvclive/metrics.json"]
-    mocked_dvc_repo.scm.add.side_effect = Exception("foo")
+    mocked_dvc_repo.scm.add.side_effect = DvcException("foo")
 
     with Live(report=None) as live:
         live.summary["foo"] = 1

tests/test_log_artifact.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 
 import pytest
+from dvc.exceptions import DvcException
 
 from dvclive import Live
 from dvclive.serialize import load_yaml
@@ -216,7 +217,7 @@ def test_log_artifact_attrs(tmp_dir, mocked_dvc_repo):
 
 def test_log_artifact_type_model_when_dvc_add_fails(tmp_dir, mocker, mocked_dvc_repo):
     (tmp_dir / "model.pth").touch()
-    mocked_dvc_repo.add.side_effect = Exception
+    mocked_dvc_repo.add.side_effect = DvcException("foo")
     with Live(save_dvc_exp=True) as live:
         live.log_artifact("model.pth", type="model")