def prepare(self, commands, environs=None, **gen_opts):
environs = environs or []
if self.num_tasks <= 0:
num_tasks_per_node = self.num_tasks_per_node or 1
min_num_tasks = (-self.num_tasks
if self.num_tasks else num_tasks_per_node)
try:
guessed_num_tasks = self.guess_num_tasks()
except NotImplementedError as e:
raise JobError('flexible task allocation is not supported by '
'this backend') from e
if guessed_num_tasks < min_num_tasks:
nodes_required = min_num_tasks // num_tasks_per_node
nodes_found = guessed_num_tasks // num_tasks_per_node
raise JobError('could not find enough nodes: '
'required %s, found %s' %
(nodes_required, nodes_found))
self._num_tasks = guessed_num_tasks
getlogger().debug('flex_alloc_tasks: setting num_tasks to %s' %
self._num_tasks)
with shell.generate_script(self.script_filename,
**gen_opts) as builder:
builder.write_prolog(self.emit_preamble())
for e in environs:
builder.write(e.emit_load_commands())
for c in commands:
builder.write_body(c)
def prepare(self, commands, environs=None, **gen_opts):
environs = environs or []
if self.num_tasks <= 0:
num_tasks_per_node = self.num_tasks_per_node or 1
min_num_tasks = (-self.num_tasks
if self.num_tasks else num_tasks_per_node)
try:
guessed_num_tasks = self.guess_num_tasks()
except NotImplementedError as e:
raise JobError('flexible node allocation is not supported by '
'this backend') from e
if guessed_num_tasks < min_num_tasks:
raise JobError(
'could not satisfy the minimum task requirement: '
'required %s, found %s' %
(min_num_tasks, guessed_num_tasks))
self.num_tasks = guessed_num_tasks
getlogger().debug('flex_alloc_nodes: setting num_tasks to %s' %
self.num_tasks)
with shell.generate_script(self.script_filename,
**gen_opts) as builder:
builder.write_prolog(self.scheduler.emit_preamble(self))
builder.write(runtime.emit_loadenv_commands(*environs))
for c in commands:
builder.write_body(c)
def test_generate(self):
with shell.generate_script(self.script_file.name) as gen:
gen.write_prolog('# this is a test script')
gen.write_prolog('# another comment')
gen.write('v1=10')
gen.write_prolog('export v2=20')
gen.write_body('((v3 = v1 + v2))')
gen.write_body('echo hello $v3')
gen.write('unset v2')
gen.write_epilog('echo foo')
gen.write_epilog('unset v1')
expected_output = '''#!/bin/bash
# this is a test script
# another comment
export v2=20
v1=10
((v3 = v1 + v2))
echo hello $v3
unset v2
echo foo
unset v1
'''
with open(self.script_file.name) as fp:
assert expected_output == fp.read()
def test_trap_exit(script_file):
with shell.generate_script(script_file, trap_exit=True) as gen:
gen.write('echo hello')
completed = osext.run_command(str(script_file), check=True)
assert 'hello' in completed.stdout
assert 0 == completed.returncode
assert '-reframe: script exiting with exit code: 0' in completed.stdout
def test_generate_login(self):
with shell.generate_script(self.script_file.name, login=True) as gen:
gen.write('echo hello')
expected_output = '''#!/bin/bash -l
echo hello
'''
with open(self.script_file.name) as fp:
assert expected_output == fp.read()
def test_trap_exit(self):
with shell.generate_script(self.script_file.name,
trap_exit=True) as gen:
gen.write('echo hello')
completed = os_ext.run_command(self.script_file.name, check=True)
self.assertIn('hello', completed.stdout)
self.assertEqual(0, completed.returncode)
self.assertIn("-reframe: script exiting with exit code: 0",
completed.stdout)
def test_trap_error(script_file):
with shell.generate_script(script_file, trap_errors=True) as gen:
gen.write('false')
gen.write('echo hello')
with pytest.raises(SpawnedProcessError) as cm:
osext.run_command(str(script_file), check=True)
exc = cm.value
assert 'hello' not in exc.stdout
assert 1 == exc.exitcode
assert "-reframe: command `false' failed (exit code: 1)" in exc.stdout
def test_trap_error(self):
with shell.generate_script(self.script_file.name,
trap_errors=True) as gen:
gen.write('false')
gen.write('echo hello')
with self.assertRaises(SpawnedProcessError) as cm:
os_ext.run_command(self.script_file.name, check=True)
exc = cm.exception
self.assertNotIn('hello', exc.stdout)
self.assertEqual(1, exc.exitcode)
self.assertIn("-reframe: command `false' failed (exit code: 1)",
exc.stdout)
def test_write_types(self):
class C:
def __str__(self):
return 'echo "C()"'
with shell.generate_script(self.script_file.name) as gen:
gen.write(['echo foo', 'echo hello'])
gen.write('echo bar')
gen.write(C())
expected_output = '''#!/bin/bash
echo foo
echo hello
echo bar
echo "C()"
'''
with open(self.script_file.name) as fp:
assert expected_output == fp.read()
def test_trap_signal(script_file):
with shell.generate_script(script_file, trap_signals=True) as gen:
gen.write('sleep 10')
gen.write('echo hello')
f_stdout = tempfile.NamedTemporaryFile(mode='w+', delete=False)
proc = osext.run_command_async(str(script_file), stdout=f_stdout,
start_new_session=True)
# Yield for some time to allow the script to start
time.sleep(1)
# We kill the whole spawned process group (we are launching a shell)
os.killpg(proc.pid, 15)
proc.wait()
f_stdout.flush()
f_stdout.seek(0)
stdout = f_stdout.read()
assert 'hello' not in stdout
assert 143 == proc.returncode
assert '-reframe: script caught signal: 15' in stdout
def test_generate_custom_filemode(script_file):
with shell.generate_script(script_file, stat.S_IREAD) as gen:
gen.write('x=1')
assert stat.S_IMODE(os.stat(script_file).st_mode) == stat.S_IREAD
