Commit 4a99331f authored by Benjamin Murauer's avatar Benjamin Murauer
Browse files

Resolve "Specify SLURM arguments"

parent 93d885bd
......@@ -10,8 +10,8 @@ from dbispipeline.core import Core
from dbispipeline.storage_handlers import PostGresOutputHandler
from dbispipeline.storage_handlers import PrintHandler
from dbispipeline.utils import LOGGER
from dbispipeline.utils import prepare_slurm_job
from dbispipeline.utils import restore_backup
from dbispipeline.utils import write_slurm_job_file
@click.command(help='Uses the plan file specified in PLAN to run the dbis '
......@@ -19,7 +19,9 @@ from dbispipeline.utils import restore_backup
@click.option('--dryrun', is_flag=True, help='Don\'t store results into DB')
@click.option('--force', is_flag=True, help='Run even if git is dirty')
@click.option('-v', '--verbose', is_flag=True, help='increase logging')
@click.option('--slurm', is_flag=True, help='create slurm file and submit')
@click.option('--slurm', is_flag=True, help='create slurm file')
@click.option('--slurm-and-run', is_flag=True,
help='create slurm job file and submit using sbatch')
@click.option('--restore', type=str, help='result file to be restored')
@click.option(
'--mail',
......@@ -30,7 +32,7 @@ from dbispipeline.utils import restore_backup
' mail will be sent for each run. If set to \'total\', one mail will '
'be sent after the entire pipeline is complete.')
@click.argument('plan', type=click.Path(exists=True))
def main(dryrun, force, verbose, slurm, restore, mail, plan):
def main(dryrun, force, verbose, slurm, slurm_and_run, restore, mail, plan):
"""Entry point that executes the pipeline given a configuration."""
if verbose:
LOGGER.setLevel(logging.DEBUG)
......@@ -50,10 +52,12 @@ def main(dryrun, force, verbose, slurm, restore, mail, plan):
except git.GitError:
pass
if slurm:
jobfile = prepare_slurm_job(dryrun, force, verbose, restore, mail,
plan)
call(['sbatch', jobfile])
if slurm or slurm_and_run:
job = write_slurm_job_file(dryrun, force, verbose, restore, mail,
plan)
if slurm_and_run:
LOGGER.info('starting slurm job: %s', job)
call(['sbatch', job])
else:
Core(plan, dryrun=dryrun, mail=mail).run()
......
......@@ -311,7 +311,7 @@ def _notify(message, subject):
s.sendmail(sender, [recipient], msg.as_string())
def prepare_slurm_job(dryrun, force, verbose, restore, mail, plan):
def write_slurm_job_file(dryrun, force, verbose, restore, mail, plan):
"""
Writes a slurm job to 'slurmjobs/<timestamp>_<jobname>.job'.
......@@ -344,8 +344,7 @@ def prepare_slurm_job(dryrun, force, verbose, restore, mail, plan):
content = '#!/bin/bash -l\n'
content += '\n'.join(['#SBATCH ' + x for x in options])
venv_manager = config.get('slurm',
'virtual_env_manager',
venv_manager = config.get('slurm', 'virtual_env_manager',
fallback='pipenv')
content += f'\nsrun {venv_manager} run python -m dbispipeline {plan} '
......@@ -366,6 +365,7 @@ def prepare_slurm_job(dryrun, force, verbose, restore, mail, plan):
jobfile = join('slurmjobs', job_filename)
with open(jobfile, 'w') as o_f:
o_f.write(content)
LOGGER.info('slurm job written to %s', jobfile)
return jobfile
......
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