Config - Support customized env for all modes (#295)

Support customized env for all modes in configuration.

docs/superbench-config.mdx

@@ -384,7 +384,7 @@ Some attributes may only be suitable for specific mode.
 | `proc_num` |    ✓    |          ✓          |   ✓   |
 | `node_num` |    ✘    |          ✓          |   ✘   |
 | `prefix`   |    ✓    |          ✘          |   ✘   |
-| `env`      |    ✘    |          ✘          |   ✓   |
+| `env`      |    ✓    |          ✓          |   ✓   |
 | `mca`      |    ✘    |          ✘          |   ✓   |
 | `parallel` |    ✓    |          ✘          |   ✘   |
 
@@ -421,8 +421,12 @@ So `prefix: CUDA_VISIBLE_DEVICES={proc_rank}` will be expressed as `CUDA_VISIBLE
 
 ### `env`
 
-Environment variables to use in the mode,
-in a flatten key-value dictionary.
+Environment variables to use in the mode, in a flatten key-value dictionary.
+The value needs to be string, any integer or boolean values need to be enclosed in quotes.
+
+Formatted string is also supported in value, available variables include:
++ `proc_rank`
++ `proc_num`
 
 ### `mca`
 

superbench/runner/runner.py

@@ -68,6 +68,8 @@ def __validate_sb_config(self):    # noqa: C901
             if not self._sb_benchmarks[name].modes:
                 self._sb_benchmarks[name].modes = []
             for idx, mode in enumerate(self._sb_benchmarks[name].modes):
+                if not mode.env:
+                    self._sb_benchmarks[name].modes[idx].env = {}
                 if mode.name == 'local':
                     if not mode.proc_num:
                         self._sb_benchmarks[name].modes[idx].proc_num = 1
@@ -84,8 +86,6 @@ def __validate_sb_config(self):    # noqa: C901
                             'btl_tcp_if_exclude': 'lo,docker0',
                             'coll_hcoll_enable': 0,
                         }
-                    if not mode.env:
-                        self._sb_benchmarks[name].modes[idx].env = {}
                     for key in ['PATH', 'LD_LIBRARY_PATH', 'SB_MICRO_PATH']:
                         self._sb_benchmarks[name].modes[idx].env.setdefault(key, None)
 
@@ -150,7 +150,10 @@ def __get_mode_command(self, benchmark_name, mode, timeout=None):
             ).format(
                 proc_num=mode.proc_num,
                 mca_list=' '.join(f'-mca {k} {v}' for k, v in mode.mca.items()),
-                env_list=' '.join(f'-x {k}={v}' if v else f'-x {k}' for k, v in mode.env.items()),
+                env_list=' '.join(
+                    f'-x {k}={str(v).format(proc_rank=mode.proc_rank, proc_num=mode.proc_num)}'
+                    if isinstance(v, str) else f'-x {k}' for k, v in mode.env.items()
+                ),
                 command=exec_command,
             )
         else:
@@ -384,11 +387,14 @@ def _run_proc(self, benchmark_name, mode, vars):
         logger.info('Runner is going to run %s in %s mode, proc rank %d.', benchmark_name, mode.name, mode.proc_rank)
 
         timeout = self._sb_benchmarks[benchmark_name].timeout
+        env_list = '--env-file sb.env'
+        for k, v in mode.env.items():
+            if isinstance(v, str):
+                env_list += f' -e {k}={str(v).format(proc_rank=mode.proc_rank, proc_num=mode.proc_num)}'
         ansible_runner_config = self._ansible_client.get_shell_config(
-            (
-                'docker exec sb-workspace bash -c '
-                "'set -o allexport && source sb.env && set +o allexport && {command}'"
-            ).format(command=self.__get_mode_command(benchmark_name, mode, timeout))
+            "docker exec {env_list} sb-workspace bash -c '{command}'".format(
+                env_list=env_list, command=self.__get_mode_command(benchmark_name, mode, timeout)
+            )
         )
         if mode.name == 'mpi':
             ansible_runner_config = self._ansible_client.update_mpi_config(ansible_runner_config)

tests/runner/test_runner.py

@@ -153,11 +153,13 @@ def test_get_mode_command(self):
                     'env': {
                         'SB_MICRO_PATH': '/sb',
                         'FOO': 'BAR',
+                        'RANK': '{proc_rank}',
+                        'NUM': '{proc_num}',
                     },
                 },
                 'expected_command': (
                     'mpirun -tag-output -allow-run-as-root -hostfile hostfile -map-by ppr:8:node -bind-to numa '
-                    '-mca coll_hcoll_enable 0 -x SB_MICRO_PATH=/sb -x FOO=BAR '
+                    '-mca coll_hcoll_enable 0 -x SB_MICRO_PATH=/sb -x FOO=BAR -x RANK=2 -x NUM=8 '
                     f'sb exec --output-dir {self.sb_output_dir} -c sb.config.yaml -C superbench.enable=foo'
                 ),
             },