update 3.0 args in lightning

iterative · Oct 18, 2023 · 66ef8b7 · 66ef8b7
1 parent f2924bd
commit 66ef8b7
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/src/dvclive/lightning.py b/src/dvclive/lightning.py
@@ -56,12 +56,13 @@ def __init__(  # noqa: PLR0913
         prefix="",
         log_model: Union[str, bool] = False,
         experiment=None,
-        dir: Optional[str] = None,  # noqa: A002
+        dir: str = "dvclive",  # noqa: A002
         resume: bool = False,
         report: Optional[str] = None,
-        save_dvc_exp: bool = False,
-        dvcyaml: bool = True,
+        save_dvc_exp: bool = True,
+        dvcyaml: Union[str, bool] = True,
         cache_images: bool = False,
+        exp_message: Optional[str] = None,
     ):
         super().__init__()
         self._prefix = prefix

diff --git a/tests/frameworks/test_lightning.py b/tests/frameworks/test_lightning.py
@@ -260,7 +260,7 @@ def validation_step(self, *args, **kwargs):
         return loss
 
 
-def test_lightning_val_udpates_to_studio(tmp_dir, mocked_dvc_repo, mocked_studio_post):
+def test_lightning_val_updates_to_studio(tmp_dir, mocked_dvc_repo, mocked_studio_post):
     """Test the `self.experiment._latest_studio_step -= 1` logic."""
     mocked_post, _ = mocked_studio_post
 
@@ -281,7 +281,7 @@ def test_lightning_val_udpates_to_studio(tmp_dir, mocked_dvc_repo, mocked_studio
     # 2: update_train_step_metrics
     # 3: log_eval_end_metrics
     plots = calls[3][1]["json"]["plots"]
-    val_loss = plots["dvclive/dvc.yaml::dvclive/plots/metrics/val/loss.tsv"]
+    val_loss = plots["dvclive/plots/metrics/val/loss.tsv"]
     # Without `self.experiment._latest_studio_step -= 1`
     # This would be empty
     assert len(val_loss["data"]) == 1