Fix bug in MPITaskScheduler where tasks with identical priorities fail with TypeError (#3794)

# Description

The `MPITaskScheduler` uses Python's PriorityQueue to prioritize tasks
based on the number of nodes requested.
When items with identical priorities are are submitted to the
PriorityQueue, they attempt to sort based on the task dict which fails
with TypeError unhashable type: dict. This PR adds a new
`PrioritizedTask` dataclass that sets the task element to
`field(compare=False)`.

I'm splitting changes in #3783 to keep the PR concise. This is split 1
of 3.


# Changed Behaviour

Fixes the bug described above.

## Type of change

Choose which options apply, and delete the ones which do not apply.

- Bug fix

Author: yadudoc

parsl/executors/high_throughput/mpi_resource_management.py

@@ -4,6 +4,7 @@
 import pickle
 import queue
 import subprocess
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Dict, List, Optional
 
@@ -69,6 +70,14 @@ def __str__(self):
         return f"MPINodesUnavailable(requested={self.requested} available={self.available})"
 
 
+@dataclass(order=True)
+class PrioritizedTask:
+    # Comparing dict will fail since they are unhashable
+    # This dataclass limits comparison to the priority field
+    priority: int
+    task: Dict = field(compare=False)
+
+
 class TaskScheduler:
     """Default TaskScheduler that does no taskscheduling
 
@@ -111,7 +120,7 @@ def __init__(
         super().__init__(pending_task_q, pending_result_q)
         self.scheduler = identify_scheduler()
         # PriorityQueue is threadsafe
-        self._backlog_queue: queue.PriorityQueue = queue.PriorityQueue()
+        self._backlog_queue: queue.PriorityQueue[PrioritizedTask] = queue.PriorityQueue()
         self._map_tasks_to_nodes: Dict[str, List[str]] = {}
         self.available_nodes = get_nodes_in_batchjob(self.scheduler)
         self._free_node_counter = SpawnContext.Value("i", len(self.available_nodes))
@@ -169,7 +178,7 @@ def put_task(self, task_package: dict):
                 allocated_nodes = self._get_nodes(nodes_needed)
             except MPINodesUnavailable:
                 logger.info(f"Not enough resources, placing task {tid} into backlog")
-                self._backlog_queue.put((nodes_needed, task_package))
+                self._backlog_queue.put(PrioritizedTask(nodes_needed, task_package))
                 return
             else:
                 resource_spec["MPI_NODELIST"] = ",".join(allocated_nodes)
@@ -183,8 +192,8 @@ def put_task(self, task_package: dict):
     def _schedule_backlog_tasks(self):
         """Attempt to schedule backlogged tasks"""
         try:
-            _nodes_requested, task_package = self._backlog_queue.get(block=False)
-            self.put_task(task_package)
+            prioritized_task = self._backlog_queue.get(block=False)
+            self.put_task(prioritized_task.task)
         except queue.Empty:
             return
         else:

parsl/tests/test_mpi_apps/test_mpi_scheduler.py

@@ -161,3 +161,28 @@ def test_MPISched_contention():
     assert task_on_worker_side['task_id'] == 2
     _, _, _, resource_spec = unpack_res_spec_apply_message(task_on_worker_side['buffer'])
     assert len(resource_spec['MPI_NODELIST'].split(',')) == 8
+
+
+@pytest.mark.local
+def test_hashable_backlog_queue():
+    """Run multiple large tasks that to force entry into backlog_queue
+    where queue.PriorityQueue expects hashability/comparability
+    """
+
+    task_q, result_q = SpawnContext.Queue(), SpawnContext.Queue()
+    scheduler = MPITaskScheduler(task_q, result_q)
+
+    assert scheduler.available_nodes
+    assert len(scheduler.available_nodes) == 8
+
+    assert scheduler._free_node_counter.value == 8
+
+    for i in range(3):
+        mock_task_buffer = pack_res_spec_apply_message("func", "args", "kwargs",
+                                                       resource_specification={
+                                                           "num_nodes": 8,
+                                                           "ranks_per_node": 2
+                                                       })
+        task_package = {"task_id": i, "buffer": mock_task_buffer}
+        scheduler.put_task(task_package)
+    assert scheduler._backlog_queue.qsize() == 2, "Expected 2 backlogged tasks"