在 Python 脚本中动态检查和分配 SLURM 节点

在 Python 脚本中动态检查和分配 SLURM 节点

我有一个计算量很大的模拟函数,希望将其分布在多节点集群中。代码如下所示:

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

我正在从具有高计算能力的节点运行代码,并希望将函数输入分发到具有不同计算能力的许多节点。最高计算能力的节点应优先考虑,并且如果它们空闲,则始终接手下一个任务。下面是我希望执行的操作的伪代码。

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 
   

我对分布式计算和 SLURM 使用还比较陌生,因此我不确定如何检查特定节点当前是否正在使用(在 Python 脚本的上下文中)。我想要一种方法来维护一个动态列表/堆,以维护集群上当前未使用的节点,以便我可以使用它来执行我的所有任务。有办法做到这一点吗?

相关内容