[WIP] Qos on database (#103)

mcucchi9 · francesconazzaro · web-flow · commit f60e9c46a540 · 2024-04-22T16:04:09.000+02:00
* create qos_rules table

* various fixes

* fix qos status update on can_run and notify

* fix alembic configuration

* fix alembic

* line saving

* update get_qos_status_from_request

* update unit-test

* fix

* revert

* add conclusion_value to qosrule definition

* tests

* add conclusion info

* qa and tests

---------

Co-authored-by: Francesco Nazzaro &lt;francesconazzaro89@gmail.com&gt;
diff --git a/alembic/versions/8deb52d20c05_add_qos_rules_table.py b/alembic/versions/8deb52d20c05_add_qos_rules_table.py
@@ -0,0 +1,48 @@
+"""add qos_rules table.
+
+Revision ID: 8deb52d20c05
+Revises: 6ee20703d353
+Create Date: 2024-04-09 17:16:01.559118
+
+"""
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "8deb52d20c05"
+down_revision = "6ee20703d353"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "qos_rules",
+        sa.Column("uid", sa.Text, primary_key=True),
+        sa.Column("name", sa.Text),
+        sa.Column("info", sa.Text),
+        sa.Column("condition", sa.Text),
+        sa.Column("conclusion", sa.Text),
+        sa.Column("conclusion_value", sa.Text),
+        sa.Column("queued", sa.Integer),
+        sa.Column("running", sa.Integer),
+    )
+    op.create_table(
+        "system_request_qos_rule",
+        sa.Column(
+            "request_uid",
+            sa.dialects.postgresql.UUID(False),
+            sa.ForeignKey("system_requests.request_uid"),
+            primary_key=True,
+        ),
+        sa.Column(
+            "rule_uid", sa.Text, sa.ForeignKey("qos_rules.uid"), primary_key=True
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("system_request_qos_rule")
+    op.drop_table("qos_rules")
diff --git a/cads_broker/database.py b/cads_broker/database.py
@@ -37,6 +37,41 @@ class InvalidRequestID(Exception):
     pass
 
 
+class QoSRule(BaseModel):
+    """QoS Rule ORM model."""
+
+    __tablename__ = "qos_rules"
+
+    uid = sa.Column(sa.Text, primary_key=True)
+    name = sa.Column(sa.Text)
+    info = sa.Column(sa.Text)
+    condition = sa.Column(sa.Text)
+    conclusion = sa.Column(sa.Text)
+    conclusion_value = sa.Column(sa.Text)
+    queued = sa.Column(sa.Integer)
+    running = sa.Column(sa.Integer)
+
+    system_requests: sa.orm.Mapped[list["SystemRequest"]] = sa.orm.relationship(
+        "SystemRequest",
+        secondary="system_request_qos_rule",
+        back_populates="qos_rules",
+        uselist=True,
+    )
+
+
+class SystemRequestQoSRule(BaseModel):
+    """Association table for SystemRequest and QoSRule."""
+
+    __tablename__ = "system_request_qos_rule"
+
+    request_uid = sa.Column(
+        sa.dialects.postgresql.UUID(False),
+        sa.ForeignKey("system_requests.request_uid"),
+        primary_key=True,
+    )
+    rule_uid = sa.Column(sa.Text, sa.ForeignKey("qos_rules.uid"), primary_key=True)
+
+
 class Events(BaseModel):
     """Events ORM model."""
 
@@ -102,9 +137,21 @@ class SystemRequest(BaseModel):
     __mapper_args__ = {"eager_defaults": False}
 
     # joined is temporary
-    cache_entry = sa.orm.relationship(cacholote.database.CacheEntry, lazy="joined")
-    adaptor_properties = sa.orm.relationship(AdaptorProperties, lazy="select")
-    events = sa.orm.relationship(Events, lazy="select", passive_deletes=True)
+    cache_entry: sa.orm.Mapped["cacholote.database.CacheEntry"] = sa.orm.relationship(
+        cacholote.database.CacheEntry, lazy="joined"
+    )
+    adaptor_properties: sa.orm.Mapped["AdaptorProperties"] = sa.orm.relationship(
+        AdaptorProperties, lazy="select"
+    )
+    events: sa.orm.Mapped[list["Events"]] = sa.orm.relationship(
+        Events, lazy="select", passive_deletes=True
+    )
+    qos_rules: sa.orm.Mapped[list["QoSRule"]] = sa.orm.relationship(
+        "QoSRule",
+        secondary="system_request_qos_rule",
+        back_populates="system_requests",
+        uselist=True,
+    )
 
     @property
     def age(self):
@@ -376,40 +423,108 @@ def get_events_from_request(
     return events
 
 
-def get_qos_status_from_request(
-    request: SystemRequest,
-) -> dict[str, list[tuple[str, str]]]:
-    ret_value: dict[str, list[str]] = {}
-    for rule_name, rules in request.qos_status.items():
-        ret_value[rule_name] = []
-        for rule in rules.values():
-            ret_value[rule_name].append(
-                (rule.get("info", ""), rule.get("conclusion", ""))
-            )
-    return ret_value
+def reset_qos_rules(session: sa.orm.Session):
+    """Delete all QoS rules."""
+    for rule in session.scalars(sa.select(QoSRule)):
+        rule.system_requests = []
+        session.delete(rule)
+    session.commit()
 
 
-def set_request_qos_rule(
-    request: SystemRequest,
+def get_qos_rule(uid: str, session: sa.orm.Session):
+    """Get a QoS rule."""
+    statement = sa.select(QoSRule).where(QoSRule.uid == uid)
+    return session.scalars(statement).one()
+
+
+def add_qos_rule(
     rule,
     session: sa.orm.Session,
+    queued: int = 0,
+    running: int = 0,
 ):
-    qos_status = request.qos_status
-    old_rules = qos_status.get(rule.name, {})
-    rule_uid = rule.get_uid(request)
-    if rule_uid in old_rules:
-        return
-    old_rules[rule_uid] = {
-        "conclusion": str(rule.evaluate(request)),
-        "info": str(rule.info).replace('"', ""),
-        "condition": str(rule.condition),
-    }
-    qos_status[rule.name] = old_rules
-    session.execute(
-        sa.update(SystemRequest)
-        .filter_by(request_uid=request.request_uid)
-        .values(qos_status=qos_status)
+    """Add a QoS rule."""
+    qos_rule = QoSRule(
+        uid=str(rule.__hash__()),
+        name=str(rule.name),
+        info=str(rule.info),
+        condition=str(rule.condition),
+        conclusion=str(rule.conclusion),
+        # conclusion_value may change over time, this case is not handled
+        conclusion_value=str(rule.evaluate(request=None)),
+        queued=queued,
+        running=running,
     )
+    session.add(qos_rule)
+    session.commit()
+    return qos_rule
+
+
+def increment_qos_rule_running(rules: list, session: sa.orm.Session):
+    """Increment the running counter of a QoS rule."""
+    for rule in rules:
+        try:
+            qos_rule = get_qos_rule(str(rule.__hash__()), session)
+        except sqlalchemy.orm.exc.NoResultFound:
+            qos_rule = add_qos_rule(rule=rule, session=session)
+        qos_rule.running += 1
+    session.commit()
+
+
+def decrement_qos_rule_running(rules: list, session: sa.orm.Session):
+    """Increment the running counter of a QoS rule."""
+    for rule in rules:
+        qos_rule = get_qos_rule(str(rule.__hash__()), session)
+        qos_rule.running -= 1
+    session.commit()
+
+
+def delete_request_qos_status(request_uid: str, rules: list, session: sa.orm.Session):
+    """Delete all QoS rules from a request."""
+    request = get_request(request_uid, session)
+    for rule in rules:
+        try:
+            qos_rule = get_qos_rule(str(rule.__hash__()), session)
+        except sqlalchemy.orm.exc.NoResultFound:
+            qos_rule = add_qos_rule(rule=rule, session=session)
+        if qos_rule in request.qos_rules:
+            request.qos_rules.remove(qos_rule)
+            qos_rule.queued -= 1
+        qos_rule.running += 1
+    session.commit()
+
+
+def add_request_qos_status(request_uid: str, rules: list, session: sa.orm.Session):
+    request = get_request(request_uid, session)
+    for rule in rules:
+        try:
+            qos_rule = get_qos_rule(str(rule.__hash__()), session)
+        except sqlalchemy.orm.exc.NoResultFound:
+            qos_rule = add_qos_rule(rule=rule, session=session)
+        if qos_rule not in request.qos_rules:
+            qos_rule.queued += 1
+            request.qos_rules.append(qos_rule)
+    session.commit()
+
+
+def get_qos_status_from_request(
+    request: SystemRequest,
+) -> dict[str, list[dict[str, str]]]:
+    ret_value: dict[str, list[dict[str, str]]] = {}
+    rules = request.qos_rules
+    for rule in rules:
+        rule_name = rule.name
+        rule_summary = {
+            "info": rule.info,
+            "queued": rule.queued,
+            "running": rule.running,
+            "conclusion": rule.conclusion_value,
+        }
+        if rule_name not in ret_value:
+            ret_value[rule_name] = [rule_summary]
+        else:
+            ret_value[rule_name].append(rule_summary)
+    return ret_value
 
 
 def requeue_request(
@@ -428,7 +543,7 @@ def requeue_request(
         request.request_metadata = metadata
         request.status = "accepted"
         session.commit()
-        logger.info("--------- requeueing request", **logger_kwargs(request=request))
+        logger.info("requeueing request", **logger_kwargs(request=request))
         return request
     else:
         return
diff --git a/cads_broker/dispatcher.py b/cads_broker/dispatcher.py
@@ -2,6 +2,7 @@
 import io
 import operator
 import os
+import sched
 import time
 import traceback
 from typing import Any
@@ -118,6 +119,7 @@ class Broker:
         repr=lambda futures: " ".join(futures.keys()),
     )
     running_requests: int = 0
+    internal_scheduler: sched.scheduler = sched.scheduler(time.time, time.sleep)
 
     @classmethod
     def from_address(
@@ -131,6 +133,8 @@ 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)
         self = cls(
             client=client,
@@ -171,7 +175,9 @@ def sync_database(self, session: sa.orm.Session) -> None:
             # 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)
+                self.qos.notify_end_of_request(
+                    request, session, scheduler=self.internal_scheduler
+                )
                 continue
             # if request is in futures, go on
             if request.request_uid in self.futures:
@@ -239,7 +245,9 @@ def on_future_done(self, future: distributed.Future) -> None:
                 )
                 return
             self.futures.pop(future.key)
-            self.qos.notify_end_of_request(request, session)
+            self.qos.notify_end_of_request(
+                request, session, scheduler=self.internal_scheduler
+            )
             logger.info(
                 "job has finished",
                 dask_status=future.status,
@@ -259,7 +267,9 @@ def submit_requests(
         requests_counter = 0
         for request in queue:
             with self.session_maker_write() as session_write:
-                if self.qos.can_run(request, session=session_write):
+                if self.qos.can_run(
+                    request, session=session_write, scheduler=self.internal_scheduler
+                ):
                     self.submit_request(request, session=session_write)
                     requests_counter += 1
                     if requests_counter == int(number_of_requests * WORKERS_MULTIPLIER):
@@ -271,7 +281,9 @@ def submit_request(
         request = db.set_request_status(
             request_uid=request.request_uid, status="running", session=session
         )
-        self.qos.notify_start_of_request(request, session)
+        self.qos.notify_start_of_request(
+            request, session, scheduler=self.internal_scheduler
+        )
         future = self.client.submit(
             worker.submit_workflow,
             key=request.request_uid,
@@ -305,6 +317,11 @@ def run(self) -> None:
                 self.qos.environment.set_session(session_read)
                 with self.session_maker_write() as session_write:
                     self.sync_database(session=session_write)
+
+                if len(self.internal_scheduler.queue) > int(
+                    os.getenv("MAX_SCHEDULER_QUEUE", 0)
+                ):
+                    self.internal_scheduler.run()
                 self.running_requests = len(
                     [
                         future
diff --git a/cads_broker/entry_points.py b/cads_broker/entry_points.py
@@ -104,7 +104,9 @@ def delete_running_requests(
             database.logger.info(f"Deleting {number_of_requests} requests.")
         for request in requests:
             database.logger.info(f"deleting {request.request_uid}...")
-            database.set_request_status(request.request_uid, "dismissed", session=session)
+            database.set_request_status(
+                request.request_uid, "dismissed", session=session
+            )
 
         session.commit()
 
diff --git a/cads_broker/qos/QoS.py b/cads_broker/qos/QoS.py
diff --git a/tests/test_02_database.py b/tests/test_02_database.py
