Add template for mpi_cmd_prefix

easybuild/framework/easyconfig/easyconfig.py

@@ -1459,17 +1459,32 @@ def generate_template_values(self):
 
     def _generate_template_values(self, ignore=None):
         """Actual code to generate the template values"""
-        if self.template_values is None:
-            self.template_values = {}
 
         # step 0. self.template_values can/should be updated from outside easyconfig
-        # (eg the run_setp code in EasyBlock)
+        # (eg the run_step code in EasyBlock)
 
         # step 1-3 work with easyconfig.templates constants
         # disable templating with creating dict with template values to avoid looping back to here via __getitem__
         prev_enable_templating = self.enable_templating
+
+        self.enable_templating = False
+
+        if self.template_values is None:
+            # if no template values are set yet, initiate with a minimal set of template values;
+            # this is important for easyconfig that use %(version_minor)s to define 'toolchain',
+            # which is a pretty weird use case, but fine...
+            self.template_values = template_constant_dict(self, ignore=ignore)
+
+        self.enable_templating = prev_enable_templating
+
+        # grab toolchain instance with templating support enabled,
+        # which is important in case the Toolchain instance was not created yet
+        toolchain = self.toolchain
+
+        # get updated set of template values, now with toolchain instance
+        # (which is used to define the %(mpi_cmd_prefix)s template)
         self.enable_templating = False
-        template_values = template_constant_dict(self, ignore=ignore)
+        template_values = template_constant_dict(self, ignore=ignore, toolchain=toolchain)
         self.enable_templating = prev_enable_templating
 
         # update the template_values dict

easybuild/framework/easyconfig/templates.py

@@ -145,7 +145,7 @@
 # versionmajor, versionminor, versionmajorminor (eg '.'.join(version.split('.')[:2])) )
 
 
-def template_constant_dict(config, ignore=None, skip_lower=None):
+def template_constant_dict(config, ignore=None, skip_lower=None, toolchain=None):
     """Create a dict for templating the values in the easyconfigs.
         - config is a dict with the structure of EasyConfig._config
     """
@@ -257,6 +257,17 @@ def template_constant_dict(config, ignore=None, skip_lower=None):
         except Exception:
             _log.warning("Failed to get .lower() for name %s value %s (type %s)", name, value, type(value))
 
+    # step 5. add additional conditional templates
+    if toolchain is not None and hasattr(toolchain, 'mpi_cmd_prefix'):
+        try:
+            # get prefix for commands to be run with mpi runtime using default number of ranks
+            mpi_cmd_prefix = toolchain.mpi_cmd_prefix()
+            if mpi_cmd_prefix is not None:
+                template_values['mpi_cmd_prefix'] = mpi_cmd_prefix
+        except EasyBuildError as err:
+            # don't fail just because we couldn't resolve this template
+            _log.warning("Failed to create mpi_cmd_prefix template, error was:\n%s", err)
+
     return template_values
 
 

easybuild/tools/toolchain/mpi.py

@@ -166,6 +166,22 @@ def mpi_family(self):
         else:
             raise EasyBuildError("mpi_family: MPI_FAMILY is undefined.")
 
+    def mpi_cmd_prefix(self, nr_ranks=1):
+        """Construct an MPI command prefix to precede an executable"""
+
+        # Verify that the command appears at the end of mpi_cmd_for
+        test_cmd = 'xxx_command_xxx'
+        mpi_cmd = self.mpi_cmd_for(test_cmd, nr_ranks)
+        if mpi_cmd.rstrip().endswith(test_cmd):
+            result = mpi_cmd.replace(test_cmd, '').rstrip()
+        else:
+            warning_msg = "mpi_cmd_for cannot be used by mpi_cmd_prefix, "
+            warning_msg += "requires that %(cmd)s template appears at the end"
+            self.log.warning(warning_msg)
+            result = None
+
+        return result
+
     def mpi_cmd_for(self, cmd, nr_ranks):
         """Construct an MPI command for the given command and number of ranks."""
 
@@ -180,10 +196,10 @@ def mpi_cmd_for(self, cmd, nr_ranks):
             self.log.info("Using specified template for MPI commands: %s", mpi_cmd_template)
         else:
             # different known mpirun commands
