From 5e70a9e6ebe213f48553398ccf2e0adc149c4b93 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Sun, 6 Dec 2020 11:25:20 +0200
Subject: [PATCH] Fix flush (wait) on auxiliary task (obtained using
 Task.get_task()) should wait on all upload events

---
 trains/backend_interface/metrics/reporter.py | 8 ++++++++
 trains/task.py                               | 3 +++
 2 files changed, 11 insertions(+)

diff --git a/trains/backend_interface/metrics/reporter.py b/trains/backend_interface/metrics/reporter.py
index cacee174..949f6f09 100644
--- a/trains/backend_interface/metrics/reporter.py
+++ b/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
diff --git a/trains/task.py b/trains/task.py
index 67d1fd57..3390fc48 100644
--- a/trains/task.py
+++ b/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