frameworks: Call live.end at the end of training.

Ensure `post_to_studio` inside `live.end` is only called once.

Author: daavoo

src/dvclive/catalyst.py

@@ -27,3 +27,6 @@ def on_epoch_end(self, runner) -> None:
             )
             utils.save_checkpoint(checkpoint, self.model_file)
         self.live.next_step()
+
+    def on_experiment_end(self, runner):  # pylint: disable=unused-argument
+        self.live.end()

src/dvclive/fastai.py

@@ -23,3 +23,6 @@ def after_epoch(self):
         if self.model_file:
             self.learn.save(self.model_file)
         self.live.next_step()
+
+    def after_fit(self):
+        self.live.end()

src/dvclive/huggingface.py

@@ -42,3 +42,12 @@ def on_epoch_end(
             tokenizer = kwargs.get("tokenizer")
             if tokenizer:
                 tokenizer.save_pretrained(self.model_file)
+
+    def on_train_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs
+    ):
+        self.live.end()

src/dvclive/keras.py

@@ -54,3 +54,8 @@ def on_epoch_end(
             else:
                 self.model.save(self.model_file)
         self.live.next_step()
+
+    def on_train_end(
+        self, logs: Optional[Dict] = None
+    ):  # pylint: disable=unused-argument
+        self.live.end()

src/dvclive/lightning.py

@@ -69,3 +69,7 @@ def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None):
             metric_name = standardize_metric_name(metric_name, __name__)
             self.experiment.log_metric(name=metric_name, val=metric_val)
         self.experiment.next_step()
+
+    @rank_zero_only
+    def finalize(self, status: str) -> None:
+        self.experiment.end()

src/dvclive/live.py

@@ -43,7 +43,7 @@ def __init__(
     ):
         self._dir: str = dir
         self._resume: bool = resume or env2bool(env.DVCLIVE_RESUME)
-
+        self._ended: bool = False
         self.studio_url = os.getenv(env.STUDIO_REPO_URL, None)
         self.studio_token = os.getenv(env.STUDIO_TOKEN, None)
         self.rev = None
@@ -243,8 +243,10 @@ def make_report(self):
     def end(self):
         self.make_summary()
         if self.report_mode == "studio":
-            if not post_to_studio(self, "done", logger):
-                logger.warning("`post_to_studio` `done` event failed.")
+            if not self._ended:
+                if not post_to_studio(self, "done", logger):
+                    logger.warning("`post_to_studio` `done` event failed.")
+                self._ended = True
         else:
             self.make_report()
 

src/dvclive/xgb.py

@@ -26,3 +26,7 @@ def after_iteration(self, model, epoch, evals_log):
         if self.model_file:
             model.save_model(self.model_file)
         self.live.next_step()
+
+    def after_training(self, model):
+        self.live.end()
+        return model

tests/test_catalyst.py

@@ -49,13 +49,17 @@ def runner_params():
     }
 
 
-def test_catalyst_callback(tmp_dir, runner, runner_params):
+def test_catalyst_callback(tmp_dir, runner, runner_params, mocker):
+    callback = DVCLiveCallback()
+    live = callback.live
+    spy = mocker.spy(live, "end")
+
     runner.train(
         **runner_params,
         num_epochs=2,
         callbacks=[
             dl.AccuracyCallback(input_key="logits", target_key="targets"),
-            DVCLiveCallback(),
+            callback,
         ],
         logdir="./logs",
         valid_loader="valid",
@@ -64,6 +68,7 @@ def test_catalyst_callback(tmp_dir, runner, runner_params):
         verbose=True,
         load_best_on_end=True,
     )
+    spy.assert_called_once()
 
     assert os.path.exists("dvclive")
 

tests/test_fastai.py

@@ -38,13 +38,16 @@ def data_loader():
     return xor_loader
 
 
-def test_fastai_callback(tmp_dir, data_loader):
+def test_fastai_callback(tmp_dir, data_loader, mocker):
     learn = tabular_learner(data_loader, metrics=accuracy)
     learn.remove_cb(ProgressCallback)
     learn.model_dir = os.path.abspath("./")
     callback = DVCLiveCallback("model")
     live = callback.live
+
+    spy = mocker.spy(live, "end")
     learn.fit_one_cycle(2, cbs=[callback])
+    spy.assert_called_once()
 
     assert os.path.exists(live.dir)
 

tests/test_huggingface.py

@@ -101,7 +101,7 @@ def args():
     )
 
 
-def test_huggingface_integration(tmp_dir, model, args, data):
+def test_huggingface_integration(tmp_dir, model, args, data, mocker):
     trainer = Trainer(
         model,
         args,
@@ -110,8 +110,11 @@ def test_huggingface_integration(tmp_dir, model, args, data):
         compute_metrics=compute_metrics,
     )
     callback = DVCLiveCallback()
+    live = callback.live
+    spy = mocker.spy(live, "end")
     trainer.add_callback(callback)
     trainer.train()
+    spy.assert_called_once()
 
     live = callback.live
     assert os.path.exists(live.dir)