Merge pull request #5702 from aldbr/cherry-pick-2-0e31badcd-integration

[sweep:integration] Multi-node allocations: Embed executableFile in the srun Wrapper

Author: fstagni

src/DIRAC/Resources/Computing/BatchSystems/SLURM.py

@@ -11,6 +11,7 @@
 import re
 import subprocess
 import shlex
+import random
 
 __RCSID__ = "$Id$"
 
@@ -51,6 +52,10 @@ def submitJob(self, **kwargs):
         errFile = os.path.expandvars(errFile)
         executable = os.path.expandvars(executable)
 
+        # There are more than 1 node, we have to run the executable in parallel on different nodes using srun
+        if numberOfNodes != "1":
+            executable = self._generateSrunWrapper(executable)
+
         jobIDs = []
         for _i in range(nJobs):
             jid = ""
@@ -95,6 +100,10 @@ def submitJob(self, **kwargs):
             jid = jid.strip()
             jobIDs.append(jid)
 
+        # Delete the srun wrapper
+        if numberOfNodes != "1":
+            os.remove(executable)
+
         if jobIDs:
             resultDict["Status"] = 0
             resultDict["Jobs"] = jobIDs
@@ -103,6 +112,35 @@ def submitJob(self, **kwargs):
             resultDict["Message"] = error
         return resultDict
 
+    def _generateSrunWrapper(self, executableFile):
+        """
+        Associate the executable with srun, to execute the same command in parallel on multiple nodes.
+        Wrap it in a new executable file
+
+        :param str executableFile: name of the executable file to wrap
+        :return str: name of the wrapper that runs the executable via srun
+        """
+        suffix = random.randrange(1, 99999)
+        wrapper = "srunExec_%s.sh" % suffix
+        with open(executableFile, "r") as f:
+            content = f.read()
+
+        # Build the script to run the executable in parallel multiple times
+        # - Embed the content of executableFile inside the parallel library wrapper script
+        # - srun is the command to execute a task multiple time in parallel
+        #   -l option: add the task ID to the output
+        #   -k option: do not kill the slurm job if one of the nodes is broken
+        cmd = """#!/bin/bash
+cat > %(wrapper)s << EOFEXEC
+%(content)s
+EOFEXEC
+chmod 755 %(wrapper)s
+srun -l -k %(wrapper)s
+""" % dict(
+            wrapper=wrapper, content=content
+        )
+        return cmd
+
     def killJob(self, **kwargs):
         """Delete a job from SLURM batch scheduler. Input: list of jobs output: int"""
 
@@ -267,3 +305,86 @@ def getCEStatus(self, **kwargs):
         resultDict["Waiting"] = waitingJobs
         resultDict["Running"] = runningJobs
         return resultDict
+
+    def getJobOutputFiles(self, **kwargs):
+        """Get output file names and templates for the specific CE
+
+        Reorder the content of the output files according to the node identifier
+        if multiple nodes were involved.
+
+        From:
+        >>> 1: line1
+        >>> 2: line1
+        >>> 1: line2
+        To:
+        >>> # On node 1
+        >>>   line1
+        >>>   line2
+        >>> # On node 2
+        >>>   line1
+        """
+        resultDict = {}
+
+        MANDATORY_PARAMETERS = ["JobIDList", "OutputDir", "ErrorDir"]
+        for argument in MANDATORY_PARAMETERS:
+            if argument not in kwargs:
+                resultDict["Status"] = -1
+                resultDict["Message"] = "No %s" % argument
+                return resultDict
+
+        outputDir = kwargs["OutputDir"]
+        errorDir = kwargs["ErrorDir"]
+        jobIDList = kwargs["JobIDList"]
+        numberOfNodes = kwargs.get("NumberOfNodes", "1")
+
+        jobDict = {}
+        for jobID in jobIDList:
+            output = "%s/%s.out" % (outputDir, jobID)
+            error = "%s/%s.err" % (errorDir, jobID)
+
+            if numberOfNodes != "1":
+                self._openFileAndSortOutput(output)
+                self._openFileAndSortOutput(error)
+
+            jobDict[jobID] = {}
+            jobDict[jobID]["Output"] = output
+            jobDict[jobID]["Error"] = error
+
+        resultDict["Status"] = 0
+        resultDict["Jobs"] = jobDict
+        return resultDict
+
+    def _openFileAndSortOutput(self, outputFile):
+        """
+        Open a file, get its content and reorder it according to the node identifiers
+
+        :param str outputFile: name of the file to sort
+        """
+        with open(outputFile, "r") as f:
+            outputContent = f.read()
+
+        sortedContent = self._sortOutput(outputContent)
+
+        with open(outputFile, "w") as f:
+            f.write(sortedContent)
+
+    def _sortOutput(self, outputContent):
+        """
+        Reorder the content of the output file according to the node identifiers
+
+        :param str outputContent: content to sort
+        :return str: content sorted
+        """
+        outputLines = outputContent.split("\n")
+        nodes = {}
+        for line in outputLines:
+            node, line_content = line.split(":", 1)
+            if node not in nodes:
+                nodes[node] = []
+            nodes[node].append(line_content)
+
+        content = ""
+        for node, lines in nodes.items():
+            content += "# On node %s\n\n" % node
+            content += "\n".join(lines) + "\n"
+        return content

