Retrieve results asynchronized vatlab/sos-notebook#303

vatlab · Dec 1, 2020 · 7013c75 · 7013c75
1 parent 7b2fef4
commit 7013c75
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/sos/task_engines.py b/src/sos/task_engines.py
@@ -552,7 +552,9 @@ def update_task_status(self, task_id, status):
                     self.running_pending_tasks.pop(task_id)
                 if status == 'completed':
                     # status changed to completed
-                    self.task_results[task_id] = self.agent.receive_result(task_id)
+                    self.task_results[task_id] = self._thread_workers.submit(
+                        self.agent.receive_result,
+                        task_id)
             # for running pending tasks
             if status == 'aborted' and task_id in self.running_pending_tasks:
                 self.pending_tasks.append(task_id)
@@ -561,8 +563,9 @@ def update_task_status(self, task_id, status):
 
     def get_results(self, task_ids):
         while True:
-            if all(x in self.task_status for x in task_ids):
-                return {x: self.task_results[x] for x in task_ids}
+            with threading.Lock():
+                if all(x in self.task_results and not self.task_results[x].running() for x in task_ids):
+                    return {x: self.task_results[x].result() for x in task_ids}
             time.sleep(0.1)
 
     def query_tasks(self,