Make MPITaskScheduler prioritize large tasks

parsl/executors/high_throughput/mpi_resource_management.py

@@ -196,7 +196,8 @@ def put_task(self, task_package: dict):
         _f, _args, _kwargs, resource_spec = unpack_res_spec_apply_message(task_package["buffer"])
 
         nodes_needed = resource_spec.get("num_nodes")
-        prioritized_task = PrioritizedTask(priority=nodes_needed,
+        # Prioritize large jobs
+        prioritized_task = PrioritizedTask(priority=-1 * nodes_needed,
                                            task=task_package,
                                            unpacked_task=(_f, _args, _kwargs, resource_spec),
                                            nodes_needed=nodes_needed)

parsl/tests/test_mpi_apps/test_mpi_scheduler.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import pickle
+import random
 from unittest import mock
 
 import pytest
@@ -218,3 +219,38 @@ def test_tiny_large_loop():
         got_result = scheduler.get_result(True, 1)
 
     assert got_result == result_pkl
+
+
+@pytest.mark.local
+def test_larger_jobs_prioritized():
+    """Larger jobs should be scheduled first"""
+
+    task_q, result_q = SpawnContext.Queue(), SpawnContext.Queue()
+    scheduler = MPITaskScheduler(task_q, result_q)
+
+    max_nodes = len(scheduler.available_nodes)
+
+    # The first task will get scheduled with all the nodes,
+    # and the remainder hits the backlog queue.
+    node_request_list = [max_nodes] + [random.randint(1, 4) for _i in range(8)]
+
+    for task_id, num_nodes in enumerate(node_request_list):
+        mock_task_buffer = pack_res_spec_apply_message("func", "args", "kwargs",
+                                                       resource_specification={
+                                                           "num_nodes": num_nodes,
+                                                           "ranks_per_node": 2
+                                                       })
+        task_package = {"task_id": task_id, "buffer": mock_task_buffer}
+        scheduler.put_task(task_package)
+
+    # Confirm that the tasks are sorted in decreasing order
+    output_priority = []
+    for i in range(len(node_request_list) - 1):
+        p_task = scheduler._backlog_queue.get()
+        output_priority.append(p_task.nodes_needed)
+
+    # Remove the first large job that blocks the nodes and forces following
+    # tasks into backlog
+    expected_priority = node_request_list[1:]
+    expected_priority.sort(reverse=True)
+    assert expected_priority == output_priority, "Expected nodes in decreasing sorted order"