src/DIRAC/Resources/Computing/BatchSystems/test/Test_SLURM.py

@@ -0,0 +1,156 @@
+""" Test SLURM
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import pytest
+import os
+from six.moves import reload_module
+
+from DIRAC import S_OK, gLogger
+from DIRAC.Resources.Computing.BatchSystems.SLURM import SLURM
+
+
+gLogger.setLevel("DEBUG")
+
+executableContent = """
+#!/bin/bash
+
+echo "hello world"
+"""
+
+expectedContent = """#!/bin/bash
+cat > srunExec_1.sh << EOFEXEC
+
+#!/bin/bash
+
+echo "hello world"
+
+EOFEXEC
+chmod 755 srunExec_1.sh
+srun -l -k srunExec_1.sh
+"""
+
+srunOutput = """
+1: line1
+1: line2
+2: line1
+1: line3
+3: line1
+2: line2
+3: line2
+2: line3
+3: line3
+"""
+
+srunExpected1 = """
+# On node 1
+
+ line1
+ line2
+ line3
+"""
+
+
+srunExpected2 = """
+# On node 3
+
+ line1
+ line2
+ line3
+"""
+
+
+srunExpected3 = """
+# On node 2
+
+ line1
+ line2
+ line3
+"""
+
+
+srunExpected = [srunExpected1, srunExpected2, srunExpected3]
+
+
+normalOutput = """
+line1
+line2
+line3
+"""
+
+
+normalExpected = """
+line1
+line2
+line3
+"""
+
+
+@pytest.mark.parametrize(
+    "expectedContent",
+    [
+        (expectedContent),
+    ],
+)
+def test_generateWrapper(mocker, expectedContent):
+    """Test generateWrapper()"""
+    mocker.patch("DIRAC.Resources.Computing.BatchSystems.SLURM.random.randrange", return_value=1)
+    slurm = SLURM()
+
+    executableFile = "executableFile.sh"
+    with open(executableFile, "w") as f:
+        f.write(executableContent)
+
+    res = slurm._generateSrunWrapper(executableFile)
+    # Make sure a wrapper file has been generated and is executable
+    assert res == expectedContent
+
+    os.remove(executableFile)
+
+
+@pytest.mark.parametrize(
+    "numberOfNodes, outputContent, expectedContent",
+    [
+        ("3-5", srunOutput, srunExpected),
+        ("1", normalOutput, normalExpected),
+    ],
+)
+def test_getJobOutputFiles(numberOfNodes, outputContent, expectedContent):
+    """Test getJobOutputFiles()"""
+    slurm = SLURM()
+
+    # We remove the '\n' at the beginning/end of the file because there are not present in reality
+    outputContent = outputContent.strip()
+    # We only remove the '\n' at the beginning because processOutput adds a '\n' at the end
+    expectedContent = [i.lstrip() for i in expectedContent]
+
+    outputFile = "./1234.out"
+    with open(outputFile, "w") as f:
+        f.write(outputContent)
+
+    errorFile = "./1234.err"
+    with open(errorFile, "w") as f:
+        f.write(outputContent)
+
+    batchDict = {
+        "JobIDList": ["1234"],
+        "OutputDir": ".",
+        "ErrorDir": ".",
+        "NumberOfNodes": numberOfNodes,
+    }
+    result = slurm.getJobOutputFiles(**batchDict)
+    assert result["Status"] == 0
+
+    output = result["Jobs"]["1234"]["Output"]
+    error = result["Jobs"]["1234"]["Error"]
+    assert output == outputFile
+    assert error == errorFile
+
+    with open(outputFile, "r") as f:
+        wrapperContent = f.read()
+    for srunLines in expectedContent:
+        assert srunLines in wrapperContent
+
+    os.remove(outputFile)
+    os.remove(errorFile)

src/DIRAC/Resources/Computing/ParallelLibraries/__init__.py renamed to src/DIRAC/Resources/Computing/BatchSystems/test/__init__.py

@@ -79,7 +79,6 @@ def __init__(self, ceName):
         self.valid = None
         self.mandatoryParameters = []
         self.batchSystem = None
-        self.parallelLibrary = None
         self.taskResults = {}
         self.minProxyTime = gConfig.getValue("/Registry/MinProxyLifeTime", 10800)  # secs
         self.defaultProxyTime = gConfig.getValue("/Registry/DefaultProxyLifeTime", 43200)  # secs
@@ -195,52 +194,19 @@ def loadBatchSystem(self, batchSystemName):
 
         :param str batchSystemName: name of the batch system
         """
