Fix flush (wait) on auxiliary task (obtained using Task.get_task()) should wait on all upload events

trains/backend_interface/metrics/reporter.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import math
+from time import time
 
 try:
     from collections.abc import Iterable  # noqa
@@ -135,6 +136,13 @@ class Reporter(InterfaceBase, AbstractContextManager, SetupUploadMixin, AsyncMan
         self._flush_event.set()
         self._thread.join()
 
+    def wait_for_events(self, timeout=None, step=2.0):
+        tic = time()
+        while self._events or self.get_num_results():
+            self.wait_for_results(timeout=step)
+            if timeout and time() - tic >= timeout:
+                break
+
     def report_scalar(self, title, series, value, iter):
         """
         Report a scalar value

trains/task.py

@@ -1215,6 +1215,9 @@ class Task(_Task):
             self._logger._flush_stdout_handler()
         if self.__reporter:
             self.__reporter.flush()
+            if wait_for_uploads:
+                self.__reporter.wait_for_events()
+
         LoggerRoot.flush()
 
         return True