[9.1] Aggregate DataStore Accounting Requests

docs/source/AdministratorGuide/Systems/RequestManagement/concepts.rst

@@ -33,3 +33,4 @@ CleanReqDBAgent
 ---------------
 
 Because the database can grow very large, the :ref:`CleanReqDBAgent` is in charge of removing old `Requests` in a final state.
+It is also capable of canceling requests if they are there for too long, or can group Requests together.

src/DIRAC/RequestManagementSystem/Agent/CleanReqDBAgent.py

@@ -1,8 +1,3 @@
-########################################################################
-# File: CleanReqDBAgent.py
-# Author: Krzysztof.Ciba@NOSPAMgmail.com
-# Date: 2013/05/17 08:31:26
-########################################################################
 """Cleaning the RequestDB from obsolete records and kicking assigned requests
 
 .. literalinclude:: ../ConfigTemplate.cfg
@@ -11,15 +6,8 @@
   :dedent: 2
   :caption: CleanReqDBAgent options
 
-.. moduleauthor:: Krzysztof.Ciba@NOSPAMgmail.com
-
 """
 
-# #
-# @file CleanReqDBAgent.py
-# @author Krzysztof.Ciba@NOSPAMgmail.com
-# @date 2013/05/17 08:32:08
-# @brief Definition of CleanReqDBAgent class.
 
 # # imports
 import datetime
@@ -54,6 +42,12 @@ class CleanReqDBAgent(AgentModule):
     # # remove failed requests flag
     DEL_FAILED = False
 
+    # # Number of Accounting requests to fetch to batch
+    # # 0 to disable
+    ACCOUNTING_BATCH_MAX_REQUESTS = 0
+    # # Number of operations to batch in a single requests.
+    ACCOUNTING_BATCH_SIZE = 100
+
     # # request db
     __requestDB = None
 
@@ -72,6 +66,14 @@ def initialize(self):
         self.KICK_LIMIT = self.am_getOption("KickLimit", self.KICK_LIMIT)
         self.log.info(f"Kick limit = {self.KICK_LIMIT} request/cycle")
 
+        self.ACCOUNTING_BATCH_MAX_REQUESTS = self.am_getOption(
+            "AccountingBatchMaxRequests", self.ACCOUNTING_BATCH_MAX_REQUESTS
+        )
+        self.log.info(f"Accounting max requests = {self.ACCOUNTING_BATCH_MAX_REQUESTS} request/cycle")
+
+        self.ACCOUNTING_BATCH_SIZE = self.am_getOption("AccountingBatchSize", self.ACCOUNTING_BATCH_SIZE)
+        self.log.info(f"Accouting batch size = {self.ACCOUNTING_BATCH_SIZE} requests")
+
         if self.cancelGraceDays >= self.DEL_GRACE_DAYS:
             self.cancelGraceDays = self.DEL_GRACE_DAYS - 1
             self.log.warn("Cancelled jobs grace period > delete period, capping to %u days" % self.cancelGraceDays)
@@ -86,6 +88,7 @@ def execute(self):
         now = datetime.datetime.utcnow()
         kickTime = now - datetime.timedelta(hours=self.KICK_GRACE_HOURS)
         rmTime = now - datetime.timedelta(days=self.DEL_GRACE_DAYS)
+        batched = None
 
         # # kick
         statusList = ["Assigned"]
@@ -161,9 +164,23 @@ def execute(self):
                         continue
                     cancelled += 1
 
+        if self.ACCOUNTING_BATCH_MAX_REQUESTS:
+            res = self.__requestDB.aggregateAccountingDataStoreRequests(
+                batch_size=self.ACCOUNTING_BATCH_SIZE, max_requests=self.ACCOUNTING_BATCH_MAX_REQUESTS
+            )
+            if not res["OK"]:
+                self.log.error("Failed to natch accounting requests requests:", res["Message"])
+                return res
+            batched = res["Value"]
+
         self.log.info("execute: kicked assigned requests", str(kicked))
         self.log.info("execute: deleted finished requests", str(deleted))
         if self.cancelGraceDays > 0:
             self.log.info("execute: cancelled overdue requests", str(cancelled))
 
+        if batched:
+            self.log.info(
+                "Batched accounting requests",
+                f"created : {len(batched['created_requests'])}, canceled: {len(batched['canceled_requests'])}",
+            )
         return S_OK()

src/DIRAC/RequestManagementSystem/ConfigTemplate.cfg

@@ -140,6 +140,12 @@ Agents
      # regardless of State
      # if set to 0 (default) Requests are never cancelled
      CancelGraceDays = 0
