Use HeartBeatTable to trace host state

amat27 · amat27 · commit 6e53b5263f6c · 2018-03-13T16:51:19.000+08:00
diff --git a/heartbeat_table.py b/heartbeat_table.py
@@ -0,0 +1,98 @@
+import threading
+from collections import namedtuple
+from datetime import datetime
+from datetime import timedelta
+import logging_aux
+
+
+class HeartBeatTable(object):
+
+    def __init__(self, provisioning_timeout=timedelta(minutes=15), heartbeat_timeout=timedelta(minutes=3)):
+        self._table = {}
+        self.logger = logging_aux.init_logger_aux(
+            "hpcframework.heartbeat", "hpcframework.heartbeat.log")
+        self.on_host_running = []
+        self._table_lock = threading.Lock()
+        self._provisioning_timeout = provisioning_timeout
+        self._heartbeat_timeout = heartbeat_timeout
+
+    def add_slaveinfo(self, hostname, agent_id, task_id, cpus, last_heartbeat=datetime.utcnow()):
+        if hostname in self._table and self._table[hostname].state != HpcState.Closed:
+            self.logger.warn("Heart beat entry of {} existed. old value: {}.".format(
+                hostname, str(self._table[hostname])))
+        slaveinfo = SlaveInfo(hostname, agent_id, task_id,
+                              cpus, last_heartbeat, HpcState.Provisioning)
+        self._table[hostname] = slaveinfo
+        self.logger.info("Heart beat entry added: {}".format(str(slaveinfo)))
+
+    def on_slave_heartbeat(self, hostname):
+        if hostname in self._table:
+            self._table[hostname].last_heartbeat = datetime.utcnow()
+            self.logger.info("Heatbeat from host {}".format(hostname))
+            if self._table[hostname].state == HpcState.Provisioning:
+                with self._table_lock:  # to ensure we only run running callback once per entry
+                    if self._table[hostname].state == HpcState.Provisioning:
+                        self._table[hostname].state = HpcState.Running
+                        self.__exec_callback(self.on_host_running)
+                        self.logger.info(
+                            "Host {} start running".format(hostname))
+        else:
+            self.logger.error(
+                "Host {} is not recognized. Heartbeat ignored.".format(hostname))
+
+    def on_slave_close(self, hostname):
+        if hostname in self._table:
+            self._table[hostname].state = HpcState.Closed
+            self.logger.info("Host {} closed".format(hostname))
+        else:
+            self.logger.error(
+                "Host {} is not recognized. Close event ignored.".format(hostname))
+
+    def get_task_info(self, hostname):
+        if hostname in self._table:
+            entry = self._table[hostname]
+            return (entry.task_id, entry.agent_id)
+        else:
+            self.logger.error(
+                "Host {} is not recognized. Failed to get task info.".format(hostname))
+
+    def __exec_callback(self, callbacks):
+        for callback in callbacks:
+            try:
+                self.logger.debug(
+                    'Callback %s on %s' % (callback.__name__)
+                )
+                callback()
+            except Exception as e:
+                self.logger.exception(
+                    'Error in %s callback: %s' % (callback.__name__, str(e))
+                )
+
+    def check_timeout(self, now=datetime.utcnow()):
+        provision_timeout_list = []
+        heartbeat_timeout_list = []
+        running_list = []
+        for host in dict(self._table):
+            if host.state == HpcState.Provisioning and now - host.heartbeat_timeout >= self._provisioning_timeout:
+                provision_timeout_list.append(host)
+            elif host.state == HpcState.Running:
+                if now - host.heartbeat_timeout >= self._heartbeat_timeout:
+                    heartbeat_timeout_list.append(host)
+                else:
+                    running_list.append(host)
+        return (provision_timeout_list, heartbeat_timeout_list, running_list)
+
+    def get_cores_in_provisioning(self):
+        cores = 0.0
+        for host in dict(self._table):
+            if host.state == HpcState.Provisioning:
+                cores += host.cpus
+        return cores
+
+
+SlaveInfo = namedtuple(
+    "SlaveInfo", "hostname agent_id task_id cpus last_heartbeat state")
+
+
+class HpcState:
+    Provisioning, Running, Closed = range(3)
diff --git a/hpcframework.py b/hpcframework.py
@@ -1,23 +1,25 @@
+import base64
+import codecs
 import json
