def _run_sbatch(rundir, commands, cwd, use_srun, use_mpi_config):
sbatchstr = _generate_sbatch(rundir, cwd, use_srun, use_mpi_config)
log.debug('Generated sbatch file', sbatchstr)
with open(_sbatch_file(rundir), 'w') as sbatch:
sbatch.write(sbatchstr)
runscriptstr = _generate_runscript(commands)
log.debug('Generated runscript file', runscriptstr)
with open(_runscript_file(rundir), 'w') as runscript:
runscript.write(runscriptstr)
command = ['sbatch', '--wait', _sbatch_file(rundir)]
log.info('Invoking sbatch', ' '.join(command))
start = time.time()
result = subprocess.run(command,
env=env.env,
stderr=subprocess.PIPE,
stdout=subprocess.PIPE)
end = time.time()
log.info(f'sbatch finished in {end - start:.2f}s')
if result.returncode != 0 and result.stderr:
log.error(
f'sbatch finished with exit code '
f'{result.returncode} and message', result.stderr.decode())
raise RuntimeError(f'Job submission failed: {result.stderr.decode()}')
m = re.match(r'Submitted batch job (\d+)', result.stdout.decode())
if not m:
log.error(f'Failed parsing sbatch output', result.stdout.decode())
raise RuntimeError('Job submission failed; sbatch output: ' +
result.stdout.decode())
return int(m.group(1))
async def _run_async(command, log_output, **kwargs):
process = await asyncio.create_subprocess_exec(
*command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env.env,
**kwargs)
async def read_output(stream):
buffer = io.StringIO()
async for line in stream:
line = line.decode()
buffer.write(line)
log_output(command[0], line.strip('\n'))
buffer.seek(0)
return buffer.read()
returncode, stdout, stderr = await asyncio.gather(
process.wait(), read_output(process.stdout),
read_output(process.stderr))
if returncode != 0:
commstr = ' '.join(f'"{c}"' for c in command)
log.error(
f'{commstr} finished with exit code {returncode} and message',
stderr)
raise RuntimeError(f'{commstr} failed with message "{stderr}"')
return stdout
def run(command, **kwargs):
if not command:
raise ValueError('No command provided')
log.info('Invoking', ' '.join(command))
start = time.time()
try:
output = subprocess.check_output(command,
env=env.env,
stderr=subprocess.STDOUT,
**kwargs)
except subprocess.CalledProcessError as e:
log.error(f'{command[0]} failed with output', e.output.decode())
raise e
end = time.time()
log.info(f'{command[0]} finished in {end - start:.2f}s')
output = output.decode().strip()
log.debug(f'{command[0]} output', output)
return output