Batch queued requests (#134)

* add limit on get_accepted_requests

* test log

* comment out reset queue

* re-introduce check queue

* cache requests qos properties and manage permission error in 1 commit

* move commits

* fix permission error

* copy values of the dictionary

* add perf logger

* perf logger and optimisation

* fix

* logs

* add logs

* add response_error

* delete log messages

* optional commit for add_event

* sessionflush every 2000

* fix permission error when properties are not cached

* commit instead of flush

* introduce limit on dismissed requests management

* move check permission outside of properties

* fix permission error

* improve performance

* qa

* qa

Author: francesconazzaro

cads_broker/config.py

@@ -39,6 +39,7 @@ class BrokerConfig(pydantic_settings.BaseSettings):
     broker_requeue_on_lost_requests: bool = True
     broker_requeue_limit: int = 3
     broker_max_internal_scheduler_tasks: int = 500
+    broker_max_accepted_requests: int = 2000
 
 
 class SqlalchemySettings(pydantic_settings.BaseSettings):

cads_broker/database.py

@@ -221,14 +221,17 @@ def get_running_requests(
 def get_accepted_requests(
     session: sa.orm.Session,
     last_created_at: datetime.datetime | None = None,
+    limit: int | None = None,
 ):
-    """Get all accepted requests."""
+    """Get all accepted requests after 'last_created_at'."""
     statement = sa.select(SystemRequest)
     if last_created_at:
         statement = statement.where(SystemRequest.created_at >= last_created_at)
     statement = statement.where(SystemRequest.status == "accepted").order_by(
         SystemRequest.created_at
     )
+    if limit:
+        statement = statement.limit(limit)
     return session.scalars(statement).all()
 
 
@@ -321,8 +324,12 @@ def count_users(status: str, entry_point: str, session: sa.orm.Session) -> int:
     )
 
 
-def get_dismissed_requests(session: sa.orm.Session) -> Iterable[SystemRequest]:
+def get_dismissed_requests(
+    session: sa.orm.Session, limit: int | None
+) -> Iterable[SystemRequest]:
     stmt_dismissed = sa.select(SystemRequest).where(SystemRequest.status == "dismissed")
+    if limit:
+        stmt_dismissed = stmt_dismissed.limit(limit)
     return session.scalars(stmt_dismissed).fetchall()
 
 
@@ -708,10 +715,12 @@ def add_event(
     request_uid: str,
     message: str,
     session: sa.orm.Session,
+    commit: bool = True,
 ):
     event = Events(event_type=event_type, request_uid=request_uid, message=message)
     session.add(event)
-    session.commit()
+    if commit:
+        session.commit()
 
 
 def dictify_request(request: SystemRequest) -> dict[str, Any]:

cads_broker/dispatcher.py

@@ -335,6 +335,7 @@ def manage_dismissed_request(self, request, session):
             request_uid=request.request_uid,
             message=dismission_metadata.get("message", ""),
             session=session,
+            commit=False,
         )
         previous_status = dismission_metadata.get("previous_status", "accepted")
         if dismission_metadata.get("reason", "DismissedRequest") == "PermissionError":
