Source code for idmtools_platform_slurm.utils.slurm_job
"""idmtools SlurmPlatform SlurmJob utils.Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved."""importosimportsysimportsubprocessfrompathlibimportPathfromtypingimportTYPE_CHECKING,NoReturnifTYPE_CHECKING:# pragma: no coverfromidmtools_platform_slurm.slurm_platformimportSlurmPlatformINDICATOR_VARIABLE='RUN_ON_SLURM'
[docs]defremove_slurm_indicator()->NoReturn:""" Remove the environment variable. Returns: None """os.environ.pop(INDICATOR_VARIABLE,None)
[docs]defcheck_slurm_indicator()->bool:""" Check if the environment set to '1'. Returns: True/False """returnos.environ.get(INDICATOR_VARIABLE,'0')=='1'
[docs]defslurm_installed()->bool:""" Check if Slurm system is installed or available. Returns: True/False """try:subprocess.check_output(["sinfo","-V"])returnTrueexcept:returnFalse
[docs]defrun_script_on_slurm(platform:'SlurmPlatform',run_on_slurm:bool=False,cleanup:bool=True)->bool:""" This is a utility tool which wraps the SlurmJob creation and run. Args: platform: idmtools Platform run_on_slurm: True/False cleanup: True/False to delete the generated slurm job related files Returns: True/False """fromidmtools_platform_slurm.utils.slurm_job.slurm_jobimportSlurmJobfromidmtools_platform_slurm.slurm_platformimportSlurmPlatform# Double make sure it is Slurm Platformifnotisinstance(platform,SlurmPlatform):returnFalseifrun_on_slurmandnotcheck_slurm_indicator():# Locate the script# Wrong path due to emod_malaria bug:# script = os.path.abspath(sys.argv[0])# Workaround: manually build full pathscript=Path(sys.path[0]).joinpath(Path(sys.argv[0]).name)# Collect script input parametersscript_params=sys.argv[1:]# Run script as Slurm jobsj=SlurmJob(script_path=script,platform=platform,script_params=script_params,cleanup=cleanup)# Kick off Slurm jobsj.run()returnTrueelse:returnFalse