Fix QoS status in case of workers refresh (#151)

* add debug message

* fix

* debug

* try to fix qos status bug

* fix dask scheduler race condition

* try to fix released race condition

* remove debug messages

* remove debug print statement from QoS class

Author: francesconazzaro

cads_broker/dispatcher.py

@@ -477,6 +477,10 @@ def sync_database(self, session: sa.orm.Session) -> None:
         for request in requests:
             # if request is in futures, go on
             if request.request_uid in self.futures:
+                # notify start of request if it is not already notified
+                self.qos.notify_start_of_request(
+                    request, scheduler=self.internal_scheduler
+                )
                 continue
             elif task := scheduler_tasks.get(request.request_uid, None):
                 if (state := task["state"]) in ("memory", "erred"):
@@ -507,11 +511,27 @@ def sync_database(self, session: sa.orm.Session) -> None:
                             **db.logger_kwargs(request=finished_request),
                         )
                 # if the task is in processing, it means that the task is still running
-                if state == "processing":
+                elif state == "processing":
+                    # notify start of request if it is not already notified
+                    self.qos.notify_start_of_request(
+                        request, scheduler=self.internal_scheduler
+                    )
+                    continue
+                elif state == "released":
+                    # notify start of request if it is not already notified
+                    queued_request = db.requeue_request(
+                        request=request, session=session
+                    )
+                    if queued_request:
+                        self.queue.add(queued_request.request_uid, request)
+                        self.qos.notify_end_of_request(
+                            request, scheduler=self.internal_scheduler
+                        )
                     continue
             # if it doesn't find the request: re-queue it
             else:
                 request = db.get_request(request.request_uid, session=session)
+                # if the broker finds the cache_id it means that the job has finished
                 if request.cache_id:
                     successful_request = db.set_successful_request(
                         request_uid=request.request_uid,

cads_broker/qos/QoS.py

@@ -339,18 +339,20 @@ def notify_start_of_request(self, request, scheduler):
         """
         limits_list = []
         for limit in self.limits_for(request):
-            limit.increment(request.request_uid)
-            limits_list.append(limit)
-        scheduler.append(
-            {
-                "function": database.delete_request_qos_status,
-                "kwargs": {
-                    "rules": limits_list,
-                    "request_uid": request.request_uid,
-                },
-            }
-        )
-        # Keep track of the running request. This is needed by reconfigure(self)
+            if request.request_uid not in limit.running:
+                limit.increment(request.request_uid)
+                limits_list.append(limit)
+        if limits_list:
+            scheduler.append(
+                {
+                    "function": database.delete_request_qos_status,
+                    "kwargs": {
+                        "rules": limits_list,
+                        "request_uid": request.request_uid,
+                    },
+                }
+            )
+            # Keep track of the running request. This is needed by reconfigure(self)
 
     @locked
     def notify_end_of_request(self, request, scheduler):
@@ -361,7 +363,7 @@ def notify_end_of_request(self, request, scheduler):
         """
         limits_list = []
         for limit in self.limits_for(request):
-            limit.decrement()
+            limit.decrement(request.request_uid)
             limits_list.append(limit)
 
         scheduler.append(

cads_broker/qos/Rule.py

@@ -134,28 +134,31 @@ class Limit(QoSRule):
 
     def __init__(self, environment, info, condition, conclusion):
         super().__init__(environment, info, condition, conclusion)
-        # running requests
-        self.value = 0
         self.queued = set()
+        self.running = set()
 
     def increment(self, request_uid):
         self.remove_from_queue(request_uid)
-        self.value += 1
+        self.running.add(request_uid)
 
     def remove_from_queue(self, request_uid):
         if request_uid in self.queued:
             self.queued.remove(request_uid)
 
-    def decrement(self):
-        if self.value > 0:
-            self.value -= 1
+    def decrement(self, request_uid):
+        if request_uid in self.running:
+            self.running.remove(request_uid)
 
     def queue(self, request_uid):
         self.queued.add(request_uid)
 
     def capacity(self, request):
         return self.evaluate(request)
 
+    @property
+    def value(self):
+        return len(self.running)
+
     def full(self, request):
         # NOTE: the self.value can be greater than the limit capacity after a
         # reconfiguration of the QoS