Fixing finite generator race condition (keras-team#8162)

datumbox · fchollet · commit 9c79e01a61b3 · 2017-10-19T12:19:00.000-07:00
* Fixing race condition

Check again for empty queue after detecting that all threads have finished.

* Updating the test to check for all set items
diff --git a/keras/utils/data_utils.py b/keras/utils/data_utils.py
@@ -654,7 +654,7 @@ def get(self):
                     yield inputs
             else:
                 all_finished = all([not thread.is_alive() for thread in self._threads])
-                if all_finished:
+                if all_finished and self.queue.empty():
                     raise StopIteration()
                 else:
                     time.sleep(self.wait_time)
diff --git a/tests/keras/utils/data_utils_test.py b/tests/keras/utils/data_utils_test.py
@@ -264,7 +264,7 @@ def test_finite_generator_enqueuer_threads():
     acc = []
     for output in gen_output:
         acc.append(int(output[0, 0, 0, 0]))
-    assert len(set(acc) - set(range(100))) == 0, "Output is not the same"
+    assert set(acc) == set(range(100)), "Output is not the same"
     enqueuer.stop()
 
 
