#!/usr/bin/env python
# SunGridEngineClusterDriver.py
#
#   Copyright (C) 2006 CCLRC, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is 
#   included in the root directory of this package.
#
# 20th June 2006
# 
# A Driver implementation to work with sun grid engine clusters via the 
# "qsub" shell command. This is based on ScriptDriver. This works like...
#
# Inherited from DefaultClusterDriver

import os
import subprocess

from DefaultClusterDriver import DefaultClusterDriver

class SunGridEngineClusterDriver(DefaultClusterDriver):

    def __init__(self):

        DefaultClusterDriver.__init__(self)
        self._job_id = 0

        return

    def _check_sge_errors(self, sge_stderr_list):
        '''Check through the standard error from Sun Grid Engine
        for indications that something went wrong... called from
        cleanup()'''

        for o in sge_stderr_list:
            if 'command not found' in o:
                missing_program = o.split(':')[2].strip()
                raise RuntimeError, 'executable "%s" missing' % \
                      missing_program

    def submit(self):
        '''This is where most of the work will be done - in here is
        where the script itself gets written and run, and the output
        file channel opened when the process has finished...'''

        # this will return almost instantly, once the job is in
        # the queue
        pipe = subprocess.Popen(['qsub', '-V', '-cwd',
                                 '%s.sh' % self._script_name],
                                cwd = self._working_directory,
                                stdout = subprocess.PIPE,
                                stderr = subprocess.PIPE)

        # this will get all of the output as a tuple (stdout, stderr)
        stdout, stderr = pipe.communicate()

        # check the standard error
        if len(stderr) > 0:
            # something probably went wrong
            if 'error opening' in stderr:
                raise RuntimeError, 'executable "%s" does not exist' % \
                      stdout.split('\n')[0].split(':')[0].replace(
                    'error opening ', '')

        # probably everything is ok then

        # the job id etc go to the standard output
        self._job_id = stdout.split('\n')[0].split()[2]

        return

    def cleanup(self):
        '''Cleanup and close-out.'''
        
        sge_stdout = os.path.join(self._working_directory,
                                  '%s.sh.o%s' % (self._script_name,
                                                 self._job_id))

        sge_stderr = os.path.join(self._working_directory,
                                  '%s.sh.e%s' % (self._script_name,
                                                 self._job_id))

        # check the standard error file for any indications that
        # something went wrong running this job...

        error_output = open(sge_stderr, 'r').readlines()
        self._check_sge_errors(error_output)

        # check the stderr records for errors...
        self.check_for_error_text(error_output)

        # at this stage I should delete the sge specific files defined
        # above to be tidy...

        try:
            os.remove(sge_stdout)
            os.remove(sge_stderr)
        except:
            # something wrong with this deletion?
            pass

        return

    def kill(self):
        '''This is meaningless...'''

        pass

    
if __name__ == '__main__':
    # Then run a simple test

    d = SunGridEngineClusterDriver()

    d.set_executable('ExampleProgram')
    d.start()
    d.close()

    while True:
        line = d.output()

        if not line:
            break

        print line.strip()

