nordugrid-arc-nagios-plugins 1.9.1-1 / usr / lib / python2.7 / dist-packages / arcnagios / jobutils.py

import arc, os, pipes, shutil, time
import nagutils, vomsutils
from arcnagios import arcutils, persistence
from arcnagios.rescheduler import Rescheduler
from arcnagios.jobplugins import load_jobplugin
from arcnagios.utils import lazy_property
from genshi.template.loader import TemplateLoader, TemplateNotFound
from glob import glob

class JobDescription(object):
    _required_attributes = ["job_name", "application_name", "script_path"]
    def __init__(self, job_name = None, application_name = None, logdir = None,
		 script_path = None, script_args = [],
		 output = 'stdout.txt', error = 'stderr.txt',
		 wall_time_limit = None, memory_limit = None,
		 staged_inputs = [], staged_outputs = [],
		 runtime_environments = [],
		 queue_name = None, template = 'default.jsdl'):
	self.template = template
	self.job_name = job_name or 'ARC Probe'
	self.application_name = application_name
	self.logdir = logdir
	self.script_path = script_path
	self.script_name = os.path.basename(script_path)
	self.script_args = script_args
	self.output = output
	self.error = error
	self.wall_time_limit = wall_time_limit
	self.memory_limit = memory_limit
	self.staged_inputs = staged_inputs
	self.staged_outputs = staged_outputs
	self.runtime_environments = runtime_environments
	self.queue_name = queue_name

    def verify(self):
	for attr in self._required_attributes:
	    if getattr(self, attr) is None:
		raise nagutils.ServiceUNKNOWN('Missing %s for job description.'
					      % attr)

class JobInfo(persistence.PersistentObject):
    persistent_attributes = {
	'submission_time':	persistence.pt_int,
	'host':			persistence.pt_str,
	'job_tag':		persistence.pt_str_opt,
	'progress_service':	persistence.pt_str_opt,
	'termination_service':	persistence.pt_str_opt,
	'job_id':		persistence.pt_str,
	'job_state':		persistence.pt_t(arcutils.jobstate_of_str),
	'job_specific_state':	persistence.pt_str_opt,
	'job_state_time':	persistence.pt_int_opt,
	'job_state_alert':	persistence.pt_int_opt,
	'check_time':		persistence.pt_int_opt,
	'check_attempts':	persistence.pt_int_opt,
	'fetch_attempts':	persistence.pt_int_opt,
	'stored_urls':		persistence.pt_str_list,
	'tests':		persistence.pt_str_list,
    }

    @property
    def host_and_tag(self):
	if self.job_tag:
	    return '%s#%s' % (self.host, self.job_tag)
	else:
	    return self.host

    def __cmp__(a, b):
	return cmp(a.host_and_tag, b.host_and_tag)

def parse_staging(spec):
    """Parse the arguments to the --stage-input and --stage-output options."""

    if ';' in spec:
	xs = spec.split(';')
	spec, urloptions = xs[0], xs[1:]
    else:
	urloptions = []
    if '=' in spec:
	filename, url = spec.split('=', 1)
    else:
	filename, url = os.path.basename(spec), spec
    return (filename, url, urloptions)

def key_value(s):
    kv = s.split('=', 1)
    if len(kv) != 2:
	raise ValueError('Expecting an argument of the form KEY=VALUE.')
    return kv

class JobNagiosPlugin(nagutils.NagiosPlugin, vomsutils.NagiosPluginVomsMixin):
    """Nagios probe to test ARC CEs.  The probe has two sub-commands
    implemented by `check_submit` and `check_monitor`.  The former is run on
    all CEs, while the latter is run to collect submitted jobs."""

    probe_name = 'ARCCE'
    main_config_section = ['arcce']

    JSDL_FILENAME = 'job.jsdl'
    JOB_SCRIPT_FILENAME = 'job.sh'
    JOBID_FILENAME = 'active.jobid'
    ACTIVE_JOB_FILENAME = 'active.map'
    JOB_OUTPUT_DIRNAME = 'job_output'

    _archive_filenames = [
	JOBID_FILENAME, JSDL_FILENAME, JOB_SCRIPT_FILENAME, ACTIVE_JOB_FILENAME
    ]

    _arc_bindir = None

    prev_status = None

    # Timeout in seconds for cleaner tasks.
    cleaner_arcrm_timeout = 5
    cleaner_arcclean_timeout = 5
    cleaner_arckill_timeout = 5

    def __init__(self, **kwargs):
	default_template_dirs = glob(os.path.join(self.config_dir(),'*.d'))
	default_template_dirs.sort()
	nagutils.NagiosPlugin.__init__(self, **kwargs)
	ap = self.argparser
	ap.add_argument('--fqan', dest = 'fqan')
	ap.add_argument('--top-workdir', dest = 'top_workdir',
		help = 'Parent directory of per-VO probe working directories '
		       ' (deprecated)')
	ap.add_argument('-O', dest = 'jobtest_options',
		action = 'append', type = key_value, default = [],
		help = 'Given a value of the form VAR=VALUE, binds VAR to '
		       'VALUE in the environment of the job tests.')
	ap.add_argument('--template-dir', dest = 'template_dirs',
		action = 'append', default = default_template_dirs,
		help = 'Add a directory from which job description templates '
		       'can be loaded.')
	self._user_config = arc.UserConfig()

    def parse_args(self, args):
	nagutils.NagiosPlugin.parse_args(self, args)
	self.opts.template_dirs.reverse()

    def top_vopath(self, suffix):
	return os.path.join(self.opts.arcnagios_spooldir,
			    self.voms_suffixed('ce') + '-' + suffix)

    @lazy_property
    def top_workdir(self):
	return self.top_vopath('state')

    def workdir_for(self, host, job_tag):
	if job_tag:
	    return os.path.join(self.top_workdir, host + '#' + job_tag)
	else:
	    return os.path.join(self.top_workdir, host)

    def archivedir_for(self, host, job_tag):
	return os.path.join(self.top_vopath('archive'),
			    time.strftime('%Y-%m/%Y-%m-%d/%H:%M:%S-')
			    + host + (job_tag and ('#' + job_tag) or ''))

    @lazy_property
    def template_loader(self):
	return TemplateLoader(self.opts.template_dirs)

    def write_jsdl(self, out_path, jobdesc):
	jobdesc.verify()
	try:
	    templ = self.template_loader.load(jobdesc.template)
	except TemplateNotFound, xc:
	    raise nagutils.ServiceUNKNOWN('%s. The searched included %s.'
		    % (xc, ', '.join(self.opts.template_dirs)))
	r = templ.generate(jd = jobdesc)
	fd = open(out_path, 'w')
	fd.write(r.render())
	fd.close()

    def _arcrm(self, url, n_attempts):
        cmd = ['arcrm', url, '-t', str(self.cleaner_arcrm_timeout)]
	if n_attempts > 8: cmd.append('-f')
        rc, output = self.run_arc_cmd(*cmd)
	if rc == 0:
	    self.log.info('Removed test file %s.' % url)
	    return True

        self.log.warn('Failed to remove %s.' % url)
        return False

    def _arcclean(self, job_id, n_attempts):
        cmd = ['arcclean', job_id, '-t', str(self.cleaner_arcclean_timeout)]
	if n_attempts > 8: cmd.append('-f')
        rc, output = self.run_arc_cmd(*cmd)
	if rc == 0:
	    self.log.info('Removed job %s' % job_id)
	    return True

        self.log.warning('Failed to clean %s' % job_id)
        return False

    def _arckill(self, job_id, n_attempts):
        cmd = ['arckill', job_id, '-t', str(self.cleaner_arckill_timeout)]
	rc, kill_output = self.run_arc_cmd(*cmd)
	if rc == 0:
	    return True

        cmd = ['arcclean', job_id, '-t', str(self.cleaner_arcclean_timeout)]
	if n_attempts > 8: cmd.append('-f')
	rc, clean_output = self.run_arc_cmd(*cmd)
	if rc == 0:
	    return True

	self.log.warning('Failed to kill %s: %s' % (job_id, kill_output))
	return False

    @lazy_property
    def cleaner(self):
	cleaner = Rescheduler(self.top_vopath('state.db'), 'cleaner',
			      log = self.log)
	cleaner.register('arcrm', self._arcrm)
	cleaner.register('arcclean', self._arcclean)
	cleaner.register('arckill', self._arckill)
	self.at_exit(cleaner.close)
	return cleaner

    def cleanup_job_files(self, host, job_tag, archive = False):
	self.log.debug('Cleaning up job files for %s.'%host)
	workdir = self.workdir_for(host, job_tag)
#	archdir = os.path.join(workdir,
#			       time.strftime('archive/%Y-%m-%d/%H:%M:%S'))
	archdir = self.archivedir_for(host, job_tag)
	archdir_created = False
	for filename in self._archive_filenames:
	    try:
		if archive:
		    if os.path.exists(os.path.join(workdir, filename)):
			if not archdir_created:
			    os.makedirs(archdir)
			    archdir_created = True
			os.rename(os.path.join(workdir, filename),
				  os.path.join(archdir, filename))
		else:
		    os.unlink(os.path.join(workdir, filename))
	    except StandardError:
		pass
	try:
	    job_output_dir = os.path.join(workdir, self.JOB_OUTPUT_DIRNAME)
	    if os.path.exists(job_output_dir) and os.listdir(job_output_dir):
		if archive:
		    if not archdir_created:
			os.makedirs(archdir)
			archdir_created = True
		    os.rename(job_output_dir,
			      os.path.join(archdir, self.JOB_OUTPUT_DIRNAME))
		else:
		    last_dir = job_output_dir + '.LAST'
		    shutil.rmtree(last_dir, ignore_errors = True)
		    os.rename(job_output_dir, last_dir)
	except StandardError, xc:
	    self.log.warn('Error clearing %s: %s'%(job_output_dir, xc))

    def run_arc_cmd(self, prog, *args, **kwargs):
	if self._arc_bindir:
	    prog = os.path.join(self._arc_bindir, prog)
	cmd = prog + ' ' + ' '.join([pipes.quote(str(arg)) for arg in args])
	self.log.debug('Exec: %s'%cmd)
	fh = os.popen(cmd + ' 2>&1')
	output = fh.read()
	err = fh.close()
	return err or 0, output

    def require_voms_proxy(self):
	proxy_path = vomsutils.NagiosPluginVomsMixin.require_voms_proxy(self)
#	self._user_config.KeyPath(self.opts.user_key)
#	self._user_config.CertificatePath(self.opts.user_cert)
	if proxy_path:
	    self._user_config.ProxyPath(proxy_path)
	try:
	    self._user_config.InitializeCredentials() # old API
	except TypeError:
	    self._user_config.InitializeCredentials(
		    arc.initializeCredentialsType(
			arc.initializeCredentialsType.RequireCredentials))

    def load_active_job(self, host, job_tag):
	"""Load information about the current job on `host : str` tagged with
	`job_tag : str`, or `None` if no information is found."""

	workdir = self.workdir_for(host, job_tag)
	ajf = os.path.join(workdir, self.ACTIVE_JOB_FILENAME)
	if os.path.exists(ajf):
	    self.log.debug('Loading job info from %s.'%ajf)
	    jobinfo = JobInfo()
	    jobinfo.persistent_load(ajf)
	    return jobinfo

    def save_active_job(self, jobinfo, host, job_tag):
	"""Save information about the current job running on `host : str`
	tagged with `job_tag : str`."""

	workdir = self.workdir_for(host, job_tag)
	ajf = os.path.join(workdir, self.ACTIVE_JOB_FILENAME)
	self.log.debug('Saving active job info.')
	# FIXME: Lock.
	jobinfo.persistent_save(ajf)

    def collect_active_jobids(self):
	jobids = set()
	for fp in glob(os.path.join(self.top_workdir, '*', 'active.jobid')):
	    try:
		fh = open(fp)
		jobids.add(fh.readline().strip())
	    except IOError, xc:
		self.log.error('Failed to read %s: %s' % (fp, xc))
	return jobids

    def discard_stored_urls(self, jobinfo):
	for url in jobinfo.stored_urls:
	    if not url.startswith('file:'):
		self.cleaner.call('arcrm', url)

    def cleanup_job_tests(self, jobinfo):
	for test_name in jobinfo.tests:
	    try:
		test = self.load_jobtest(test_name, hostname = jobinfo.host)
		test.cleanup(jobinfo.job_state)
	    except Exception, xc:
		self.log.error('Error in cleanup %s for %s: %s'
			       % (test_name, jobinfo.job_id, xc))

    def cleanup_job(self, jobinfo, archive = False):
	"""Clean up job state from a fetched job."""

	self.discard_stored_urls(jobinfo)
	self.cleanup_job_tests(jobinfo)
	self.cleanup_job_files(jobinfo.host, jobinfo.job_tag, archive = archive)

    def discard_job(self, jobinfo, archive = False):
	"""Discard the job described by `jobinfo : JobInfo`."""

	if jobinfo.job_state.is_final():
	    self.cleaner.call('arcclean', jobinfo.job_id)
	else:
	    self.cleaner.call('arckill', jobinfo.job_id)
	self.cleanup_job(jobinfo, archive = archive)

    def load_jobtest(self, jobtest_name, **env):
	"""Load a plugin-based job-test from the section of the configuration
	specified by `jobtest_name`.  The result is an instance of `JobPlugin`
	subclass specified by the ``jobplugin`` variable of the given
	section."""

	env['config_dir'] = self.config_dir()
	if self.voms:
	    env['voms'] = self.voms
	env.update(self.opts.jobtest_options)

	jobplugin_section = 'arcce.%s'%jobtest_name
	if not self.config.has_section(jobplugin_section):
	    if self.config.has_section('arc-ce.%s'%jobtest_name):
		self.log.warn('The section arc-ce.%s is deprecated, please use %s.'
			      % (jobtest_name, jobplugin_section))
		jobplugin_section = 'arc-ce.%s'%jobtest_name
	    else:
		raise nagutils.ServiceUNKNOWN(
			'Missing configuration section %s for '
			'job-plugin test.' % jobplugin_section)
	jobplugin_name = self.config.get(jobplugin_section, 'jobplugin')
	jobplugin_cls = load_jobplugin(jobplugin_name)
	return jobplugin_cls(jobplugin_name, self.config, jobplugin_section,
			     self.log, env)