From b987e667df18f8a65470b6bd72fd401c6f904470 Mon Sep 17 00:00:00 2001
From: Ivan Shcheklein <shcheklein@gmail.com>
Date: Sun, 9 Jun 2024 17:37:46 -0700
Subject: [PATCH] fix(studio): wait for studio metrics publish to complete on
 end (#827)

---
 src/dvclive/live.py          |  7 +++++++
 tests/test_post_to_studio.py | 25 +++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/src/dvclive/live.py b/src/dvclive/live.py
index 67d0855..8f49fc9 100644
--- a/src/dvclive/live.py
+++ b/src/dvclive/live.py
@@ -905,6 +905,11 @@ def worker():
 
         self._studio_queue.put(self)
 
+    def _wait_for_studio_updates_posted(self):
+        if self._studio_queue:
+            logger.debug("Waiting for studio updates to be posted")
+            self._studio_queue.join()
+
     def end(self):
         """
         Signals that the current experiment has ended.
@@ -946,6 +951,8 @@ def end(self):
 
         self.save_dvc_exp()
 
+        self._wait_for_studio_updates_posted()
+
         # Mark experiment as done
         post_to_studio(self, "done")
 
diff --git a/tests/test_post_to_studio.py b/tests/test_post_to_studio.py
index 6aee074..a9ac838 100644
--- a/tests/test_post_to_studio.py
+++ b/tests/test_post_to_studio.py
@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import pytest
+import time
 from dvc.env import DVC_EXP_GIT_REMOTE
 from dvc_studio_client import DEFAULT_STUDIO_URL
 from dvc_studio_client.env import DVC_STUDIO_REPO_URL, DVC_STUDIO_TOKEN
@@ -186,6 +187,30 @@ def test_post_to_studio_done_only_once(tmp_dir, mocked_dvc_repo, mocked_studio_p
     assert expected_done_calls == actual_done_calls
 
 
+def test_studio_updates_posted_on_end(tmp_path, mocked_dvc_repo, mocked_studio_post):
+    mocked_post, valid_response = mocked_studio_post
+    metrics_file = tmp_path / "metrics.json"
+    metrics_content = "metrics"
+
+    def long_post(*args, **kwargs):
+        # in case of `data` `long_post` should be called from a separate thread,
+        # meanwhile main thread go forward without slowing down, so if there is no
+        # some kind of wait in the Live main thread, then it will complete before
+        # we even can have a chance to write the file below
+        if kwargs["json"]["type"] == "data":
+            time.sleep(1)
+            metrics_file.write_text(metrics_content)
+
+        return valid_response
+
+    mocked_post.side_effect = long_post
+
+    with Live() as live:
+        live.log_metric("foo", 1)
+
+    assert metrics_file.read_text() == metrics_content
+
+
 @pytest.mark.studio()
 def test_post_to_studio_skip_start_and_done_on_env_var(
     tmp_dir, mocked_dvc_repo, mocked_studio_post, monkeypatch