diff --git a/meshroom/core/desc/__init__.py b/meshroom/core/desc/__init__.py index c9c11a48ec..234de83721 100644 --- a/meshroom/core/desc/__init__.py +++ b/meshroom/core/desc/__init__.py @@ -49,4 +49,6 @@ Node, NodeVersionType, NodeVersionTypeEnum, + SubmissionSettings, + StageSettings, ) diff --git a/meshroom/core/desc/node.py b/meshroom/core/desc/node.py index b27463aa85..b48a3bc688 100644 --- a/meshroom/core/desc/node.py +++ b/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. diff --git a/meshroom/core/submitter.py b/meshroom/core/submitter.py index 888c9c07b6..16f3ce4735 100644 --- a/meshroom/core/submitter.py +++ b/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