feat (TS): RequestTasks use getBulkRequestStatus

DIRACGrid · Aug 12, 2024 · 7d9bfd6 · 7d9bfd6
1 parent b903955
commit 7d9bfd6
Showing 1 changed file with 39 additions and 16 deletions.
diff --git a/src/DIRAC/TransformationSystem/Client/RequestTasks.py b/src/DIRAC/TransformationSystem/Client/RequestTasks.py
@@ -317,26 +317,49 @@ def getSubmittedTaskStatus(self, taskDicts):
         Check if tasks changed status, and return a list of tasks per new status
         """
         updateDict = {}
-        badRequestID = 0
+        externalIDs = [
+            int(taskDict["ExternalID"])
+            for taskDict in taskDicts
+            if taskDict["ExternalID"] and int(taskDict["ExternalID"])
+        ]
+
+        # Count how many tasks don't have an valid external ID
+        badRequestID = len(taskDicts) - len(externalIDs)
+
+        res = self.requestClient.getBulkRequestStatus(externalIDs)
+        if not res["OK"]:
+            # We need a transformationID for the log, and although we expect a single one,
+            # do things ~ properly
+            tids = list({taskDict["TransformationID"] for taskDict in taskDicts})
+            try:
+                tid = tids[0]
+            except IndexError:
+                tid = 0
+
+            self._logWarn(
+                "getSubmittedTaskStatus: Failed to get bulk requestIDs",
+                res["Message"],
+                transID=tid,
+            )
+            return S_OK({})
+        new_statuses = res["Value"]
+
         for taskDict in taskDicts:
             oldStatus = taskDict["ExternalStatus"]
             # ExternalID is normally a string
-            if taskDict["ExternalID"] and int(taskDict["ExternalID"]):
-                newStatus = self.requestClient.getRequestStatus(taskDict["ExternalID"])
-                if not newStatus["OK"]:
-                    log = self._logVerbose if "not exist" in newStatus["Message"] else self._logWarn
-                    log(
-                        "getSubmittedTaskStatus: Failed to get requestID for request",
-                        newStatus["Message"],
-                        transID=taskDict["TransformationID"],
-                    )
-                else:
-                    newStatus = newStatus["Value"]
-                    # We don't care updating the tasks to Assigned while the request is being processed
-                    if newStatus != oldStatus and newStatus != "Assigned":
-                        updateDict.setdefault(newStatus, []).append(taskDict["TaskID"])
+
+            newStatus = new_statuses.get(taskDict["ExternalID"])
+            if not newStatus:
+                self._logVerbose(
+                    "getSubmittedTaskStatus: Failed to get requestID for request",
+                    f"No such RequestID {taskDict['ExternalID']}",
+                    transID=taskDict["TransformationID"],
+                )
             else:
-                badRequestID += 1
+                # We don't care updating the tasks to Assigned while the request is being processed
+                if newStatus != oldStatus and newStatus != "Assigned":
+                    updateDict.setdefault(newStatus, []).append(taskDict["TaskID"])
+
         if badRequestID:
             self._logWarn("%d requests have identifier 0" % badRequestID)
         return S_OK(updateDict)