diff --git a/meshroom/tractorSubmitter/api/base.py b/meshroom/tractorSubmitter/api/base.py index 9702529..400254d 100644 --- a/meshroom/tractorSubmitter/api/base.py +++ b/meshroom/tractorSubmitter/api/base.py @@ -11,22 +11,20 @@ """ import os -import sys -import re import json import getpass -import logging -import shlex -import shutil -from collections import namedtuple import tempfile +from collections import namedtuple from typing import Optional, Tuple +from tractorSubmitter.rezUtils import CommandArgsBuilder + +from meshroom.core.node import BaseNode + TRACTOR_JOB_URL = "http://tractor-engine/tv/#jid={jid}" Chunk = namedtuple("chunk", ["iteration", "start", "end"]) -REZ_DELIMITER_PATTERN = re.compile(r"(-|==|>=|>|<=|<)") LICENSES_MAP = { 'mtoa': 'arnold', 'houdiniE': 'houdinie', @@ -47,84 +45,6 @@ def createTmpFolder(create=False): return tmpFolder -def getResolvedVersionsDict(): - """ Get a dict {packageName: version} corresponding to the current context """ - resolvedPackages = os.environ.get('REZ_RESOLVE', '').split() - resolvedVersions = {} - for r in resolvedPackages: - if r.startswith('~'): # remove implicit packages - continue - v = r.split('-') - if len(v) == 2: - resolvedVersions[v[0]] = v[1] - elif len(v) > 2: # Handle case with multiple hyphen-minus - resolvedVersions[v[0]] = "-".join(v[1:]) - return resolvedVersions - - -def getRequestPackages(packagesDelimiter="=="): - """ - Get list of packages required for the job - Depends on env var and current rez context - - By default we use the "==" delimiter to make sure we have the same version - in the job that the one we have in the env where meshroom is launched - """ - reqPackages = set() - if 'REZ_REQUEST' in os.environ: - # Get the names of the packages that have been requested - requestedPackages = os.environ.get('REZ_USED_REQUEST', '').split() - usedPackages = set() # Use set to remove duplicates - for p in requestedPackages: - if p.startswith('~') or p.startswith("!"): - continue - v = REZ_DELIMITER_PATTERN.split(p) - usedPackages.add(v[0]) - # Add requested packages to the reqPackages set - resolvedVersions = getResolvedVersionsDict() - for p in usedPackages: - reqPackages.add(packagesDelimiter.join([p, resolvedVersions[p]])) - logging.debug(f"TractorSubmitter: REZ Packages: {str(reqPackages)}") - elif 'REZ_MESHROOM_VERSION' in os.environ: - reqPackages.add(f"meshroom{packagesDelimiter}{os.environ.get('REZ_MESHROOM_VERSION', '')}") - return list(reqPackages) - - -def rezWrapCommand(cmd, useCurrentContext=False, useRequestedContext=True, - otherRezPkg: list[str] = None): - """ Wrap command to be runned using rez - :param cmd: command to run - :type cmd: bool - :param useCurrentContext: use current rez context to retrieve a list of rez packages - :type useCurrentContext: bool - :param useRequestedContext: use rez packages that have been requested (not the full context) # TODO : remove it - :type useRequestedContext: bool - :param otherRezPkg: Additionnal rez packages - :type otherRezPkg: list[str] - """ - packages = set() - if useCurrentContext: - # In this case we want to use the full context - packages.update([p for p in os.environ.get('REZ_RESOLVE', '').split(" ") if p]) - elif useRequestedContext: - # In this case we want to use only packages in the rez request - packages.update(getRequestPackages()) - # Add additional packages - if otherRezPkg: - packages.update(otherRezPkg) - packagesStr = " ".join([p for p in packages if p]) - if packagesStr: - rezBin = "rez" - if "REZ_BIN" in os.environ and os.environ["REZ_BIN"]: - rezBin = os.environ["REZ_BIN"] - elif "REZ_PACKAGES_ROOT" in os.environ and os.environ["REZ_PACKAGES_ROOT"]: - 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}" - return cmd - - def toTractorEnv(environment): """ Format env for Tractor """ return [f"setenv {k}={v}" for k, v in environment.items()] @@ -163,12 +83,11 @@ def toTractorEnv(environment): # class JobInfo: - def __init__(self, name, share=None, service=None, environment=None, tags=None, user=None, + def __init__(self, name, share=None, serviceKey=None, environment=None, tags=None, user=None, comment="", paused=False): self.name = name self.share = self.getShare(share) - self.requirements = service or {} - self.service = service or os.environ.get("DEFAULT_TRACTOR_SERVICE", "") + self.serviceKey = serviceKey or os.environ.get("DEFAULT_TRACTOR_SERVICE", "") self.tags = tags or {} self.paused = paused self.comment = comment @@ -193,7 +112,7 @@ def cook(self): env = self.environment.copy() return { "title": self.name, - "service": self.service, + "service": self.serviceKey, "metadata": json.dumps(tags), "envkey": toTractorEnv(env), "paused": self.paused, @@ -206,12 +125,12 @@ def cook(self): class TaskInfo: def __init__(self, name: str, + node: BaseNode, cmdArgs: str, - nodeUid: str, cacheFolder: str="", environment: dict=None, reqPackages: list=None, - service: str=None, + config: str = None, licenses=None, taskType:Optional[Tuple]=None, tags=None): @@ -220,29 +139,29 @@ def __init__(self, Args: name: name of the task (usually the node name). For the final title we add the task type (chunk index) + node: Node cmdArgs: Command to execute - nodeUid: Node UID cacheFolder: Folder containing the node cache. environment: Environment to set. Dict with key:value. reqPackages: List of requested packages. - service: Service key expression (used for machine targeting). licenses: Eequired licenses. taskType: Task type and iteration if needed. Tuple[task type, iteration] tags: Additional metadata to set on the task. """ + self.node = node self.name = name - self.uid = nodeUid self.taskCommandArgs = cmdArgs + self.config = config # Env self.environment = environment or {} # Requested packages self.reqPackages = reqPackages or [] # self.limits - self.service = service or os.environ.get("DEFAULT_TRACTOR_SERVICE", "") self.limits = self.getLimits(licenses) # Tags self.tags = tags or {} - self.tags["nodeUid"] = nodeUid + if node: + self.tags["nodeUid"] = node._uid # Expanding / Chunks taskType_, iteration_ = taskType or ("placeholder", None) @@ -251,8 +170,19 @@ def __init__(self, self.preprocessTask = (taskType_ == "preprocess") self.postprocessTask = (taskType_ == "postprocess") self.chunkTask = (taskType_ == "chunk") - self.iteration = iteration_ + + # Submitter settings + self.taskSubmitterSettings = self._getTaskSubmitterSettings() + + def _getTaskSubmitterSettings(self): + nodeSubmitSettings = self.node.nodeDesc.getSubmitSettings(self.node) + if self.preprocessTask: + return nodeSubmitSettings.preprocess + elif self.postprocessTask: + return nodeSubmitSettings.postprocess + else: + return nodeSubmitSettings.process @staticmethod def getLimits(licenses=None): @@ -262,6 +192,34 @@ def getLimits(licenses=None): taskLimits.append(os.environ['DEFAULT_TRACTOR_LIMIT']) return taskLimits + @property + def service(self): + """ Get the service key for the task """ + if self.placeholderTask: + return "" + if self.expandingTask: + return self.config.GLOBAL_KEY + # Get submitter settings + settings = self.taskSubmitterSettings + # If an explicit service key is set use it + if "service_key" in settings: + return settings.service_key + # Else try to build from the config file + kwargs = {} + if "cuda_tag" in settings: + kwargs["cuda_tag"] = settings.cuda_tag + if "excluded_hosts" in settings: + hosts = settings.excluded_hosts + hosts = hosts if isinstance(hosts, list) else [hosts] + kwargs["excluded_hosts"] = hosts + service = self.config.get_config( + cpu=settings.cpu, + ram=settings.ram, + gpu=settings.gpu, + **kwargs + ) + return service + def _setExpandingTaskFile(self, cacheFolder): """ Doesn't work with current python API ! It should be possible starting Tractor 1.7 to give a file path to cmd.expand @@ -283,18 +241,20 @@ def _setExpandingTaskFile(self, cacheFolder): @property def envkey(self): - return toTractorEnv(self.environment) - - def cook(self): + settings = self.taskSubmitterSettings + env = self.environment + print("env", env) + if "env" in settings: + env.update(settings.env) + return toTractorEnv(env) + + def get_kwargs(self): title = f"{self.name}" tags = self.tags - cmd = self.taskCommandArgs if self.preprocessTask: - cmd += f" --preprocess" title += "_preprocess" tags["iteration"] = "preprocess" elif self.postprocessTask: - cmd += f" --postprocess" title += "_postprocess" tags["iteration"] = "postprocess" elif self.chunkTask: @@ -302,24 +262,49 @@ def cook(self): title += f"_{self.iteration}" else: title += f"_0" - cmd += f" --iteration {self.iteration}" tags["iteration"] = self.iteration - - if self.expandingTask: - cmd = rezWrapCommand(cmd, otherRezPkg=self.reqPackages) - # Wrap with tractor wrapper (will redirect stdout to stderr) - # to make sure stdout only has the - wrapperModule = "tractorExpander.py" - wrapperPath = os.path.join(os.environ["MR_SUBMITTERS_SCRITPS"], wrapperModule) - cmd = f"{sys.executable} {wrapperPath} {cmd}" - elif self.placeholderTask: - cmd = None - else: - cmd = rezWrapCommand(cmd, otherRezPkg=self.reqPackages) - return { "title": title, - "argv": shlex.split(cmd) if cmd else cmd, "service": self.service, "metadata": json.dumps(self.tags) } + + def get_commands(self): + """Build list of commands + + Note: + seatup/teardown commands do not work for tractor expand commands + """ + cmd = self.taskCommandArgs + if self.placeholderTask or not cmd: + return [] + settings = self.taskSubmitterSettings + commands = [] + + # Setup + if "setup_command" in settings and not self.expandingTask: + commands.append(settings.setup_command) + + # Build process command + processCommand = CommandArgsBuilder(cmd) + processCommand.setRequiredPackages(self.reqPackages) + processCommand.setSubmissionSettings(settings) + if self.expandingTask: + wrapperModule = "tractorExpander.py" + wrapperPath = os.path.join(os.environ["MR_SUBMITTERS_SCRITPS"], wrapperModule) + processCommand.setTractorWrapper(wrapperPath) + commands.append(processCommand.getWrappedCommand()) + else: + if self.preprocessTask: + processCommand.cmd += f" --preprocess" + elif self.postprocessTask: + processCommand.cmd += f" --postprocess" + elif self.chunkTask: + processCommand.cmd += f" --iteration {self.iteration}" + commands.append(processCommand.getWrappedCommand()) + + # Teardown + if "teardown_command" in settings and not self.expandingTask: + commands.append(settings.teardown_command) + + return commands diff --git a/meshroom/tractorSubmitter/api/subtaskCreator.py b/meshroom/tractorSubmitter/api/subtaskCreator.py index 6536968..52897c0 100644 --- a/meshroom/tractorSubmitter/api/subtaskCreator.py +++ b/meshroom/tractorSubmitter/api/subtaskCreator.py @@ -82,13 +82,13 @@ def sendTractorCmd(task_def): f.write("\n" + task_def + "\n") -def queueSubtask(title, argv, service="", limits=None, metadata=None, envkey=None): +def queueSubtask(commands, title, service="", limits=None, metadata=None, envkey=None): """ Queue a subtask to be created in Tractor. Args: + commands (list): Commands to run title (str): Task title - cmd (str or list): Command to run (string or argv list) service (str): Tractor service key limits (list): Limit tags (e.g. ["blender", "nuke"]) metadata (dict): Metadata as key:value pairs @@ -98,22 +98,14 @@ def queueSubtask(title, argv, service="", limits=None, metadata=None, envkey=Non Example: queueSubtask( - title="render_frame_0001", cmd="render --frame 1 scene.ma", + title="render_frame_0001", service="mikrosRender", limits=["blender"], metadata={'user': 'john', 'iteration': '1', 'prod': 'mvg'} ) """ - # Parse command - if isinstance(argv, str): - cmd_argv = shlex.split(argv) - else: - cmd_argv = list(argv) - - cmd_str = " ".join(cmd_argv) - # Build tags string tags_str = "" if limits: @@ -131,12 +123,21 @@ def queueSubtask(title, argv, service="", limits=None, metadata=None, envkey=Non # Build service string service_str = f"-service {{{service}}}" if service else "" + + # Build list of commands + tclCmds = [] + for argv in commands: + cmd_str = " ".join(argv) + tclCmds.append(f"""RemoteCmd {{{cmd_str}}} {service_str} {tags_str} {envkey_str}""") # Write Alfred task definition # TODO : we can use tractor API to convert a Task into alf (asTcl) task_def = f""" -Task -title {{{title}}} {service_str} {metadata_str} -cmds {{ - RemoteCmd {{{cmd_str}}} {service_str} {tags_str} {envkey_str} +Task -title {{{title}}} {service_str} {metadata_str} -cmds {{""" + for cmd in tclCmds: + task_def += f""" + {cmd}""" + task_def += f""" }} """ print(task_def) @@ -144,7 +145,7 @@ def queueSubtask(title, argv, service="", limits=None, metadata=None, envkey=Non log(f"Queued subtask: {title}") -def queueChunkTask(node, cmdArgs, service, tags=None, reqPackages=None, environment=None): +def queueChunkTask(node, taskCommand, config, tags=None, reqPackages=None, environment=None): blockSize, fullSize, nbBlocks = node.nodeDesc.parallelization.getSizes(node) if nbBlocks <= 0: return @@ -152,19 +153,21 @@ def queueChunkTask(node, cmdArgs, service, tags=None, reqPackages=None, environm for iteration in range(nbBlocks): taskInfo = TaskInfo( - name=node.name, - cmdArgs=cmdArgs, - nodeUid=node._uid, + name=node.name, + node=node, + cmdArgs=taskCommand, environment=environment, reqPackages=reqPackages, - service=service, + config=config, licenses=licenses, taskType=("chunk", iteration), tags=tags.copy() if tags else None, ) # title, argv, service, metadata - taskArgs = taskInfo.cook() + taskArgs = taskInfo.get_kwargs() # limits, envkey taskArgs['limits'] = taskInfo.limits taskArgs['envkey'] = taskInfo.envkey - queueSubtask(**taskArgs) + # Commands + commands = taskInfo.get_commands() + queueSubtask(commands, **taskArgs) diff --git a/meshroom/tractorSubmitter/rezUtils.py b/meshroom/tractorSubmitter/rezUtils.py new file mode 100644 index 0000000..c610a62 --- /dev/null +++ b/meshroom/tractorSubmitter/rezUtils.py @@ -0,0 +1,149 @@ +import os +import re +import logging +import shlex +import shutil +import sys +from pathlib import Path + + +REZ_DELIMITER_PATTERN = re.compile(r"(-|==|>=|>|<=|<)") + + +def getResolvedVersionsDict(): + """ Get a dict {packageName: version} corresponding to the current context """ + resolvedPackages = os.environ.get('REZ_RESOLVE', '').split() + resolvedVersions = {} + for r in resolvedPackages: + if r.startswith('~'): # remove implicit packages + continue + v = r.split('-') + if len(v) == 2: + resolvedVersions[v[0]] = v[1] + elif len(v) > 2: # Handle case with multiple hyphen-minus + resolvedVersions[v[0]] = "-".join(v[1:]) + return resolvedVersions + + +def getRequestPackages(packagesDelimiter="=="): + """ + Get list of packages required for the job + Depends on env var and current rez context + + By default we use the "==" delimiter to make sure we have the same version + in the job that the one we have in the env where meshroom is launched + """ + reqPackages = set() + if 'REZ_REQUEST' in os.environ: + # Get the names of the packages that have been requested + requestedPackages = os.environ.get('REZ_USED_REQUEST', '').split() + usedPackages = set() # Use set to remove duplicates + for p in requestedPackages: + if p.startswith('~') or p.startswith("!"): + continue + v = REZ_DELIMITER_PATTERN.split(p) + usedPackages.add(v[0]) + # Add requested packages to the reqPackages set + resolvedVersions = getResolvedVersionsDict() + for p in usedPackages: + reqPackages.add(packagesDelimiter.join([p, resolvedVersions[p]])) + logging.debug(f"TractorSubmitter: REZ Packages: {str(reqPackages)}") + elif 'REZ_MESHROOM_VERSION' in os.environ: + reqPackages.add(f"meshroom{packagesDelimiter}{os.environ.get('REZ_MESHROOM_VERSION', '')}") + return list(reqPackages) + + +class CommandArgsBuilder: + def __init__(self, cmd): + self.cmd = cmd + self.settings = None + self.packages = [] + self.useCurrentContext=False + self.useRequestedContext=True + self.rezPkgDelimiter="==" + self.tractorWrapper = None + + def setSubmissionSettings(self, settings): + self.settings = settings + + def setRequiredPackages(self, packages: list[str] = None): + self.packages = packages + + def setRezSettings(self, + useCurrentContext: bool = False, + useRequestedContext: bool = True, + rezPkgDelimiter: str = "=="): + """Set additional settings for rez + + Args: + useCurrentContext: Use current rez context to retrieve a list of rez packages. + useRequestedContext: Use rez packages that have been requested (not the full context) + rezPkgDelimiter: Delimiter used for the request. + """ + self.useCurrentContext = useCurrentContext + self.useRequestedContext = useRequestedContext + self.rezPkgDelimiter = rezPkgDelimiter + + def setTractorWrapper(self, wrapperPath): + """ Sets a python script wrapper to wrap the command executed on farm + + It needs to be used on several occasions for example on the expanding tasks, + it is wrapping the whodl process so that we only write tractor TCL commands + on the output. + + Example + cmd = "rez env PKGS -- meshroom_createChunks ARGS" + -> cmd = "python tractorWrapper.py rez env PKGS -- meshroom_createChunks ARGS" + """ + self.tractorWrapper = wrapperPath + + def getRezPackages(self): + """ Get list of packages depending on current environment and rez settings. """ + packages = set() + if self.useCurrentContext: + packages.update([p for p in os.environ.get('REZ_RESOLVE', '').split(" ") if p]) + elif self.useRequestedContext: + packages.update(getRequestPackages(packagesDelimiter=self.rezPkgDelimiter)) + if self.packages: + packages.update(self.packages) + return [p for p in packages if p] + + def getRezExecutable(self) -> str: + """ Find path to rez executable. If not found, use the alias "rez". """ + rezBin = "rez" + if "REZ_BIN" in os.environ and os.environ["REZ_BIN"]: + rezBin = os.environ["REZ_BIN"] + elif "REZ_PACKAGES_ROOT" in os.environ and os.environ["REZ_PACKAGES_ROOT"]: + rezBin = os.path.join(os.environ["REZ_PACKAGES_ROOT"], "bin/rez") + elif shutil.which("rez"): + rezBin = shutil.which("rez") + if Path(rezBin).exists(): + return str(Path(rezBin).resolve()) + return rezBin + + def getWrappedCommand(self) -> list[str]: + """ Wraps the rez command. + If a "rezWrapper" is found on the settings, call it to build the command. + """ + # First split the command to execute + if "target_os" in self.settings and self.settings.target_os == "windows": + args = shlex.split(self.cmd, posix=False) + else: + args = shlex.split(self.cmd) + # Get rez executable and packages + rez_bin = self.getRezExecutable() + rez_packages = self.getRezPackages() + # Use the task-specific wrapper if we find one + if "rezWrapper" in self.settings: + args = self.settings.rezWrapper( + rez_bin = rez_bin, + rez_packages=rez_packages, + args=args, + tractor_wrapper=self.tractorWrapper + ) + else: + # Default : "rez env PKGS -- CMD" + args = [rez_bin, "env"] + rez_packages + ["--"] + args + if self.tractorWrapper: + args = [sys.executable, self.tractorWrapper] + args + return args diff --git a/meshroom/tractorSubmitter/tractorSubmitter.py b/meshroom/tractorSubmitter/tractorSubmitter.py index c6cb6dc..f75e451 100644 --- a/meshroom/tractorSubmitter/tractorSubmitter.py +++ b/meshroom/tractorSubmitter/tractorSubmitter.py @@ -9,8 +9,8 @@ # ========== Tractor ========== from tractor.api import author as tractorAuthor +from tractorSubmitter.rezUtils import getRequestPackages from tractorSubmitter.api.base import ( - getRequestPackages, TaskInfo, JobInfo, TRACTOR_JOB_URL, PRIORITY_DICT ) @@ -40,36 +40,48 @@ def wrapMeshroomBin(_bin): class Task: - def __init__(self, name, command, uid=None, nodeCache="", tags=None, - reqPackages=None, service=None, licenses=None, taskType=None): + def __init__(self, node, command, defaultName="", nodeCache="", tags=None, + reqPackages=None, config=None, licenses=None, taskType=None): self.taskInfos = TaskInfo( - name=name, + name=node.name if node else defaultName, + node=node, cmdArgs=command, - nodeUid=uid, cacheFolder=nodeCache, reqPackages=reqPackages, - service=service, + config=config, licenses=licenses, taskType=taskType, tags=tags.copy() if tags else None, ) - taskKwargs = self.taskInfos.cook() + # Create task + taskKwargs = self.taskInfos.get_kwargs() + logging.info(f"Task {self.taskInfos.name} -> kwargs: {taskKwargs}") + logging.info(f"Task {self.taskInfos.name} -> envkey: {self.taskInfos.envkey}") self.tractorTask: tractorAuthor.Task = tractorAuthor.Task(**taskKwargs) - for cmd in self.tractorTask.cmds: - cmd.tags = self.taskInfos.limits - cmd.envkey = self.taskInfos.envkey - cmd.expand = self.taskInfos.expandingTask - # If we use a file for expanding task instead we could use this : - # if taskInfos.expandingTask: - # cmd.expand = taskInfos.expandingFile + # Add commands + for i, cmd in enumerate(self.taskInfos.get_commands()): + # All attrs: + # msg, tags, service, metrics, id, refersto, expand, atleast, atmost, + # minrunsecs, maxrunsecs, samehost, envkey, retryrc, when, resumewhile, + # resumepin, metadata, + logging.info(f"Task {self.taskInfos.name} -> Command {i}: {cmd}") + self.tractorTask.newCommand( + argv=cmd, + service=taskKwargs.get("service"), + envkey=self.taskInfos.envkey, + tags=self.taskInfos.limits, + expand=self.taskInfos.expandingTask, + # If we use a file for expanding task instead we could use this : + # expand = taskInfos.expandingFile + ) class Job: - def __init__(self, name, tags=None, requirements=None, environment=None, user=None, comment="", paused=False): + def __init__(self, name, tags=None, serviceKey=None, environment=None, user=None, comment="", paused=False): self.jobInfo = JobInfo( name, share="", - service=requirements, + serviceKey=serviceKey, environment=environment, tags=tags, user=user, @@ -319,11 +331,11 @@ def retrieveJob(self, jid) -> TractorJob: def createTask(self, meshroomFile: str, orderedTask: OrderedTask, createdTasks: Dict[OrderedTask, Task], **kwargs) -> Task: node = orderedTask.node if orderedTask.taskType == OrderedTaskType.PLACEHOLDER: - defaultName = kwargs.get("jobName", "") - defaultName += " (placeholder)" + defaultName = kwargs.get("jobName", "") + " (placeholder)" return Task( - name=orderedTask.node.name if orderedTask.node else defaultName, - command="", + node=node, + defaultName=defaultName, + command="", ) if orderedTask.taskType == OrderedTaskType.CHUNK: @@ -341,12 +353,11 @@ def createTask(self, meshroomFile: str, orderedTask: OrderedTask, createdTasks: tags['prod'] = self.prod taskParams = { - "name": node.name, - "uid": node._uid, # Provide unicity info + "node": node, "nodeCache": node._internalFolder, "tags": tags, "reqPackages": self.reqPackages, - "service": self.getTaskService(node), + "config": self.config, "licenses": node.nodeDesc._licenses, "taskType": taskType } @@ -414,12 +425,12 @@ def createChunkTask(self, node, graphFile, environment=None, **kwargs): # Environment environment = environment or {} # Command - cmdArgs = f"meshroom_compute --node {node.name} \"{graphFile}\" --extern" + taskCommand = f"meshroom_compute --node {node.name} \"{graphFile}\" --extern" # Add task to the queue queueChunkTask( node=node, - cmdArgs=cmdArgs, - service=self.getTaskService(node), + taskCommand=taskCommand, + config=self.config, tags=taskTags, reqPackages=self.reqPackages, environment=environment