Merge pull request #33 from crowdresearch/develop

Travis setup + Websocket graceful closure

Author: shirishgoyal

.travis.yml

@@ -18,7 +18,7 @@ before_script:
     flake8 .
 
 # command to run tests
-script: pytest
+script: python -m pytest
 
 deploy:
   provider: pypi

daemo/api.py

@@ -48,14 +48,14 @@ def get_task_results_by_taskworker_id(self, taskworker_id):
             return None
 
     def update_approval_status(self, task):
-        log.debug(msg="updating status for task %d" % task["task_id"])
+        log.debug(msg="updating status for task worker %d" % task["taskworker_id"])
 
         STATUS_ACCEPTED = 3
         STATUS_REJECTED = 4
 
         data = {
             "status": STATUS_ACCEPTED if task["accept"] else STATUS_REJECTED,
-            "workers": [task["id"]]
+            "workers": [task["taskworker_id"]]
         }
 
         response = self.client.post(self.route.update_task_status, data=json.dumps(data))

daemo/channel.py

@@ -28,6 +28,8 @@ def __init__(self, queue, api_client, url):
         self.url = url
         self.api_client = api_client
         self.lock = threading.Lock()
+        # defer.setDebugging(True)
+
         self.clientDisconnected = defer.Deferred()
 
         access_token = self.api_client.get_auth_token()
@@ -38,9 +40,7 @@ def __init__(self, queue, api_client, url):
 
         self.factory = ClientFactory(self.url, headers=headers)
 
-        self.factory.onConnectionLost = self.clientDisconnected
-        self.clientDisconnected.addCallback(self.stop_reactor)
-
+        self.factory.force_close = False
         self.factory.protocol = ClientProtocol
         self.factory.queue = self.queue
 
@@ -50,33 +50,30 @@ def return_name(self):
     def run(self):
         log.debug(msg="opening channel...")
 
-        self.connector = connectWS(self.factory)
-
         if self.state == 0:
             self.state = 1
+
+            self.factory.onConnectionLost = self.clientDisconnected
+            self.clientDisconnected.addCallback(self.on_client_disconnected)
+
+            self.connector = connectWS(self.factory)
             self.connector.reactor.run()
 
-    def stop(self):
-        log.info(msg="closing channel...")
-        # print self.clientDisconnected
+    def stop(self, forced_closure):
+        if forced_closure:
+            self.factory.force_close = forced_closure
+            self.factory.stopTrying()
+            self.connector.disconnect()
 
-        self.factory.stopTrying()
-        self.connector.disconnect()
         return defer.gatherResults([self.clientDisconnected])
 
-    def stop_reactor(self, protocol):
-        if not self.factory.continueTrying > 0:
-
+    def on_client_disconnected(self, protocol):
+        if not (self.factory.continueTrying > 0):
             time.sleep(2)
 
-            # print "channel: stop_reactor"
-
             with self.lock:
                 if self.state > 0:
                     self.state = 0
                     self.connector.reactor.stop()
-        else:
-            print "setup factory callback"
-
-            # self.factory.onConnectionLost = self.clientDisconnected
-            # self.clientDisconnected.addCallback(self.stop_reactor)
+        # else:
+        #     log.warning("continue trying")

daemo/client.py

