clustercontrolsge.py

#!/usr/bin/env python3
 
 
 
# std
import logging
import os
import subprocess
import stat
import shutil
 
# ours
import validationfunctions
from validationscript import Script
 
 
class Cluster:
    """
    A class that provides the controls for running jobs on a (remote)
    Sun Grid Engine cluster. It provides two methods:
    - is_job_finished(job): Returns True or False, depending on whether the job
        has finished execution
    - execute(job): Takes a job and executes it by sending it to the cluster
    """
 
    @staticmethod
    def is_supported():
        """
        Check if qsub is available
        """
        return shutil.which("qsub") is not None
 
    @staticmethod
    def name():
        """
        Returns name of this job control
        """
        return "cluster-sge"
 
    @staticmethod
    def description():
        """
        Returns description of this job control
        """
        return "Batch submission via command line to Grid Engine"
 
    def __init__(self):
        """!
        The default constructor.
        - Holds the current working directory, which is also the location of
          the shellscripts that are being sent to the cluster.
        - Initializes a logger which writes to validate_basf2.py's log.
        - Finds the revision of basf2 that will be set up on the cluster.
        """
 
        
        self.submit_command = (
            "qsub -cwd -l h_vmem={requirement_vmem}G,"
            "h_fsize={requirement_storage}G "
            "-oo {logfile} -q {queuename} -V"
        )
 
        
        self.requirement_vmem = 4
 
        
        self.requirement_storage = 50
 
        
        self.queuename = "short.q"
 
        
        self.path = os.getcwd()
 
        
        self.logger = logging.getLogger("validate_basf2")
 
        
 
        
        self.tools = self.adjust_path(os.environ["BELLE2_TOOLS"])
        belle2_release_dir = os.environ.get("BELLE2_RELEASE_DIR", None)
        belle2_local_dir = os.environ.get("BELLE2_LOCAL_DIR", None)
 
        
        self.b2setup = "b2setup"
        if belle2_release_dir is not None:
            self.b2setup += " " + belle2_release_dir.split("/")[-1]
        if belle2_local_dir is not None:
            self.b2setup = (
                "MY_BELLE2_DIR="
                + self.adjust_path(belle2_local_dir)
                + " "
                + self.b2setup
            )
        if os.environ.get("BELLE2_OPTION") != "debug":
            self.b2setup += "; b2code-option " + os.environ.get("BELLE2_OPTION")
 
        # Write to log which revision we are using
        self.logger.debug(f"Setting up the following release: {self.b2setup}")
 
        # Define the folder in which the log of the cluster messages will be
        # stored (same folder like the log for validate_basf2.py)
        clusterlog_dir = "./html/logs/__general__/"
        if not os.path.exists(clusterlog_dir):
            os.makedirs(clusterlog_dir)
 
        
        self.clusterlog = open(clusterlog_dir + "clusterlog.log", "w+")
 
    # noinspection PyMethodMayBeStatic
    def adjust_path(self, path: str):
        """!
        This method can be used if path names are different on submission
        and execution hosts.
        @param path: The past that needs to be adjusted
        @return: The adjusted path
        """
 
        return path
 
    # noinspection PyMethodMayBeStatic
    def available(self):
        """!
        The cluster should always be available to accept new jobs.
        @return: Will always return True if the function can be called
        """
 
        return True
 
    def execute(self, job: Script, options="", dry=False, tag="current"):
        """!
        Takes a Script object and a string with options and runs it on the
        cluster, either with ROOT or with basf2, depending on the file type.
 
        @param job: The steering file object that should be executed
        @param options: Options that will be given to the basf2 command
        @param dry: Whether to perform a dry run or not
        @param tag: The folder within the results directory
        @return: None
        """
 
        # Define the folder in which the results (= the ROOT files) should be
        # created. This is where the files containing plots will end up. By
        # convention, data files will be stored in the parent dir.
        # Then make sure the folder exists (create if it does not exist) and
        # change to cwd to this folder.
        output_dir = os.path.abspath(f"./results/{tag}/{job.package}")
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
 
        # Path where log file is supposed to be created
        log_file = output_dir + "/" + os.path.basename(job.path) + ".log"
 
        # Remove any left over done files
        donefile_path = f"{self.path}/script_{job.name}.done"
        if os.path.isfile(donefile_path):
            os.remove(donefile_path)
 
        # Now we need to distinguish between .py and .C files:
        extension = os.path.splitext(job.path)[1]
        if extension == ".C":
            # .c files are executed with root
            command = "root -b -q " + job.path
        else:
            # .py files are executed with basf2
            # 'options' contains an option-string for basf2, e.g. '-n 100'
            params = validationfunctions.basf2_command_builder(
                job.path, options.split()
            )
            command = subprocess.list2cmdline(params)
 
        # Create a helpfile-shellscript, which contains all the commands that
        # need to be executed by the cluster.
        # First, set up the basf2 tools and perform b2setup with the correct
        # revision. The execute the command (i.e. run basf2 or ROOT on a
        # steering file). Write the return code of that into a *.done file.
        # Delete the helpfile-shellscript.
        tmp_name = self.path + "/" + "script_" + job.name + ".sh"
        with open(tmp_name, "w+") as tmp_file:
            tmp_file.write(
                "#!/bin/bash \n\n"
                + "BELLE2_NO_TOOLS_CHECK=1 \n"
                + f"source {self.tools}/b2setup \n"
                + f"cd {self.adjust_path(output_dir)} \n"
                + f"{command} \n"
                + f"echo $? > {self.path}/script_{job.name}.done \n"
                + f"rm {tmp_name} \n"
            )
 
        # Make the helpfile-shellscript executable
        st = os.stat(tmp_name)
        os.chmod(tmp_name, st.st_mode | stat.S_IEXEC)
 
        # Prepare the command line command for submission to the cluster
        params = self.submit_command.format(
            queuename=self.queuename,
            requirement_storage=self.requirement_storage,
            requirement_vmem=self.requirement_vmem,
            logfile=log_file,
        ).split() + [tmp_name]
 
        # Log the command we are about the execute
        self.logger.debug(subprocess.list2cmdline(params))
 
        # Submit it to the cluster. The steering
        # file output will be written to 'log_file' (see above).
        # If we are performing a dry run, don't send anything to the cluster
        # and just create the *.done file right away and delete the *.sh file.
        if not dry:
            process = subprocess.Popen(
                params, stdout=self.clusterlog, stderr=subprocess.STDOUT
            )
 
            # Check whether the submission succeeded
            if process.wait() != 0:
                job.status = "failed"
        else:
            os.system(f"echo 0 > {self.path}/script_{job.name}.done")
            os.unlink(tmp_name)
 
    def is_job_finished(self, job: Script):
        """!
        Checks whether the '.done'-file has been created for a job. If so, it
        returns True, else it returns False.
        Also deletes the .done-File once it has returned True.
 
        @param job: The job of which we want to know if it finished
        @return: True if the job has finished, otherwise False
        """
 
        # If there is a file indicating the job is done, that is its name:
        donefile_path = f"{self.path}/script_{job.name}.done"
 
        # Check if such a file exists. If so, this means that the job has
        # finished.
        if os.path.isfile(donefile_path):
 
            # Read the returncode/exit_status for the job from the *.done-file
            with open(donefile_path) as f:
                try:
                    returncode = int(f.read().strip())
                except ValueError:
                    returncode = -654
 
            # Delete the *.done file
            os.remove(donefile_path)
 
            # Return that the job is finished + the return code for it
            return [True, returncode]
 
        # If no such file exists, the job has not yet finished
        else:
            return [False, 0]
 
    # noinspection PyMethodMayBeStatic
    def terminate(self, job: Script):
        """!
        Terminate a running job, not support with this backend so ignore the
        call.
        """
        self.logger.error("Script termination not supported.")