Use advanced task settings

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,43 +241,70 @@ 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:
             if self.iteration >= 0:
                 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