@@ -73,6 +73,8 @@ def __init__(self, credentials_path='credentials.json', rerun_key=None, multi_th
         if host is not None:
             self.host = host
 
+        self.pid = os.getpid()
+
         self.api_client = ApiClient(self.credentials_path, self.host, self.http_proto)
 
         self.store = Store()
@@ -566,9 +568,11 @@ def _process_task(self, payload):
                 self.store.batches[batch_index]["submissions"][task_group_id] += 1
 
                 if stream:
-                    self._stream_response(batch_index, task_id, task_group_id, task_data, approve, completed)
+                    self._stream_response(batch_index, task_id, task_group_id, taskworker_id, task_data, approve,
+                                          completed)
                 else:
-                    self._aggregate_responses(batch_index, task_id, task_group_id, task_data, approve, completed)
+                    self._aggregate_responses(batch_index, task_id, task_group_id, taskworker_id, task_data, approve,
+                                              completed)
 
                 self.check_for_pending_tasks_reviews()
         else:
@@ -598,17 +602,17 @@ def check_for_pending_tasks_reviews(self):
     def _review_completed(self, project_key, ratings, ignore_history=True):
         self.rate(project_key, ratings, ignore_history=ignore_history)
 
-    def _stream_response(self, batch_index, task_id, task_group_id, task_data, approve, completed):
+    def _stream_response(self, batch_index, task_id, task_group_id, taskworker_id, task_data, approve, completed):
         log.info(msg="streaming responses...")
 
         log.info(msg="calling approve callback...")
 
         if approve([task_data]):
             task_data["accept"] = True
-            log.info(msg="task %d approved" % task_id)
+            log.info(msg="task worker %d approved" % taskworker_id)
         else:
             task_data["accept"] = False
-            log.info(msg="task %d rejected" % task_id)
+            log.info(msg="task worker %d rejected" % taskworker_id)
 
             # reverse increment as rejection will create another task
             self.store.batches[batch_index]["submissions"][task_group_id] -= 1
@@ -624,11 +628,11 @@ def _stream_response(self, batch_index, task_id, task_group_id, task_data, appro
         if is_done:
             self.store.mark_task_completed(batch_index, task_id, task_group_id)
 
-    def _aggregate_responses(self, batch_index, task_id, task_group_id, task_data, approve, completed):
+    def _aggregate_responses(self, batch_index, task_id, task_group_id, taskworker_id, task_data, approve, completed):
         log.info(msg="aggregating responses...")
 
         # store it for aggregation (stream = False)
-        self.store.aggregate(batch_index, task_id, task_group_id, task_data)
+        self.store.aggregate(batch_index, task_id, task_group_id, taskworker_id, task_data)
 
         is_done = self.store.is_task_complete(batch_index, task_id, task_group_id)
 
@@ -638,44 +642,52 @@ def _aggregate_responses(self, batch_index, task_id, task_group_id, task_data, a
             is_done = self.store.is_batch_complete(batch_index)
 
             if is_done:
-                self.store.mark_batch_completed(batch_index)
+                self._on_batch_complete(batch_index, approve, completed)
 
-                tasks_data = self.store.get_aggregated(batch_index)
+    def _on_batch_complete(self, batch_index, approve, completed):
+        self.store.mark_batch_completed(batch_index)
 
-                log.info(msg="calling approve callback...")
-                approvals = approve(tasks_data)
+        tasks_data = self.store.get_aggregated(batch_index)
 
-                tasks_approvals = zip(tasks_data, approvals)
+        log.info(msg="calling approve callback...")
+        approvals = approve(tasks_data)
 
-                for task_approval in tasks_approvals:
-                    task_data = task_approval[0]
-                    approval = task_approval[1]
+        tasks_approvals = zip(tasks_data, approvals)
 
-                    task_data["accept"] = approval
+        for task_approval in tasks_approvals:
+            task_data = task_approval[0]
+            approval = task_approval[1]
 
-                    if approval:
-                        log.info(msg="task %d approved" % task_data.get("task_id"))
-                    else:
-                        log.info(msg="task %d rejected" % task_data.get("task_id"))
-                        self.store.batches[batch_index]["submissions"][task_group_id] -= 1
-                        self.store.mark_task_incomplete(batch_index, task_group_id)
-                        self.store.mark_batch_incomplete(batch_index)
+            task_data["accept"] = approval
 
-                    self.api_client.update_approval_status(task_data)
+            if approval:
+                log.info(msg="task worker %d approved" % task_data.get("taskworker_id"))
+            else:
+                log.info(msg="task worker %d rejected" % task_data.get("taskworker_id"))
+                self.store.mark_task_incomplete(
+                    batch_index,
+                    task_data.get("task_id"),
+                    task_data.get("task_group_id")
+                )
+                self.store.mark_batch_incomplete(batch_index)
 
-                is_done = self.store.is_batch_complete(batch_index)
+            self.api_client.update_approval_status(task_data)
 
-                if is_done:
-                    approved_tasks = [x[0] for x in zip(tasks_data, approvals) if x[1]]
+        is_done = self.store.is_batch_complete(batch_index)
 
-                    log.info(msg="calling completed callback...")
-                    completed(approved_tasks)
+        if is_done:
+            approved_tasks = [x[0] for x in zip(tasks_data, approvals) if x[1]]
+
+            log.info(msg="calling completed callback...")
+            completed(approved_tasks)
 
     def _fetch_task(self, task_id):
         data = self.api_client.fetch_task(task_id)
         return transform_task(data)
 
     def _open_channel(self):
+        signal.signal(signal.SIGINT, self._handler)
+
         # shared queue between main process and channel for message passing
         self.queue = multiprocessing.Queue()
 
@@ -687,23 +699,29 @@ def _open_channel(self):
         thread = callback_thread(name='signal monitor', target=signal.pause)
         thread.start()
 
-        signal.signal(signal.SIGINT, self._handler)
-
         subscribe_url = self.websock_proto + self.host + self.api_client.route.subscribe
 
-        print "starting channel..."
         self.channel = Channel(self.queue, self.api_client, subscribe_url)
         self.channel.start()
 
     def _handler(self, signum, frame):
+        forced_closure = signum in [signal.SIGINT, signal.SIGTERM]
+
         # call this handler to stop the processes definitively
-        if signum in [signal.SIGINT, signal.SIGTERM, signal.SIGABRT] and os.getpid() == self.channel.pid:
-            self.channel.stop()
+        if signum in [signal.SIGINT, signal.SIGTERM, signal.SIGABRT]:
+
+            if self.channel is not None and os.getpid() == self.channel.pid:
+                # log.warn(msg="closing channel thread")
+
+                self.channel.stop(forced_closure)
+
+                if self.queue is not None:
+                    self.queue.put(None)
+                    self.queue = None
 
-            if self.queue is not None:
-                self.queue.put(None)
-                self.queue = None
+            # if self.pid is not None and os.getpid() == self.pid:
+            #     log.warning(msg="client:closing main thread")
 
     def _stop(self):
-        log.info(msg="disconnecting...")
+        log.warn(msg="disconnecting...")
         os.kill(int(self.channel.pid), signal.SIGINT)

daemo/client_factory.py

@@ -8,19 +8,33 @@
 
 class ClientFactory(WebSocketClientFactory, ReconnectingClientFactory):
     maxRetries = 5
+    maxDelay = 300  # 5 min
 
     def clientConnectionFailed(self, connector, reason):
         log.warning("websocket connection failed.")
-        log.warning(reason.value)
+        # log.warning(reason.value)
 
-        if self.continueTrying > 0:
-            log.info("connecting again %d..." % self.retries)
-            self.retry(connector)
+        self.retryConnection(connector, reason, True)
 
     def clientConnectionLost(self, connector, reason):
-        # super(ClientFactory, self).clientConnectionLost(connector, reason)
         log.warning("websocket connection lost.")
 
-        if self.continueTrying > 0:
-            log.info("connecting again %d..." % self.retries)
-            self.retry(connector)
+        self.retryConnection(connector, reason, False)
+
+    def retryConnection(self, connector, reason, fail=False):
+        if self.continueTrying and self.retries < self.maxRetries:
+            self.connector = connector
+            log.info("connecting again (%d/%d)..." % (self.retries + 1, self.maxRetries))
+            self.retry(self.connector)
+        else:
+            self.stop(connector, fail)
+
+    def stop(self, connector, fail):
+        connector.disconnect()
+
+        if self.queue is not None:
+            self.queue.put(None)
+            self.queue = None
+
+        if fail:
+            connector.reactor.stop()

daemo/protocol.py

@@ -10,10 +10,13 @@
 class ClientProtocol(WebSocketClientProtocol):
     lost = False
 
-    def onConnect(self, response):
+    def connectionMade(self):
         log.info("channel connected")
         self.factory.resetDelay()
 
+    # def onConnect(self, response):
+    #     log.info("channel connected")
+
     def onOpen(self):
         log.info("channel opened")
 
@@ -31,19 +34,16 @@ def onSend(self, data):
         log.debug("<<<{}>>>".format(data))
 
     def onClose(self, wasClean, code, reason):
-        log.debug("channel closed")
+        log.warning("channel closed")
 
         if not wasClean:
             log.error(reason.value)
 
     def connectionLost(self, reason):
-        log.debug("connection lost")
-        # log.error(reason.value)
-
-        self.factory.onConnectionLost.callback(self)
+        if self.factory.force_close:
+            onConnectionLost = self.factory.onConnectionLost
 
-        # if isinstance(reason.value, ConnectionDone):
-        #     try:
-        #         self.factory.onConnectionLost.callback(self)
-        #     except Exception as e:
-        #         log.error(e)
+            # do not let callback fire again
+            if onConnectionLost is not None:
+                self.factory.onConnectionLost = None
+                onConnectionLost.callback(self)

daemo/storage.py

@@ -57,15 +57,19 @@ def is_task_complete(self, batch_index, task_id, task_group_id):
     def mark_task_completed(self, batch_index, task_id, task_group_id):
         if task_group_id in self.batches[batch_index]["status"]:
             self.batches[batch_index]["status"][task_group_id] = True
-            log.debug(msg="task %d is complete" % task_id)
+            log.debug(msg="task %d is complete" % task_group_id)
 
     def mark_task_incomplete(self, batch_index, task_id, task_group_id):
         if task_group_id in self.batches[batch_index]["status"]:
+            self.batches[batch_index]["submissions"][task_group_id] = self.batches[batch_index]["submissions"][
+                task_group_id] - 1
             self.batches[batch_index]["status"][task_group_id] = False
-            log.debug(msg="task %d is NOT complete" % task_id)
+            log.debug(msg="task %d is NOT complete" % task_group_id)
 
     def is_batch_complete(self, batch_index):
-        return all(self.batches[batch_index]["status"].values())
+        is_complete = all(self.batches[batch_index]["status"].values())
+        log.debug(msg="batch %d is %s complete" % (batch_index, ''if is_complete else 'NOT'))
+        return is_complete
 
     def mark_batch_completed(self, batch_index):
         log.debug(msg="batch %d is complete" % batch_index)
@@ -88,7 +92,8 @@ def all_batches_complete(self):
     def all_reviews_complete(self):
         return all([self.cache[match_group_id]["is_complete"] for match_group_id in self.cache.keys()])
 
-    def aggregate(self, batch_index, task_id, task_group_id, task_data):
+    def aggregate(self, batch_index, task_id, task_group_id, taskworker_id, task_data):
+        task_data["taskworker_id"] = taskworker_id
         self.batches[batch_index]["aggregated_data"].append({
             "task_id": task_id,
             "task_group_id": task_group_id,

daemo/utils.py

@@ -118,6 +118,11 @@ def transform_task_results(data):
 
     data["task_id"] = data["task"]
     data["worker_id"] = data["worker"]
+    data["taskworker_id"] = data["id"]
+
+    del data['task']
+    del data['worker']
+    del data['id']
 
     return data