<s> from __future__ import absolute_import <EOL> import os <EOL> import random <EOL> import socket <EOL> import sys <EOL> import threading <EOL> import pickle <EOL> import logging <EOL> import subprocess <EOL> import traceback <EOL> from time import sleep , time <EOL> import psutil <EOL> import mesos . interface <EOL> from bd2k . util . expando import Expando <EOL> from mesos . interface import mesos_pb2 <EOL> import mesos . native <EOL> from struct import pack <EOL> from toil . batchSystems . abstractBatchSystem import BatchSystemSupport <EOL> from toil . resource import Resource <EOL> log = logging . getLogger ( __name__ ) <EOL> class MesosExecutor ( mesos . interface . Executor ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self ) : <EOL> super ( MesosExecutor , self ) . __init__ ( ) <EOL> self . popenLock = threading . Lock ( ) <EOL> self . runningTasks = { } <EOL> self . workerCleanupInfo = None <EOL> Resource . prepareSystem ( ) <EOL> self . address = None <EOL> def registered ( self , driver , executorInfo , frameworkInfo , slaveInfo ) : <EOL> """<STR_LIT>""" <EOL> log . info ( "<STR_LIT>" ) <EOL> self . address = socket . gethostbyname ( slaveInfo . hostname ) <EOL> nodeInfoThread = threading . Thread ( target = self . _sendFrameworkMessage , args = [ driver ] ) <EOL> nodeInfoThread . daemon = True <EOL> nodeInfoThread . start ( ) <EOL> def reregistered ( self , driver , slaveInfo ) : <EOL> """<STR_LIT>""" <EOL> log . info ( "<STR_LIT>" ) <EOL> def disconnected ( self , driver ) : <EOL> """<STR_LIT>""" <EOL> log . critical ( "<STR_LIT>" ) <EOL> def killTask ( self , driver , taskId ) : <EOL> try : <EOL> pid = self . runningTasks [ taskId ] <EOL> except KeyError : <EOL> pass <EOL> else : <EOL> os . kill ( pid , <NUM_LIT:9> ) <EOL> def shutdown ( self , driver ) : <EOL> log . critical ( '<STR_LIT>' ) <EOL> for taskId in self . runningTasks . keys ( ) : <EOL> self . killTask ( driver , taskId ) <EOL> Resource . cleanSystem ( ) <EOL> BatchSystemSupport . workerCleanup ( self . workerCleanupInfo ) <EOL> log . critical ( '<STR_LIT>' ) <EOL> def error ( self , driver , message ) : <EOL> """<STR_LIT>""" <EOL> log . critical ( "<STR_LIT>" + message ) <EOL> def _sendFrameworkMessage ( self , driver ) : <EOL> message = None <EOL> while True : <EOL> if message is None : <EOL> message = Expando ( address = self . address ) <EOL> psutil . cpu_percent ( ) <EOL> else : <EOL> message . nodeInfo = dict ( cores = float ( psutil . cpu_percent ( ) ) * <NUM_LIT> , <EOL> memory = float ( psutil . virtual_memory ( ) . percent ) * <NUM_LIT> , <EOL> workers = len ( self . runningTasks ) ) <EOL> driver . sendFrameworkMessage ( repr ( message ) ) <EOL> sleep ( random . randint ( <NUM_LIT> , <NUM_LIT> ) ) <EOL> def launchTask ( self , driver , task ) : <EOL> """<STR_LIT>""" <EOL> def runTask ( ) : <EOL> log . debug ( "<STR_LIT>" , task . task_id . value ) <EOL> sendUpdate ( mesos_pb2 . TASK_RUNNING ) <EOL> taskData = pickle . loads ( task . data ) <EOL> if self . workerCleanupInfo is not None : <EOL> assert self . workerCleanupInfo == taskData . workerCleanupInfo <EOL> else : <EOL> self . workerCleanupInfo = taskData . workerCleanupInfo <EOL> startTime = time ( ) <EOL> try : <EOL> popen = runJob ( taskData ) <EOL> self . runningTasks [ task . task_id . value ] = popen . pid <EOL> try : <EOL> exitStatus = popen . wait ( ) <EOL> wallTime = time ( ) - startTime <EOL> if <NUM_LIT:0> == exitStatus : <EOL> sendUpdate ( mesos_pb2 . TASK_FINISHED , wallTime ) <EOL> elif - <NUM_LIT:9> == exitStatus : <EOL> sendUpdate ( mesos_pb2 . TASK_KILLED , wallTime ) <EOL> else : <EOL> sendUpdate ( mesos_pb2 . TASK_FAILED , wallTime , message = str ( exitStatus ) ) <EOL> finally : <EOL> del self . runningTasks [ task . task_id . value ] <EOL> except : <EOL> wallTime = time ( ) - startTime <EOL> exc_info = sys . exc_info ( ) <EOL> log . error ( '<STR_LIT>' , exc_info = exc_info ) <EOL> exc_type , exc_value , exc_trace = exc_info <EOL> sendUpdate ( mesos_pb2 . TASK_FAILED , wallTime , <EOL> message = '<STR_LIT>' . join ( traceback . format_exception_only ( exc_type , exc_value ) ) ) <EOL> def runJob ( job ) : <EOL> """<STR_LIT>""" <EOL> if job . userScript : <EOL> job . userScript . register ( ) <EOL> log . debug ( "<STR_LIT>" , job . command ) <EOL> with self . popenLock : <EOL> return subprocess . Popen ( job . command , <EOL> shell = True , env = dict ( os . environ , ** job . environment ) ) <EOL> def sendUpdate ( taskState , wallTime = None , message = '<STR_LIT>' ) : <EOL> log . debug ( '<STR_LIT>' ) <EOL> status = mesos_pb2 . TaskStatus ( ) <EOL> status . task_id . value = task . task_id . value <EOL> status . message = message <EOL> status . state = taskState <EOL> if wallTime is not None : <EOL> status . data = pack ( '<STR_LIT:d>' , wallTime ) <EOL> driver . sendStatusUpdate ( status ) <EOL> log . debug ( '<STR_LIT>' ) <EOL> thread = threading . Thread ( target = runTask ) <EOL> thread . start ( ) <EOL> def frameworkMessage ( self , driver , message ) : <EOL> """<STR_LIT>""" <EOL> log . debug ( "<STR_LIT>" . format ( message ) ) <EOL> def main ( executorClass = MesosExecutor ) : <EOL> logging . basicConfig ( level = logging . DEBUG ) <EOL> log . debug ( "<STR_LIT>" ) <EOL> executor = executorClass ( ) <EOL> driver = mesos . native . MesosExecutorDriver ( executor ) <EOL> exit_value = <NUM_LIT:0> if driver . run ( ) == mesos_pb2 . DRIVER_STOPPED else <NUM_LIT:1> <EOL> assert len ( executor . runningTasks ) == <NUM_LIT:0> <EOL> sys . exit ( exit_value ) </s>
