Advanced submitter settings

meshroom/core/desc/__init__.py

@@ -49,4 +49,6 @@
     Node,
     NodeVersionType,
     NodeVersionTypeEnum,
+    SubmissionSettings,
+    StageSettings,
 )

meshroom/core/desc/node.py

@@ -1,5 +1,3 @@
-# desc/node.py
-
 import enum
 from inspect import getfile, getattr_static
 from pathlib import Path
@@ -10,6 +8,7 @@
 import signal
 import subprocess
 from collections import OrderedDict
+from types import SimpleNamespace
 import psutil
 
 from meshroom import _MESHROOM_ROOT
@@ -62,6 +61,48 @@ class MrNodeType(enum.Enum):
     BACKDROP = enum.auto()
 
 
+class StageSettings(SimpleNamespace):
+    def __init__(self, cpu=Level.NORMAL, ram=Level.NORMAL, gpu=Level.NONE, **kwargs):
+        super().__init__(**kwargs)
+        self.cpu = cpu.value if hasattr(cpu, "value") else cpu
+        self.ram = ram.value if hasattr(cpu, "value") else cpu
+        self.gpu = gpu.value if hasattr(cpu, "value") else cpu
+
+    def __contains__(self, item):
+        return item in self.__dict__
+
+
+class SubmissionSettings:
+    def __init__(self, node):
+        """Holds infos used when we submit the node for remote computing.
+
+        Args:
+            node (BaseNode): the BaseNode instance being updated
+        """
+        # Pre/post process use the default settings
+        self.preprocess = StageSettings()
+        self.postprocess = StageSettings()
+        # Process use the values we set on the node description
+        self.process = StageSettings(cpu=node.cpu, ram=node.ram, gpu=node.gpu)
+        self.process.licenses = node.nodeDesc._licenses
+        # Retrocompatibility behaviour
+        if hasattr(node.nodeDesc, "_cuda_tag"):
+            logging.warning(f"DeprecationWarning : Node of type {node.nodeDesc} uses '_cuda_tag'. Please use SubmissionSettings instead.")
+            self.process.cuda_tag = node.nodeDesc._cuda_tag
+        if hasattr(node.nodeDesc, "_service_key"):
+            logging.warning(f"DeprecationWarning : Node of type {node.nodeDesc} uses '_service_key'. Please use SubmissionSettings instead.")
+            self.process.service_key = node.nodeDesc._service_key
+
+    def getStageSettings(self, stageName="process"):
+        stage = getattr(self, stageName, None)
+        if stage:
+            return stage
+        raise ValueError(
+             f"Unknown stageName: {stageName!r}. Expected one of: "
+             "'preprocess', 'process', 'postprocess'."
+         )
+
+
 class InternalAttributesFactory:
     BASIC = [
         StringParam(
@@ -390,6 +431,15 @@ def postUpdate(cls, node):
         """
         pass
 
+    def getSubmissionSettings(self, node) -> SubmissionSettings:
+        """ Gets invoked when we submit the node on farm.
+        SubmissionSettings contain all the settings that we can use on the submitter
+
+        Args:
+            node: The BaseNode instance about to be processed.
+        """
+        return SubmissionSettings(node)
+
     def preprocess(self, node):
         """ Gets invoked just before the processChunk method for the node.
 

meshroom/core/submitter.py

@@ -7,12 +7,15 @@
 import operator
 
 from enum import IntFlag, auto
-from typing import Optional, Dict, List
+from typing import Optional, Dict, List, TYPE_CHECKING
 from itertools import accumulate
 
 import meshroom
 from meshroom.common import BaseObject, Property
 
+if TYPE_CHECKING:
+    from meshroom.core.node import BaseNode
+
 
 logger = logging.getLogger("Submitter")
 
@@ -80,33 +83,30 @@ class OrderedTaskType(IntFlag):
 
 
 class OrderedTask:
-    _usedUids = set()
+    _lastUid = 0
 
-    def __init__(self, taskType, node = None, iteration : int = -1):
+    def __init__(self, taskType: OrderedTaskType, node: BaseNode = None, iteration : int = -1):
         self.taskType: OrderedTaskType = taskType
-        self.node = node  # BaseNode
-        self.iteration = iteration
-        self.dependencies = []
-        self.uid = self._generateUid()
-
-    @property
-    def nodeName(self):
-        return self.node.name if self.node else "NONE"
+        self.node: BaseNode = node
+        self.iteration: int = iteration
+        self.dependencies : list["OrderedTask"] = []
+        # Generate UID
+        self.uid: int = self._generateUid()
 
     @classmethod
     def _generateUid(cls) -> int:
-        nextUid = max(cls._usedUids) + 1 if len(cls._usedUids) > 0 else 0
-        cls._usedUids.add(nextUid)
-        return nextUid
+        cls._lastUid += 1
+        return cls._lastUid
 
-    def addDependency(self, otherTask: OrderedTask):
-        self.dependencies.append(otherTask)
+    @property
+    def nodeName(self) -> str:
+        return self.node.name if self.node else "NONE"
 
     @property
-    def shortName(self):
-        sn = self.nodeName if self.node else "NONE"
+    def name(self) -> str:
+        sn = self.nodeName
         if self.taskType == OrderedTaskType.CHUNK:
-            sn += f"_{self.iteration if self.iteration >=0 else 'allchunks'}"
+            sn += f"_{self.iteration if self.iteration >=0 else 'process'}"
         else:
             sn += f"_{self.taskType.name}"
         return f"{self.uid:03d} {sn}"
@@ -124,6 +124,9 @@ def __repr__(self):
         string += f" ({len(self.dependencies)} deps)>"
         return string
 
+    def addDependency(self, otherTask: OrderedTask):
+        self.dependencies.append(otherTask)
+
 
 class OrderedNode:
     """ Intermediate structure used to order tasks """
@@ -220,7 +223,7 @@ def gatherTasks(task):
             return allTasks
         tasks: list[OrderedTask] = gatherTasks(self.rootTask)
         for task in set(tasks):
-            logging.debug(f"[{task.shortName}] {task} -> depends on {[t.shortName for t in task.dependencies]}")
+            logging.debug(f"[{task.name}] {task} -> depends on {[t.name for t in task.dependencies]}")
 
     def iterOnTasks(self, current:OrderedTask=None, skipRootTask=False):
         skipCurrent = (current is None) and skipRootTask