Commit c1bc12ce authored by one's avatar one
Browse files

Use env file in docker instead of /tmp

parent c128dabb
...@@ -534,21 +534,23 @@ def _run_proc(self, benchmark_name, mode, vars): ...@@ -534,21 +534,23 @@ def _run_proc(self, benchmark_name, mode, vars):
if isinstance(timeout, int): if isinstance(timeout, int):
timeout = max(timeout, 60) timeout = max(timeout, 60)
env_list = '--env-file /tmp/sb.env'
if 'skip' not in self._docker_config: if 'skip' not in self._docker_config:
self._docker_config.skip = False self._docker_config.skip = False
base_env_cmd = 'set -o allexport && source /root/sb.env && set +o allexport'
mode_env_cmds = []
if self._docker_config.skip: if self._docker_config.skip:
env_list = 'set -o allexport && source /tmp/sb.env && set +o allexport' base_env_cmd = 'set -o allexport && source /tmp/sb.env && set +o allexport'
for k, v in mode.env.items(): for k, v in mode.env.items():
formatted_value = self.__format_mode_env_value(v, mode.proc_rank, mode.proc_num) formatted_value = self.__format_mode_env_value(v, mode.proc_rank, mode.proc_num)
if formatted_value is not None: if formatted_value is not None:
envvar = self.__quote_env_assignment(k, formatted_value) envvar = self.__quote_env_assignment(k, formatted_value)
env_list += ( mode_env_cmds.append(f'export {envvar}')
f' -e {envvar}'
if not self._docker_config.skip else f' && export {envvar}' env_list = base_env_cmd
) if mode_env_cmds:
env_list = f"{env_list} && {' && '.join(mode_env_cmds)}"
fcmd = "docker exec {env_list} sb-workspace bash -lc '{command}'" fcmd = "docker exec sb-workspace bash -lc '{env_list} && {command}'"
if self._docker_config.skip: if self._docker_config.skip:
fcmd = "bash -c '{env_list} && cd $SB_WORKSPACE && {command}'" fcmd = "bash -c '{env_list} && cd $SB_WORKSPACE && {command}'"
ansible_runner_config = self._ansible_client.get_shell_config( ansible_runner_config = self._ansible_client.get_shell_config(
......
...@@ -501,7 +501,39 @@ def fake_get_shell_config(cmd): ...@@ -501,7 +501,39 @@ def fake_get_shell_config(cmd):
self.runner._run_proc('foo', mode, {'proc_rank': 0}) self.runner._run_proc('foo', mode, {'proc_rank': 0})
self.assertIn('-e "NCCL_BUFFSIZE=4194304"', captured['cmd']) self.assertIn('docker exec sb-workspace bash -lc', captured['cmd'])
self.assertIn('-e "NCCL_RINGS=0 1 2 3|0 3 2 1"', captured['cmd']) self.assertIn('source /root/sb.env', captured['cmd'])
self.assertIn('export "NCCL_BUFFSIZE=4194304"', captured['cmd'])
self.assertIn('export "NCCL_RINGS=0 1 2 3|0 3 2 1"', captured['cmd'])
self.assertIn('-x "NCCL_BUFFSIZE=4194304"', captured['cmd']) self.assertIn('-x "NCCL_BUFFSIZE=4194304"', captured['cmd'])
self.assertIn('-x "NCCL_RINGS=0 1 2 3|0 3 2 1"', captured['cmd']) self.assertIn('-x "NCCL_RINGS=0 1 2 3|0 3 2 1"', captured['cmd'])
self.assertIn('-x PATH', captured['cmd'])
self.assertNotIn('export "PATH=', captured['cmd'])
@mock.patch('superbench.runner.ansible.AnsibleClient.run')
def test_run_proc_no_docker_keeps_tmp_env_source(self, mock_ansible_client_run):
"""Test _run_proc still sources /tmp/sb.env in no_docker mode."""
mock_ansible_client_run.return_value = 0
self.runner._sb_benchmarks = {'foo': {}}
self.runner._docker_config.skip = True
captured = {}
def fake_get_shell_config(cmd):
captured['cmd'] = cmd
return {'module_args': cmd, 'cmdline': '', 'host_pattern': 'localhost', 'module': 'shell'}
self.runner._ansible_client.get_shell_config = fake_get_shell_config
mode = OmegaConf.create({
'name': 'local',
'proc_num': 1,
'env': {
'FOO': 'a b',
},
'prefix': '',
})
self.runner._run_proc('foo', mode, {'proc_rank': 0})
self.assertIn("bash -c 'set -o allexport && source /tmp/sb.env && set +o allexport", captured['cmd'])
self.assertIn('export "FOO=a b"', captured['cmd'])
self.assertIn('cd $SB_WORKSPACE && PROC_RANK=0 sb exec', captured['cmd'])
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment