
__author__ = 'Miroslav Bauer'


import os
import zmq
import time
import util
import pickle
import operator
from logger    import Logger
from objproxy  import RpcProxy
from zmqthread import ZmqThread
from jexecutor import JobExecutor
from resource  import ResourceAgent
from job       import Job, JobParam
from resman    import ResourceManager
from constants import RES_STAT_WEIGHT, RESCHEDULE_TRIES
from constants import JOB_EXTERNAL_STATE as state

log = Logger('dmorph.scheduler')

class Scheduler(ZmqThread):
	''' Class responsible for planning jobs onto resources available '''

	def __init__(self, context):
		super(Scheduler, self).__init__(context)
		self.name       = 'scheduler'
		self._tasks     = {}
		self._resources = {}
		self.resman     = None

	def _bind(self):
		self._sockets['responder'] = self._context.socket(zmq.ROUTER)
		self._sockets['executors'] = self._context.socket(zmq.PULL)
		try:
			self._sockets['responder'].bind('inproc:///scheduler/responses')
			self._sockets['executors'].bind('inproc:///scheduler/executors')
		except IOError:
			log.err('Could not bind scheduler sockets:')
			self.terminate()
			raise

	def _connect(self):
		super(Scheduler, self)._connect()
		self.resman = RpcProxy(self._context, ResourceManager)
		self.resman.connect('resman')

	def _updateResources(self):
		for name, raProxy in self._resources.items():
			raProxy._terminate()

		self._resources.clear()

		resnames = self.resman.getOnlineResources()
		try:
			for name in resnames:
				raProxy = RpcProxy(self._context, ResourceAgent)
				raProxy.connect('resource/%s' % name)
				self._resources[name] = raProxy
		except Exception, e:
			log.err('Failed to update resources: %s' % e)
			self.terminate()
			raise

	def run(self):
		''' Polls for new jobs, picks resource and starts JobExecutor
		    thread for them. Then polls for executor results.
		'''
		self._bind()
		self._connect()
		self._waitFor(['resman'])

		responder = self._sockets['responder']
		executors = self._sockets['executors']
		pl = util.createPoller([responder, executors])
		self._signalReady()

		log.info('Scheduler started!')
		while not self.stop.is_set():
			socks = dict(pl.poll())
			if responder in socks and socks[responder] == zmq.POLLIN:
				self._handleRequest()
			elif executors in socks and socks[executors] == zmq.POLLIN:
				self._handleExecutorResponse()

		self._cleanup()

	def _handleRequest(self):
		self._updateResources()

		job = None
		msg = self._sockets['responder'].recv_multipart()
		if len(msg) != 2:
			log.err('Invalid message for scheduler: %s' % msg)
			return None
		origin, jobpick = msg
		try:
			job = self._acceptJob(origin, jobpick)
			self._scheduleJob(job)
		except RuntimeError, e:
			self._sendResponse(origin, e)
			log.err('%s:%s' % (e.__class__.__name__, e))
			if job: self._removeTask(job.jid)
		except Exception, e:
			log.err('%s:%s' % (e.__class__.__name__, e))	
			if job: self._removeTask(job.jid)
			raise

		return job
	
	def _handleExecutorResponse(self):
		jid, response = pickle.loads(self._sockets['executors'].recv())
		task = self._tasks[jid]
		if isinstance(response, Exception) and task['try'] < RESCHEDULE_TRIES:
			self._rescheduleJob(jid)
		else:
			self._sendResponse(task['owner'], response)
			self._removeTask(jid)

	def _sendResponse(self, recipient, response):
		self._sockets['responder'].send_multipart([recipient, pickle.dumps(response)])

	def _acceptJob(self, originator, jobpick):
		job = pickle.loads(jobpick)
		if job.jid in self._tasks:
			raise RuntimeError('Job #%s is already being processed' % job.jid)
		self._tasks[job.jid] = {'jobData':job, 'owner':originator, 'resource':None, 'try':0}
		self._changeJobState(job.jid, state['queued'])
		return job
	
	def _removeTask(self, jid):
		log.debug('Removing job #%s from scheduler registry' % jid)
		if jid in self._tasks: del self._tasks[jid]

	def _changeJobState(self, jid, state):
		self.notice(str(jid), 'STATE:%s' % state)
	
	def _rescheduleJob(self, jid):
		task = self._tasks[jid]
		self._tasks[jid]['try'] += 1
		badResource = task['resource']
		log.warn('Job #%s failed. Rescheduling. Tries left:'\
			' %s' % (jid, (RESCHEDULE_TRIES - task['try'])))
		self._scheduleJob(task['jobData'], [badResource])

	def _scheduleJob(self, job, excludingResources=[]):
		log.info('Job #%s needs %s' % (job.jid, job.envRequirements))
		if job.preferred and job.preferred in self._resources:
			winner = job.preferred
		else:
			resourceList = {}
			for name, res in self._resources.items():
				if name not in excludingResources:
					resourceList[name] = res
			resourceCandidates = self._filterResourcesByReqs(
				job.envRequirements, resourceList)
			if not resourceCandidates: raise RuntimeError(
				'No suitable resource found for job #%s' % job.jid)
			scoreList = { cand : 0 for cand in resourceCandidates }
			jsize = job.size
			log.debug('Job #%s approx. size is: %s kB' % (job.jid, jsize))
			for res in resourceCandidates:
				scoreList[res] = self._calculateScore(job, res, jsize)
			winner = max(scoreList.iteritems(), key=operator.itemgetter(1))[0]
			log.info('Scores for #%s: %s' % (job.jid, scoreList))

		log.info('Scheduling #%s to %s' % (job.jid, winner))
		self._changeJobState(job.jid, state['scheduled'])
		self._execJob(job, winner)

	def _execJob(self, job, resource):
		self._resources[resource].updateStat('last_job_assignment_time', time.time())
		self._tasks[job.jid]['resource'] = resource
		executor = JobExecutor(self._context, job, resource)
		self._threads[executor.name] = executor
		self._threads[executor.name].start()

	def _filterResourcesByReqs(self, requirements, resources):
		result = resources.keys()
		for req, val in requirements.items():
			for rn, ra in resources.items():
				if not ra.meetsRequirement(req, val):
					if rn in result: result.remove(rn)
					log.debug('%s doesn\'t meet requirement: %s' % (rn, req))
		log.debug('Suitable resources: %s' % result)
		return result
	
	def _calculateScore(self, job, resourcename, jobsize):
		stats = self._resources[resourcename].getStats()
		score = 0.0
		# Score for reliability
		rel = (float(stats['jobs_processed'] - stats['jobs_failed'])/
			(stats['jobs_processed']+0.001))*RES_STAT_WEIGHT['reliability']

		# Least recently used based score
		lru = (float(time.time() - stats['last_job_assignment_time'])/
			(60*15))*RES_STAT_WEIGHT['last_job_assignment_time']
		
		# Resource load and performance based score
		perf = float(RES_STAT_WEIGHT['exec_time_avg'])/(stats['exec_time_avg']+1.0)
		perf += float(RES_STAT_WEIGHT['preparing_jobs'])/(stats['preparing_jobs']+1.0)
		perf += float(RES_STAT_WEIGHT['executing_jobs'])/(stats['executing_jobs']+1.0)
		perf += float(RES_STAT_WEIGHT['unstaging_jobs'])/(stats['unstaging_jobs']+1.0)

		# Approx. filetransfer speed based score
		xfer = (RES_STAT_WEIGHT['transfer_speed_avg']*stats['transfer_speed_avg'])/(jobsize+0.01)
		log.debug('%s rel:%d lru:%d perf:%d xfer:%d' % (resourcename, rel,lru,perf,xfer))
		return int(rel + lru + perf + xfer)
	
	def terminate(self):
		''' Terminate all threads an close proxy object to resource manager '''
		super(Scheduler, self).terminate()
		if self.resman: self.resman._terminate()
