Distributing python code across nodes in slurm

Q&A
1

I have a computationally expensive simulation function I am looking to distribute accross a multi-node cluster. The code looks something like this:

input_tasks = [input_0, input_1, ..., input_n]
for i in input_tasks:
    expensive_function(i)

I am running the code from a node with high compute and I am looking to distribute the function inputs to many nodes with varying compute power. The highest compute nodes should take priority and always pick up the next task if they are free. A pseudocode of what I wish to do is written below.

input_tasks = [input_0, input_1, ..., input_n]
available_nodes_ranked_by_compute = [node_0, node_1, ..., etc]
While(input_tasks): 
   i = input_tasks.pop(0)
   #get best current node or wait for a node to free up
   node_i = available_nodes_ranked_by_compute.pop(0)
   expensive_function(i, node_i)
   #add node back to avaiable node list when its done
   available_nodes_ranked_by_compute.append(node_i)
   #re-sort available nodes by compute

I want a way to maintain a dynamic list/heap that maintains the currrently unused nodes on the cluster so I can use it to execute all my tasks. Is there a basic way to do this?

2

Use Dask, it’s made exactly for this kind of problems

3

‘heapq’ module may be used for this purpose:
([heapq — Heap queue algorithm — Python 3.11.5 documentation] - Heap queue algorithm.

image
image981×464 18.5 KB

Here, ‘available_nodes’ is managed as a priority queue using Python’s heapq module. The ‘get_node_priority’ function is used to determine the priority of each node based on their computing power.