@@ -371,8 +372,10 @@ def sync_database(self, session: sa.orm.Session) -> None:
         """
         # the retrieve API sets the status to "dismissed",
         # here the broker fixes the QoS and queue status accordingly
-        dismissed_requests = db.get_dismissed_requests(session)
-        for request in dismissed_requests:
+        dismissed_requests = db.get_dismissed_requests(
+            session, limit=CONFIG.broker_max_accepted_requests
+        )
+        for i, request in enumerate(dismissed_requests):
             if future := self.futures.pop(request.request_uid, None):
                 future.cancel()
             session = self.manage_dismissed_request(request, session)
@@ -561,6 +564,34 @@ def on_future_done(self, future: distributed.Future) -> str:
             future.release()
         return future.key
 
+    @perf_logger
+    def cache_requests_qos_properties(self, requests, session: sa.orm.Session) -> None:
+        """Cache the qos properties of the requests."""
+        # copy list of requests to avoid RuntimeError: dictionary changed size during iteration
+        for request in list(requests):
+            try:
+                self.qos._properties(request, check_permissions=True, session=session)
+            except PermissionError as exception:
+                db.add_event(
+                    event_type="user_visible_error",
+                    request_uid=request.request_uid,
+                    message=exception.args[0],
+                    session=session,
+                )
+                request = db.get_request(request.request_uid, session=session)
+                request.status = "failed"
+                request.finished_at = datetime.datetime.now()
+                request.response_error = {
+                    "reason": "PermissionError",
+                    "message": exception.args[0],
+                }
+                self.queue.pop(request.request_uid, None)
+                self.qos.notify_dismission_of_request(
+                    request, session, scheduler=self.internal_scheduler
+                )
+                logger.info("job has finished", **db.logger_kwargs(request=request))
+        session.commit()
+
     def processing_time_priority_algorithm(
         self,
         session_write: sa.orm.Session,
@@ -592,7 +623,9 @@ def processing_time_priority_algorithm(
             )
             for request in requests:
                 # need to check the limits on each request to update the qos_rules table
-                can_run = self.qos.can_run(request, session=session_write, scheduler=self.internal_scheduler)
+                can_run = self.qos.can_run(
+                    request, session=session_write, scheduler=self.internal_scheduler
+                )
                 if can_run and may_run and requests_counter < number_of_requests:
                     self.submit_request(request, session=session_write)
                     may_run = False
@@ -673,6 +706,7 @@ def run(self) -> None:
                         db.get_accepted_requests(
                             session=session_write,
                             last_created_at=self.queue.last_created_at,
+                            limit=CONFIG.broker_max_accepted_requests,
                         )
                     )
                     self.sync_qos_rules(session_write)
@@ -688,10 +722,13 @@ def run(self) -> None:
                         # if the internal queue is not in sync with the database, re-sync it
                         logger.info(
                             "re-syncing internal queue",
-                            internal_queue={queue_length},
-                            db_queue={db_queue},
+                            internal_queue=queue_length,
+                            db_queue=db_queue,
                         )
                         self.queue.reset()
+                    self.cache_requests_qos_properties(
+                        self.queue.values(), session_write
+                    )
 
                 running_requests = len(db.get_running_requests(session=session_read))
                 queue_length = self.queue.len()

cads_broker/entry_points.py

@@ -3,7 +3,6 @@
 import datetime
 import os
 import random
-import sqlite3
 import uuid
 from enum import Enum
 from typing import Any, Optional
@@ -19,42 +18,47 @@
 
 @app.command()
 def add_dummy_requests(
-    requests_db: str, number_of_requests: int = 1000, number_of_users: int = 100
+    number_of_requests: int = 1000, number_of_users: int = 100, max_length: int = 60
 ) -> None:
-    connection = sqlite3.connect(requests_db)
-    connection.row_factory = sqlite3.Row
-    cursor = connection.execute(f"SELECT * FROM broker limit {number_of_requests}")
     user_uids = [str(uuid.uuid4()) for _ in range(number_of_users)]
     with database.ensure_session_obj(None)() as session:
-        for i, row in enumerate(cursor):
-            if row["elapsed"] != "null":
-                database.ensure_adaptor_properties(
-                    hash="test",
-                    config={},
-                    form={},
-                    session=session,
-                )
-                request = database.SystemRequest(
-                    request_uid=str(uuid.uuid4()),
-                    process_id="test-adaptor-dummy",
-                    user_uid=random.choice(user_uids),
-                    status="accepted",
-                    request_body={
-                        "setup_code": None,
-                        "request": {
-                            "elapsed": str(datetime.timedelta(seconds=row["elapsed"])),
-                            "timestamp": str(datetime.datetime.now()),
-                        },
+        database.ensure_adaptor_properties(
+            hash="test",
+            config={},
+            form={},
+            session=session,
+        )
+        for i in range(number_of_requests):
+            request = database.SystemRequest(
+                request_uid=str(uuid.uuid4()),
+                process_id="test-adaptor-dummy",
+                user_uid=random.choice(user_uids),
+                # avoid using the same timestamp for all requests
+                created_at=datetime.datetime.now()
+                - datetime.timedelta(seconds=number_of_requests - i),
+                status="accepted",
+                request_body={
+                    "setup_code": None,
+                    "request": {
+                        "elapsed": str(
+                            datetime.timedelta(seconds=random.randint(0, max_length))
+                        ),
+                        "timestamp": str(datetime.datetime.now()),
                     },
-                    request_metadata={},
-                    origin="api",
-                    portal="c3s",
-                    adaptor_properties_hash="test",
-                    entry_point="cads_adaptors:DummyAdaptor",
-                )
-                session.add(request)
-                session.commit()
+                },
+                request_metadata={},
+                origin="api",
+                portal="c3s",
+                adaptor_properties_hash="test",
+                entry_point="cads_adaptors:DummyAdaptor",
+            )
+            session.add(request)
+
+            print(f"Added request {i} with request_uid {request.request_uid}")
+
+        # Commit all at the end
         session.commit()
+        print(f"Committed {number_of_requests} requests")
 
 
 @app.command()

cads_broker/qos/QoS.py

@@ -122,7 +122,7 @@ def can_run(self, request, session, scheduler):
         return not len(limits) and not len(permissions)
 
     @locked
-    def _properties(self, request, session):
+    def _properties(self, request, session, check_permissions=False):
         """Return the Properties object associated with a request.
 
         If it does not exists it is created.
@@ -135,18 +135,8 @@ def _properties(self, request, session):
 
         properties = Properties()
 