-import datetime
-import time
-import os
-import sys
-import threading
 import logging
+import os
 import signal
 import sys
+import threading
+import time
 import uuid
-import base64
-import codecs
-import restserver
-import restclient
-import logging_aux
+from collections import namedtuple
+from datetime import datetime
 
 from mesoshttp.client import MesosClient
+from mesoshttp.offers import Offer
+
+import logging_aux
+import restclient
+import restserver
 from restclient import AutoScaleRestClient
+import heartbeat_table
 
-from mesoshttp.offers import Offer
 
 class Test(object):
     class MesosFramework(threading.Thread):
@@ -39,15 +41,17 @@ def __init__(self):
         # signal.signal(signal.SIGINT, signal.SIG_IGN)
         logging.getLogger('mesoshttp').setLevel(logging.DEBUG)
 
-        self.hpc_client = AutoScaleRestClient()   
-        self.core_provisioning = 0.0     
+        self.heartbeat_table = heartbeat_table.HeartBeatTable()
+
+        self.hpc_client = AutoScaleRestClient()
+        self.core_provisioning = 0.0
         with open("setupscript.ps1") as scriptfile:
             hpc_setup_ps1 = scriptfile.read()
-        self.logger.info("Loaded HPC setup script:/n{}".format(hpc_setup_ps1))
+        self.logger.info("Loaded HPC setup script:\n{}".format(hpc_setup_ps1))
         hpc_setup_ps1_utf16 = hpc_setup_ps1.encode('utf-16')
-        hpc_setup_ps1_utf16_nobom = hpc_setup_ps1_utf16[2:] if hpc_setup_ps1_utf16[0:2] == codecs.BOM_UTF16 else hpc_setup_ps1_utf16
+        hpc_setup_ps1_utf16_nobom = hpc_setup_ps1_utf16[2:] if hpc_setup_ps1_utf16[
+            0:2] == codecs.BOM_UTF16 else hpc_setup_ps1_utf16
         self.hpc_setup_ps1_b64 = base64.b64encode(hpc_setup_ps1_utf16_nobom)
-        
 
         self.driver = None  # type: MesosClient.SchedulerDriver
         self.mesos_client = MesosClient(mesos_urls=['http://172.16.1.4:5050'])
@@ -57,6 +61,10 @@ def __init__(self):
         self.mesos_client.on(MesosClient.UPDATE, self.status_update)
         self.th = Test.MesosFramework(self.mesos_client)
         self.th.start()
+
+        self.heartbeat_server = restserver.RestServer(self.heartbeat_table ,8088)
+        self.heartbeat_server.start()
+
         while True and self.th.isAlive():
             try:
                 self.th.join(1)
@@ -65,11 +73,12 @@ def __init__(self):
                 break
 
     def shutdown(self):
-        print('Stop requested by user, stopping framework....')
+        print 'Stop requested by user, stopping framework....'
         self.logger.warn('Stop requested by user, stopping framework....')
         self.driver.tearDown()
         self.mesos_client.stop = True
         self.stop = True
+        self.heartbeat_server.stop()
 
     def subscribed(self, driver):
         self.logger.warn('SUBSCRIBED')
@@ -78,24 +87,25 @@ def subscribed(self, driver):
     def status_update(self, update):
         # if update['status']['state'] == 'TASK_RUNNING':
         #     self.driver.kill(update['status']['agent_id']['value'], update['status']['task_id']['value'])
-        print(str(update))
+        self.logger.log("Update received:\n{}".format(str(update)))
 
     def offer_received(self, offers):
         # self.logger.info('OFFER: %s' % (str(offers)))
         grow_decision = self.hpc_client.get_grow_decision()
