[submitter] Update local farm to work with preprocess/postprocess chunks

Author: Alxiice

bin/meshroom_compute

@@ -138,20 +138,21 @@ if args.node:
         print(f"InputNode: No computation to do.")
         sys.exit(0)
 
-    if args.iteration == ChunkIndex.NONE:
-        chunks = node.chunks
-    elif args.preprocess:
+    if args.preprocess:
         chunks = [node._preprocessChunk]
     elif args.postprocess:
         chunks = [node._postprocessChunk]
+    elif args.iteration == ChunkIndex.NONE:  # Warning : default value
+        chunks = node.chunks
     else:
         chunks = [node.chunks[args.iteration]]
-
+    
     if not args.forceStatus and not args.forceCompute:
         for chunk in chunks:
             if chunk.status.status in submittedStatuses:
                 # Particular case for the local isolated, the node status is set to RUNNING by the submitter directly.
-                # We ensure that no other instance has started to compute, by checking that the computeSessionUid is empty.
+                # We ensure that no other instance has start
+                # ed to compute, by checking that the computeSessionUid is empty.
                 if chunk.node.getMrNodeType() == meshroom.core.MrNodeType.NODE and \
                     not chunk.status.computeSessionUid and node._nodeStatus.submitterSessionUid:
                     continue

localfarm/localFarmBackend.py

@@ -358,6 +358,7 @@ def startTask(self, task: Task):
             with open(task.logFile, "w") as log:
                 log.write(f"# ========== Starting task {task.tid} at {task.started_at.isoformat()}"
                           f" (command=\"{task.command}\") ==========\n")
+                log.write(f"# metadata: {task.metadata}\n")
                 log.write(f"# process_env:\n")
                 log.write(f"# Additional env variables:\n")
                 for _k, _v in additional_env.items():

meshroom/core/submitter.py

@@ -299,14 +299,12 @@ def display(self, task:OrderedTask=None, level=0):
     def iterOnTasks(self, current:OrderedTask=None):
         if current is None:
             current = self.rootTask
-        items = [current]
+        yield current
         for task in current.dependencies:
-            items.extend(self.iterOnTasks(task))
-        return items
+            yield from self.iterOnTasks(task)
 
     def __iter__(self):
-        for item in self.iterOnTasks():
-            yield item
+        yield from self.iterOnTasks()
 
 
 class BaseSubmittedJob:

meshroom/submitters/localFarmSubmitter.py

@@ -78,7 +78,7 @@ def getRequestPackages(packagesDelimiter="=="):
     return list(reqPackages)
 
 
