Repository: dmwm/wmcore
Issue: Top level task must consider pileup location of subsequent tasks

**Impact of the bug**
WMAgent

**Describe the bug**
I still cannot believe this is the current buggy behaviour...
Apparently, when the agent acquires a TaskChain WQE without any input dataset, it creates the top level jobs against all the sites in the SiteWhitelist, even if the subsequent tasks require a pileup dataset.

**How to reproduce it**
Injecting a TaskChain workflow, without any input dataset and a pileup dataset in the Task2.
Assign it to a site containing the pileup and another without it.

**Expected behavior**
Top level jobs should be created against the intersection(SiteWhitelist - SiteBlacklist, PileupLocation). 

**Additional context and error message**
Problem reported in this issue: https://its.cern.ch/jira/browse/CMSCOMPPR-12155
Error message in WMStats is:
```
    Fatal Exception (Exit Code: 8029)

        An exception of category 'NoSecondaryFiles' occurred while
           [0] Constructing the EventProcessor
           [1] Constructing module: class=MixingModule label='mix'
        Exception Message:
        RootEmbeddedFileSequence no input files specified for secondary input source.
```


Changed files:
test/python/WMCore_t/WorkQueue_t/WorkQueue_t.py, src/python/WMCore/WMBS/MySQL/Locations/GetSiteInfo.py, src/python/WMCore/WorkQueue/WMBSHelper.py, test/python/WMCore_t/WorkQueue_t/WMBSHelper_t.py, src/python/WMCore/WorkQueue/WorkQueue.py

diff --git a/src/python/WMCore/WMBS/MySQL/Locations/GetSiteInfo.py b/src/python/WMCore/WMBS/MySQL/Locations/GetSiteInfo.py
index 5f3329b5d..2b4eb06c3 100644
--- a/src/python/WMCore/WMBS/MySQL/Locations/GetSiteInfo.py
+++ b/src/python/WMCore/WMBS/MySQL/Locations/GetSiteInfo.py
@@ -19,10 +19,14 @@ class GetSiteInfo(DBFormatter):
                  INNER JOIN wmbs_location_pnns wls ON wls.location = wl.id
                  INNER JOIN wmbs_pnns wpnn ON wpnn.id = wls.pnn
                  INNER JOIN wmbs_location_state wlst ON wlst.id = wl.state
-               WHERE wl.site_name = :site"""
+          """
 
     def execute(self, siteName=None, conn=None, transaction=False):
