Make file_pusher block like it promised

wandb · Aug 7, 2020 · e6cb00a · e6cb00a
1 parent b24c416
commit e6cb00a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/wandb/internal/file_pusher.py b/wandb/internal/file_pusher.py
@@ -128,6 +128,8 @@ def commit_artifact(self, artifact_id, on_commit=None):
 
     def finish(self):
         self._incoming_queue.put(step_checksum.RequestFinish())
+        while self.is_alive():
+            time.sleep(0.2)
 
     def is_alive(self):
         return (self._step_checksum.is_alive()

diff --git a/wandb/internal/sender.py b/wandb/internal/sender.py
@@ -135,9 +135,6 @@ def handle_final(self, data):
 
         if self._pusher:
             self._pusher.finish()
-            # TODO: move this into pusher.finish()
-            while self._pusher.is_alive():
-                time.sleep(0.1)
 
         if self._fs:
             # TODO(jhr): now is a good time to output pending output lines
@@ -157,9 +154,7 @@ def handle_final(self, data):
         # TODO(david): this info should be in exit_result footer?
         if self._pusher:
             self._pusher.print_status()
-
-        # Dont let anyone else touch pusher
-        self._pusher = None
+            self._pusher = None
 
     def handle_run(self, data):
         run = data.run