mirror of
https://github.com/alicevision/Meshroom.git
synced 2025-04-30 18:57:53 +02:00
94 lines
3.5 KiB
Python
94 lines
3.5 KiB
Python
#!/usr/bin/env python
|
|
# coding:utf-8
|
|
|
|
import os
|
|
import json
|
|
|
|
import simpleFarm
|
|
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(os.path.dirname(currentDir)))
|
|
|
|
class SimpleFarmSubmitter(BaseSubmitter):
|
|
if 'REZ_MESHROOM_VERSION' in os.environ:
|
|
MESHROOM_PACKAGE = "meshroom-{}".format(os.environ.get('REZ_MESHROOM_VERSION', ''))
|
|
else:
|
|
MESHROOM_PACKAGE = None
|
|
|
|
filepath = os.environ.get('SIMPLEFARMCONFIG', os.path.join(currentDir, 'simpleFarmConfig.json'))
|
|
config = json.load(open(filepath))
|
|
|
|
ENGINE = ''
|
|
DEFAULT_TAGS = {'prod': ''}
|
|
|
|
def __init__(self, parent=None):
|
|
super(SimpleFarmSubmitter, self).__init__(name='SimpleFarm', parent=parent)
|
|
self.engine = os.environ.get('MESHROOM_SIMPLEFARM_ENGINE', 'tractor')
|
|
self.share = os.environ.get('MESHROOM_SIMPLEFARM_SHARE', 'vfx')
|
|
self.prod = os.environ.get('PROD', 'mvg')
|
|
|
|
def createTask(self, meshroomFile, node):
|
|
tags = self.DEFAULT_TAGS.copy() # copy to not modify default tags
|
|
nbFrames = node.size
|
|
arguments = {}
|
|
parallelArgs = ''
|
|
print('node: ', node.name)
|
|
if node.isParallelized:
|
|
blockSize, fullSize, nbBlocks = node.nodeDesc.parallelization.getSizes(node)
|
|
parallelArgs = ' --iteration @start'
|
|
arguments.update({'start': 0, 'end': nbBlocks - 1, 'step': 1})
|
|
|
|
tags['nbFrames'] = nbFrames
|
|
tags['prod'] = self.prod
|
|
allRequirements = list()
|
|
allRequirements.extend(self.config['CPU'].get(node.nodeDesc.cpu.name, []))
|
|
allRequirements.extend(self.config['RAM'].get(node.nodeDesc.ram.name, []))
|
|
allRequirements.extend(self.config['GPU'].get(node.nodeDesc.gpu.name, []))
|
|
|
|
task = simpleFarm.Task(
|
|
name=node.nodeType,
|
|
command='{exe} --node {nodeName} "{meshroomFile}" {parallelArgs} --extern'.format(
|
|
exe='meshroom_compute' if self.MESHROOM_PACKAGE else os.path.join(binDir, 'meshroom_compute'),
|
|
nodeName=node.name, meshroomFile=meshroomFile, parallelArgs=parallelArgs),
|
|
tags=tags,
|
|
rezPackages=[self.MESHROOM_PACKAGE] if self.MESHROOM_PACKAGE else None,
|
|
requirements={'service': str(','.join(allRequirements))},
|
|
**arguments)
|
|
return task
|
|
|
|
def submit(self, nodes, edges, filepath):
|
|
name = os.path.splitext(os.path.basename(filepath))[0] + ' [Meshroom]'
|
|
comment = filepath
|
|
nbFrames = max([node.size for node in nodes])
|
|
|
|
mainTags = {
|
|
'prod': self.prod,
|
|
'nbFrames': str(nbFrames),
|
|
'comment': comment,
|
|
}
|
|
allRequirements = list(self.config.get('BASE', []))
|
|
|
|
# Create Job Graph
|
|
job = simpleFarm.Job(name,
|
|
tags=mainTags,
|
|
requirements={'service': str(','.join(allRequirements))},
|
|
)
|
|
|
|
nodeNameToTask = {}
|
|
|
|
for node in nodes:
|
|
task = self.createTask(filepath, node)
|
|
job.addTask(task)
|
|
nodeNameToTask[node.name] = task
|
|
|
|
for u, v in edges:
|
|
nodeNameToTask[u.name].dependsOn(nodeNameToTask[v.name])
|
|
|
|
if self.engine == 'tractor-dummy':
|
|
job.submit(share=self.share, engine='tractor', execute=True)
|
|
return True
|
|
else:
|
|
res = job.submit(share=self.share, engine=self.engine)
|
|
return len(res) > 0
|