-        results = self.dbi.processData(self.sql, {'site': siteName},
+        if not siteName:
+            results = self.dbi.processData(self.sql, conn=conn,
+                                           transaction=transaction)
+        else:
+            sql = self.sql + " WHERE wl.site_name = :site"
+            results = self.dbi.processData(sql, {'site': siteName},
                                        conn=conn, transaction=transaction)
-        formatted = self.formatDict(results)
-        return formatted
+        return self.formatDict(results)
diff --git a/src/python/WMCore/WorkQueue/WMBSHelper.py b/src/python/WMCore/WorkQueue/WMBSHelper.py
index 2235b9343..b2cb2ec34 100644
--- a/src/python/WMCore/WorkQueue/WMBSHelper.py
+++ b/src/python/WMCore/WorkQueue/WMBSHelper.py
@@ -153,7 +153,8 @@ class WMBSHelper(WMConnectionBase):
     Interface between the WorkQueue and WMBS.
     """
 
-    def __init__(self, wmSpec, taskName, blockName=None, mask=None, cachepath='.'):
+    def __init__(self, wmSpec, taskName, blockName=None, mask=None,
+                 cachepath='.', commonLocation=None):
         """
         _init_
 
@@ -169,6 +170,7 @@ class WMBSHelper(WMConnectionBase):
         self.topLevelFileset = None
         self.topLevelSubscription = None
         self.topLevelTaskDBSBufferId = None
+        self.commonLocation = commonLocation or []
 
         self.mergeOutputMapping = {}
 
@@ -185,7 +187,6 @@ class WMBSHelper(WMConnectionBase):
         self.setFileAddChecksum = self.daofactory(classname="Files.AddChecksumByLFN")
         self.addFileAction = self.daofactory(classname="Files.Add")
         self.addToFileset = self.daofactory(classname="Files.AddDupsToFileset")
-        self.getLocations = self.daofactory(classname="Locations.ListSites")
         self.getLocationInfo = self.daofactory(classname="Locations.GetSiteInfo")
 
         # DAOs from DBSBuffer
@@ -348,18 +349,14 @@ class WMBSHelper(WMConnectionBase):
             if self.mask and self.mask.get(key) is None:
                 msg = 'Invalid value "%s" for %s' % (self.mask.get(key), key)
                 raise WorkQueueWMBSException(msg)
+
         locations = set()
-        for site in self.getLocations.execute(conn=self.getDBConn(),
-                                              transaction=self.existingTransaction()):
-            try:
-                siteInfo = self.getLocationInfo.execute(site, conn=self.getDBConn(),
-                                                        transaction=self.existingTransaction())
-                if not siteInfo:
-                    self.logger.info('Skipping MonteCarlo injection to site "%s" as unknown to wmbs' % site)
-                    continue
-                locations.add(siteInfo[0]['pnn'])
-            except Exception as ex:
-                self.logger.error('Error getting storage element for "%s": %s' % (site, str(ex)))
+        siteInfo = self.getLocationInfo.execute(conn=self.getDBConn(),
+                                                transaction=self.existingTransaction())
+        for site in siteInfo:
+            if site['pnn'] in self.commonLocation:
+                locations.add(site['pnn'])
+
         if not locations:
             msg = 'No locations to inject Monte Carlo work to, unable to proceed'
             raise WorkQueueWMBSException(msg)
diff --git a/src/python/WMCore/WorkQueue/WorkQueue.py b/src/python/WMCore/WorkQueue/WorkQueue.py
index ae42ea0d1..71c4ad32d 100644
--- a/src/python/WMCore/WorkQueue/WorkQueue.py
+++ b/src/python/WMCore/WorkQueue/WorkQueue.py
@@ -31,6 +31,7 @@ from WMCore.Services.WorkQueue.WorkQueue import WorkQueue as WorkQueueDS
 from WMCore.WMSpec.WMWorkload import WMWorkloadHelper, getWorkloadFromTask
 from WMCore.WorkQueue.DataLocationMapper import WorkQueueDataLocationMapper
 from WMCore.WorkQueue.DataStructs.ACDCBlock import ACDCBlock
+from WMCore.WorkQueue.DataStructs.WorkQueueElement import possibleSites
 from WMCore.WorkQueue.DataStructs.WorkQueueElementsSummary import getGlobalSiteStatusSummary
 from WMCore.WorkQueue.Policy.End import endPolicy
 from WMCore.WorkQueue.Policy.Start import startPolicy
@@ -427,11 +428,16 @@ class WorkQueue(WorkQueueBase):
         from WMCore.WorkQueue.WMBSHelper import WMBSHelper
         # the parent element (from local couch) can be fetch via:
         # curl -ks -X GET 'http://localhost:5984/workqueue/<ParentQueueId>'
-        msg = "Running WMBS preparation for %s with ParentQueueId %s"
-        self.logger.info(msg, match['RequestName'], match['ParentQueueId'])
+
+        # Keep in mind that WQE contains sites, wmbs location contains pnns
+        commonSites = possibleSites(match)
+        commonLocation = self.cric.PSNstoPNNs(commonSites)
+        msg = "Running WMBS preparation for %s with ParentQueueId %s,\\n  with common location %s"
+        self.logger.info(msg, match['RequestName'], match['ParentQueueId'], commonLocation)
 
         mask = match['Mask']
-        wmbsHelper = WMBSHelper(wmspec, match['TaskName'], blockName, mask, self.params['CacheDir'])
+        wmbsHelper = WMBSHelper(wmspec, match['TaskName'], blockName, mask,
+                                self.params['CacheDir'], commonLocation)
 
         sub, match['NumOfFilesAdded'] = wmbsHelper.createSubscriptionAndAddFiles(block=dbsBlock)
         self.logger.info("Created top level subscription %s for %s with %s files",
diff --git a/test/python/WMCore_t/WorkQueue_t/WMBSHelper_t.py b/test/python/WMCore_t/WorkQueue_t/WMBSHelper_t.py
index 6591240dc..3a0ba37f5 100644
--- a/test/python/WMCore_t/WorkQueue_t/WMBSHelper_t.py
+++ b/test/python/WMCore_t/WorkQueue_t/WMBSHelper_t.py
@@ -529,12 +529,13 @@ class WMBSHelperTest(EmulatedUnitTestCase):
         # dbsDict = {self.inputDataset.dbsurl : self.dbs}
         return dbs
 
-    def createWMBSHelperWithTopTask(self, wmspec, block, mask=None,
-                                    parentFlag=False, detail=False):
+    def createWMBSHelperWithTopTask(self, wmspec, block, mask=None, parentFlag=False,
+                                    detail=False, commonLocation=None):
 
         topLevelTask = getFirstTask(wmspec)
 
-        wmbs = WMBSHelper(wmspec, topLevelTask.name(), block, mask, cachepath=self.workDir)
+        wmbs = WMBSHelper(wmspec, topLevelTask.name(), block, mask, cachepath=self.workDir,
+                          commonLocation=commonLocation)
         if block:
             if parentFlag:
                 block = self.dbs.getFileBlockWithParents(block)[block]
@@ -891,6 +892,7 @@ class WMBSHelperTest(EmulatedUnitTestCase):
 
     def testDuplicateSubscription(self):
         """Can't duplicate subscriptions"""
+        siteWhitelist = ["T2_XX_SiteA", "T2_XX_SiteB"]
         # using default wmspec
         block = self.dataset + "#" + BLOCK1
         wmbs = self.createWMBSHelperWithTopTask(self.wmspec, block)
@@ -918,7 +920,8 @@ class WMBSHelperTest(EmulatedUnitTestCase):
         self.setupMCWMSpec()
         mask = Mask(FirstRun=12, FirstLumi=1234, FirstEvent=12345,
                     LastEvent=999995, LastLumi=12345, LastRun=12)
-        wmbs = self.createWMBSHelperWithTopTask(self.wmspec, None, mask)
+        wmbs = self.createWMBSHelperWithTopTask(self.wmspec, None, mask,
+                                                commonLocation=siteWhitelist)
         wmbs.topLevelFileset.loadData()
         numOfFiles = len(wmbs.topLevelFileset.files)
         filesetId = wmbs.topLevelFileset.id
@@ -932,7 +935,8 @@ class WMBSHelperTest(EmulatedUnitTestCase):
         self.assertEqual(numOfFiles, numDbsFiles)
 
         # reinsert subscription - shouldn't create anything new
-        wmbs = self.createWMBSHelperWithTopTask(self.wmspec, None, mask)
+        wmbs = self.createWMBSHelperWithTopTask(self.wmspec, None, mask,
+                                                commonLocation=siteWhitelist)
         wmbs.topLevelFileset.loadData()
         self.assertEqual(numOfFiles, len(wmbs.topLevelFileset.files))
         self.assertEqual(filesetId, wmbs.topLevelFileset.id)
@@ -977,13 +981,15 @@ class WMBSHelperTest(EmulatedUnitTestCase):
         # in BasicProductionWorkload.getProdArgs() but changing it to
         # "reqmgr_config_cache_t" from StdBase test arguments does not fix the
         # situation. testDuplicateSubscription probably has the same issue
+        siteWhitelist = ["T2_XX_SiteA", "T2_XX_SiteB"]
 
         self.setupMCWMSpec()
 
         mask = Mask(FirstRun=12, FirstLumi=1234, FirstEvent=12345,
                     LastEvent=999995, LastLumi=12345, LastRun=12)
 
-        wmbs = self.createWMBSHelperWithTopTask(self.wmspec, None, mask)
+        wmbs = self.createWMBSHelperWithTopTask(self.wmspec, None, mask,
+                                                commonLocation=siteWhitelist)
         subscription = wmbs.topLevelSubscription
         self.assertEqual(1, subscription.exists())
         fileset = subscription['fileset']
diff --git a/test/python/WMCore_t/WorkQueue_t/WorkQueue_t.py b/test/python/WMCore_t/WorkQueue_t/WorkQueue_t.py
index f1f02560c..0559d4735 100644
--- a/test/python/WMCore_t/WorkQueue_t/WorkQueue_t.py
+++ b/test/python/WMCore_t/WorkQueue_t/WorkQueue_t.py
@@ -266,7 +266,7 @@ class WorkQueueTest(WorkQueueTestCase):
 
         # create relevant sites in wmbs
         rc = ResourceControl()
-        site_se_mapping = {'T2_XX_SiteA': 'a.example.com', 'T2_XX_SiteB': 'b.example.com'}
+        site_se_mapping = {'T2_XX_SiteA': 'T2_XX_SiteA', 'T2_XX_SiteB': 'T2_XX_SiteB'}
         for site, se in site_se_mapping.iteritems():
             rc.insertSite(site, 100, 200, se, cmsName=site, plugin="MockPlugin")
             daofactory = DAOFactory(package="WMCore.WMBS",
@@ -874,22 +874,22 @@ class WorkQueueTest(WorkQueueTestCase):
         self.queue.queueWork(processingSpec.specUrl())
         elements = len(self.queue)
         self.queue.updateLocationInfo()
+        self.assertEqual(len(self.queue.status()), NBLOCKS_HICOMM)
         work = self.queue.getWork({'T2_XX_SiteA': 1000, 'T2_XX_SiteB': 1000}, {})
         self.assertEqual(len(self.queue), 0)
         self.assertEqual(len(self.queue.status(status='Running')), elements)
         ids = [x.id for x in work]
+        self.assertEqual(len(ids), NBLOCKS_HICOMM)
         canceled = self.queue.cancelWork(ids)
         self.assertEqual(sorted(canceled), sorted(ids))
-        self.assertEqual(len(self.queue), 0)
         self.assertEqual(len(self.queue.status()), NBLOCKS_HICOMM)
+        self.assertEqual(len(self.queue.status(status='Running')), NBLOCKS_HICOMM)
         self.assertEqual(len(self.queue.statusInbox(status='Canceled')), 1)
 
-        # now cancel a request
+        # create a new request with one fake file
         self.queue.queueWork(self.spec.specUrl())
-        elements = len(self.queue)
-        work = self.queue.getWork({'T2_XX_SiteA': 1000, 'T2_XX_SiteB': 1000},
-                                  {})
-        self.assertEqual(len(self.queue), 0)
+        self.assertEqual(len(self.queue), 1)
+        work = self.queue.getWork({'T2_XX_SiteA': 1000, 'T2_XX_SiteB': 1000}, {})
         self.assertEqual(len(self.queue.status(status='Running')), len(self.queue.status()))
         ids = [x.id for x in work]
         canceled = self.queue.cancelWork(WorkflowName='testProduction')