|
11 | 11 | from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getVOForGroup |
12 | 12 | from DIRAC.Core.Base.API import API |
13 | 13 | from DIRAC.Core.Security.ProxyInfo import getProxyInfo |
| 14 | +from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader |
14 | 15 | from DIRAC.Core.Utilities.PromptUser import promptUser |
15 | 16 | from DIRAC.FrameworkSystem.Client.NotificationClient import NotificationClient |
16 | 17 | from DIRAC.FrameworkSystem.Client.ProxyManagerClient import gProxyManager |
|
19 | 20 | from DIRAC.WorkloadManagementSystem.Client.JobManagerClient import JobManagerClient |
20 | 21 | from DIRAC.WorkloadManagementSystem.Client.PilotManagerClient import PilotManagerClient |
21 | 22 | from DIRAC.WorkloadManagementSystem.Client.WMSAdministratorClient import WMSAdministratorClient |
| 23 | +from DIRAC.WorkloadManagementSystem.Service.WMSUtilities import ( |
| 24 | + killPilotsInQueues, |
| 25 | +) |
22 | 26 |
|
23 | 27 | voName = "" |
24 | 28 | ret = getProxyInfo(disableVOMS=True) |
@@ -452,8 +456,28 @@ def killPilot(self, gridReference): |
452 | 456 | if not isinstance(gridReference, (str, list)): |
453 | 457 | return self._errorReport("Expected string or list of strings for pilot reference") |
454 | 458 |
|
455 | | - result = PilotManagerClient().killPilot(gridReference) |
456 | | - return result |
| 459 | + # Make a list if it is not yet |
| 460 | + if isinstance(gridReference, str): |
| 461 | + gridReference = [gridReference] |
| 462 | + |
| 463 | + # Regroup pilots per site |
| 464 | + pilotRefDict = {} |
| 465 | + for pilotReference in gridReference: |
| 466 | + result = PilotManagerClient().getPilotInfo(pilotReference) |
| 467 | + if not result["OK"] or not result["Value"]: |
| 468 | + return S_ERROR(f"Failed to get info for pilot {pilotReference}") |
| 469 | + |
| 470 | + pilotDict = result["Value"][pilotReference] |
| 471 | + queue = "@@@".join( |
| 472 | + [pilotDict["VO"], pilotDict["GridSite"], pilotDict["DestinationSite"], pilotDict["Queue"]] |
| 473 | + ) |
| 474 | + gridType = pilotDict["GridType"] |
| 475 | + pilotRefDict.setdefault(queue, {}) |
| 476 | + pilotRefDict[queue].setdefault("PilotList", []) |
| 477 | + pilotRefDict[queue]["PilotList"].append(pilotReference) |
| 478 | + pilotRefDict[queue]["GridType"] = gridType |
| 479 | + |
| 480 | + return killPilotsInQueues(pilotRefDict) |
457 | 481 |
|
458 | 482 | ############################################################################# |
459 | 483 | def getPilotLoggingInfo(self, gridReference): |
|
0 commit comments