end: bug fix for step (#759)

* add jsonargparse[signatures] req to lightning

* end: bug fix for step

pyproject.toml

@@ -65,7 +65,7 @@ lgbm = ["lightgbm"]
 huggingface = ["transformers", "datasets"]
 catalyst = ["catalyst>22"]
 fastai = ["fastai"]
-lightning = ["lightning>=2.0", "torch"]
+lightning = ["lightning>=2.0", "torch", "jsonargparse[signatures]>=4.26.1"]
 optuna = ["optuna"]
 all = [
    "dvclive[image,mmcv,tf,xgb,lgbm,huggingface,catalyst,fastai,lightning,optuna,plots,markdown]"

src/dvclive/live.py

@@ -586,7 +586,7 @@ def end(self):
             self.cache(images_path)
 
         # If next_step called before end, don't want to update step number
-        if self._step is not None:
+        if "step" in self.summary:
             self.step = self.summary["step"]
         self.sync()
 

tests/test_step.py

@@ -86,3 +86,23 @@ def test_get_step_control_flow(tmp_dir):
     steps, values = read_history(dvclive, "i")
     assert steps == list(range(10))
     assert values == [float(x) for x in range(10)]
+
+
+def test_set_step_only(tmp_dir):
+    dvclive = Live()
+    dvclive.step = 1
+    dvclive.end()
+
+    assert dvclive.read_latest() == {"step": 1}
+    assert not os.path.exists(os.path.join(tmp_dir, "dvclive", "plots"))
+
+
+def test_step_on_end(tmp_dir):
+    dvclive = Live()
+    for metric in range(3):
+        dvclive.log_metric("m", metric)
+        dvclive.next_step()
+    dvclive.end()
+    assert dvclive.step == metric
+
+    assert dvclive.read_latest() == {"step": metric, "m": metric}