mirror of
https://github.com/alicevision/Meshroom.git
synced 2025-06-06 12:51:57 +02:00
122 lines
3.4 KiB
Python
122 lines
3.4 KiB
Python
|
|
#!/usr/bin/env python
|
|
# coding:utf-8
|
|
|
|
import os
|
|
import json
|
|
import subprocess
|
|
|
|
from meshroom.core.desc import Level
|
|
from meshroom.core.submitter import BaseSubmitter
|
|
|
|
currentDir = os.path.dirname(os.path.realpath(__file__))
|
|
binDir = os.path.dirname(os.path.dirname(currentDir))
|
|
|
|
class Vertex:
|
|
|
|
def __init__(self, node):
|
|
self.pid = -1
|
|
self.node = node
|
|
self.depends = list()
|
|
|
|
class SlurmSubmitter(BaseSubmitter):
|
|
|
|
def __init__(self, parent=None):
|
|
super(SlurmSubmitter, self).__init__(name='Slurm', parent=parent)
|
|
self.vertices = list()
|
|
self.verticesByNode = dict()
|
|
|
|
def start(self, node, dependencies, meshroomFile):
|
|
|
|
nbFrames = node.size
|
|
arguments = {}
|
|
parallelArgs = ''
|
|
|
|
binary = os.path.join(binDir, 'bin/meshroom_compute')
|
|
bash_content = os.path.join(currentDir, 'slurmCommand.sh')
|
|
jobNameParameter = "--job-name=%s" % node.name
|
|
dependenciesString = ""
|
|
|
|
print(node.name)
|
|
|
|
command = [
|
|
"sbatch",
|
|
"--parsable",
|
|
jobNameParameter,
|
|
]
|
|
|
|
if (len(dependencies) > 0):
|
|
dependenciesString = "--dependency=afterok"
|
|
for dependency in dependencies:
|
|
dependenciesString = "%s:%d" % (dependenciesString, dependency.pid)
|
|
command.append(dependenciesString)
|
|
|
|
if node.nodeDesc.gpu.name != "NONE":
|
|
command.append("--partition=gpu")
|
|
command.append("--gres=gpu:1")
|
|
|
|
countCores=24
|
|
if node.isParallelized:
|
|
blockSize, fullSize, nbBlocks = node.nodeDesc.parallelization.getSizes(node)
|
|
|
|
if nbBlocks > 1:
|
|
command.append("--array=1-%d" % nbBlocks)
|
|
countCores = 1
|
|
|
|
command.append("--cpus-per-task=%d" % countCores)
|
|
|
|
additional = [
|
|
bash_content,
|
|
binary,
|
|
node.name,
|
|
meshroomFile,
|
|
str(countCores)
|
|
]
|
|
|
|
returnString = subprocess.check_output(command + additional)
|
|
returnId = int(returnString)
|
|
|
|
return returnId
|
|
|
|
def submit(self, nodes, edges, filepath):
|
|
|
|
#build a list of nodes with dependencies
|
|
for item in nodes:
|
|
v = Vertex(item)
|
|
self.vertices.append(v)
|
|
self.verticesByNode[item] = v
|
|
|
|
#set dependencies
|
|
for edge in edges:
|
|
node = self.verticesByNode[edge[0]]
|
|
dependency = self.verticesByNode[edge[1]]
|
|
node.depends.append(dependency)
|
|
|
|
#While all the graph is not submitted
|
|
while True:
|
|
|
|
#Try to find a node with all its dependencies submitted
|
|
somethingFound = False
|
|
|
|
for vertex in self.vertices:
|
|
|
|
#check if vertex was already launched
|
|
if vertex.pid >= 0:
|
|
continue
|
|
|
|
#check That all dependency were launched
|
|
valid = True
|
|
for dependency in vertex.depends:
|
|
if dependency.pid < 0:
|
|
valid = False
|
|
break
|
|
|
|
#launch
|
|
if valid:
|
|
vertex.pid = self.start(vertex.node, vertex.depends, filepath)
|
|
somethingFound = True
|
|
|
|
if somethingFound == False:
|
|
break
|
|
|
|
return True
|