-            mpirun_n_cmd = "mpirun -n %(nr_ranks)d %(cmd)s"
+            mpirun_n_cmd = "mpirun -n %(nr_ranks)s %(cmd)s"
             mpi_cmds = {
                 toolchain.OPENMPI: mpirun_n_cmd,
-                toolchain.QLOGICMPI: "mpirun -H localhost -np %(nr_ranks)d %(cmd)s",
+                toolchain.QLOGICMPI: "mpirun -H localhost -np %(nr_ranks)s %(cmd)s",
                 toolchain.INTELMPI: mpirun_n_cmd,
                 toolchain.MVAPICH2: mpirun_n_cmd,
                 toolchain.MPICH: mpirun_n_cmd,
@@ -201,7 +217,7 @@ def mpi_cmd_for(self, cmd, nr_ranks):
                 impi_ver = self.get_software_version(self.MPI_MODULE_NAME)[0]
                 if LooseVersion(impi_ver) <= LooseVersion('4.1'):
 
-                    mpi_cmds[toolchain.INTELMPI] = "mpirun %(mpdbf)s %(nodesfile)s -np %(nr_ranks)d %(cmd)s"
+                    mpi_cmds[toolchain.INTELMPI] = "mpirun %(mpdbf)s %(nodesfile)s -np %(nr_ranks)s %(cmd)s"
 
                     # set temporary dir for MPD
                     # note: this needs to be kept *short*,
@@ -230,7 +246,7 @@ def mpi_cmd_for(self, cmd, nr_ranks):
 
                     # create nodes file
                     nodes = os.path.join(tmpdir, 'nodes')
-                    write_file(nodes, "localhost\n" * nr_ranks)
+                    write_file(nodes, "localhost\n" * int(nr_ranks))
 
                     params.update({'nodesfile': "-machinefile %s" % nodes})
 
@@ -240,9 +256,19 @@ def mpi_cmd_for(self, cmd, nr_ranks):
             else:
                 raise EasyBuildError("Don't know which template MPI command to use for MPI family '%s'", mpi_family)
 
+        missing = []
+        for key in sorted(params.keys()):
+            tmpl = '%(' + key + ')s'
+            if tmpl not in mpi_cmd_template:
+                missing.append(tmpl)
+        if missing:
+            raise EasyBuildError("Missing templates in mpi-cmd-template value '%s': %s",
+                                 mpi_cmd_template, ', '.join(missing))
+
         try:
             res = mpi_cmd_template % params
         except KeyError as err:
-            raise EasyBuildError("Failed to complete MPI cmd template '%s' with %s: %s", mpi_cmd_template, params, err)
+            raise EasyBuildError("Failed to complete MPI cmd template '%s' with %s: KeyError %s",
+                                 mpi_cmd_template, params, err)
 
         return res

test/framework/easyconfig.py

@@ -1008,6 +1008,19 @@ def test_templating(self):
         eb['description'] = "test easyconfig % %% %s% %%% %(name)s %%(name)s %%%(name)s %%%%(name)s"
         self.assertEqual(eb['description'], "test easyconfig % %% %s% %%% PI %(name)s %PI %%(name)s")
 
+        # test use of %(mpi_cmd_prefix)s template
+        test_ecs_dir = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'easyconfigs', 'test_ecs')
+        gompi_ec = os.path.join(test_ecs_dir, 't', 'toy', 'toy-0.0-gompi-2018a.eb')
+        test_ec = os.path.join(self.test_prefix, 'test.eb')
+        write_file(test_ec, read_file(gompi_ec) + "\nsanity_check_commands = ['%(mpi_cmd_prefix)s toy']")
+
+        ec = EasyConfig(test_ec)
+        self.assertEqual(ec['sanity_check_commands'], ['mpirun -n 1 toy'])
+
+        init_config(build_options={'mpi_cmd_template': "mpiexec -np %(nr_ranks)s -- %(cmd)s  "})
+        ec = EasyConfig(test_ec)
+        self.assertEqual(ec['sanity_check_commands'], ['mpiexec -np 1 -- toy'])
+
     def test_templating_doc(self):
         """test templating documentation"""
         doc = avail_easyconfig_templates()

test/framework/robot.py

@@ -424,14 +424,14 @@ def test_resolve_dependencies_minimal(self):
             # to test resolving of dependencies with minimal toolchain
             # for each of these, we know test easyconfigs are available (which are required here)
             "dependencies = [",
-            "   ('OpenMPI', '2.1.2'),",  # available with GCC/6.4.0-2.28
+            # the use of %(version_minor)s here is mainly to check if templates are being handled correctly
+            # (it doesn't make much sense, but it serves the purpose)
+            "   ('OpenMPI', '%(version_minor)s.1.2'),",  # available with GCC/6.4.0-2.28
             "   ('OpenBLAS', '0.2.20'),",  # available with GCC/6.4.0-2.28
             "   ('ScaLAPACK', '2.0.2', '-OpenBLAS-0.2.20'),",  # available with gompi/2018a
             "   ('SQLite', '3.8.10.2'),",
             "]",
             # toolchain as list line, for easy modification later;
-            # the use of %(version_minor)s here is mainly to check if templates are being handled correctly
-            # (it doesn't make much sense, but it serves the purpose)
             "toolchain = {'name': 'foss', 'version': '%(version_minor)s018a'}",
         ]
         write_file(barec, '\n'.join(barec_lines))

