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