-        
-        if(grow_decision.cores_to_grow - self.core_provisioning > 0):
-            for offer in offers: # type: Offer                
-                self.logger.info("offer_received: {}".format(
-                    (str(offer.get_offer()))))
+
+        if grow_decision.cores_to_grow - self.core_provisioning > 0:
+            for offer in offers:  # type: Offer
                 mesos_offer = offer.get_offer()
+                self.logger.info("offer_received: {}".format(
+                    (str(mesos_offer))))
                 if 'attributes' in mesos_offer:
                     attributes = mesos_offer['attributes']
                     if self.get_text(attributes, 'os') != 'windows_server':
                         offer.decline()
                     else:
                         cores = self.get_scalar(attributes, 'cores')
-                        cpus = self.get_scalar(mesos_offer['resources'], 'cpus')
+                        cpus = self.get_scalar(
+                            mesos_offer['resources'], 'cpus')
 
                         if cores == cpus:
                             self.accept_offer(offer)
@@ -105,44 +115,44 @@ def offer_received(self, offers):
                     offer.decline()
         else:
             for offer in offers:
-                offer.decline()                
+                offer.decline()
 
     def accept_offer(self, offer):
-        self.logger.info("Offer %s meets hpc's requiremnt" %
+        self.logger.info("Offer %s meets HPC's requirement" %
                          offer.get_offer()['id']['value'])
         self.run_job(offer)
-
-        # i = 0
-        # for offer in offers:
-        #     if i == 0:
-        #         self.run_job(offer)
-        #     else:
-        #         offer.decline()
-        #     i+=1
-    def get_scalar(self, dict, name):
-        for i in dict:
+       
+    def get_scalar(self, collection, name):
+        for i in collection:
             if i['name'] == name:
                 return i['scalar']['value']
         return 0.0
 
-    def get_text(self, dict, name):
-        for i in dict:
+    def get_text(self, collection, name):
+        for i in collection:
             if i['name'] == name:
                 return i['text']['value']
         return ""
 
     def run_job(self, mesos_offer):
         offer = mesos_offer.get_offer()
         self.logger.info("Accepting offer: {}".format(str(offer)))
+
+        agent_id = offer['agent_id']['value']
+        hostname = offer['hostname']
+        task_id = uuid.uuid4().hex
+        cpus = self.get_scalar(offer['resources'], 'cpus')
+
         task = {
             'name': 'sample test',
-            'task_id': {'value': uuid.uuid4().hex},
-            'agent_id': {'value': offer['agent_id']['value']},
+            'task_id': {'value': task_id},
+            'agent_id': {'value': agent_id},
             'resources': [
                 {
                     'name': 'cpus',
                     'type': 'SCALAR',
-                    'scalar': {'value': self.get_scalar(offer['resources'], 'cpus') - 0.1}
+                    # work around of MESOS-8631
+                    'scalar': {'value': cpus - 0.1}
                 },
                 {
                     'name': 'mem',
@@ -152,14 +162,10 @@ def run_job(self, mesos_offer):
             ],
             'command': {'value': 'powershell -EncodedCommand ' + self.hpc_setup_ps1_b64}
         }
-        self.logger.debug("Sending command:/n{}".format(task['command']['value']))
+        self.logger.debug(
+            "Sending command:\n{}".format(task['command']['value']))
         mesos_offer.accept([task])
-
+        self.heartbeat_table.add_slaveinfo(hostname, agent_id, task, cpus)        
 
 if __name__ == "__main__":
-    rest_server = restserver.RestServer(8088)
-    server_thread = threading.Thread(target=rest_server.run)
-    server_thread.start()
     test_mesos = Test()
-    rest_server.stop()
-    server_thread.join()
diff --git a/restserver.py b/restserver.py
@@ -1,41 +1,64 @@
 from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
 import json
+import threading
+from heartbeat_table import HeartBeatTable
+import logging_aux
+import logging
 
-class RestServer(object):
-    class S(BaseHTTPRequestHandler):
-        def _set_headers(self):
-            self.send_response(200)
-            self.send_header('Content-type', 'text/html')
-            self.end_headers()
-
-        def do_GET(self):
-            self._set_headers()
-            self.wfile.write("<html><body><h1>hi!</h1></body></html>")
-
-        def do_HEAD(self):
-            self._set_headers()
-            
-        def do_POST(self):
-            # Doesn't do anything with posted data
-            content_length = int(self.headers['Content-Length']) # <--- Gets the size of data
-            post_data = self.rfile.read(content_length) # <--- Gets the data itself
-            self._set_headers()
-            json_obj = json.loads(post_data)        
-            self.wfile.write("<html><body><h1>POST!</h1><pre>" + str(json_obj) + "</pre></body></html>")
-
-    def __init__(self, port = 80):
+class RestServer(object): # TODO: replace this implementation with twisted based implementation
+    def __init__(self, heartbeat_table, port = 80):
+        self.logger = logging_aux.init_logger_aux(
+            "hpcframework.heatbeat_server", "hpcframework.heatbeat_server.log")
+        self._heartbeat_table = heartbeat_table # type: HeartBeatTable
         self._server_address = ('', port)
-        self._server_class = HTTPServer
-        self._handler_class = self.S
+        self._server_class = HTTPServer        
+        self._handler_class = HeartBeatHandler
         self._port = port
         self._httpd = self._server_class(self._server_address, self._handler_class)
+        self._server_thread = threading.Thread(target=self.run)    
+        HeartBeatHandler.logger = self.logger
+        HeartBeatHandler.heartbeat_table = self._heartbeat_table
 
     def run(self):        
-        print 'Starting httpd...'
+        self.logger.debug('Starting httpd...')
         self._httpd.serve_forever()            
     
     def stop(self):
         self._httpd.shutdown()   
+        self._server_thread.join()
+
+    def start(self):
+        self._server_thread.start()
+
+class HeartBeatHandler(BaseHTTPRequestHandler):
+    logger = None # type: logging.Logger
+    heartbeat_table = None # type: HeartBeatTable
+
+    def _set_headers(self):
+        self.send_response(200)
+        self.send_header('Content-type', 'text/html')
+        self.end_headers()
+
+    def do_GET(self):
+        self._set_headers()
+        self.wfile.write("<html><body><h1>hi from thread {}!</h1></body></html>".format(threading._get_ident()))
+
+    def do_HEAD(self):
+        self._set_headers()
+        
+    def do_POST(self):
+        # Doesn't do anything with posted data
+        content_length = int(self.headers['Content-Length']) # <--- Gets the size of data
+        post_data = self.rfile.read(content_length) # <--- Gets the data itself
+        self._set_headers()
+        json_obj = json.loads(post_data)        
+        # self.wfile.write("<html><body><h1>POST!</h1><pre>" + str(json_obj) + "</pre></body></html>")
+        self.logger.debug("Received heartbeat object {}".format(str(json_obj)))
+        try:
+            self.heartbeat_table.on_slave_heartbeat(json_obj['hostname'])
+        except Exception as ex:
+            self.logger.exception(ex)               
+ 
 
 # if __name__ == "__main__":
 #     from sys import argv
diff --git a/setupscript.ps1 b/setupscript.ps1
@@ -15,7 +15,7 @@ while (!$broughtOnline -and ($retryCount -lt 120)) {
     catch {
         Write-host "Wait for 5 secs and then retry"
         ++$retryCount
-        sleep 5
+        Start-Sleep 5
     }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ while (!$broughtOnline -and ($retryCount -lt 120)) {`
`15`	`15`	`catch {`
`16`	`16`	`Write-host "Wait for 5 secs and then retry"`
`17`	`17`	`++$retryCount`
`18`		`- sleep 5`
	`18`	`+ Start-Sleep 5`
`19`	`19`	`}`
`20`	`20`	`}`
`21`	`21`