Reduce scheduler memory (#113)

* try to reduce scheduler memory

* add cache

* fix

* add performance logger to submit_requests

* add perf logger

* add perf logger

* add perflogger

* block adding new queued requests if the table is too big

* revert

* try to improve performance

* truncate instead of delete

* fix truncate

* fix

* improve performance of delete_request_qos_status

* remove wrong queue update

* fix

* adding cascading rules

* style

* remove comments

* qa

---------

Co-authored-by: Alessio Siniscalchi <[email protected]>

Author: francesconazzaro

alembic/versions/ca178571bdc5_cascading_rules_for_requests_rules.py

@@ -0,0 +1,63 @@
+"""cascading rules for requests-rules.
+
+Revision ID: ca178571bdc5
+Revises: 8deb52d20c05
+Create Date: 2024-06-10 17:36:06.551396
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "ca178571bdc5"
+down_revision = "8deb52d20c05"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.drop_constraint(
+        "system_request_qos_rule_rule_uid_fkey", "system_request_qos_rule"
+    )
+    op.drop_constraint(
+        "system_request_qos_rule_request_uid_fkey", "system_request_qos_rule"
+    )
+    op.create_foreign_key(
+        "system_request_qos_rule_rule_uid_fkey",
+        "system_request_qos_rule",
+        "qos_rules",
+        ["rule_uid"],
+        ["uid"],
+        ondelete="CASCADE",
+    )
+    op.create_foreign_key(
+        "system_request_qos_rule_request_uid_fkey",
+        "system_request_qos_rule",
+        "system_requests",
+        ["request_uid"],
+        ["request_uid"],
+        ondelete="CASCADE",
+    )
+
+
+def downgrade() -> None:
+    op.drop_constraint(
+        "system_request_qos_rule_rule_uid_fkey", "system_request_qos_rule"
+    )
+    op.drop_constraint(
+        "system_request_qos_rule_request_uid_fkey", "system_request_qos_rule"
+    )
+    op.create_foreign_key(
+        "system_request_qos_rule_rule_uid_fkey",
+        "system_request_qos_rule",
+        "qos_rules",
+        ["rule_uid"],
+        ["uid"],
+    )
+    op.create_foreign_key(
+        "system_request_qos_rule_request_uid_fkey",
+        "system_request_qos_rule",
+        "system_requests",
+        ["request_uid"],
+        ["request_uid"],
+    )

cads_broker/database.py

@@ -70,10 +70,12 @@ class SystemRequestQoSRule(BaseModel):
 
     request_uid = sa.Column(
         sa.dialects.postgresql.UUID(False),
-        sa.ForeignKey("system_requests.request_uid"),
+        sa.ForeignKey("system_requests.request_uid", ondelete="CASCADE"),
         primary_key=True,
     )
-    rule_uid = sa.Column(sa.Text, sa.ForeignKey("qos_rules.uid"), primary_key=True)
+    rule_uid = sa.Column(
+        sa.Text, sa.ForeignKey("qos_rules.uid", ondelete="CASCADE"), primary_key=True
+    )
 
 
 class Events(BaseModel):
@@ -467,9 +469,11 @@ def get_events_from_request(
 
 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)
+    session.execute(sa.text("truncate qos_rules cascade"))
+    # 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
@@ -484,6 +488,11 @@ def reset_qos_rules(session: sa.orm.Session, qos):
     session.commit()
 
 
+def count_system_request_qos_rule(session: sa.orm.Session) -> int:
+    """Count the number of rows in system_request_qos_rule."""
+    return session.query(SystemRequestQoSRule).count()
+
+
 def get_qos_rule(uid: str, session: sa.orm.Session):
     """Get a QoS rule."""
     statement = sa.select(QoSRule).where(QoSRule.uid == uid)
@@ -557,11 +566,10 @@ def delete_request_qos_status(
             except sqlalchemy.orm.exc.NoResultFound:
                 qos_rule = add_qos_rule(rule=rule, session=session)
                 created_rules[qos_rule.uid] = qos_rule
-        if qos_rule.uid in [r.uid for r in request.qos_rules]:
-            request.qos_rules.remove(qos_rule)
         qos_rule.queued = len(rule.queued)
         qos_rule.running = rule.value
-    return request, created_rules
+    request.qos_rules = []
+    return None, created_rules
 
 
 def add_request_qos_status(

cads_broker/dispatcher.py

@@ -79,6 +79,7 @@ def get_tasks_from_scheduler(client: distributed.Client) -> Any:
 
     This function is executed on the scheduler pod.
     """
+
     def get_tasks_on_scheduler(dask_scheduler: distributed.Scheduler) -> dict[str, Any]:
         tasks = {}
         for task_id, task in dask_scheduler.tasks.items():
@@ -177,6 +178,7 @@ def len(self) -> int:
         with self._lock:
             return len(self.queue_dict)
 
+    @cachetools.cachedmethod(lambda self: cachetools.TTLCache(maxsize=1024, ttl=60))
     def reset(self) -> None:
         with self._lock:
             self.queue_dict = dict()
@@ -235,8 +237,8 @@ def from_address(
             rules_hash=rules_hash,
         )
         with session_maker_write() as session:
-            qos.reload_rules(session=session)
-            db.reset_qos_rules(session, qos)
+            perf_logger(qos.reload_rules)(session=session)
+            perf_logger(db.reset_qos_rules)(session, qos)
         self = cls(
             client=client,
             session_maker_read=session_maker_read,
@@ -263,7 +265,7 @@ def set_request_error_status(
 
         If the error reason is "KilledWorker":
             - if the worker has been killed by the Nanny for memory usage, it add the event for the user
-            - if the worker is killed for unknown reasons, it re-queues the request 
+            - if the worker is killed for unknown reasons, it re-queues the request
               if the requeue limit is not reached. This is configurable with the environment variable
         """
         error_message = "".join(traceback.format_exception(exception))
@@ -327,10 +329,10 @@ def sync_database(self, session: sa.orm.Session) -> None:
 
         - If the task is in the futures list it does nothing.
         - If the task is not in the futures list but it is in the scheduler:
-            - If the task is in memory (it is successful but it has been lost by the broker), 
+            - If the task is in memory (it is successful but it has been lost by the broker),
               it is set to successful.
             - If the task is in error, it is set to failed.
-        - If the task is not in the dask scheduler, it is re-queued. 
+        - If the task is not in the dask scheduler, it is re-queued.
           This behaviour can be changed with an environment variable.
         """
         # the retrieve API sets the status to "dismissed", here the broker deletes the request
@@ -444,15 +446,15 @@ def sync_qos_rules(self, session_write) -> None:
             - the requests from the self.queue.
               If a request is updated the relative self.queue entry is updated too
         """
-        qos_rules = db.get_qos_rules(session=session_write)
+        qos_rules = perf_logger(db.get_qos_rules)(session=session_write)
         if tasks_number := len(self.internal_scheduler.queue):
             logger.info("performance", tasks_number=tasks_number)
         for task in list(self.internal_scheduler.queue)[
             : int(os.getenv("BROKER_MAX_INTERNAL_SCHEDULER_TASKS", 500))
         ]:
             # the internal scheduler is used to asynchronously add qos rules to database
             # it returns a new qos rule if a new qos rule is added to database
-            request, new_qos_rules = task["function"](
+            request, new_qos_rules = perf_logger(task["function"])(
                 session=session_write,
                 request=self.queue.get(task["kwargs"].get("request_uid")),
                 rules_in_db=qos_rules,
@@ -527,6 +529,7 @@ def on_future_done(self, future: distributed.Future) -> str:
             future.release()
         return future.key
 
+    @perf_logger
     def submit_requests(
         self,
         session_write: sa.orm.Session,
@@ -568,10 +571,7 @@ def submit_request(
                 request_uid=request.request_uid,
                 user_uid=request.user_uid,
                 hostname=os.getenv("CDS_PROJECT_URL"),
-                **request.adaptor_properties.config,
             ),
-            form=request.adaptor_properties.form,
-            request=request.request_body.get("request", {}),
             resources=request.request_metadata.get("resources", {}),
             metadata=request.request_metadata,
         )