Source code for idmtools_platform_slurm.utils.slurm_job

idmtools SlurmPlatform SlurmJob utils.

Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
import os
import sys
import subprocess
from pathlib import Path
from typing import TYPE_CHECKING, NoReturn

if TYPE_CHECKING:  # pragma: no cover
    from idmtools_platform_slurm.slurm_platform import SlurmPlatform


[docs]def create_slurm_indicator() -> NoReturn: """ Add environment variable. Returns: None """ os.environ[INDICATOR_VARIABLE] = '1'
[docs]def remove_slurm_indicator() -> NoReturn: """ Remove the environment variable. Returns: None """ os.environ.pop(INDICATOR_VARIABLE, None)
[docs]def check_slurm_indicator() -> bool: """ Check if the environment set to '1'. Returns: True/False """ return os.environ.get(INDICATOR_VARIABLE, '0') == '1'
[docs]def slurm_installed() -> bool: """ Check if Slurm system is installed or available. Returns: True/False """ try: subprocess.check_output(["sinfo", "-V"]) return True except: return False
[docs]def run_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 """ from idmtools_platform_slurm.utils.slurm_job.slurm_job import SlurmJob from idmtools_platform_slurm.slurm_platform import SlurmPlatform # Double make sure it is Slurm Platform if not isinstance(platform, SlurmPlatform): return False if run_on_slurm and not check_slurm_indicator(): # Locate the script # Wrong path due to emod_malaria bug: # script = os.path.abspath(sys.argv[0]) # Workaround: manually build full path script = Path(sys.path[0]).joinpath(Path(sys.argv[0]).name) # Collect script input parameters script_params = sys.argv[1:] # Run script as Slurm job sj = SlurmJob(script_path=script, platform=platform, script_params=script_params, cleanup=cleanup) # Kick off Slurm job return True else: return False