-        # First check permissions
-        for rule in self.rules.permissions:
-            if rule.match(request):
-                properties.permissions.append(rule)
-                if not rule.evaluate(request):
-                    database.set_dismissed_request(
-                        request_uid=request.request_uid,
-                        session=session,
-                        message=rule.info.evaluate(Context(request, self.environment)),
-                        reason="PermissionError",
-                    )
-                    break
+        if check_permissions:
+            self.check_permissions_for(request, properties)
 
         # Add general limits
         for rule in self.rules.global_limits:
@@ -217,7 +207,7 @@ def _status(self, request, session, out):
             out("    {}".format(priority))
 
         out("Permissions rules:")
-        for permission in self.permissions_for(request, session):
+        for permission in self.check_permissions_for(request, session):
             out("    {}".format(permission))
 
     @locked
@@ -229,12 +219,23 @@ def limits_for(self, request, session):
         return self._properties(request, session).limits
 
     @locked
-    def permissions_for(self, request, session):
+    def check_permissions_for(self, request, properties):
         """Return the permission rules that applies to a request.
 
         Ensure that the properties cache is created if needed.
         """
-        return self._properties(request, session).permissions
+        # check permissions
+        for rule in self.rules.permissions:
+            if rule.match(request):
+                properties.permissions.append(rule)
+                if not rule.evaluate(request):
+                    # Store in cache with empty properties
+                    self.requests_properties_cache[request.request_uid] = properties
+                    raise PermissionError(
+                        rule.info.evaluate(Context(request, self.environment))
+                    )
+
+        return properties.permissions
 
     @locked
     def priorities_for(self, request, session):
@@ -303,15 +304,16 @@ def notify_dismission_of_request(self, request, session, scheduler):
         for limit in self.limits_for(request, session):
             limit.remove_from_queue(request.request_uid)
             limits_list.append(limit)
-        scheduler.append(
-            {
-                "function": database.delete_request_qos_status,
-                "kwargs": {
-                    "rules": limits_list,
-                    "request_uid": request.request_uid,
-                },
-            }
-        )
+        if limits_list:
+            scheduler.append(
+                {
+                    "function": database.delete_request_qos_status,
+                    "kwargs": {
+                        "rules": limits_list,
+                        "request_uid": request.request_uid,
+                    },
+                }
+            )
 
     @locked
     def notify_start_of_request(self, request, session, scheduler):

tests/test_02_database.py

@@ -100,18 +100,25 @@ def test_get_accepted_requests(session_obj: sa.orm.sessionmaker) -> None:
     successful_request = mock_system_request(
         status="running", adaptor_properties_hash=adaptor_properties.hash
     )
-    accepted_request = mock_system_request(
+    accepted_request1 = mock_system_request(
         status="accepted", adaptor_properties_hash=adaptor_properties.hash
     )
-    accepted_request_uid = accepted_request.request_uid
+    accepted_request2 = mock_system_request(
+        status="accepted", adaptor_properties_hash=adaptor_properties.hash
+    )
+    accepted_request_uid = accepted_request1.request_uid
     with session_obj() as session:
         session.add(adaptor_properties)
         session.add(successful_request)
-        session.add(accepted_request)
+        session.add(accepted_request1)
+        session.add(accepted_request2)
         session.commit()
-        requests = db.get_accepted_requests(session=session)
-    assert len(requests) == 1
-    assert requests[0].request_uid == accepted_request_uid
+        all_requests = db.get_accepted_requests(session=session)
+        one_request = db.get_accepted_requests(session=session, limit=1)
+    assert len(all_requests) == 2
+    assert len(one_request) == 1
+    assert min(all_requests, key=lambda x: x.created_at) == one_request[0]
+    assert all_requests[0].request_uid == accepted_request_uid
 
 
 def test_count_finished_requests_per_user(session_obj: sa.orm.sessionmaker) -> None:
@@ -863,4 +870,4 @@ def test_count_running_users(session_obj: sa.orm.sessionmaker) -> None:
         )
         assert 1 == db.count_users(
             session=session, entry_point="foobar", status="running"
-        )
+        )

tests/test_20_dispatcher.py

@@ -135,7 +135,6 @@ def test_processing_time_priority_algorithm(
     request1 = mock_system_request(user_uid="user1", request_uid="request1")
     request2 = mock_system_request(user_uid="user1", request_uid="request2")
 
-
     def mock_get_users_queue_from_processing_time():
         return {"user1": [request1, request2]}
 
@@ -154,9 +153,7 @@ def mock_submit_request(self, request, session):
             if candidate.request_uid == request.request_uid:
                 candidates.remove(candidate)
 
-    mocker.patch(
-        "cads_broker.dispatcher.Broker.submit_request", mock_submit_request
-    )
+    mocker.patch("cads_broker.dispatcher.Broker.submit_request", mock_submit_request)
 
     environment = Environment.Environment()
     qos = QoS.QoS(