Improve dismissed requests (#109)

* track killed worker log event from nanny

* add configurability to requeue

* request not found

* rename requeue env var

* change env variable

* check all qos status

* qa

* manage dismissed requests

* fix

* improve reset_qos_rules

* fix reset_qos_rules

* remove warning to restart the broker. it is not needed anymore

* qa

* tests

Author: francesconazzaro

cads_broker/database.py

@@ -14,6 +14,7 @@
 import sqlalchemy_utils
 import structlog
 from sqlalchemy.dialects.postgresql import JSONB
+from typing_extensions import Iterable
 
 import alembic.command
 import alembic.config
@@ -27,6 +28,9 @@
 status_enum = sa.Enum(
     "accepted", "running", "failed", "successful", "dismissed", name="status"
 )
+DISMISSED_MESSAGE = os.getenv(
+    "DISMISSED_MESSAGE", "The request has been dismissed by the system."
+)
 
 
 class NoResultFound(Exception):
@@ -421,6 +425,28 @@ def count_users(status: str, entry_point: str, session: sa.orm.Session) -> int:
     )
 
 
+def update_dismissed_requests(session: sa.orm.Session) -> Iterable[str]:
+    stmt_dismissed = (
+        sa.update(SystemRequest)
+        .where(SystemRequest.status == "dismissed")
+        .returning(SystemRequest.request_uid)
+        .values(status="failed", response_error={"reason": "dismissed request"})
+    )
+    dismissed_uids = session.scalars(stmt_dismissed).fetchall()
+    session.execute(  # type: ignore
+        sa.insert(Events),
+        map(
+            lambda x: {
+                "request_uid": x,
+                "message": DISMISSED_MESSAGE,
+                "event_type": "user_visible_error",
+            },
+            dismissed_uids,
+        ),
+    )
+    return dismissed_uids
+
+
 def get_events_from_request(
     request_uid: str,
     session: sa.orm.Session,
@@ -439,11 +465,23 @@ def get_events_from_request(
     return events
 
 
-def reset_qos_rules(session: sa.orm.Session):
+def reset_qos_rules(session: sa.orm.Session, qos):
     """Delete all QoS rules."""
     for rule in session.scalars(sa.select(QoSRule)):
         rule.system_requests = []
         session.delete(rule)
+    cached_rules: dict[str, Any] = {}
+    for request in get_running_requests(session):
+        # Recompute the limits
+        limits = qos.limits_for(request, session)
+        cached_rules.update(
+            delete_request_qos_status(
+                request_uid=request.request_uid,
+                rules=limits,
+                session=session,
+                rules_in_db=cached_rules,
+            )
+        )
     session.commit()
 
 
@@ -482,24 +520,6 @@ def add_qos_rule(
     return qos_rule
 
 
-def increment_qos_rule_running(
-    rules: list, session: sa.orm.Session, rules_in_db: dict[str, QoSRule] = {}, **kwargs
-):
-    """Increment the running counter of a QoS rule."""
-    created_rules: dict = {}
-    for rule in rules:
-        if (rule_uid := str(rule.__hash__())) in rules_in_db:
-            qos_rule = rules_in_db[rule_uid]
-        else:
-            try:
-                qos_rule = get_qos_rule(rule_uid, session)
-            except sqlalchemy.orm.exc.NoResultFound:
-                qos_rule = add_qos_rule(rule=rule, session=session)
-                created_rules[qos_rule.uid] = qos_rule
-        qos_rule.running += 1
-    return created_rules
-
-
 def decrement_qos_rule_running(
     rules: list, session: sa.orm.Session, rules_in_db: dict[str, QoSRule] = {}, **kwargs
 ):

cads_broker/dispatcher.py

@@ -212,19 +212,20 @@ def from_address(
         factory.register_functions()
         session_maker_read = db.ensure_session_obj(session_maker_read, mode="r")
         session_maker_write = db.ensure_session_obj(session_maker_write, mode="w")
-        with session_maker_write() as session:
-            db.reset_qos_rules(session)
         rules_hash = get_rules_hash(qos_config.rules_path)
+        qos = QoS.QoS(
+            qos_config.rules,
+            qos_config.environment,
+            rules_hash=rules_hash,
+        )
+        with session_maker_write() as session:
+            db.reset_qos_rules(session, qos)
         self = cls(
             client=client,
             session_maker_read=session_maker_read,
             session_maker_write=session_maker_write,
             environment=qos_config.environment,
-            qos=QoS.QoS(
-                qos_config.rules,
-                qos_config.environment,
-                rules_hash=rules_hash,
-            ),
+            qos=qos,
             address=address,
         )
         return self
@@ -244,21 +245,22 @@ def sync_database(self, session: sa.orm.Session) -> None:
 
         If the task is not in the dask scheduler, it is re-queued.
         """
+        # the retrieve API sets the status to "dismissed", here the broker deletes the request
+        # this is to better control the status of the QoS
+        dismissed_uids = db.update_dismissed_requests(session)
+        for uid in dismissed_uids:
+            if future := self.futures.pop(uid, None):
+                future.cancel()
+        if dismissed_uids:
+            self.queue.reset()
+            self.qos.reload_rules(session)
+            db.reset_qos_rules(session, self.qos)
+
         statement = sa.select(db.SystemRequest).where(
             db.SystemRequest.status.in_(("running", "dismissed"))
         )
         dask_tasks = get_tasks(self.client)
         for request in session.scalars(statement):
-            # the retrieve API set the status to "dismissed", here the broker deletes the request
-            # this is to better control the status of the QoS
-            if request.status == "dismissed":
-                db.delete_request(request=request, session=session)
-                self.qos.notify_end_of_request(
-                    request, session, scheduler=self.internal_scheduler
-                )
-                if future := self.futures.get(request.request_uid):
-                    future.cancel()
-                continue
             # if request is in futures, go on
             if request.request_uid in self.futures:
                 continue

cads_broker/entry_points.py

@@ -159,7 +159,6 @@ def delete_requests(
         typer.echo(
             f"Status set to 'dismissed' for {number_of_requests} requests in the broker database."
         )
-        typer.echo("Please restart the broker.")
 
 
 @app.command()

tests/test_02_database.py

@@ -516,25 +516,6 @@ def test_add_qos_rule(session_obj: sa.orm.sessionmaker) -> None:
         assert db.get_qos_rule(str(rule.__hash__()), session=session).name == rule.name
 
 
-def test_increment_qos_rule_running(session_obj: sa.orm.sessionmaker) -> None:
-    rule1 = MockRule("name1", "conclusion1", "info1", "condition1")
-    rule2 = MockRule("name2", "conclusion2", "info2", "condition2")
-    with session_obj() as session:
-        db.add_qos_rule(rule1, session=session)
-    with session_obj() as session:
-        db.increment_qos_rule_running([rule1, rule2], rules_in_db={}, session=session)
-        session.commit()
-    with session_obj() as session:
-        assert db.get_qos_rule(str(rule1.__hash__()), session=session).running == 1
-        rule = db.get_qos_rule(str(rule2.__hash__()), session=session)
-        assert rule.running == 1
-        rules_in_db = {rule.uid: rule}
-        db.increment_qos_rule_running([rule2], rules_in_db=rules_in_db, session=session)
-        session.commit()
-    with session_obj() as session:
-        assert db.get_qos_rule(str(rule2.__hash__()), session=session).running == 2
-
-
 def test_add_request_qos_status(session_obj: sa.orm.sessionmaker) -> None:
     rule1 = MockRule("name1", "conclusion1", "info1", "condition1")
     rule2 = MockRule("name2", "conclusion2", "info2", "condition2")
@@ -629,30 +610,6 @@ def test_decrement_qos_rule_running(session_obj: sa.orm.sessionmaker) -> None:
         )
 
 
-def test_reset_qos_rules(session_obj: sa.orm.sessionmaker) -> None:
-    rule1 = MockRule("name1", "conclusion1", "info1", "condition1")
-    rule2 = MockRule("name2", "conclusion2", "info2", "condition2")
-    adaptor_properties = mock_config()
-    request = mock_system_request(adaptor_properties_hash=adaptor_properties.hash)
-    request_uid = request.request_uid
-    with session_obj() as session:
-        rule1_db = db.add_qos_rule(rule1, session=session)
-        rule2_db = db.add_qos_rule(rule2, session=session)
-        rules_in_db = {rule1_db.uid: rule1_db, rule2_db.uid: rule2_db}
-        session.add(adaptor_properties)
-        session.add(request)
-        session.commit()
-
-        db.add_request_qos_status(
-            request, [rule1, rule2], session=session, rules_in_db=rules_in_db
-        )
-        session.commit()
-
-        db.reset_qos_rules(session=session)
-        request = db.get_request(request_uid, session=session)
-        assert db.get_qos_status_from_request(request) == {}
-
-
 def test_get_events_from_request(session_obj: sa.orm.sessionmaker) -> None:
     adaptor_properties = mock_config()
     request = mock_system_request(adaptor_properties_hash=adaptor_properties.hash)

tests/test_20_dispatcher.py

@@ -59,7 +59,7 @@ def test_broker_sync_database(
     in_futures_request_uid = str(uuid.uuid4())
     in_dask_request_uid = str(uuid.uuid4())
     lost_request_uid = str(uuid.uuid4())
-    dismissed_request_uid = str(uuid.uuid4())
+    # dismissed_request_uid = str(uuid.uuid4())
     adaptor_metadata = mock_config()
     in_futures_request = mock_system_request(
         request_uid=in_futures_request_uid,
@@ -76,17 +76,17 @@ def test_broker_sync_database(
         status="running",
         adaptor_properties_hash=adaptor_metadata.hash,
     )
-    dismissed_request = mock_system_request(
-        request_uid=dismissed_request_uid,
-        status="dismissed",
-        adaptor_properties_hash=adaptor_metadata.hash,
-    )
+    # dismissed_request = mock_system_request(
+    #     request_uid=dismissed_request_uid,
+    #     status="dismissed",
+    #     adaptor_properties_hash=adaptor_metadata.hash,
+    # )
     with session_obj() as session:
         session.add(adaptor_metadata)
         session.add(in_futures_request)
         session.add(in_dask_request)
         session.add(lost_request)
-        session.add(dismissed_request)
+        # session.add(dismissed_request)
         session.commit()
 
     def mock_get_tasks() -> dict[str, str]:
@@ -115,9 +115,9 @@ def mock_get_tasks() -> dict[str, str]:
             db.SystemRequest.request_uid == lost_request_uid
         )
         output_request = session.scalars(statement).first()
-        assert output_request.status == "accepted"
-        assert output_request.request_metadata.get("resubmit_number") == 1
+        assert output_request.status == "failed"
+        assert output_request.request_metadata.get("resubmit_number") is None
 
-        with pytest.raises(db.NoResultFound):
-            with session_obj() as session:
-                db.get_request(dismissed_request_uid, session=session)
+        # with pytest.raises(db.NoResultFound):
+        #     with session_obj() as session:
+        #         db.get_request(dismissed_request_uid, session=session)