-def rezWrapCommand(cmd, useCurrentContext=False, useRequestedContext=True, otherRezPkg: list[str] = None):
+def rezWrapCommand(cmd, useCurrentContext=False, useRequestedContext=True, otherRezPkg: list[str] = None, additionalEnv: dict=None):
     """ Wrap command to be runned using rez.
     :param cmd: command to run
     :type cmd: bool
@@ -108,7 +108,11 @@ def rezWrapCommand(cmd, useCurrentContext=False, useRequestedContext=True, other
             rezBin = os.path.join(os.environ["REZ_PACKAGES_ROOT"], "bin/rez")
         elif shutil.which("rez"):
             rezBin = shutil.which("rez")
-        return f"{rezBin} env {packagesStr} -- {cmd}"
+        addEnvCmd = ""
+        if additionalEnv:
+            envVars = " ".join([f'{k}="{v}"' for k, v in additionalEnv.items()])
+            addEnvCmd = f"env {envVars} "
+        return f"{rezBin} env {packagesStr} -- {addEnvCmd}{cmd}"
     return cmd
 
 
@@ -236,33 +240,44 @@ def getChunks(chunkParams) -> list[Chunk]:
                 it = [Chunk(i, item[0], item[-1]) for i, item in enumerate(slices) if i not in ignoreIterations]
         return it
 
-    @staticmethod
-    def getExpandWrappedCmd(cmdArgs, rezPackages):
+    def getExpandWrappedCmd(self, cmdArgs, rezPackages):
         # Wrap with create_chunks
         cmdBin = wrapMeshroomBin("meshroom_createChunks")
         cmd = f"{cmdBin} --submitter LocalFarm {cmdArgs}"
         # Wrap with rez
-        cmd = rezWrapCommand(cmd, otherRezPkg=rezPackages)
+        cmd = rezWrapCommand(cmd, otherRezPkg=rezPackages, additionalEnv=self.jobEnv)
         return cmd
 
     def createFarmTask(self, meshroomFile: str, orderedTask: OrderedTask, createdTasks: Dict[OrderedTask, Task]) -> Task:
         metadata = dict()
         if orderedTask.node:
-            metadata = {"nodeUid": orderedTask.node._uid, "iteration": orderedTask.iteration}
+            metadata = {"nodeUid": orderedTask.node._uid}
 
+        if orderedTask.iteration >= 0:
+            metadata["iteration"] = orderedTask.iteration
+        elif orderedTask.taskType == OrderedTaskType.PREPROCESS:
+            metadata["iteration"] = "preprocess"
+        elif orderedTask.taskType == OrderedTaskType.POSTPROCESS:
+            metadata["iteration"] = "postprocess"
+
         if orderedTask.taskType == OrderedTaskType.PLACEHOLDER:
             return Task(name=orderedTask.node.name if orderedTask.node else "", command="", metadata=metadata)
 
         cmdArgs = f"--node {orderedTask.node.name} \"{meshroomFile}\" --extern"
-        metadata = {"nodeUid": orderedTask.node._uid, "iteration": orderedTask.iteration}
 
         if orderedTask.taskType == OrderedTaskType.EXPANDING:
             cmd = self.getExpandWrappedCmd(cmdArgs, self.reqPackages)
             task = Task(name=orderedTask.node.name, command=cmd, metadata=metadata, env=self.jobEnv)
         else:
             cmdBin = wrapMeshroomBin("meshroom_compute")
-            cmd = f"{cmdBin} {cmdArgs} --iteration {orderedTask.iteration}"
-            cmd = rezWrapCommand(cmd, otherRezPkg=self.reqPackages)
+            cmd = f"{cmdBin} {cmdArgs}"
+            if orderedTask.taskType == OrderedTaskType.PREPROCESS:
+                cmd += f" --preprocess"
+            elif orderedTask.taskType == OrderedTaskType.POSTPROCESS:
+                cmd += f" --postprocess"
+            elif orderedTask.taskType == OrderedTaskType.CHUNK:
+                cmd += f" --iteration {orderedTask.iteration}"
+            cmd = rezWrapCommand(cmd, otherRezPkg=self.reqPackages, additionalEnv=self.jobEnv)
             task = Task(name=orderedTask.node.name, command=cmd, metadata=metadata, env=self.jobEnv)
 
         return task
@@ -303,24 +318,25 @@ def createJob(self, orderedTasks, filepath, submitLabel="{projectName}") -> Loca
         name = submitLabel.format(projectName=projectName)
         # Create job
         job = Job(name)
-        
+
         # Create tasks
         orderedTasks.display()
         createdTasks: Dict[OrderedTask, Task] = dict()
         for taskToCreate in orderedTasks.iterOnTasks():
             if taskToCreate in createdTasks.keys():
                 continue
             createdTask = self.createFarmTask(filepath, taskToCreate, createdTasks)
+            job.addTask(createdTask)
             createdTasks[taskToCreate] = createdTask
-        
+
         for orderedTask, task in createdTasks.items():
             print(orderedTask, "->", task)
-        
+
         for orderedTask, task in createdTasks.items():
             deps = [createdTasks.get(t) for t in orderedTask.dependencies]
             for dependency in deps:
-                job.addTaskDependency(dependency, task)
-        
+                job.addTaskDependency(task, dependency)
+    
         # Submit job
         engine = LocalFarmEngine(self.farmPath)
         res = job.submit(engine)
@@ -359,7 +375,7 @@ def createChunkTask(self, node, graphFile, **kwargs):
             metadata = {"nodeUid": node._uid, "iteration": chunk.iteration}
             cmdBin = wrapMeshroomBin("meshroom_compute")
             cmd = f"{cmdBin} {cmdArgs} --iteration {chunk.iteration}"
-            cmd = rezWrapCommand(cmd, otherRezPkg=self.reqPackages)
+            cmd = rezWrapCommand(cmd, otherRezPkg=self.reqPackages, additionalEnv=self.jobEnv)
             print("Additional chunk task command: ", cmd)
             task = Task(name=name, command=cmd, metadata=metadata, env=taskEnv)
             engine.create_additional_task(currentJid, currentTid, task)

tests/test_submit.py

@@ -52,18 +52,25 @@ def waitForNodeCompletion(job: LocalFarmJob, node: Node, timeout=25):
     print(f"Waiting for node {node.name} to complete...")
     startTime = time.time()
     while True:
+        time.sleep(1)
+        if time.time() - startTime > timeout:
+            raise TimeoutError(f"Node {node.name} did not complete within {timeout} seconds")
+        # Check for job error
+        err = job.getJobErrors()
+        if err:
+            raise RuntimeError(f"Job encountered an error: {err}")
+        # Check that all tasks are finished
+        for task in job.localfarmTasks.values():
+            if task.get("status") not in (Status.NONE.name, Status.SUCCESS.name, Status.STOPPED.name, Status.ERROR.name):
+                continue
+        break
+        # Stop if the node switched to done
         node.updateStatusFromCache()
         nodeStatus = node.getGlobalStatus()
         if nodeStatus not in (Status.SUBMITTED, Status.RUNNING):
             print(f"Node status switched to {nodeStatus}")
             return
-        # Check for job error
-        err = job.getJobErrors()
-        if err:
-            raise RuntimeError(f"Job encountered an error: {err}")
-        if time.time() - startTime > timeout:
-            raise TimeoutError(f"Node {node.name} did not complete within {timeout} seconds")
-        time.sleep(1)
+
 
 def processSubmit(node: Node, graph, tmp_path):
     """