+
+    # Number of Accounting requests to fetch to batch
+    # 0 to disable (default)
+     AccountingBatchMaxRequests = 0
+     # Number of operations to batch in a single requests.
+     AccountingBatchSize = 100
   }
   ##END
 }

src/DIRAC/RequestManagementSystem/DB/RequestDB.py

@@ -16,6 +16,7 @@
 import datetime
 import errno
 import random
+import uuid
 from urllib.parse import quote_plus
 
 from sqlalchemy import (
@@ -25,13 +26,15 @@
     DateTime,
     Enum,
     ForeignKey,
+    Index,
     Integer,
     MetaData,
     String,
     Table,
     create_engine,
     distinct,
     func,
+    insert,
     inspect,
 )
 from sqlalchemy.exc import SQLAlchemyError
@@ -154,6 +157,7 @@
     Column("RequestID", Integer, primary_key=True),
     Column("SourceComponent", String(255)),
     Column("NotBefore", DateTime),
+    Index("idx_request_status_notbefore_lastupdate", "Status", "NotBefore", "LastUpdate"),
     mysql_engine="InnoDB",
 )
 
@@ -624,6 +628,114 @@ def deleteRequest(self, requestID):
 
         return S_OK()
 
+    def aggregateAccountingDataStoreRequests(self, batch_size=100, max_requests=10000):
+        """Aggregate Accounting.DataStore requests with ForwardDISET operations into batches.
+
+        Groups requests by Owner/OwnerGroup and processes up to batch_size operations per new request.
+        ForwardDISET operations are copied to the new aggregated request. Original requests are set to
+        "Assigned" then "Canceled", with guarded state transitions.
+
+        :param int batch_size: Number of operations to batch together (default: 100)
+        :param int max_requests: Maximum number of source requests to process per execution (default: 10000)
+        :return: S_OK({"created_requests": [new_request_ids], "canceled_requests": [old_request_ids]}) or S_ERROR
+        """
+        created_requests = []
+        canceled_requests = []
+        session = self.DBSession()
+        try:
+            now = datetime.datetime.utcnow()
+
+            rows = (
+                session.query(
+                    requestTable.c.RequestID.label("old_request_id"),
+                    requestTable.c.Owner,
+                    requestTable.c.OwnerGroup,
+                    operationTable.c.OperationID,
+                    operationTable.c.TargetSE,
+                    operationTable.c.CreationTime,
+                    operationTable.c.SourceSE,
+                    operationTable.c.Arguments,
+                    operationTable.c.Error,
+                    operationTable.c.Type,
+                    operationTable.c.Status,
+                    operationTable.c.LastUpdate,
+                    operationTable.c.SubmitTime,
+                    operationTable.c.Catalog,
+                )
+                .join(operationTable, requestTable.c.RequestID == operationTable.c.RequestID)
+                .filter(requestTable.c.RequestName.like("Accounting.DataStore%"))
+                .filter(requestTable.c.Status == "Waiting")
+                .order_by(requestTable.c.Owner, requestTable.c.OwnerGroup, requestTable.c.RequestID)
+                .with_for_update(skip_locked=True)
+                .limit(max_requests)
+                .all()
+            )
+            if not rows:
+                return S_OK({"created_requests": [], "canceled_requests": []})
+
+            grouped_rows = {}
+            for row in rows:
+                grouped_rows.setdefault((row.Owner, row.OwnerGroup), []).append(row)
+            for (owner, owner_group), group_rows in grouped_rows.items():
+                for i in range(0, len(group_rows), batch_size):
+                    batch = group_rows[i : i + batch_size]
+                    old_request_ids = sorted({row.old_request_id for row in batch})
+
+                    new_request_id = session.execute(
+                        insert(requestTable).values(
+                            RequestName=f"Aggregated_Accounting.DataStore_{uuid.uuid4()}",
+                            Status="Waiting",
+                            Owner=owner,
+                            OwnerGroup=owner_group,
+                            CreationTime=now,
+                            LastUpdate=now,
+                            SubmitTime=now,
+                            NotBefore=now,
+                            SourceComponent="Aggregation",
+                        )
+                    ).inserted_primary_key[0]
+
+                    operation_rows = [
+                        {
+                            "TargetSE": row.TargetSE,
+                            "CreationTime": row.CreationTime,
+                            "SourceSE": row.SourceSE,
+                            "Arguments": row.Arguments,
+                            "Error": row.Error,
+                            "Type": row.Type,
+                            "Order": order,
+                            "Status": row.Status,
+                            "LastUpdate": row.LastUpdate,
+                            "SubmitTime": row.SubmitTime,
+                            "Catalog": row.Catalog,
+                            "RequestID": new_request_id,
+                        }
+                        for order, row in enumerate(batch)
+                    ]
+                    session.execute(insert(operationTable), operation_rows)
+
+                    canceled_result = session.execute(
+                        update(requestTable)
+                        .where(requestTable.c.RequestID.in_(old_request_ids))
+                        .values(Status="Canceled", LastUpdate=now)
+                    )
+
+                    if canceled_result.rowcount != len(old_request_ids):
+                        raise RuntimeError("State transition mismatch while canceling requests")
+
+                    created_requests.append(new_request_id)
+                    canceled_requests.extend(old_request_ids)
+
+            session.commit()
+        except Exception as e:
+            session.rollback()
+            self.log.exception("aggregateAccountingDataStoreRequests: unexpected exception", lException=e)
+            return S_ERROR(f"aggregateAccountingDataStoreRequests: unexpected exception {e}")
+        finally:
+            session.close()
+
+        return S_OK({"created_requests": created_requests, "canceled_requests": canceled_requests})
+
     def getDBSummary(self):
         """get db summary"""
         # # this will be returned