-        result = self._loadCEAttribute("BatchSystem", batchSystemName, "Resources.Computing.BatchSystems")
-        if not result["OK"]:
-            return result
-        self.batchSystem = result["Value"]
-        return S_OK()
-
-    def loadParallelLibrary(self, parallelLibraryName, workingDirectory="."):
-        """Instantiate object representing the parallel library that will generate a script to wrap the executable
-
-        :param str parallelLibraryName: name of the parallel library
-        """
-        result = self._loadCEAttribute(
-            "ParallelLibrary",
-            parallelLibraryName,
-            "Resources.Computing.ParallelLibraries",
-            workingDirectory=workingDirectory,
-        )
-        if not result["OK"]:
-            return result
-        self.parallelLibrary = result["Value"]
-        return S_OK()
-
-    def _loadCEAttribute(self, ceAttribute, typeAttribute, moduleLocation, **kwargs):
-        """Instantiate specific CE attributes
-
-        :param str ceAttribute: name of the CE attribute
-        :param str typeAttribute: class of the CE attribute
-        :param str moduleLocation: location of the class definition in the source code
-        :param dict kwargs: parameters to pass to the constructor of the ceAttribute instance
-        """
-        if typeAttribute is None:
-            typeAttribute = self.ceParameters[ceAttribute]
+        if batchSystemName is None:
+            batchSystemName = self.ceParameters["BatchSystem"]
 
         objectLoader = ObjectLoader()
-        moduleLocation += ".%s" % typeAttribute
-        result = objectLoader.loadObject(moduleLocation, typeAttribute)
+        result = objectLoader.loadObject("Resources.Computing.BatchSystems.%s" % batchSystemName, batchSystemName)
         if not result["OK"]:
-            self.log.error("Failed to load %s object: %s" % (ceAttribute, result["Message"]))
+            self.log.error("Failed to load batch object: %s" % result["Message"])
             return result
-        ceAttributeClass = result["Value"]
-        ceAttributeModuleFile = result["ModuleFile"]
-
-        ceAttributeObject = ceAttributeClass(**kwargs)
-        self.log.info("Class from module: ", ceAttributeModuleFile)
-
-        return S_OK(ceAttributeObject)
+        batchClass = result["Value"]
+        batchModuleFile = result["ModuleFile"]
+        self.batchSystem = batchClass()
+        self.log.info("Batch system class from module: ", batchModuleFile)
+        return S_OK()
 
     def setParameters(self, ceOptions):
         """Add parameters from the given dictionary overriding the previous values