Faster overwrite for NanoAODFromDAS method

core/python/Sample.py

@@ -253,9 +253,10 @@ def fromDirectory(cls, name, directory, treeName = "Events", normalization = Non
     @classmethod
     def nanoAODfromDAS(cls, name, DASname, instance = 'global', redirector='root://hephyse.oeaw.ac.at/', dbFile=None, overwrite=False, treeName = "Events", maxN = None, \
             selectionString = None, weightString = None, xSection=-1,
-            isData = False, color = 0, texName = None, multithreading=True, genWeight='genWeight', json=None):
+            isData = False, color = 0, texName = None, multithreading=True, genWeight='genWeight', json=None, localSite='T2_AT_Vienna'):
         '''
         get nanoAOD from DAS and make a local copy on afs 
+        if overwrite is true, old entries will be overwritten, no matter what the old entry contains. if overwrite=='update', file-list and normalization are checked, and only if they potentially changed the old entry is overwritten.
         '''
         from multiprocessing import Pool
         from RootTools.fwlite.Database import Database
@@ -272,16 +273,23 @@ def nanoAODfromDAS(cls, name, DASname, instance = 'global', redirector='root://h
         else:
             cache = None
 
+        # first check if there are already files in the cache
+        normalizationFromCache = 0.
+        if n_cache_files:
+            filesFromCache          = [ f["value"] for f in cache.getDicts({'name':name, 'DAS':DASname}) ]
+            normalizationFromCache  = cache.getDicts({'name':name, 'DAS':DASname})[0]["normalization"]
+        else:
+            filesFromCache = []
+
+        # if we don't want to overwrite, and there's a filelist in the cache we're already done
         if n_cache_files and not overwrite:
-            files = [ f["value"] for f in cache.getDicts({'name':name, 'DAS':DASname}) ]
-            normalization = cache.getDicts({'name':name, 'DAS':DASname})[0]["normalization"]
+            files           = filesFromCache
+            normalization   = normalizationFromCache
 
             logger.info('Found sample %s in cache %s, return %i files.', name, dbFile, len(files))
 
         else:
-            if overwrite:
-                cache.removeObjects({"name":name, 'DAS':DASname})
-
+            # only entered if overwrite is not set or sample not in the cache yet
             def _dasPopen(dbs):
                 if 'LSB_JOBID' in os.environ:
                     raise RuntimeError, "Trying to do a DAS query while in a LXBatch job (env variable LSB_JOBID defined)\nquery was: %s" % dbs
@@ -302,23 +310,62 @@ def _dasPopen(dbs):
                     filename = redirector+'/'+line
                     files.append(filename)
 
-            if DASname.endswith('SIM') or not 'Run20' in DASname:
-                # need to read the proper normalization for MC
-                logger.info("Reading normalization. This is slow, so grab a coffee.")
-                tmp_sample = cls(name=name, files=files, treeName = treeName, selectionString = selectionString, weightString = weightString,
-                    isData = isData, color=color, texName = texName, xSection = xSection, normalization=1)
-                normalization = tmp_sample.getYieldFromDraw('(1)', genWeight)['val']
-                logger.info("Got normalization %s", normalization)
+            if (sorted(files) == sorted(filesFromCache)) and float(normalizationFromCache) > 0.0 and overwrite=='update':
+                # if the files didn't change we don't need to read the normalization again (slowest part!). If the norm was 0 previously, also get it again.
+                logger.info("File list for %s didn't change. Skipping.", name)
+                normalization = normalizationFromCache
+                logger.info('Sample %s from cache %s returned %i files.', name, dbFile, len(files))
+
             else:
-                # for data, we can just use the number of events, although no normalization is needed anyway.
-                dbs='dasgoclient -query="summary %s=%s instance=prod/%s" --format=json'%(qwhat,query, instance)
-                jdata = json.load(_dasPopen(dbs))['data'][0]['summary'][0]
-                normalization = int(jdata['nevents'])
-
-        if overwrite or n_cache_files<1:
-            for f in files:
-                if cache is not None:
-                    cache.add({"name":name, 'DAS':DASname, 'normalization':str(normalization)}, f, save=True)
+                if overwrite:
+                    # remove old entry
+                    cache.removeObjects({"name":name, 'DAS':DASname})
+                    logger.info("Removed old DB entry.")
+
+                if instance == 'global':
+                    # check if dataset is available in local site, otherwise don't read a normalization
+                    dbs='dasgoclient -query="site %s=%s instance=prod/%s" --format=json'%(qwhat,query, instance)
+                    jdata = json.load(_dasPopen(dbs))
+
+                    filesOnLocalT2 = False
+                    for d in jdata['data']:
+                        if d['site'][0]['name'] == localSite and d['site'][0].has_key('replica_fraction'):
+                            fraction = d['site'][0]['replica_fraction']
+                            if float(str(fraction).replace('%','')) < 100.:
+                                filesOnLocalT2 = False
+                                break
+                            else:
+                                filesOnLocalT2 = True
+                else:
+                    # if we produced the samples ourselves we don't need to check this
+                    filesOnLocalT2 = True
+
+                if filesOnLocalT2:
+                    logger.info("Files are available at %s", localSite)
+
+                if DASname.endswith('SIM') or not 'Run20' in DASname:
+                    # need to read the proper normalization for MC
+                    if filesOnLocalT2:
+                        logger.info("Reading normalization. This is slow, so grab a coffee.")
+                        tmp_sample = cls(name=name, files=files, treeName = treeName, selectionString = selectionString, weightString = weightString,
+                            isData = isData, color=color, texName = texName, xSection = xSection, normalization=1)
+                        normalization = tmp_sample.getYieldFromDraw('(1)', genWeight)['val']
+                        logger.info("Got normalization %s", normalization)
+                    else:
+                        logger.info("Files only available on the grid. Can't read a normalization.")
+                        normalization = 0
+                else:
+                    # for data, we can just use the number of events, although no normalization is needed anyway.
+                    dbs='dasgoclient -query="summary %s=%s instance=prod/%s" --format=json'%(qwhat,query, instance)
+                    jdata = json.load(_dasPopen(dbs))['data'][0]['summary'][0]
+                    normalization = int(jdata['nevents'])
+
+                for f in files:
+                    if cache is not None:
+                        cache.add({"name":name, 'DAS':DASname, 'normalization':str(normalization)}, f, save=True)
+
+                logger.info('Found sample %s in cache %s, return %i files.', name, dbFile, len(files))
+
 
         if limit>0: files=files[:limit]
         sample = cls(name=name, files=files, treeName = treeName, selectionString = selectionString, weightString = weightString,

plot/python/plotting.py

@@ -462,7 +462,7 @@ def draw(plot, \
 
     for o in drawObjects:
         if o:
-            if type(o) in [ ROOT.TF1 ]:
+            if type(o) in [ ROOT.TF1, ROOT.TGraph ]:
                 o.Draw('same')
             else:
                 o.Draw()