@@ -796,7 +908,7 @@ def getRequestCountersWeb(self, groupingAttribute, selectDict):
         resultDict = {}
 
         session = self.DBSession()
-
+        breakpoint()
         try:
             if groupingAttribute == "Type":
                 groupingColumn = self._get_column("Operation", "Type")

src/DIRAC/RequestManagementSystem/DB/test/Test_RequestDB.py

@@ -1,5 +1,6 @@
-""" This runs the RMS scenari as unit test for the DB. For that,
-it replaces the normal MySQL connection with an inmemory SQLite db
+"""This runs the RMS scenari as unit test for the DB.
+
+For that, it replaces the normal MySQL connection with an in-memory SQLite db.
 """
 
 # pylint: disable=invalid-name,wrong-import-position
@@ -76,3 +77,51 @@ def test_web_queries_reject_unknown_attributes(reqDB):
     invalid_distinct = reqDB.getDistinctValues("Request", "__class__")
     assert not invalid_distinct["OK"], invalid_distinct
     assert invalid_distinct["Message"] == "Unknown Request attribute '__class__'"
+
+
+def _create_accounting_request(reqDB, suffix):
+    request = Request({"RequestName": f"Accounting.DataStore.{suffix}", "Owner": "owner", "OwnerGroup": "group"})
+    operation = Operation({"Type": "ForwardDISET", "Arguments": f"payload-{suffix}", "TargetSE": "CERN-USER"})
+    request += operation
+
+    result = reqDB.putRequest(request)
+    assert result["OK"], result
+    return result["Value"]
+
+
+def test_aggregate_accounting_requests_copy_forwarddiset_operations(reqDB):
+    original_request_id = _create_accounting_request(reqDB, "direct-db")
+
+    result = reqDB.aggregateAccountingDataStoreRequests(batch_size=100, max_requests=10000)
+    assert result["OK"], result
+    assert len(result["Value"]["created_requests"]) == 1, result
+    assert result["Value"]["canceled_requests"] == [original_request_id], result
+
+    aggregated_request_id = result["Value"]["created_requests"][0]
+
+    session = reqDB.DBSession()
+    try:
+        requests = session.query(Request).order_by(Request.RequestID).all()  # pylint: disable=no-member
+        assert len(requests) == 2
+
+        original_request = next(req for req in requests if req.RequestID == original_request_id)
+        aggregated_request = next(req for req in requests if req.RequestID == aggregated_request_id)
+
+        assert original_request._Status == "Canceled"
+        assert aggregated_request._Status == "Waiting"
+        assert len(original_request.__operations__) == 1
+        assert len(aggregated_request.__operations__) == 1
+
+        original_operation = original_request.__operations__[0]
+        aggregated_operation = aggregated_request.__operations__[0]
+
+        assert original_operation.OperationID != aggregated_operation.OperationID
+        assert original_operation.RequestID == original_request_id
+        assert aggregated_operation.RequestID == aggregated_request_id
+        assert original_operation.Type == aggregated_operation.Type == "ForwardDISET"
+        assert original_operation.Arguments == aggregated_operation.Arguments
+        assert original_operation.TargetSE == aggregated_operation.TargetSE
+        assert len(original_operation.__files__) == 0
+        assert len(aggregated_operation.__files__) == 0
+    finally:
+        session.close()