Add submit_high_priority for high priority jobs

loky/process_executor.py

@@ -533,8 +533,8 @@ def weakref_cb(_,
         # A ctx.SimpleQueue of _ResultItems generated by the process workers.
         self.result_queue = executor._result_queue
 
-        # A queue.Queue of work ids e.g. Queue([5, 6, ...]).
-        self.work_ids_queue = executor._work_ids
+        # A _WorkTracker of work ids.
+        self.work_ids_tracker = executor._work_ids
 
         # A dict mapping work ids to _WorkItems e.g.
         #     {5: <_WorkItem...>, 6: <_WorkItem...>, ...}
@@ -584,7 +584,7 @@ def add_call_item_to_queue(self):
             if self.call_queue.full():
                 return
             try:
-                work_id = self.work_ids_queue.get(block=False)
+                work_id = self.work_ids_tracker.get()
             except queue.Empty:
                 return
             else:
@@ -930,6 +930,26 @@ class ShutdownExecutorError(RuntimeError):
     """
 
 
+class _WorkTracker(object):
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._work_ids = []
+
+    def get(self):
+        with self._lock:
+            if not len(self._work_ids):
+                raise queue.Empty()
+
+            return self._work_ids.pop(0)
+
+    def put(self, work_id, high_priority=False):
+        with self._lock:
+            if high_priority:
+                self._work_ids.insert(0, work_id)
+            else:
+                self._work_ids.append(work_id)
+
+
 class ProcessPoolExecutor(_base.Executor):
 
     _at_exit = None
@@ -1000,7 +1020,7 @@ def __init__(self, max_workers=None, job_reducers=None,
         self._queue_count = 0
         self._pending_work_items = {}
         self._running_work_items = []
-        self._work_ids = queue.Queue()
+        self._work_ids = _WorkTracker()
         self._processes_management_lock = self._context.Lock()
         self._executor_manager_thread = None
         self._shutdown_lock = threading.Lock()
@@ -1109,6 +1129,16 @@ def _ensure_executor_running(self):
             self._start_executor_manager_thread()
 
     def submit(self, fn, *args, **kwargs):
+        """ Submits a job to be executed. """
+        return self._submit(fn, high_priority=False, args=args, kwargs=kwargs)
+
+    def submit_high_priority(self, fn, *args, **kwargs):
+        """ Submits a high priority job to be executed. The job will be placed at the top
+            of the queue.
+        """
+        return self._submit(fn, high_priority=True, args=args, kwargs=kwargs)
+
+    def _submit(self, fn, high_priority, args, kwargs):
         with self._flags.shutdown_lock:
             if self._flags.broken is not None:
                 raise self._flags.broken
@@ -1126,13 +1156,15 @@ def submit(self, fn, *args, **kwargs):
             w = _WorkItem(f, fn, args, kwargs)
 
             self._pending_work_items[self._queue_count] = w
-            self._work_ids.put(self._queue_count)
+            self._work_ids.put(self._queue_count, high_priority=high_priority)
             self._queue_count += 1
+
             # Wake up queue management thread
             self._executor_manager_thread_wakeup.wakeup()
 
             self._ensure_executor_running()
             return f
+
     submit.__doc__ = _base.Executor.submit.__doc__
 
     def map(self, fn, *iterables, **kwargs):

tests/test_reusable_executor.py

@@ -6,7 +6,7 @@
 import pytest
 import warnings
 import threading
-from time import sleep
+from time import sleep, time
 from multiprocessing import util, current_process
 from pickle import PicklingError, UnpicklingError
 from distutils.version import LooseVersion
@@ -827,3 +827,24 @@ def test_reusable_initializer(self):
         executor = get_reusable_executor(max_workers=4)
         for x in executor.map(self._test_initializer, delay=.1):
             assert x == 'uninitialized'
+
+
+class TestExecutorHighPrioritySubmit(ReusableExecutorMixin):
+    def _now(self):
+        return time()
+
+    def test_high_priority(self):
+        executor = get_reusable_executor(max_workers=1)
+
+        # Submit a bunch of jobs.
+        later_job_futures = []
+        for _ in range(0, 10):
+            later_job_futures.append(executor.submit(self._now))
+
+        high_priority_job = executor.submit_high_priority(self._now)
+
+        high_priority_result = high_priority_job.result()
+        last_job_result = later_job_futures[-1].result()
+
+        # Ensure the high priority job ran before the last submitted job.
+        assert high_priority_result < last_job_result