test/framework/toolchain.py

@@ -948,6 +948,48 @@ def test_nosuchtoolchain(self):
         tc = self.get_toolchain('intel', version='1970.01')
         self.assertErrorRegex(EasyBuildError, "No module found for toolchain", tc.prepare)
 
+    def test_mpi_cmd_prefix(self):
+        """Test mpi_exec_nranks function."""
+        self.modtool.prepend_module_path(self.test_prefix)
+
+        tc = self.get_toolchain('gompi', version='2018a')
+        tc.prepare()
+        self.assertEqual(tc.mpi_cmd_prefix(nr_ranks=2), "mpirun -n 2")
+        self.assertEqual(tc.mpi_cmd_prefix(nr_ranks='2'), "mpirun -n 2")
+        self.assertEqual(tc.mpi_cmd_prefix(), "mpirun -n 1")
+        self.modtool.purge()
+
+        self.setup_sandbox_for_intel_fftw(self.test_prefix)
+        tc = self.get_toolchain('intel', version='2018a')
+        tc.prepare()
+        self.assertEqual(tc.mpi_cmd_prefix(nr_ranks=2), "mpirun -n 2")
+        self.assertEqual(tc.mpi_cmd_prefix(nr_ranks='2'), "mpirun -n 2")
+        self.assertEqual(tc.mpi_cmd_prefix(), "mpirun -n 1")
+        self.modtool.purge()
+
+        self.setup_sandbox_for_intel_fftw(self.test_prefix, imklver='10.2.6.038')
+        tc = self.get_toolchain('intel', version='2012a')
+        tc.prepare()
+
+        mpi_exec_nranks_re = re.compile("^mpirun --file=.*/mpdboot -machinefile .*/nodes -np 4")
+        self.assertTrue(mpi_exec_nranks_re.match(tc.mpi_cmd_prefix(nr_ranks=4)))
+        mpi_exec_nranks_re = re.compile("^mpirun --file=.*/mpdboot -machinefile .*/nodes -np 1")
+        self.assertTrue(mpi_exec_nranks_re.match(tc.mpi_cmd_prefix()))
+
+        # test specifying custom template for MPI commands
+        init_config(build_options={'mpi_cmd_template': "mpiexec -np %(nr_ranks)s -- %(cmd)s", 'silent': True})
+        self.assertEqual(tc.mpi_cmd_prefix(nr_ranks="7"), "mpiexec -np 7 --")
+        self.assertEqual(tc.mpi_cmd_prefix(), "mpiexec -np 1 --")
+
+        # check that we return None when command does not appear at the end of the template
+        init_config(build_options={'mpi_cmd_template': "mpiexec -np %(nr_ranks)s -- %(cmd)s option", 'silent': True})
+        self.assertEqual(tc.mpi_cmd_prefix(nr_ranks="7"), None)
+        self.assertEqual(tc.mpi_cmd_prefix(), None)
+
+        # template with extra spaces at the end if fine though
+        init_config(build_options={'mpi_cmd_template': "mpirun -np %(nr_ranks)s %(cmd)s  ", 'silent': True})
+        self.assertEqual(tc.mpi_cmd_prefix(), "mpirun -np 1")
+
     def test_mpi_cmd_for(self):
         """Test mpi_cmd_for function."""
         self.modtool.prepend_module_path(self.test_prefix)
@@ -974,6 +1016,17 @@ def test_mpi_cmd_for(self):
         init_config(build_options={'mpi_cmd_template': "mpiexec -np %(nr_ranks)s -- %(cmd)s", 'silent': True})
         self.assertEqual(tc.mpi_cmd_for('test123', '7'), "mpiexec -np 7 -- test123")
 
+        # check whether expected error is raised when a template with missing keys is used;
+        # %(ranks)s should be %(nr_ranks)s
+        init_config(build_options={'mpi_cmd_template': "mpiexec -np %(ranks)s -- %(cmd)s", 'silent': True})
+        error_pattern = \
+            r"Missing templates in mpi-cmd-template value 'mpiexec -np %\(ranks\)s -- %\(cmd\)s': %\(nr_ranks\)s"
+        self.assertErrorRegex(EasyBuildError, error_pattern, tc.mpi_cmd_for, 'test', 1)
+
+        init_config(build_options={'mpi_cmd_template': "mpirun %(foo)s -np %(nr_ranks)s %(cmd)s", 'silent': True})
+        error_pattern = "Failed to complete MPI cmd template .* with .*: KeyError 'foo'"
+        self.assertErrorRegex(EasyBuildError, error_pattern, tc.mpi_cmd_for, 'test', 1)
+
     def test_prepare_deps(self):
         """Test preparing for a toolchain when dependencies are involved."""
         tc = self.get_toolchain('GCC', version='6.4.0-2.28')